use std::sync::Arc;
use duckdb::types::Value as DuckValue;
use tokio::sync::{Mutex, Semaphore};
use tokio::task::JoinHandle;
use crate::foundation::{DbError, DbResult};
#[derive(Debug, Clone, PartialEq)]
pub struct DuckDbRow {
pub columns: Vec<(String, DuckValue)>,
}
impl DuckDbRow {
pub fn get(&self, column_name: &str) -> Option<&DuckValue> {
self.columns
.iter()
.find(|(name, _)| name == column_name)
.map(|(_, value)| value)
}
pub fn column_count(&self) -> usize {
self.columns.len()
}
}
#[derive(Debug, Clone)]
pub struct DuckDbExecResult {
pub rows_affected: usize,
}
const DEFAULT_POOL_SIZE: usize = 4;
#[derive(Clone)]
pub struct DuckDbConnection {
pool: Arc<Mutex<Vec<duckdb::Connection>>>,
pool_size: usize,
spawn_permit: Arc<Semaphore>,
}
impl DuckDbConnection {
pub fn new(url: &str) -> Result<Self, DbError> {
Self::with_pool_size(url, DEFAULT_POOL_SIZE)
}
pub fn with_pool_size(url: &str, pool_size: usize) -> Result<Self, DbError> {
let pool_size = pool_size.max(1);
let db_path = Self::parse_url(url);
let primary = duckdb::Connection::open(&db_path)
.map_err(|e| DbError::Connection(sea_orm::DbErr::Custom(format!("DuckDB connection failed: {e}"))))?;
let mut pool = Vec::with_capacity(pool_size);
pool.push(primary);
for i in 1..pool_size {
let cloned = pool[0].try_clone().map_err(|e| {
DbError::Connection(sea_orm::DbErr::Custom(format!(
"DuckDB try_clone failed for connection {}: {e}",
i + 1
)))
})?;
pool.push(cloned);
}
Ok(Self {
pool: Arc::new(Mutex::new(pool)),
pool_size,
spawn_permit: Arc::new(Semaphore::new(pool_size)),
})
}
fn parse_url(url: &str) -> String {
let lower = url.to_lowercase();
if lower == ":memory:" || lower == "duckdb::memory:" {
return ":memory:".to_string();
}
if let Some(rest) = url.strip_prefix("duckdb:") {
return rest.trim_start_matches('/').to_string();
}
url.to_string()
}
pub async fn execute(&self, sql: &str) -> DbResult<DuckDbExecResult> {
let permit = self.acquire_permit().await?;
let conn = {
let mut pool = self.pool.lock().await;
pool.pop().ok_or_else(|| {
DbError::Connection(sea_orm::DbErr::Custom(
"DuckDB pool exhausted: no connection available".to_string(),
))
})?
};
let sql_owned = sql.to_string();
let handle: JoinHandle<DbResult<(duckdb::Connection, DuckDbExecResult)>> =
tokio::task::spawn_blocking(move || {
let rows_affected = conn
.execute(&sql_owned, [])
.map_err(|e| DbError::Connection(sea_orm::DbErr::Custom(format!("DuckDB execute failed: {e}"))))?;
Ok((conn, DuckDbExecResult { rows_affected }))
});
let result = handle
.await
.map_err(|e| DbError::Connection(sea_orm::DbErr::Custom(format!("spawn_blocking join failed: {e}"))))?;
drop(permit);
let (conn, exec_result) = result?;
{
let mut pool = self.pool.lock().await;
pool.push(conn);
}
Ok(exec_result)
}
pub async fn query(&self, sql: &str) -> DbResult<Vec<DuckDbRow>> {
let permit = self.acquire_permit().await?;
let conn = {
let mut pool = self.pool.lock().await;
pool.pop().ok_or_else(|| {
DbError::Connection(sea_orm::DbErr::Custom(
"DuckDB pool exhausted: no connection available".to_string(),
))
})?
};
let sql_owned = sql.to_string();
let handle: JoinHandle<DbResult<(duckdb::Connection, Vec<DuckDbRow>)>> =
tokio::task::spawn_blocking(move || {
let mut stmt = conn
.prepare(&sql_owned)
.map_err(|e| DbError::Connection(sea_orm::DbErr::Custom(format!("DuckDB prepare failed: {e}"))))?;
let rows = stmt
.query_map([], |row| {
let stmt_ref = row.as_ref();
let column_count = stmt_ref.column_count();
let column_names: Vec<String> = (0..column_count)
.map(|i| stmt_ref.column_name(i).ok().map(|s| s.to_string()).unwrap_or_default())
.collect();
let mut columns = Vec::with_capacity(column_count);
for (i, name) in column_names.iter().enumerate() {
let value: DuckValue = row.get(i).unwrap_or(DuckValue::Null);
columns.push((name.clone(), value));
}
Ok(DuckDbRow { columns })
})
.map_err(|e| DbError::Connection(sea_orm::DbErr::Custom(format!("DuckDB query failed: {e}"))))?;
let mut result = Vec::new();
for row_result in rows {
let row = row_result.map_err(|e| {
DbError::Connection(sea_orm::DbErr::Custom(format!("DuckDB row fetch failed: {e}")))
})?;
result.push(row);
}
drop(stmt);
Ok((conn, result))
});
let result = handle
.await
.map_err(|e| DbError::Connection(sea_orm::DbErr::Custom(format!("spawn_blocking join failed: {e}"))))?;
drop(permit);
let (conn, rows) = result?;
{
let mut pool = self.pool.lock().await;
pool.push(conn);
}
Ok(rows)
}
pub async fn health_check(&self) -> DbResult<()> {
let rows = self.query("SELECT 1 AS health").await?;
if rows.is_empty() {
return Err(DbError::Connection(sea_orm::DbErr::Custom(
"DuckDB health check returned no rows".to_string(),
)));
}
Ok(())
}
pub fn pool_size(&self) -> usize {
self.pool_size
}
async fn acquire_permit(&self) -> DbResult<tokio::sync::SemaphorePermit<'_>> {
self.spawn_permit
.acquire()
.await
.map_err(|_| DbError::Connection(sea_orm::DbErr::Custom("Semaphore closed".to_string())))
}
}
impl std::fmt::Debug for DuckDbConnection {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("DuckDbConnection")
.field("pool_size", &self.pool_size)
.field("max_concurrency", &self.pool_size)
.finish()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn test_duckdb_connection_create_memory() {
let conn = DuckDbConnection::new(":memory:").expect("Failed to create memory connection");
assert_eq!(conn.pool_size(), DEFAULT_POOL_SIZE);
assert_eq!(DEFAULT_POOL_SIZE, 4);
let _ = conn;
}
#[tokio::test]
async fn test_duckdb_connection_create_via_url() {
let conn = DuckDbConnection::new("duckdb::memory:").expect("Failed to create connection via URL");
let _ = conn;
}
#[tokio::test]
async fn test_duckdb_execute_create_table() {
let conn = DuckDbConnection::new(":memory:").expect("Failed to create connection");
let result = conn
.execute("CREATE TABLE test_table (id INTEGER PRIMARY KEY, name VARCHAR)")
.await
.expect("Failed to create table");
assert_eq!(result.rows_affected, 0);
}
#[tokio::test]
async fn test_duckdb_execute_insert_and_query() {
let conn = DuckDbConnection::new(":memory:").expect("Failed to create connection");
conn.execute("CREATE TABLE users (id INTEGER PRIMARY KEY, name VARCHAR)")
.await
.expect("Failed to create table");
conn.execute("INSERT INTO users VALUES (1, 'Alice')")
.await
.expect("Failed to insert");
conn.execute("INSERT INTO users VALUES (2, 'Bob')")
.await
.expect("Failed to insert");
let rows = conn
.query("SELECT id, name FROM users ORDER BY id")
.await
.expect("Failed to query");
assert_eq!(rows.len(), 2);
assert_eq!(rows[0].column_count(), 2);
let name = rows[0].get("name").expect("Failed to get name column");
if let DuckValue::Text(s) = name {
assert_eq!(s, "Alice");
} else {
panic!("Expected Text value, got {:?}", name);
}
}
#[tokio::test]
async fn test_duckdb_health_check() {
let conn = DuckDbConnection::new(":memory:").expect("Failed to create connection");
conn.health_check().await.expect("Health check should pass");
}
#[tokio::test]
async fn test_duckdb_parse_url_variants() {
assert_eq!(DuckDbConnection::parse_url(":memory:"), ":memory:");
assert_eq!(DuckDbConnection::parse_url("duckdb::memory:"), ":memory:");
assert_eq!(DuckDbConnection::parse_url("duckdb:test.db"), "test.db");
assert_eq!(
DuckDbConnection::parse_url("duckdb://path/to/file.db"),
"path/to/file.db"
);
assert_eq!(DuckDbConnection::parse_url("/absolute/path.db"), "/absolute/path.db");
}
#[tokio::test]
async fn test_duckdb_concurrent_execute_respects_semaphore() {
let conn = Arc::new(DuckDbConnection::new(":memory:").expect("Failed to create connection"));
conn.execute("CREATE TABLE concurrent_test (id INTEGER)")
.await
.expect("Failed to create table");
let mut handles = Vec::new();
for i in 0..8 {
let conn_clone = conn.clone();
handles.push(tokio::spawn(async move {
conn_clone
.execute(&format!("INSERT INTO concurrent_test VALUES ({i})"))
.await
}));
}
for handle in handles {
let result = handle.await.expect("Task panicked");
assert!(result.is_ok(), "Concurrent insert should succeed");
}
let rows = conn
.query("SELECT COUNT(*) AS cnt FROM concurrent_test")
.await
.expect("Failed to count");
assert_eq!(rows.len(), 1);
let count = rows[0].get("cnt").expect("Failed to get count");
if let DuckValue::BigInt(n) = count {
assert_eq!(*n, 8);
} else {
panic!("Expected BigInt, got {:?}", count);
}
}
#[tokio::test]
async fn test_duckdb_pool_shares_memory_database() {
let conn = DuckDbConnection::new(":memory:").expect("Failed to create connection");
conn.execute("CREATE TABLE shared_test (id INTEGER PRIMARY KEY, val VARCHAR)")
.await
.expect("Failed to create table");
conn.execute("INSERT INTO shared_test VALUES (1, 'hello')")
.await
.expect("Failed to insert");
let rows = conn
.query("SELECT val FROM shared_test WHERE id = 1")
.await
.expect("Failed to query");
assert_eq!(rows.len(), 1);
let val = rows[0].get("val").expect("Failed to get val");
if let DuckValue::Text(s) = val {
assert_eq!(s, "hello");
} else {
panic!("Expected Text, got {:?}", val);
}
}
#[tokio::test]
async fn test_duckdb_custom_pool_size() {
let conn =
DuckDbConnection::with_pool_size(":memory:", 2).expect("Failed to create connection with pool size 2");
assert_eq!(conn.pool_size(), 2);
conn.execute("CREATE TABLE custom_pool_test (id INTEGER)")
.await
.expect("Failed to create table");
conn.execute("INSERT INTO custom_pool_test VALUES (42)")
.await
.expect("Failed to insert");
let rows = conn
.query("SELECT id FROM custom_pool_test")
.await
.expect("Failed to query");
assert_eq!(rows.len(), 1);
}
#[tokio::test]
async fn test_duckdb_pool_concurrent_queries_use_different_connections() {
let conn = Arc::new(
DuckDbConnection::with_pool_size(":memory:", 4).expect("Failed to create connection with pool size 4"),
);
conn.execute("CREATE TABLE parallel_test (id INTEGER, thread_id INTEGER)")
.await
.expect("Failed to create table");
let mut handles = Vec::new();
for i in 0..4 {
let conn_clone = conn.clone();
handles.push(tokio::spawn(async move {
conn_clone
.execute(&format!("INSERT INTO parallel_test VALUES ({i}, {i})"))
.await
}));
}
for (i, handle) in handles.into_iter().enumerate() {
let result = handle.await.expect("Task panicked");
assert!(result.is_ok(), "Task {} should succeed: {:?}", i, result);
}
let rows = conn
.query("SELECT COUNT(*) AS cnt FROM parallel_test")
.await
.expect("Failed to count");
let count = rows[0].get("cnt").expect("Failed to get count");
if let DuckValue::BigInt(n) = count {
assert_eq!(*n, 4, "All 4 concurrent inserts should succeed");
} else {
panic!("Expected BigInt, got {:?}", count);
}
}
#[tokio::test]
async fn test_duckdb_data_types_boolean() {
let conn = DuckDbConnection::new(":memory:").expect("Failed to create connection");
conn.execute("CREATE TABLE bool_test (id INTEGER, flag BOOLEAN)")
.await
.expect("create table");
conn.execute("INSERT INTO bool_test VALUES (1, true), (2, false)")
.await
.expect("insert");
let rows = conn
.query("SELECT flag FROM bool_test ORDER BY id")
.await
.expect("query");
assert_eq!(rows.len(), 2);
match &rows[0].get("flag").expect("should have flag") {
DuckValue::Boolean(b) => assert!(*b, "first row should be true"),
other => panic!("Expected Boolean, got {:?}", other),
}
match &rows[1].get("flag").expect("should have flag") {
DuckValue::Boolean(b) => assert!(!*b, "second row should be false"),
other => panic!("Expected Boolean, got {:?}", other),
}
}
#[tokio::test]
async fn test_duckdb_data_types_double() {
let conn = DuckDbConnection::new(":memory:").expect("Failed to create connection");
conn.execute("CREATE TABLE double_test (val DOUBLE)")
.await
.expect("create table");
conn.execute("INSERT INTO double_test VALUES (3.14), (-0.001)")
.await
.expect("insert");
let rows = conn
.query("SELECT val FROM double_test ORDER BY val")
.await
.expect("query");
assert_eq!(rows.len(), 2);
match &rows[0].get("val").expect("should have val") {
DuckValue::Double(f) => assert!((*f - (-0.001)).abs() < 1e-10, "first should be -0.001"),
other => panic!("Expected Double, got {:?}", other),
}
}
#[tokio::test]
async fn test_duckdb_null_handling() {
let conn = DuckDbConnection::new(":memory:").expect("Failed to create connection");
conn.execute("CREATE TABLE null_test (id INTEGER, name VARCHAR)")
.await
.expect("create table");
conn.execute("INSERT INTO null_test VALUES (1, NULL), (2, 'hello')")
.await
.expect("insert");
let rows = conn
.query("SELECT name FROM null_test ORDER BY id")
.await
.expect("query");
assert_eq!(rows.len(), 2);
match &rows[0].get("name").expect("should have name column") {
DuckValue::Null => {} other => panic!("Expected Null, got {:?}", other),
}
match &rows[1].get("name").expect("should have name column") {
DuckValue::Text(s) => assert_eq!(s, "hello"),
other => panic!("Expected Text, got {:?}", other),
}
}
#[tokio::test]
async fn test_duckdb_syntax_error_returns_error() {
let conn = DuckDbConnection::new(":memory:").expect("Failed to create connection");
let result = conn.execute("CREAT TABL broken (id INTEGER)").await;
assert!(result.is_err(), "syntax error should return error");
let err = result.unwrap_err();
let msg = format!("{err}");
assert!(msg.contains("DuckDB"), "error should mention DuckDB: {msg}");
}
#[tokio::test]
async fn test_duckdb_table_not_exists_returns_error() {
let conn = DuckDbConnection::new(":memory:").expect("Failed to create connection");
let result = conn.query("SELECT * FROM nonexistent_table").await;
assert!(result.is_err(), "query on nonexistent table should error");
}
#[tokio::test]
async fn test_duckdb_insert_duplicate_pk_returns_error() {
let conn = DuckDbConnection::new(":memory:").expect("Failed to create connection");
conn.execute("CREATE TABLE pk_test (id INTEGER PRIMARY KEY)")
.await
.expect("create table");
conn.execute("INSERT INTO pk_test VALUES (1)")
.await
.expect("first insert");
let result = conn.execute("INSERT INTO pk_test VALUES (1)").await;
assert!(result.is_err(), "duplicate PK should return error");
}
#[tokio::test]
async fn test_duckdb_row_get_nonexistent_column_returns_none() {
let conn = DuckDbConnection::new(":memory:").expect("Failed to create connection");
conn.execute("CREATE TABLE row_test (id INTEGER)")
.await
.expect("create table");
conn.execute("INSERT INTO row_test VALUES (42)").await.expect("insert");
let rows = conn.query("SELECT id FROM row_test").await.expect("query");
assert_eq!(rows.len(), 1);
assert!(
rows[0].get("nonexistent").is_none(),
"get() with unknown column should return None"
);
assert!(rows[0].get("").is_none(), "get() with empty string should return None");
}
#[tokio::test]
async fn test_duckdb_row_column_count_empty_result() {
let conn = DuckDbConnection::new(":memory:").expect("Failed to create connection");
conn.execute("CREATE TABLE empty_test (a INTEGER, b VARCHAR, c DOUBLE)")
.await
.expect("create table");
let rows = conn.query("SELECT a, b, c FROM empty_test").await.expect("query");
assert_eq!(rows.len(), 0, "empty table should return 0 rows");
}
#[tokio::test]
async fn test_duckdb_connection_debug_format() {
let conn = DuckDbConnection::with_pool_size(":memory:", 3).expect("Failed to create connection");
let debug_str = format!("{:?}", conn);
assert!(
debug_str.contains("DuckDbConnection"),
"Debug should contain struct name"
);
assert!(debug_str.contains("pool_size: 3"), "Debug should contain pool_size");
}
#[tokio::test]
async fn test_duckdb_file_database_persistence() {
let mut db_path = std::env::temp_dir();
db_path.push(format!("dbnexus_duckdb_test_{}.db", std::process::id()));
if let Some(parent) = db_path.parent() {
let _ = std::fs::create_dir_all(parent);
}
let url = format!("duckdb:{}", db_path.display());
{
let conn = DuckDbConnection::new(&url).expect("Failed to create file connection");
conn.execute("CREATE TABLE persist_test (id INTEGER, val VARCHAR)")
.await
.expect("create table");
conn.execute("INSERT INTO persist_test VALUES (1, 'durable')")
.await
.expect("insert");
}
let conn2 = DuckDbConnection::new(&url).expect("Failed to reopen file connection");
let rows = conn2
.query("SELECT val FROM persist_test WHERE id = 1")
.await
.expect("query after reopen");
assert_eq!(rows.len(), 1, "data should persist across connections");
match &rows[0].get("val").expect("should have val") {
DuckValue::Text(s) => assert_eq!(s, "durable"),
other => panic!("Expected Text, got {:?}", other),
}
drop(conn2);
let _ = std::fs::remove_file(&db_path);
}
#[tokio::test]
async fn test_duckdb_pool_exhaustion_queues_requests() {
let conn = Arc::new(DuckDbConnection::with_pool_size(":memory:", 1).expect("Failed to create connection"));
conn.execute("CREATE TABLE queue_test (id INTEGER)")
.await
.expect("create table");
for i in 0..5 {
conn.execute(&format!("INSERT INTO queue_test VALUES ({i})"))
.await
.expect("insert should succeed after queuing");
}
let rows = conn
.query("SELECT COUNT(*) AS cnt FROM queue_test")
.await
.expect("count");
let count = rows[0].get("cnt").expect("should have cnt");
if let DuckValue::BigInt(n) = count {
assert_eq!(*n, 5, "all 5 inserts should succeed with queuing");
} else {
panic!("Expected BigInt, got {:?}", count);
}
}
#[tokio::test]
async fn test_duckdb_parse_url_case_insensitive() {
assert_eq!(DuckDbConnection::parse_url(":MEMORY:"), ":memory:");
assert_eq!(DuckDbConnection::parse_url("DuckDB::memory:"), ":memory:");
assert_eq!(DuckDbConnection::parse_url("duckdb:test.db"), "test.db");
}
}