feat: support loading toml config from file

cli/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "bendsql"
-version = "0.1.1"
+version = "0.1.2"
 edition = "2021"
 license = "Apache-2.0"
 description = "Databend Native Cli tool"
@@ -34,12 +34,16 @@ tokio = { version = "1.26", features = [
 async-trait = "0.1.68"
 clap = { version = "4.1.0", features = ["derive"] }
 comfy-table = "6.1.4"
+humantime-serde = "1.1.1"
 indicatif = "0.17.3"
+itertools = "0.10.5"
 logos = "0.12.1"
 serde = { version = "1.0.159", features = ["derive"] }
 serde_json = "1.0.95"
+sqlformat = "0.2.1"
 strum = "0.24"
 strum_macros = "0.24"
+toml = "0.7.3"
 tonic = { version = "0.8", default-features = false, features = [
     "transport",
     "codegen",

cli/src/ast/mod.rs

@@ -0,0 +1,27 @@
+// Copyright 2023 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.
+
+mod tokenizer;
+
+use sqlformat::{Indent, QueryParams};
+pub use tokenizer::*;
+
+pub fn format_query(query: &str) -> String {
+    let options = sqlformat::FormatOptions {
+        indent: Indent::Spaces(2),
+        uppercase: true,
+        lines_between_queries: 1,
+    };
+    sqlformat::format(query, &QueryParams::None, options)
+}

cli/src/config.rs

@@ -0,0 +1,94 @@
+// Copyright 2023 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.
+
+// Loading from `$HOME/.config/bendsql/config.toml`
+
+use std::{path::Path, time::Duration};
+
+use serde::Deserialize;
+
+#[derive(Clone, Debug, Default, Deserialize)]
+pub struct Config {
+    pub connection: Connection,
+    pub settings: Settings,
+}
+
+#[derive(Clone, Debug, Deserialize)]
+#[serde(default)]
+pub struct Settings {
+    pub display_pretty_sql: bool,
+    pub prompt: String,
+}
+
+#[derive(Clone, Debug, Deserialize)]
+#[serde(default)]
+pub struct Connection {
+    #[serde(with = "humantime_serde")]
+    pub connect_timeout: Duration,
+    #[serde(with = "humantime_serde")]
+    pub query_timeout: Duration,
+    pub tcp_nodelay: bool,
+    #[serde(with = "humantime_serde")]
+    pub tcp_keepalive: Option<Duration>,
+    #[serde(with = "humantime_serde")]
+    pub http2_keep_alive_interval: Duration,
+    #[serde(with = "humantime_serde")]
+    pub keep_alive_timeout: Duration,
+    pub keep_alive_while_idle: bool,
+}
+
+impl Config {
+    pub fn load() -> Self {
+        let path = format!(
+            "{}/.config/bendsql/config.toml",
+            std::env::var("HOME").unwrap_or_else(|_| ".".to_string())
+        );
+
+        let path = Path::new(&path);
+        if !path.exists() {
+            return Self::default();
+        }
+
+        match toml::from_str(&std::fs::read_to_string(path).unwrap()) {
+            Ok(config) => config,
+            Err(e) => {
+                eprintln!("Failed to load config: {}, will use default config", e);
+                Self::default()
+            }
+        }
+    }
+}
+
+impl Default for Settings {
+    fn default() -> Self {
+        Settings {
+            display_pretty_sql: true,
+            prompt: "bendsql :) ".to_string(),
+        }
+    }
+}
+
+impl Default for Connection {
+    fn default() -> Self {
+        Connection {
+            connect_timeout: Duration::from_secs(20),
+            query_timeout: Duration::from_secs(60),
+            tcp_nodelay: true,
+            tcp_keepalive: Some(Duration::from_secs(3600)),
+            http2_keep_alive_interval: Duration::from_secs(300),
+            keep_alive_timeout: Duration::from_secs(20),
+            keep_alive_while_idle: true,
+        }
+    }
+}

cli/src/display.rs

@@ -26,30 +26,43 @@ use comfy_table::{Cell, CellAlignment, Table};
 
 use arrow_cast::display::{ArrayFormatter, FormatOptions};
 use futures::StreamExt;
+use rustyline::highlight::Highlighter;
 use serde::{Deserialize, Serialize};
 use tokio::time::Instant;
 use tonic::Streaming;
 
 use indicatif::{HumanBytes, ProgressBar, ProgressState, ProgressStyle};
 
+use crate::{ast::format_query, config::Config, helper::CliHelper};
+
 #[async_trait::async_trait]
 pub trait ChunkDisplay {
     async fn display(&mut self) -> Result<(), ArrowError>;
     fn total_rows(&self) -> usize;
 }
 
-pub struct ReplDisplay {
-    schema: Schema,
+pub struct ReplDisplay<'a> {
+    config: &'a Config,
+    query: &'a str,
+    schema: &'a Schema,
     stream: Streaming<FlightData>,
 
     rows: usize,
     progress: Option<ProgressBar>,
     start: Instant,
 }
 
-impl ReplDisplay {
-    pub fn new(schema: Schema, start: Instant, stream: Streaming<FlightData>) -> Self {
+impl<'a> ReplDisplay<'a> {
+    pub fn new(
+        config: &'a Config,
+        query: &'a str,
+        schema: &'a Schema,
+        start: Instant,
+        stream: Streaming<FlightData>,
+    ) -> Self {
         Self {
+            config,
+            query,
             schema,
             stream,
             rows: 0,
@@ -60,11 +73,17 @@ impl ReplDisplay {
 }
 
 #[async_trait::async_trait]
-impl ChunkDisplay for ReplDisplay {
+impl<'a> ChunkDisplay for ReplDisplay<'a> {
     async fn display(&mut self) -> Result<(), ArrowError> {
         let mut batches = Vec::new();
         let mut progress = ProgressValue::default();
 
+        if self.config.settings.display_pretty_sql {
+            let format_sql = format_query(self.query);
+            let format_sql = CliHelper::new().highlight(&format_sql, format_sql.len());
+            println!("\n{}\n", format_sql);
+        }
+
         while let Some(datum) = self.stream.next().await {
             match datum {
                 Ok(datum) => {

cli/src/helper.rs

@@ -27,9 +27,9 @@ use rustyline::Context;
 use rustyline::Helper;
 use rustyline::Result;
 
-use crate::token::all_reserved_keywords;
-use crate::token::tokenize_sql;
-use crate::token::TokenKind;
+use crate::ast::all_reserved_keywords;
+use crate::ast::tokenize_sql;
+use crate::ast::TokenKind;
 
 pub struct CliHelper {
     completer: FilenameCompleter,

cli/src/main.rs

@@ -14,16 +14,16 @@
 
 #![allow(clippy::upper_case_acronyms)]
 
+mod ast;
+mod config;
 mod display;
 mod helper;
 mod session;
-mod token;
-
-use std::time::Duration;
 
 use arrow::error::ArrowError;
 
 use clap::Parser;
+use config::{Config, Connection};
 use tonic::transport::{ClientTlsConfig, Endpoint};
 
 #[derive(Debug, Parser, PartialEq)]
@@ -65,14 +65,16 @@ pub async fn main() -> Result<(), ArrowError> {
         return Ok(());
     }
 
+    let config = Config::load();
+
     let protocol = if args.tls { "https" } else { "http" };
     // Authenticate
     let url = format!("{protocol}://{}:{}", args.host, args.port);
-    let endpoint = endpoint(&args, url)?;
+    let endpoint = endpoint(&config.connection, &args, url)?;
     let is_repl = atty::is(atty::Stream::Stdin);
 
     let mut session =
-        session::Session::try_new(endpoint, &args.user, &args.password, is_repl).await?;
+        session::Session::try_new(config, endpoint, &args.user, &args.password, is_repl).await?;
 
     session.handle().await;
     Ok(())
@@ -83,16 +85,16 @@ fn print_usage() {
     println!("{}", msg);
 }
 
-fn endpoint(args: &Args, addr: String) -> Result<Endpoint, ArrowError> {
+fn endpoint(conn: &Connection, args: &Args, addr: String) -> Result<Endpoint, ArrowError> {
     let mut endpoint = Endpoint::new(addr)
         .map_err(|_| ArrowError::IoError("Cannot create endpoint".to_string()))?
-        .connect_timeout(Duration::from_secs(20))
-        .timeout(Duration::from_secs(20))
-        .tcp_nodelay(true) // Disable Nagle's Algorithm since we don't want packets to wait
-        .tcp_keepalive(Option::Some(Duration::from_secs(3600)))
-        .http2_keep_alive_interval(Duration::from_secs(300))
-        .keep_alive_timeout(Duration::from_secs(20))
-        .keep_alive_while_idle(true);
+        .connect_timeout(conn.connect_timeout)
+        .timeout(conn.query_timeout)
+        .tcp_nodelay(conn.tcp_nodelay) // Disable Nagle's Algorithm since we don't want packets to wait
+        .tcp_keepalive(conn.tcp_keepalive)
+        .http2_keep_alive_interval(conn.http2_keep_alive_interval)
+        .keep_alive_timeout(conn.keep_alive_timeout)
+        .keep_alive_while_idle(conn.keep_alive_while_idle);
 
     if args.tls {
         let tls_config = ClientTlsConfig::new();

cli/src/session.rs

@@ -27,18 +27,21 @@ use std::io::BufRead;
 use tokio::time::Instant;
 use tonic::transport::Endpoint;
 
+use crate::ast::{TokenKind, Tokenizer};
+use crate::config::Config;
 use crate::display::{format_error, ChunkDisplay, FormatDisplay, ReplDisplay};
 use crate::helper::CliHelper;
-use crate::token::{TokenKind, Tokenizer};
 
 pub struct Session {
     client: FlightSqlServiceClient,
     is_repl: bool,
+    config: Config,
     prompt: String,
 }
 
 impl Session {
     pub async fn try_new(
+        config: Config,
         endpoint: Endpoint,
         user: &str,
         password: &str,
@@ -61,8 +64,15 @@ impl Session {
         client.set_header("bendsql", "1");
         let _token = client.handshake(user, password).await.unwrap();
 
-        let prompt = format!("{} :) ", endpoint.uri().host().unwrap());
+        let mut prompt = config.settings.prompt.clone();
+
+        {
+            prompt = prompt.replace("{host}", endpoint.uri().host().unwrap());
+            prompt = prompt.replace("{user}", user);
+        }
+
         Ok(Self {
+            config,
             client,
             is_repl,
             prompt,
@@ -141,11 +151,8 @@ impl Session {
     }
 
     pub async fn handle_query(&mut self, is_repl: bool, query: &str) -> Result<bool, ArrowError> {
-        if is_repl {
-            if query == "exit" || query == "quit" {
-                return Ok(true);
-            }
-            println!("\n{}\n", query);
+        if is_repl && (query == "exit" || query == "quit") {
+            return Ok(true);
         }
 
         let start = Instant::now();
@@ -184,7 +191,7 @@ impl Session {
         let schema = fb_to_schema(ipc_schema);
 
         if is_repl {
-            let mut displayer = ReplDisplay::new(schema, start, flight_data);
+            let mut displayer = ReplDisplay::new(&self.config, query, &schema, start, flight_data);
             displayer.display().await?;
         } else {
             let mut displayer = FormatDisplay::new(schema, flight_data);