Skip to content

Commit

Permalink
fix: args override config and dsn (#416)
Browse files Browse the repository at this point in the history
  • Loading branch information
sundy-li committed May 6, 2024
1 parent f6865a8 commit cff0252
Show file tree
Hide file tree
Showing 5 changed files with 55 additions and 44 deletions.
5 changes: 0 additions & 5 deletions cli/src/display.rs
Original file line number Diff line number Diff line change
Expand Up @@ -35,7 +35,6 @@ use crate::{
#[async_trait::async_trait]
pub trait ChunkDisplay {
async fn display(&mut self) -> Result<ServerStats>;
fn total_rows(&self) -> usize;
}

pub struct FormatDisplay<'a> {
Expand Down Expand Up @@ -342,10 +341,6 @@ impl<'a> ChunkDisplay for FormatDisplay<'a> {
let stats = self.stats.take().unwrap_or_default();
Ok(stats)
}

fn total_rows(&self) -> usize {
self.rows
}
}

fn format_read_progress(ss: &ServerStats, elapsed: f64) -> String {
Expand Down
77 changes: 40 additions & 37 deletions cli/src/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -227,7 +227,7 @@ where

#[tokio::main]
pub async fn main() -> Result<()> {
let mut config = Config::load();
let config = Config::load();

let args = Args::parse();
let mut cmd = Args::command();
Expand All @@ -237,7 +237,7 @@ pub async fn main() -> Result<()> {
}

let mut conn_args = match args.dsn {
Some(dsn) => {
Some(ref dsn) => {
if args.host.is_some() {
eprintln!("warning: --host is ignored when --dsn is set");
}
Expand All @@ -256,51 +256,54 @@ pub async fn main() -> Result<()> {
ConnectionArgs::from_dsn(dsn.inner())?
}
None => {
if let Some(host) = args.host {
config.connection.host = host;
}
if let Some(port) = args.port {
config.connection.port = Some(port);
}
if !args.tls {
config
.connection
.args
.insert("sslmode".to_string(), "disable".to_string());
let host = args.host.unwrap_or_else(|| config.connection.host.clone());
let mut port = config.connection.port;
if args.port.is_some() {
port = args.port;
}

let user = args.user.unwrap_or_else(|| config.connection.user.clone());
let password = args.password.unwrap_or_else(|| SensitiveString::from(""));

ConnectionArgs {
host: config.connection.host.clone(),
port: config.connection.port,
user: config.connection.user.clone(),
password: SensitiveString::from(""),
host,
port,
user,
password,
database: config.connection.database.clone(),
flight: args.flight,
args: config.connection.args.clone(),
}
}
};
// override database if specified in command line
if args.database.is_some() {
conn_args.database = args.database;
}
// override user if specified in command line
if let Some(user) = args.user {
config.connection.user = user;
}
// override password if specified in command line
if let Some(password) = args.password {
conn_args.password = password;
}
// override role if specified in command line
if let Some(role) = args.role {
config.connection.args.insert("role".to_string(), role);
}
// override args if specified in command line
for (k, v) in args.set {
config.connection.args.insert(k, v);

// Override connection args with command line options
{
if args.database.is_some() {
conn_args.database.clone_from(&args.database);
}

// override only if args.dsn is none
if args.dsn.is_none() {
if !args.tls {
conn_args
.args
.insert("sslmode".to_string(), "disable".to_string());
}

// override args if specified in command line
for (k, v) in args.set {
conn_args.args.insert(k, v);
}
}

// override role if specified in command line
if let Some(role) = args.role {
conn_args.args.insert("role".to_string(), role);
}
}
let dsn = conn_args.get_dsn()?;

let dsn = conn_args.get_dsn()?;
let mut settings = Settings::default();
let is_terminal = stdin().is_terminal();
let is_repl = is_terminal && !args.non_interactive && !args.check && args.query.is_none();
Expand Down
14 changes: 13 additions & 1 deletion cli/test.sh
Original file line number Diff line number Diff line change
Expand Up @@ -33,8 +33,20 @@ case $TEST_HANDLER in
;;
"http")
echo "==> Testing REST API handler"
export BENDSQL_DSN="databend+http://${DATABEND_USER}:${DATABEND_PASSWORD}@${DATABEND_HOST}:8000/?sslmode=disable&presign=on"
export BENDSQL="${CARGO_TARGET_DIR}/debug/bendsql"

echo "create user if not exists databend identified by 'databend'" | $BENDSQL -dsn="databend+http://${DATABEND_USER}:${DATABEND_PASSWORD}@${DATABEND_HOST}:8000/?sslmode=disable&presign=on"
echo "grant all on *.* to databend" | $BENDSQL -dsn="databend+http://${DATABEND_USER}:${DATABEND_PASSWORD}@${DATABEND_HOST}:8000/?sslmode=disable&presign=on"

export BENDSQL_NEW="${BENDSQL} --user databend --password databend --host ${DATABEND_HOST} --port 8000"

$BENDSQL_NEW --query="select 1 from numbers(10) where number > 1000"
$BENDSQL_NEW --query="create database if not exists aaa"
$BENDSQL_NEW -D aaa --query="create table if not exists bbb(a int)"
$BENDSQL_NEW -D aaa --query="drop table bbb"
$BENDSQL_NEW -D default --query="drop database aaa"

export BENDSQL_DSN="databend+http://${DATABEND_USER}:${DATABEND_PASSWORD}@${DATABEND_HOST}:8000/?sslmode=disable&presign=on"
;;
*)
echo "Usage: $0 [flight|http]"
Expand Down
2 changes: 1 addition & 1 deletion driver/src/rest_api.rs
Original file line number Diff line number Diff line change
Expand Up @@ -124,7 +124,7 @@ impl Connection for RestAPIConnection {
let mut data = data;
let mut size = size;

if file_format_options.get("compression").is_none() {
if !file_format_options.contains_key("compression") {
let mut buffer = Vec::new();
let real_size = data.read_to_end(&mut buffer).await?;
if real_size != size as usize && size != 0 {
Expand Down
1 change: 1 addition & 0 deletions sql/src/cursor_ext/cursor_read_bytes_ext.rs
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@ use std::io::Cursor;
use std::io::ErrorKind;
use std::io::Result;

#[allow(dead_code)]
pub trait ReadBytesExt {
fn peek(&mut self) -> Option<char>;
fn peek_byte(&mut self) -> Option<u8>;
Expand Down

0 comments on commit cff0252

Please sign in to comment.