Skip to content

Commit

Permalink
feat(cli): use sled for autocomplete keywords (#472)
Browse files Browse the repository at this point in the history 
* feat(cli): use sled for autocomplete keywords

* fix: apply batch in batch
  • Loading branch information
@everpcpc
everpcpc authored Aug 9, 2024
1 parent 39f690e commit 2339917
Show file tree
Hide file tree
Showing 5 changed files with 100 additions and 85 deletions.
1 change: 1 addition & 0 deletions cli/Cargo.toml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down
28 changes: 27 additions & 1 deletion cli/src/ast/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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
}
6 changes: 2 additions & 4 deletions cli/src/display.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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,
};

Expand Down Expand Up @@ -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();
Expand Down
104 changes: 35 additions & 69 deletions cli/src/helper.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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,
Expand All @@ -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)
}

Expand Down Expand Up @@ -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,
}
}
}
Expand All @@ -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)
}
Expand All @@ -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)
Expand Down
46 changes: 35 additions & 11 deletions cli/src/session.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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");
Expand Down Expand Up @@ -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 {
Expand All @@ -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 {
Expand All @@ -121,7 +145,7 @@ impl Session {
is_repl,
settings,
query: String::new(),
keywords: Arc::new(keywords),
keywords,
})
}

Expand Down Expand Up @@ -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 {
Expand Down

0 comments on commit 2339917

Please sign in to comment.