From 81de2a45422e96ee8828541d4c2f2b86742e0dfd Mon Sep 17 00:00:00 2001 From: sundyli <543950155@qq.com> Date: Tue, 4 Apr 2023 11:13:45 +0800 Subject: [PATCH 1/3] feat: support loading toml config from file --- cli/Cargo.toml | 4 ++ cli/src/ast/mod.rs | 27 ++++++++ cli/src/{token.rs => ast/tokenizer.rs} | 0 cli/src/config.rs | 95 ++++++++++++++++++++++++++ cli/src/display.rs | 29 ++++++-- cli/src/helper.rs | 6 +- cli/src/main.rs | 28 ++++---- cli/src/session.rs | 23 ++++--- 8 files changed, 183 insertions(+), 29 deletions(-) create mode 100644 cli/src/ast/mod.rs rename cli/src/{token.rs => ast/tokenizer.rs} (100%) create mode 100644 cli/src/config.rs diff --git a/cli/Cargo.toml b/cli/Cargo.toml index 4b12082d5..782eaf72e 100644 --- a/cli/Cargo.toml +++ b/cli/Cargo.toml @@ -46,6 +46,10 @@ tonic = { version = "0.8", default-features = false, features = [ "tls", "prost", ] } +toml = "0.7.3" +humantime-serde = "1.1.1" +itertools = "0.10.5" +sqlformat = "0.2.1" [[bin]] name = "bendsql" diff --git a/cli/src/ast/mod.rs b/cli/src/ast/mod.rs new file mode 100644 index 000000000..95143c8d1 --- /dev/null +++ b/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) +} diff --git a/cli/src/token.rs b/cli/src/ast/tokenizer.rs similarity index 100% rename from cli/src/token.rs rename to cli/src/ast/tokenizer.rs diff --git a/cli/src/config.rs b/cli/src/config.rs new file mode 100644 index 000000000..bcdeacec7 --- /dev/null +++ b/cli/src/config.rs @@ -0,0 +1,95 @@ +// 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, + #[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(); + } + + let config = 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() + } + }; + config + } +} + +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, + } + } +} diff --git a/cli/src/display.rs b/cli/src/display.rs index 18e411a41..207da8d47 100644 --- a/cli/src/display.rs +++ b/cli/src/display.rs @@ -26,20 +26,25 @@ 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, rows: usize, @@ -47,9 +52,17 @@ pub struct ReplDisplay { start: Instant, } -impl ReplDisplay { - pub fn new(schema: Schema, start: Instant, stream: Streaming) -> Self { +impl<'a> ReplDisplay<'a> { + pub fn new( + config: &'a Config, + query: &'a str, + schema: &'a Schema, + start: Instant, + stream: Streaming, + ) -> 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) => { diff --git a/cli/src/helper.rs b/cli/src/helper.rs index 6cc321666..00c88f03f 100644 --- a/cli/src/helper.rs +++ b/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, diff --git a/cli/src/main.rs b/cli/src/main.rs index 2a08be53f..4dfb1cb69 100644 --- a/cli/src/main.rs +++ b/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 { +fn endpoint(conn: &Connection, args: &Args, addr: String) -> Result { 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(); diff --git a/cli/src/session.rs b/cli/src/session.rs index 9a39ccd7e..66779aa67 100644 --- a/cli/src/session.rs +++ b/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 { - 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); From aefd22262c20ce16581b9929117e42b1c8830d9e Mon Sep 17 00:00:00 2001 From: sundyli <543950155@qq.com> Date: Tue, 4 Apr 2023 11:14:12 +0800 Subject: [PATCH 2/3] feat: bump --- cli/Cargo.toml | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/cli/Cargo.toml b/cli/Cargo.toml index 782eaf72e..ee7ebd1fb 100644 --- a/cli/Cargo.toml +++ b/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,22 +34,22 @@ 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", "tls", "prost", ] } -toml = "0.7.3" -humantime-serde = "1.1.1" -itertools = "0.10.5" -sqlformat = "0.2.1" [[bin]] name = "bendsql" From 8d83e6e3b45699da2a5c2b7a6a33dac877d3bcf1 Mon Sep 17 00:00:00 2001 From: sundyli <543950155@qq.com> Date: Tue, 4 Apr 2023 11:20:41 +0800 Subject: [PATCH 3/3] feat: update --- cli/src/config.rs | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/cli/src/config.rs b/cli/src/config.rs index bcdeacec7..c1935d46d 100644 --- a/cli/src/config.rs +++ b/cli/src/config.rs @@ -60,14 +60,13 @@ impl Config { return Self::default(); } - let config = match toml::from_str(&std::fs::read_to_string(path).unwrap()) { + 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() } - }; - config + } } }