chore(cli): refactor pretty sql

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",

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()