feat(cli): support setting database -D with DSN

cli/src/args.rs

@@ -0,0 +1,134 @@
+// Copyright 2021 Datafuse Labs
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use std::collections::BTreeMap;
+
+use anyhow::{anyhow, Result};
+
+#[derive(Debug, Clone, PartialEq, Default)]
+pub struct ConnectionArgs {
+    pub host: String,
+    pub port: Option<u16>,
+    pub user: String,
+    pub password: Option<String>,
+    pub database: Option<String>,
+    pub flight: bool,
+    pub args: BTreeMap<String, String>,
+}
+
+impl ConnectionArgs {
+    pub fn get_dsn(self) -> Result<String> {
+        let mut dsn = url::Url::parse("databend://")?;
+        dsn.set_host(Some(&self.host))?;
+        _ = dsn.set_port(self.port);
+        _ = dsn.set_username(&self.user);
+        if let Some(password) = self.password {
+            _ = dsn.set_password(Some(&password))
+        };
+        if let Some(database) = self.database {
+            dsn.set_path(&database);
+        }
+        if self.flight {
+            _ = dsn.set_scheme("databend+flight");
+        }
+        let mut query = url::form_urlencoded::Serializer::new(String::new());
+        if !self.args.is_empty() {
+            for (k, v) in self.args {
+                query.append_pair(&k, &v);
+            }
+        }
+        dsn.set_query(Some(&query.finish()));
+        Ok(dsn.to_string())
+    }
+
+    pub fn from_dsn(dsn: &str) -> Result<Self> {
+        let u = url::Url::parse(dsn)?;
+        let mut args = BTreeMap::new();
+        if let Some(query) = u.query() {
+            for (k, v) in url::form_urlencoded::parse(query.as_bytes()) {
+                args.insert(k.to_string(), v.to_string());
+            }
+        }
+        let flight = matches!(u.scheme(), "databend+flight");
+        let host = u.host_str().ok_or(anyhow!("missing host"))?.to_string();
+        let port = u.port();
+        let user = u.username().to_string();
+        let password = u.password().map(|s| s.to_string());
+        let database = u.path().strip_prefix('/').map(|s| s.to_string());
+        Ok(Self {
+            host,
+            port,
+            user,
+            password,
+            database,
+            flight,
+            args,
+        })
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use super::*;
+
+    #[test]
+    fn parse_dsn() -> Result<()> {
+        let dsn = "databend://username:3a%40SC(nYE1k%3D%7B%7BR@app.databend.com/test?wait_time_secs=10&max_rows_in_buffer=5000000&max_rows_per_page=10000&warehouse=wh&sslmode=disable";
+        let args = ConnectionArgs::from_dsn(dsn)?;
+        let expected = ConnectionArgs {
+            host: "app.databend.com".to_string(),
+            port: None,
+            user: "username".to_string(),
+            password: Some("3a@SC(nYE1k={{R".to_string()),
+            database: Some("test".to_string()),
+            flight: false,
+            args: {
+                let mut args = BTreeMap::new();
+                args.insert("wait_time_secs".to_string(), "10".to_string());
+                args.insert("max_rows_in_buffer".to_string(), "5000000".to_string());
+                args.insert("max_rows_per_page".to_string(), "10000".to_string());
+                args.insert("warehouse".to_string(), "wh".to_string());
+                args.insert("sslmode".to_string(), "disable".to_string());
+                args
+            },
+        };
+        assert_eq!(args, expected);
+        Ok(())
+    }
+
+    #[test]
+    fn format_dsn() -> Result<()> {
+        let args = ConnectionArgs {
+            host: "app.databend.com".to_string(),
+            port: Some(443),
+            user: "username".to_string(),
+            password: Some("3a@SC(nYE1k={{R".to_string()),
+            database: Some("test".to_string()),
+            flight: false,
+            args: {
+                let mut args = BTreeMap::new();
+                args.insert("wait_time_secs".to_string(), "10".to_string());
+                args.insert("max_rows_in_buffer".to_string(), "5000000".to_string());
+                args.insert("max_rows_per_page".to_string(), "10000".to_string());
+                args.insert("warehouse".to_string(), "wh".to_string());
+                args.insert("sslmode".to_string(), "disable".to_string());
+                args
+            },
+        };
+        let dsn = args.get_dsn()?;
+        let expected = "databend://username:3a%40SC(nYE1k%3D%7B%7BR@app.databend.com:443/test?max_rows_in_buffer=5000000&max_rows_per_page=10000&sslmode=disable&wait_time_secs=10&warehouse=wh";
+        assert_eq!(dsn, expected);
+        Ok(())
+    }
+}

cli/src/config.rs

@@ -193,7 +193,7 @@ impl Settings {
 #[serde(default)]
 pub struct ConnectionConfig {
     pub host: String,
-    pub port: u16,
+    pub port: Option<u16>,
     pub user: String,
     pub database: Option<String>,
     pub args: BTreeMap<String, String>,
@@ -254,7 +254,7 @@ impl Default for ConnectionConfig {
     fn default() -> Self {
         Self {
             host: "localhost".to_string(),
-            port: 8000,
+            port: Some(8000),
             user: "root".to_string(),
             database: None,
             args: BTreeMap::new(),

cli/src/main.rs

@@ -14,6 +14,7 @@
 
 #![allow(clippy::upper_case_acronyms)]
 
+mod args;
 mod ast;
 mod config;
 mod display;
@@ -26,6 +27,7 @@ use std::{
     io::{stdin, IsTerminal},
 };
 
+use crate::args::ConnectionArgs;
 use crate::config::OutputQuoteStyle;
 use anyhow::{anyhow, Result};
 use clap::{ArgAction, CommandFactory, Parser, ValueEnum};
@@ -198,50 +200,6 @@ where
     Ok((s[..pos].parse()?, s[pos + 1..].parse()?))
 }
 
-struct ConnectionArgs {
-    host: String,
-    port: u16,
-    user: String,
-    password: Option<String>,
-    database: Option<String>,
-    tls: bool,
-    flight: bool,
-    args: BTreeMap<String, String>,
-}
-
-impl ConnectionArgs {
-    fn get_dsn(self) -> Result<String> {
-        let mut dsn = url::Url::parse("databend://")?;
-        dsn.set_host(Some(&self.host))?;
-        _ = dsn.set_port(Some(self.port));
-        _ = dsn.set_username(&self.user);
-        if let Some(password) = self.password {
-            _ = dsn.set_password(Some(&password))
-        };
-        if let Some(database) = self.database {
-            dsn.set_path(&database);
-        }
-        if self.flight {
-            _ = dsn.set_scheme("databend+flight");
-        } else if self.tls {
-            _ = dsn.set_scheme("databend+https");
-        }
-        let mut query = url::form_urlencoded::Serializer::new(String::new());
-        if !self.args.is_empty() {
-            for (k, v) in self.args {
-                query.append_pair(&k, &v);
-            }
-        }
-        if self.tls {
-            query.append_pair("sslmode", "require");
-        } else {
-            query.append_pair("sslmode", "disable");
-        }
-        dsn.set_query(Some(&query.finish()));
-        Ok(dsn.to_string())
-    }
-}
-
 #[tokio::main]
 pub async fn main() -> Result<()> {
     let mut config = Config::load();
@@ -252,7 +210,8 @@ pub async fn main() -> Result<()> {
         cmd.print_help()?;
         return Ok(());
     }
-    let dsn = match args.dsn {
+
+    let mut conn_args = match args.dsn {
         Some(dsn) => {
             if args.host.is_some() {
                 eprintln!("warning: --host is ignored when --dsn is set");
@@ -266,9 +225,6 @@ pub async fn main() -> Result<()> {
             if args.password.is_some() {
                 eprintln!("warning: --password is ignored when --dsn is set");
             }
-            if args.database.is_some() {
-                eprintln!("warning: --database is ignored when --dsn is set");
-            }
             if args.role.is_some() {
                 eprintln!("warning: --role is ignored when --dsn is set");
             }
@@ -281,40 +237,47 @@ pub async fn main() -> Result<()> {
             if args.flight {
                 eprintln!("warning: --flight is ignored when --dsn is set");
             }
-            dsn
+            ConnectionArgs::from_dsn(&dsn)?
         }
         None => {
             if let Some(host) = args.host {
                 config.connection.host = host;
             }
             if let Some(port) = args.port {
-                config.connection.port = port;
+                config.connection.port = Some(port);
             }
             if let Some(user) = args.user {
                 config.connection.user = user;
             }
-            if args.database.is_some() {
-                config.connection.database = args.database;
-            }
             for (k, v) in args.set {
                 config.connection.args.insert(k, v);
             }
+            if !args.tls {
+                config
+                    .connection
+                    .args
+                    .insert("sslmode".to_string(), "disable".to_string());
+            }
             if let Some(role) = args.role {
                 config.connection.args.insert("role".to_string(), role);
             }
-            let conn_args = ConnectionArgs {
+            ConnectionArgs {
                 host: config.connection.host.clone(),
                 port: config.connection.port,
                 user: config.connection.user.clone(),
                 password: args.password,
                 database: config.connection.database.clone(),
-                tls: args.tls,
                 flight: args.flight,
                 args: config.connection.args.clone(),
-            };
-            conn_args.get_dsn()?
+            }
         }
     };
+    // override database if specified in command line
+    if args.database.is_some() {
+        conn_args.database = args.database;
+    }
+    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.query.is_none();

cli/tests/00-base.sql

@@ -39,3 +39,4 @@ select a[1], b['k1'], c:x, c:y from test_nested;
 select 'bye';
 drop table test;
 drop table test_decimal;
+drop table test_nested;

cli/tests/01-put.sh

@@ -1,11 +1,11 @@
 #!/bin/bash
 
-echo "DROP STAGE IF EXISTS ss_temp" | ${BENDSQL}
-echo "CREATE STAGE ss_temp" | ${BENDSQL}
+echo "DROP STAGE IF EXISTS ss_01" | ${BENDSQL}
+echo "CREATE STAGE ss_01" | ${BENDSQL}
 
 cat <<SQL | ${BENDSQL}
-DROP TABLE IF EXISTS books;
-CREATE TABLE books
+DROP TABLE IF EXISTS books_01;
+CREATE TABLE books_01
 (
     title VARCHAR,
     author VARCHAR,
@@ -14,19 +14,22 @@ CREATE TABLE books
 SQL
 
 cat <<SQL | ${BENDSQL}
-SELECT * FROM books;
+SELECT * FROM books_01;
 SQL
 
 mkdir -p /tmp/abc
 cp "${PWD}/cli/tests/data/books.parquet" /tmp/abc/books.parquet
 
-echo "put fs:///tmp/abc/b*.parquet @ss_temp/abc/" | ${BENDSQL}
-echo 'get @ss_temp/abc fs:///tmp/edf' | ${BENDSQL}
+echo "put fs:///tmp/abc/b*.parquet @ss_01/abc/" | ${BENDSQL}
+echo 'get @ss_01/abc fs:///tmp/edf' | ${BENDSQL}
 
 cat <<SQL | ${BENDSQL}
-COPY INTO books FROM @ss_temp/abc/ files=('books.parquet') FILE_FORMAT = (TYPE = PARQUET);
+COPY INTO books_01 FROM @ss_01/abc/ files=('books.parquet') FILE_FORMAT = (TYPE = PARQUET);
 SQL
 
 cat <<SQL | ${BENDSQL}
-SELECT * FROM books LIMIT 2;
+SELECT * FROM books_01 LIMIT 2;
 SQL
+
+echo "DROP STAGE IF EXISTS ss_01" | ${BENDSQL}
+echo "DROP TABLE IF EXISTS books_01" | ${BENDSQL}

cli/tests/data/ontime.sql

@@ -1,4 +1,4 @@
-CREATE TABLE test_ontime (
+CREATE TABLE http_ontime_03 (
     Year SMALLINT UNSIGNED,
     Quarter TINYINT UNSIGNED,
     Month TINYINT UNSIGNED,

cli/tests/http/01-load_stdin.sh

@@ -1,17 +1,17 @@
 #!/bin/bash
 
 cat <<SQL | ${BENDSQL}
-DROP TABLE IF EXISTS test_books;
+DROP TABLE IF EXISTS http_books_01;
 SQL
 
 cat <<SQL | ${BENDSQL}
-CREATE TABLE test_books (title VARCHAR NULL, author VARCHAR NULL, date VARCHAR NULL, publish_time TIMESTAMP NULL);
+CREATE TABLE http_books_01 (title VARCHAR NULL, author VARCHAR NULL, date VARCHAR NULL, publish_time TIMESTAMP NULL);
 SQL
 
-${BENDSQL} --query='INSERT INTO test_books VALUES;' --format=csv --data=@- <cli/tests/data/books.csv
+${BENDSQL} --query='INSERT INTO http_books_01 VALUES;' --format=csv --data=@- <cli/tests/data/books.csv
 
-${BENDSQL} --query='SELECT * FROM test_books LIMIT 10;' --output=tsv
+${BENDSQL} --query='SELECT * FROM http_books_01 LIMIT 10;' --output=tsv
 
 cat <<SQL | ${BENDSQL}
-DROP TABLE test_books;
+DROP TABLE http_books_01;
 SQL

cli/tests/http/02-load_file.sh

@@ -1,17 +1,17 @@
 #!/bin/bash
 
 cat <<SQL | ${BENDSQL}
-DROP TABLE IF EXISTS test_books;
+DROP TABLE IF EXISTS http_books_02;
 SQL
 
 cat <<SQL | ${BENDSQL}
-CREATE TABLE test_books (title VARCHAR NULL, author VARCHAR NULL, date VARCHAR NULL, publish_time TIMESTAMP NULL);
+CREATE TABLE http_books_02 (title VARCHAR NULL, author VARCHAR NULL, date VARCHAR NULL, publish_time TIMESTAMP NULL);
 SQL
 
-${BENDSQL} --query='INSERT INTO test_books VALUES;' --format=csv --data=@cli/tests/data/books.csv
+${BENDSQL} --query='INSERT INTO http_books_02 VALUES;' --format=csv --data=@cli/tests/data/books.csv
 
-${BENDSQL} --query='SELECT * FROM test_books LIMIT 10;' --output=tsv
+${BENDSQL} --query='SELECT * FROM http_books_02 LIMIT 10;' --output=tsv
 
 cat <<SQL | ${BENDSQL}
-DROP TABLE test_books;
+DROP TABLE http_books_02;
 SQL