fix: args override config and dsn (#416)

databendlabs · May 6, 2024 · cff0252 · cff0252
1 parent f6865a8
commit cff0252
Show file tree

Hide file tree

Showing 5 changed files with 55 additions and 44 deletions.
diff --git a/cli/src/display.rs b/cli/src/display.rs
@@ -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> {
@@ -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 {

diff --git a/cli/src/main.rs b/cli/src/main.rs
@@ -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();
@@ -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");
             }
@@ -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();

diff --git a/cli/test.sh b/cli/test.sh
@@ -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]"

diff --git a/driver/src/rest_api.rs b/driver/src/rest_api.rs
@@ -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 {

diff --git a/sql/src/cursor_ext/cursor_read_bytes_ext.rs b/sql/src/cursor_ext/cursor_read_bytes_ext.rs
@@ -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>;