feat(cli): use sled for autocomplete keywords

cli/Cargo.toml

@@ -28,6 +28,7 @@ indicatif = "0.17"
 log = "0.4"
 once_cell = "1.18"
 percent-encoding = "2.3"
+sled = "0.34"
 rustyline = "12.0"
 serde = { version = "1.0", features = ["derive"] }
 terminal_size = "0.3"

cli/src/ast/mod.rs

@@ -14,7 +14,7 @@
 
 use databend_common_ast::{
     ast::pretty_statement,
-    parser::{parse_sql, Dialect},
+    parser::{parse_sql, token::TokenKind, tokenize_sql, Dialect},
 };
 
 use crate::session::QueryKind;
@@ -31,3 +31,29 @@ pub fn format_query(query: &str) -> String {
     }
     query.to_string()
 }
+
+pub fn highlight_query(line: &str) -> String {
+    let tokens = tokenize_sql(line);
+    let mut line = line.to_owned();
+
+    if let Ok(tokens) = tokens {
+        for token in tokens.iter().rev() {
+            if TokenKind::is_keyword(&token.kind)
+                || TokenKind::is_reserved_ident(&token.kind, false)
+                || TokenKind::is_reserved_function_name(&token.kind)
+            {
+                line.replace_range(
+                    std::ops::Range::from(token.span),
+                    &format!("\x1b[1;32m{}\x1b[0m", token.text()),
+                );
+            } else if TokenKind::is_literal(&token.kind) {
+                line.replace_range(
+                    std::ops::Range::from(token.span),
+                    &format!("\x1b[1;33m{}\x1b[0m", token.text()),
+                );
+            }
+        }
+    }
+
+    line
+}

cli/src/display.rs

@@ -19,16 +19,14 @@ use anyhow::{anyhow, Result};
 use comfy_table::{Cell, CellAlignment, Table};
 use databend_driver::{Row, RowStatsIterator, RowWithStats, SchemaRef, ServerStats};
 use indicatif::{HumanBytes, ProgressBar, ProgressState, ProgressStyle};
-use rustyline::highlight::Highlighter;
 use terminal_size::{terminal_size, Width};
 use tokio::time::Instant;
 use tokio_stream::StreamExt;
 use unicode_segmentation::UnicodeSegmentation;
 
 use crate::{
-    ast::format_query,
+    ast::{format_query, highlight_query},
     config::{ExpandMode, OutputFormat, OutputQuoteStyle, Settings},
-    helper::CliHelper,
     session::QueryKind,
 };
 
@@ -93,7 +91,7 @@ impl<'a> FormatDisplay<'a> {
     async fn display_table(&mut self) -> Result<()> {
         if self.settings.display_pretty_sql {
             let format_sql = format_query(self.query);
-            let format_sql = CliHelper::new().highlight(&format_sql, format_sql.len());
+            let format_sql = highlight_query(&format_sql);
             println!("\n{}\n", format_sql);
         }
         let mut rows = Vec::new();

cli/src/helper.rs

@@ -15,9 +15,6 @@
 use std::borrow::Cow;
 use std::sync::Arc;
 
-use databend_common_ast::parser::all_reserved_keywords;
-use databend_common_ast::parser::token::TokenKind;
-use databend_common_ast::parser::tokenize_sql;
 use rustyline::completion::Completer;
 use rustyline::completion::FilenameCompleter;
 use rustyline::completion::Pair;
@@ -31,20 +28,15 @@ use rustyline::Context;
 use rustyline::Helper;
 use rustyline::Result;
 
+use crate::ast::highlight_query;
+
 pub struct CliHelper {
     completer: FilenameCompleter,
-    keywords: Arc<Vec<String>>,
+    keywords: Option<Arc<sled::Db>>,
 }
 
 impl CliHelper {
-    pub fn new() -> Self {
-        Self {
-            completer: FilenameCompleter::new(),
-            keywords: Arc::new(Vec::new()),
-        }
-    }
-
-    pub fn with_keywords(keywords: Arc<Vec<String>>) -> Self {
+    pub fn new(keywords: Option<Arc<sled::Db>>) -> Self {
         Self {
             completer: FilenameCompleter::new(),
             keywords,
@@ -54,28 +46,7 @@ impl CliHelper {
 
 impl Highlighter for CliHelper {
     fn highlight<'l>(&self, line: &'l str, _pos: usize) -> Cow<'l, str> {
-        let tokens = tokenize_sql(line);
-        let mut line = line.to_owned();
-
-        if let Ok(tokens) = tokens {
-            for token in tokens.iter().rev() {
-                if TokenKind::is_keyword(&token.kind)
-                    || TokenKind::is_reserved_ident(&token.kind, false)
-                    || TokenKind::is_reserved_function_name(&token.kind)
-                {
-                    line.replace_range(
-                        std::ops::Range::from(token.span),
-                        &format!("\x1b[1;32m{}\x1b[0m", token.text()),
-                    );
-                } else if TokenKind::is_literal(&token.kind) {
-                    line.replace_range(
-                        std::ops::Range::from(token.span),
-                        &format!("\x1b[1;33m{}\x1b[0m", token.text()),
-                    );
-                }
-            }
-        }
-
+        let line = highlight_query(line);
         Cow::Owned(line)
     }
 
@@ -118,12 +89,16 @@ impl Hinter for CliHelper {
         if last_word.is_empty() {
             return None;
         }
-
-        let (_, res) = KeyWordCompleter::complete(line, pos, &self.keywords);
-        if !res.is_empty() {
-            Some(res[0].replacement[last_word.len()..].to_owned())
-        } else {
-            None
+        match self.keywords {
+            Some(ref keywords) => {
+                let (_, res) = KeyWordCompleter::complete(line, pos, keywords);
+                if !res.is_empty() {
+                    Some(res[0].replacement[last_word.len()..].to_owned())
+                } else {
+                    None
+                }
+            }
+            None => None,
         }
     }
 }
@@ -137,9 +112,11 @@ impl Completer for CliHelper {
         pos: usize,
         ctx: &Context<'_>,
     ) -> std::result::Result<(usize, Vec<Pair>), ReadlineError> {
-        let keyword_candidates = KeyWordCompleter::complete(line, pos, self.keywords.as_ref());
-        if !keyword_candidates.1.is_empty() {
-            return Ok(keyword_candidates);
+        if let Some(ref keywords) = self.keywords {
+            let keyword_candidates = KeyWordCompleter::complete(line, pos, keywords);
+            if !keyword_candidates.1.is_empty() {
+                return Ok(keyword_candidates);
+            }
         }
         self.completer.complete(line, pos, ctx)
     }
@@ -161,35 +138,24 @@ impl Helper for CliHelper {}
 struct KeyWordCompleter {}
 
 impl KeyWordCompleter {
-    fn complete(s: &str, pos: usize, keywords: &[String]) -> (usize, Vec<Pair>) {
+    fn complete(s: &str, pos: usize, db: &sled::Db) -> (usize, Vec<Pair>) {
         let hint = s
             .split(|p: char| p.is_whitespace() || p == '.')
             .last()
-            .unwrap_or(s);
-        let all_keywords = all_reserved_keywords();
-
-        let mut results: Vec<Pair> = all_keywords
-            .iter()
-            .filter(|keyword| keyword.starts_with(&hint.to_ascii_lowercase()))
-            .map(|keyword| Pair {
-                display: keyword.to_string(),
-                replacement: keyword.to_string(),
-            })
-            .collect();
-
-        results.extend(
-            keywords
-                .iter()
-                .filter(|keyword| {
-                    keyword
-                        .to_lowercase()
-                        .starts_with(&hint.to_ascii_lowercase())
-                })
-                .map(|keyword| Pair {
-                    display: keyword.to_string(),
-                    replacement: keyword.to_string(),
-                }),
-        );
+            .unwrap_or(s)
+            .to_ascii_lowercase();
+
+        let r = db.scan_prefix(&hint);
+        let mut results = Vec::new();
+        for line in r {
+            let (w, t) = line.unwrap();
+            let word = String::from_utf8_lossy(&w);
+            let category = String::from_utf8_lossy(&t);
+            results.push(Pair {
+                display: format!("{}({})", word, category),
+                replacement: word.to_string(),
+            });
+        }
 
         if pos >= hint.len() {
             (pos - hint.len(), results)

cli/src/session.rs

@@ -21,6 +21,7 @@ use anyhow::anyhow;
 use anyhow::Result;
 use async_recursion::async_recursion;
 use chrono::NaiveDateTime;
+use databend_common_ast::parser::all_reserved_keywords;
 use databend_common_ast::parser::token::TokenKind;
 use databend_common_ast::parser::token::Tokenizer;
 use databend_driver::ServerStats;
@@ -41,7 +42,7 @@ use crate::display::{format_write_progress, ChunkDisplay, FormatDisplay};
 use crate::helper::CliHelper;
 use crate::VERSION;
 
-static PROMPT_SQL: &str = "select name from system.tables union all select name from system.columns union all select name from system.databases union all select name from system.functions limit 10000";
+static PROMPT_SQL: &str = "select name, 'f' as type from system.functions union all select name, 'd' as type from system.databases union all select name, 't' as type from system.tables union all select name, 'c' as type from system.columns limit 10000";
 
 static VERSION_SHORT: Lazy<String> = Lazy::new(|| {
     let version = option_env!("CARGO_PKG_VERSION").unwrap_or("unknown");
@@ -70,15 +71,16 @@ pub struct Session {
     settings: Settings,
     query: String,
 
-    keywords: Arc<Vec<String>>,
+    keywords: Option<Arc<sled::Db>>,
 }
 
 impl Session {
     pub async fn try_new(dsn: String, settings: Settings, is_repl: bool) -> Result<Self> {
         let client = Client::new(dsn).with_name(format!("bendsql/{}", VERSION_SHORT.as_str()));
         let conn = client.get_conn().await?;
         let info = conn.info().await;
-        let mut keywords = Vec::with_capacity(1024);
+        let mut keywords: Option<Arc<sled::Db>> = None;
+
         if is_repl {
             println!("Welcome to BendSQL {}.", VERSION.as_str());
             match info.warehouse {
@@ -97,22 +99,44 @@ impl Session {
             }
             let version = conn.version().await.unwrap_or_default();
             println!("Connected to {}", version);
-            println!();
 
+            let config = sled::Config::new().temporary(true);
+            let db = config.open()?;
+            // ast keywords
+            {
+                let keywords = all_reserved_keywords();
+                let mut batch = sled::Batch::default();
+                for word in keywords {
+                    batch.insert(word.to_ascii_lowercase().as_str(), "k")
+                }
+                db.apply_batch(batch)?;
+            }
+            // server keywords
             if !settings.no_auto_complete {
                 let rows = conn.query_iter(PROMPT_SQL).await;
                 match rows {
                     Ok(mut rows) => {
+                        let mut count = 0;
+                        let mut batch = sled::Batch::default();
                         while let Some(Ok(row)) = rows.next().await {
-                            let name: (String,) = row.try_into().unwrap();
-                            keywords.push(name.0);
+                            let (w, t): (String, String) = row.try_into().unwrap();
+                            batch.insert(w.as_str(), t.as_str());
+                            count += 1;
+                            if count % 1000 == 0 {
+                                db.apply_batch(batch)?;
+                                batch = sled::Batch::default();
+                            }
                         }
+                        db.apply_batch(batch)?;
+                        println!("Loaded {} auto complete keywords from server.", db.len());
                     }
                     Err(e) => {
-                        eprintln!("loading auto complete keywords failed: {}", e);
+                        eprintln!("WARN: loading auto complete keywords failed: {}", e);
                     }
                 }
             }
+            keywords = Some(Arc::new(db));
+            println!();
         }
 
         Ok(Self {
@@ -121,7 +145,7 @@ impl Session {
             is_repl,
             settings,
             query: String::new(),
-            keywords: Arc::new(keywords),
+            keywords,
         })
     }
 
@@ -232,12 +256,12 @@ impl Session {
 
     pub async fn handle_repl(&mut self) {
         let config = Builder::new()
-            .completion_prompt_limit(5)
-            .completion_type(CompletionType::Circular)
+            .completion_prompt_limit(10)
+            .completion_type(CompletionType::List)
             .build();
         let mut rl = Editor::<CliHelper, DefaultHistory>::with_config(config).unwrap();
 
-        rl.set_helper(Some(CliHelper::with_keywords(self.keywords.clone())));
+        rl.set_helper(Some(CliHelper::new(self.keywords.clone())));
         rl.load_history(&get_history_path()).ok();
 
         'F: loop {