利用DeepSeek实现rust调用duckdb动态链接库的duckdb CLI

Rust调用DuckDB实现CLI及功能介绍

提示词:请用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;
xyz

1 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.03702ms

echo "select sum(i)::varchar from range(10000000)t(i)" |./duckdb-cli
49999995000000
Query executed in 127.35006ms
评论
成就一亿技术人!
拼手气红包6.0元
还能输入1000个字符
 
红包 添加红包
表情包 插入表情
 条评论被折叠 查看
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值