提示词:请用rust调用duckdb-rs实现一个duckdb CLI,支持语法突出显示和计时,还支持命令行管道输入输出
Cargo.toml
[package]
name = "duckdb-cli"
version = "0.1.0"
edition = "2024"[dependencies]
duckdb = "1.3.1"
syntect = "5.0"
atty = "0.2"
rustyline = "12.0"
main.rs
use std::{error::Error,io::{self, BufRead, Write},time::Instant,
};use duckdb::{params, Connection};
use rustyline::{completion::Completer,highlight::Highlighter,hint::Hinter,validate::Validator,Helper, history::History,
};
use syntect::{easy::HighlightLines,highlighting::{Style, ThemeSet},parsing::SyntaxSet,util::{as_24_bit_terminal_escaped, LinesWithEndings},
};fn main() -> Result<(), Box<dyn Error>> {// 初始化DuckDB连接let conn = Connection::open_in_memory()?;// 检查是否是管道输入if atty::isnt(atty::Stream::Stdin) {// 管道模式:从stdin读取SQL并执行let stdin = io::stdin();let mut stdout = io::stdout();for line in stdin.lock().lines() {let sql = line?;if sql.trim().is_empty() {continue;}let start_time = Instant::now();let mut stmt = conn.prepare(&sql)?;let rows = stmt.query_map(params![], |row| Ok(row.get::<_, String>(0)?))?;for row in rows {writeln!(stdout, "{}", row?)?;}eprintln!("Query executed in {:?}", start_time.elapsed());}} else {// 交互模式:println!("DuckDB CLI (v{})", env!("CARGO_PKG_VERSION"));println!("Enter SQL commands or .exit to quit");let mut rl = rustyline::Editor::<SQLHelper, rustyline::history::DefaultHistory>::new()?;loop {match rl.readline("duckdb> ") {Ok(line) => {let line = line.trim();if line.eq_ignore_ascii_case(".exit") {break;}if line.is_empty() {continue;}// 执行并计时let start_time = Instant::now();let mut stmt = conn.prepare(line)?;let rows = stmt.query_map(params![], |row| Ok(row.get::<_, String>(0)?))?;// 输出结果let mut count = 0;for row in rows {println!("{}", row?);count += 1;}println!("\n{} row{} in {:?}",count,if count == 1 { "" } else { "s" },start_time.elapsed());rl.add_history_entry(line);}Err(_) => break,}}}Ok(())
}struct SQLHelper {ps: SyntaxSet,ts: ThemeSet,
}impl SQLHelper {fn new() -> Self {Self {ps: SyntaxSet::load_defaults_newlines(),ts: ThemeSet::load_defaults(),}}
}impl Highlighter for SQLHelper {fn highlight<'l>(&self, line: &'l str, _pos: usize) -> std::borrow::Cow<'l, str> {let syntax = self.ps.find_syntax_by_extension("sql").unwrap();let mut h = HighlightLines::new(syntax, &self.ts.themes["base16-ocean.dark"]);let mut highlighted = String::new();for line in LinesWithEndings::from(line) {let ranges = h.highlight_line(line, &self.ps).unwrap();highlighted.push_str(&as_24_bit_terminal_escaped(&ranges[..], false));}std::borrow::Cow::Owned(highlighted)}fn highlight_prompt<'b, 's: 'b, 'p: 'b>(&'s self,prompt: &'p str,_default: bool,) -> std::borrow::Cow<'b, str> {std::borrow::Cow::Borrowed(prompt)}fn highlight_hint<'h>(&self, hint: &'h str) -> std::borrow::Cow<'h, str> {std::borrow::Cow::Borrowed(hint)}fn highlight_candidate<'c>(&self,candidate: &'c str,_completion: rustyline::CompletionType,) -> std::borrow::Cow<'c, str> {std::borrow::Cow::Borrowed(candidate)}fn highlight_char(&self, _line: &str, _pos: usize) -> bool {false}
}impl Completer for SQLHelper {type Candidate = String;fn complete(&self,_line: &str,_pos: usize,_ctx: &rustyline::Context<'_>,) -> rustyline::Result<(usize, Vec<Self::Candidate>)> {Ok((0, Vec::new()))}
}impl Hinter for SQLHelper {type Hint = String;
}impl Validator for SQLHelper {}impl Helper for SQLHelper {}
目前的功能很简单,就是输入查询语句,输出结果和计时信息,以及退出。让他实现带语法高亮的REPL,但实际上没有实现。(你看到了是因为csdn的代码渲染结果)
程序还有限制,只能输出字符类型,其他类型要用::转换,否则报错。
交互模式
./duckdb-cli
DuckDB CLI (v0.1.0)
Enter SQL commands or .exit to quit
duckdb> select 1 a;
Error: InvalidColumnType(0, "a", Int)./duckdb-cli
DuckDB CLI (v0.1.0)
Enter SQL commands or .exit to quit
duckdb> select 'xyz' a;
xyz1 row in 1.01681ms
duckdb> create table t(i int);0 rows in 3.20001ms
duckdb> insert into t select 1;
Error: InvalidColumnType(0, "Count", BigInt)
命令行管道模式
echo "select 'a'" |./duckdb-cli
a
Query executed in 1.03702msecho "select sum(i)::varchar from range(10000000)t(i)" |./duckdb-cli
49999995000000
Query executed in 127.35006ms