diff --git a/cli/Cargo.toml b/cli/Cargo.toml index 4b12082d5..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,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", 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..c1935d46d --- /dev/null +++ b/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, + #[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, + } + } +} 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);