DuckDB Local UI みたいな拡張を Rust で書いてみようとしたときのメモ（v1.2.1 時点）
え、DuckDB の拡張って Rust で書けんの？？、という人は、以下の記事に状況をまとめたのでそっちを読んでください。ここではもっと細かい話をします。
DuckDB Local UI
MotherDuck の UI みたいなやつが表示できる拡張です。CALL start_ui() とやると、ウェブブラウザが立ち上がってブラウザ上からいろいろできるようになります。
私がこれを見て思ったのは、「え、拡張ってこんなことできるの？？」ということでした。上にリンクを貼ったように、DuckDB の拡張というのは、
- table function: テーブルを返す関数（例：
read_parquet()） - scalar function: 値を受け取って同じ長さの値を返す関数（例：
sqrt()） - aggregate function: 値を受け取って1つに集約した値を返す関数（例：
sum()） - replacement scan: table function を自動で使えるようにする仕組み
というようなもので、こんな飛び道具的な拡張が書けるとは思ってもいませんでした。おもしろそうなので自分でも何か作ってみたい、と思って調べたことをメモります。
拡張のタイプ
DuckDB Local UI の start_ui() といった関数は、table function として実装されています。目的は UI のサーバーを起動するという副作用ですが、形式としてはテーブルを返しています。ということで、自分で作るのもこれに従ってみることにします。
ブラウザを立ち上げる
duckdb-ui の実装を読んでみる
まず、duckdb の中からブラウザを立ち上げるのどうやるんだろう、と思ったら、ふつうにコマンドを実行してるだけでした。
Rust で実装する
Rust の場合は、std::process::Command が Windows だとうまく動かないらしく、Command::new("start") だと command not found になるんですが、 cmd を介して立ち上げるようにするとうまくいきました。それ以外は特筆すべき点は特になさそうです。
use std::process::Command;
#[cfg(target_os = "windows")]
fn open_browser(url: &str) {
// It seems new("start").arg(url) doesn't work...
Command::new("cmd")
.args(["/c", "start", url])
.spawn()
.unwrap();
}
#[cfg(target_os = "macos")]
fn open_browser(url: &str) {
Command::new("open").arg(url).spawn().unwrap();
}
#[cfg(target_os = "linux")]
fn open_browser(url: &str) {
Command::new("xdg-open").arg(url).spawn().unwrap();
}
サーバーを立ち上げる
duckdb-ui の実装を読んでみる
httplib という header-only のシンプルな実装の HTTP サーバーを使っているようです。
Rust で実装する
HTTP サーバーは何でもいいと思いますが、私は warp を使ってみました。#[tokio::main] じゃなくて、自分で tokio::runtime::Runtime を作って使うところが少し情報が少ないところかもしれません。私もあんまりわかってませんが、とりあえずこれで動きました。
（注：このままだと、関数を実行するたびにサーバーが重複して起動してしまいますが、そのあたりの考慮はとりあえず抜きにしています）
use tokio::runtime::Runtime;
use warp::Filter;
static RUNTIME: LazyLock<Runtime> = LazyLock::new(|| Runtime::new().unwrap());
struct HelloVTab;
impl VTab for HelloVTab {
...
fn init(_: &InitInfo) -> Result<Self::InitData, Box<dyn std::error::Error>> {
let get = warp::get().map(|| {
warp::reply::html("ここに HTML が入る")
});
RUNTIME.spawn(async move { warp::serve(get).run(([127, 0, 0, 1], 3030)).await });
// ほんとはサーバーの準備ができるまで待ってからの方がいいのかも
open_browser("http://127.0.0.1:3030");
Ok(HelloInitData { ... })
}
...
拡張の側から DuckDB を操作する
私が DuckDB Local UI を見て一番びっくりしたところはここで、それまで拡張というのは受動的なもの（値を受け取って返す、みたいな感じ）だというイメージだったので、拡張の側から DuckDB 内のデータを覗いたりクエリを実行したりできるのはどういう仕掛けなんだろう？と調べてみました。
duckdb-ui の実装を読んでみる
うまく追いきれなかったのですが、start_ui() の中身の関数は ClientContext というオブジェクトを受け取っています。
このオブジェクトは、いろいろな情報を保持しています。特に db に注目です。
HttpServer は、この ClientContext から db を取り出して内部に保持していて、これに対してクエリを投げたりしているようでした。
Rust で実装する
さて、では Rust でこれをどう実装するかというと...
どうも、いい方法がなさそうです。table function を実装するための trait（VTab）は以下のようになっていますが、見るとわかるように、 ClientContext に類するようなものは存在しません。そもそもインターフェースが違うみたいです。
pub trait VTab: Sized {
type InitData: Sized + Send + Sync;
type BindData: Sized + Send + Sync;
// Required methods
fn bind(bind: &BindInfo) -> Result<Self::BindData, Box<dyn Error>>;
fn init(init: &InitInfo) -> Result<Self::InitData, Box<dyn Error>>;
fn func(
func: &TableFunctionInfo<Self>,
output: &mut DataChunkHandle,
) -> Result<(), Box<dyn Error>>;
// Provided methods
fn supports_pushdown() -> bool { ... }
fn parameters() -> Option<Vec<LogicalTypeHandle>> { ... }
fn named_parameters() -> Option<Vec<(String, LogicalTypeHandle)>> { ... }
}
いちおう、↑の関数をそのセッションに読み込むための関数があって、そこでは Connection があるので、それを持っておいて使うことはできそうです。ただ、この拡張が読み込まれた後に DB が切り替わると何が起こるかとか、これが安全な方法なのかはいまいち自信が持てません。。まあ、これで動いてはいます。
static CONN: OnceLock<Mutex<Connection>> = OnceLock::new();
#[duckdb_entrypoint_c_api()]
pub unsafe fn extension_entrypoint(con: Connection) -> Result<(), Box<dyn Error>> {
con.register_table_function::<HelloVTab>(EXTENSION_NAME)
.expect("Failed to register hello table function");
let _ = CONN.set(Mutex::new(con.try_clone()?));
Ok(())
}
ちなみに、Connection はこんな感じで使えます。.pragma_query() は PRAGMA 文を実行するためのもので、通常の SELECT 文などは query_row() で実行できます。
let mut tables = vec![];
let guard = CONN.get().unwrap().lock().unwrap();
guard
.pragma_query(None, "show_tables", |row| {
let t: String = row.get(0)?;
tables.push(format!("<li>{t}</li>"));
Ok(())
})
.unwrap();
感想
ということで何の成果も得られませんでした。もうちょっと仕組みが整ってくれないとなんともなあ...という感じなので、ここで力尽きるかもしれませんが、いちおう作りかけのコードは以下のレポジトリにあります。
Discussion