Support chDB

Author: wyhaya

Cargo.lock

@@ -0,0 +1,16 @@
+use chdb::Connection;
+
+#[tokio::main]
+async fn main() {
+    let conn = Connection::connect(":memory:").await.unwrap();
+
+    conn.execute(r"create table if not exists test (id UInt32, name String)")
+        .unwrap();
+    conn.execute("insert into test (id, name) values (1, 'Alice')")
+        .unwrap();
+    conn.execute("insert into test (id, name) values (2, 'Bob')")
+        .unwrap();
+
+    let query = conn.query(r#"select * from test"#).unwrap();
+    dbg!(query);
+}

src-crates/chdb/examples/chdb.rs

@@ -0,0 +1,86 @@
+use dylib::ffi::{ErrorMessage, StringRef, TypedValue};
+use std::ffi::c_void;
+
+#[repr(C)]
+pub struct Meta {
+    pub column_count: usize,
+    pub row_count: usize,
+    pub duration: u32,
+}
+
+#[repr(C)]
+pub struct Column {
+    pub name: StringRef,
+    pub datatype: StringRef,
+}
+
+#[repr(C)]
+#[derive(Debug, PartialEq, Eq)]
+#[allow(dead_code)]
+pub enum DataKind {
+    Null,
+    Bool,
+    I8,
+    I16,
+    I32,
+    I64,
+    U8,
+    U16,
+    U32,
+    U64,
+    F32,
+    F64,
+    String,
+}
+
+#[repr(C)]
+pub union Data {
+    pub null: (),
+    pub bool: bool,
+    pub i8: i8,
+    pub i16: i16,
+    pub i32: i32,
+    pub i64: i64,
+    pub u8: u8,
+    pub u16: u16,
+    pub u32: u32,
+    pub u64: u64,
+    pub f32: f32,
+    pub f64: f64,
+    pub string: StringRef,
+}
+
+#[repr(C)]
+pub struct ConnectOptions {
+    pub path: StringRef,
+}
+
+pub const CONNECT: &[u8] = b"df_connect";
+pub type ConnectFn =
+    extern "C" fn(options: ConnectOptions, error: *mut ErrorMessage) -> *mut c_void;
+
+pub const CLOSE: &[u8] = b"df_close";
+pub type CloseFn = extern "C" fn(conn: *mut c_void);
+
+pub const EXECUTE: &[u8] = b"df_execute";
+pub type ExecuteFn = extern "C" fn(conn: *mut c_void, sql: StringRef, error: *mut ErrorMessage);
+
+pub const QUERY: &[u8] = b"df_query";
+pub type QueryFn =
+    extern "C" fn(conn: *mut c_void, sql: StringRef, error: *mut ErrorMessage) -> *mut c_void;
+
+pub const QUERY_META: &[u8] = b"df_query_meta";
+pub type QueryMetaFn = extern "C" fn(query: *mut c_void) -> Meta;
+
+pub const QUERY_COLUMN: &[u8] = b"df_query_column";
+pub type QueryColumnFn = extern "C" fn(query: *mut c_void, index: usize) -> Column;
+
+pub const QUERY_VALUE: &[u8] = b"df_query_value";
+pub type QueryValueFn =
+    extern "C" fn(query: *mut c_void, row: usize, col: usize) -> TypedValue<DataKind, Data>;
+
+pub const FREE_QUERY: &[u8] = b"df_free_query";
+pub type FreeQueryFn = extern "C" fn(query: *mut c_void);
+
+pub const FREE_ERROR: &[u8] = b"df_free_error";
+pub type FreeErrorFn = extern "C" fn(error: ErrorMessage);

src-crates/chdb/src/ffi.rs

@@ -0,0 +1,174 @@
+mod ffi;
+
+use dylib::Dylib;
+use dylib::driver::{Error, Result};
+use dylib::ffi::{ErrorMessage, StringRef};
+use ffi::*;
+use query::{Query, QueryColumn, Value};
+use std::{ffi::c_void, sync::Mutex};
+
+// NOTE:
+// Do not update manually
+// Use `node ./src-dylib/driver-update.mjs` update the sha256 values.
+
+const CHDB_DRIVER_VERSION: &str = "20260603";
+#[cfg(all(target_os = "macos", target_arch = "aarch64"))]
+const CHDB_SHA256: &str = "82af2fa737c04737b0f15f2deb802338117d7c36279f828220dc2955cfa5c8fc";
+#[cfg(all(target_os = "macos", target_arch = "x86_64"))]
+const CHDB_SHA256: &str = "ecddfd0861931a9a43c623f76882873f08a3d0a312afc6ac893318237445ca78";
+#[cfg(all(target_os = "linux", target_arch = "aarch64", target_env = "gnu"))]
+const CHDB_SHA256: &str = "b60cf82d932fe24d38144513c9669118949a1003eff0325f8252f3096f06ef57";
+#[cfg(all(target_os = "linux", target_arch = "x86_64", target_env = "gnu"))]
+const CHDB_SHA256: &str = "6e77f2d46e51c71b9740fe50122e660732d325b4507ed9530c28e1decf311384";
+#[cfg(all(target_os = "windows", target_arch = "aarch64", target_env = "msvc"))]
+const CHDB_SHA256: &str = ""; // Unsupported platform
+#[cfg(all(target_os = "windows", target_arch = "x86_64", target_env = "msvc"))]
+const CHDB_SHA256: &str = ""; // Unsupported platform
+
+#[derive(Debug)]
+pub struct Connection {
+    conn: Mutex<*mut c_void>,
+    dylib: Dylib,
+}
+
+unsafe impl Send for Connection {}
+unsafe impl Sync for Connection {}
+
+impl Drop for Connection {
+    fn drop(&mut self) {
+        let _ = self.close();
+    }
+}
+
+fn free_error(dylib: &Dylib, error: ErrorMessage) -> Result<Option<Error>, Error> {
+    if !error.is_null() {
+        let msg = error.as_str().to_string();
+        dylib.symbol::<FreeErrorFn>(FREE_ERROR)?(error);
+        return Ok(Some(Error::Message(msg)));
+    }
+    Ok(None)
+}
+
+impl Connection {
+    pub async fn connect(path: &str) -> Result<Self> {
+        if CHDB_SHA256.is_empty() {
+            return Err(Error::Message(
+                "chDB is not supported on the current platform.".into(),
+            ));
+        }
+        let mut error = ErrorMessage::null();
+        let dylib = Dylib::try_load("chdb", CHDB_DRIVER_VERSION, CHDB_SHA256).await?;
+        let options = ConnectOptions {
+            path: StringRef::new(path),
+        };
+        let conn = dylib.symbol::<ConnectFn>(CONNECT)?(options, &mut error);
+        if let Some(err) = free_error(&dylib, error)? {
+            return Err(err);
+        }
+        Ok(Self {
+            conn: Mutex::new(conn),
+            dylib,
+        })
+    }
+
+    fn close(&self) -> Result<(), Error> {
+        let conn = self.conn.lock().map_err(|_| Error::Mutex)?;
+        self.dylib.symbol::<CloseFn>(CLOSE)?(*conn);
+        Ok(())
+    }
+
+    pub fn execute(&self, sql: &str) -> Result<(), Error> {
+        let conn = self.conn.lock().map_err(|_| Error::Mutex)?;
+        let mut error = ErrorMessage::null();
+        self.dylib.symbol::<ExecuteFn>(EXECUTE)?(*conn, StringRef::new(sql), &mut error);
+        if let Some(err) = free_error(&self.dylib, error)? {
+            return Err(err);
+        }
+        Ok(())
+    }
+
+    pub fn query(&self, sql: &str) -> Result<Query, Error> {
+        let conn = self.conn.lock().map_err(|_| Error::Mutex)?;
+        let mut error = ErrorMessage::null();
+        let query = self.dylib.symbol::<QueryFn>(QUERY)?(*conn, StringRef::new(sql), &mut error);
+        if let Some(err) = free_error(&self.dylib, error)? {
+            return Err(err);
+        }
+        let meta = self.dylib.symbol::<QueryMetaFn>(QUERY_META)?(query);
+        let query_column = self.dylib.symbol::<QueryColumnFn>(QUERY_COLUMN)?;
+        let query_value = self.dylib.symbol::<QueryValueFn>(QUERY_VALUE)?;
+        let columns = (0..meta.column_count)
+            .map(|i| {
+                let col = query_column(query, i);
+                QueryColumn {
+                    name: col.name.as_str().to_string(),
+                    datatype: col.datatype.as_str().to_string(),
+                }
+            })
+            .collect::<Vec<_>>();
+        let mut rows = Vec::with_capacity(meta.row_count);
+        for y in 0..meta.row_count {
+            let mut row = Vec::with_capacity(meta.column_count);
+            for x in 0..meta.column_count {
+                let data = query_value(query, y, x);
+                row.push(unsafe {
+                    match data.kind {
+                        DataKind::Null => Value::Null,
+                        DataKind::Bool => Value::Bool(data.value.bool),
+                        DataKind::I8 => Value::I8(data.value.i8),
+                        DataKind::I16 => Value::I16(data.value.i16),
+                        DataKind::I32 => Value::I32(data.value.i32),
+                        DataKind::I64 => Value::I64(data.value.i64),
+                        DataKind::U8 => Value::U8(data.value.u8),
+                        DataKind::U16 => Value::U16(data.value.u16),
+                        DataKind::U32 => Value::U32(data.value.u32),
+                        DataKind::U64 => Value::U64(data.value.u64),
+                        DataKind::F32 => Value::F32(data.value.f32),
+                        DataKind::F64 => Value::F64(data.value.f64),
+                        DataKind::String => Value::String(data.value.string.as_str().to_string()),
+                    }
+                });
+            }
+            rows.push(row);
+        }
+        self.dylib.symbol::<FreeQueryFn>(FREE_QUERY)?(query);
+        Ok(Query {
+            columns,
+            rows,
+            // TODO
+            rows_affected: None,
+            duration: meta.duration,
+        })
+    }
+
+    pub fn select(&self, sql: &str) -> Result<Vec<Vec<Value>>, Error> {
+        self.query(sql).map(|q| q.rows)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::*;
+    use query::{QueryColumn, Value};
+
+    #[tokio::test]
+    async fn test_query() {
+        let conn = Connection::connect(":memory:").await.unwrap();
+        let mut query = conn.query("select 'hello' as hello").unwrap();
+        assert_eq!(query.columns.len(), 1);
+        assert_eq!(
+            query.columns.remove(0),
+            QueryColumn {
+                name: "hello".into(),
+                datatype: "String".into()
+            }
+        );
+        assert_eq!(query.rows.len(), 1);
+        assert_eq!(query.rows[0].len(), 1);
+        assert_eq!(
+            query.rows.remove(0).remove(0),
+            Value::String("hello".into())
+        );
+        assert_eq!(query.rows_affected, None);
+    }
+}

src-crates/chdb/src/lib.rs

@@ -18,6 +18,9 @@ pub enum ConnectionConfig {
     ManticoreSearch(ManticoreSearchConfig),
     MSSQL(MsSqlConfig),
     ClickHouse(ClickHouseConfig),
+    #[serde(rename = "chDB")]
+    #[strum(serialize = "chDB")]
+    ChDb(ChDbConfig),
     Databend(DatabendConfig),
     BigQuery(BigQueryConfig),
     Trino(TrinoConfig),
@@ -74,6 +77,8 @@ pub enum SqlDatabaseType {
     ManticoreSearch,
     MSSQL,
     ClickHouse,
+    #[serde(rename = "chDB")]
+    ChDb,
     Databend,
     BigQuery,
     Trino,
@@ -117,6 +122,13 @@ pub struct ClickHouseConfig {
     pub proxy: Option<ProxyConfig>,
 }
 
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct ChDbConfig {
+    pub path: String,
+    pub readonly: bool,
+    pub initial: Option<String>,
+}
+
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct DatabendConfig {
     pub protocol: ConnectProtocol,

src-crates/connection-config/connection_config.rs

@@ -25,6 +25,7 @@ r2sql = { path = "../r2sql" }
 mysql = { path = "../mysql" }
 pgsq = { path = "../pgsq" }
 duckdb = { path = "../duckdb" }
+chdb = { path = "../chdb" }
 clickhouse = { path = "../clickhouse" }
 databend = { path = "../databend" }
 bigquery = { path = "../bigquery" }