diff --git a/cli/Cargo.toml b/cli/Cargo.toml index 82ee94b6..c0070300 100644 --- a/cli/Cargo.toml +++ b/cli/Cargo.toml @@ -24,7 +24,7 @@ async-trait = "0.1" clap = { version = "4.4", features = ["derive", "env"] } comfy-table = "7.1" csv = "1.3" -databend-common-ast = "0.1.2" +databend-common-ast = "0.1.3" fern = { version = "0.6", features = ["colored"] } indicatif = "0.17" log = "0.4" @@ -36,6 +36,7 @@ rustyline = "12.0" serde = { version = "1.0", features = ["derive"] } serde_json = "1.0" sled = "0.34" +sqlformat = "0.3.3" terminal_size = "0.3" tokio = { version = "1.34", features = [ "macros", diff --git a/cli/src/ast/mod.rs b/cli/src/ast/mod.rs index cec12d7b..ac984484 100644 --- a/cli/src/ast/mod.rs +++ b/cli/src/ast/mod.rs @@ -12,10 +12,8 @@ // See the License for the specific language governing permissions and // limitations under the License. -use databend_common_ast::{ - ast::pretty_statement, - parser::{parse_sql, token::TokenKind, tokenize_sql, Dialect}, -}; +use databend_common_ast::parser::{parse_sql, token::TokenKind, tokenize_sql, Dialect}; +use sqlformat::{FormatOptions, QueryParams}; use crate::session::QueryKind; @@ -24,9 +22,22 @@ pub fn format_query(query: &str) -> String { if kind == QueryKind::Put || kind == QueryKind::Get { return query.to_owned(); } + if let Ok(tokens) = databend_common_ast::parser::tokenize_sql(query) { if let Ok((stmt, _)) = parse_sql(&tokens, Dialect::Experimental) { - return pretty_statement(stmt, 80).unwrap(); + let options = FormatOptions::default(); + + let pretty_sql = sqlformat::format(query, &QueryParams::None, &options); + // if pretty sql could be parsed into same stmt, return pretty sql + if let Ok(pretty_tokens) = databend_common_ast::parser::tokenize_sql(&pretty_sql) { + if let Ok((pretty_stmt, _)) = parse_sql(&pretty_tokens, Dialect::Experimental) { + if stmt.to_string() == pretty_stmt.to_string() { + return pretty_sql; + } + } + } + + return stmt.to_string(); } } query.to_string()