238 lines
7.1 KiB
Rust
238 lines
7.1 KiB
Rust
//! DuckDB 连接管理
|
||
//!
|
||
//! 提供线程安全的数据库连接管理
|
||
|
||
use crate::error::{DuckDbError, DuckDbResult};
|
||
use duckdb::Connection;
|
||
use parking_lot::Mutex;
|
||
use std::sync::Arc;
|
||
|
||
/// 线程安全的 DuckDB 连接包装器
|
||
///
|
||
/// 使用 `Arc<Mutex<Connection>>` 实现线程安全访问
|
||
#[derive(Clone)]
|
||
pub struct DuckDbConnection {
|
||
inner: Arc<Mutex<Connection>>,
|
||
}
|
||
|
||
impl DuckDbConnection {
|
||
/// DuckDB 内存限制默认值（2GB）
|
||
const DEFAULT_MEMORY_LIMIT: &str = "2GB";
|
||
|
||
/// 创建内存模式的数据库连接
|
||
pub fn open_in_memory() -> DuckDbResult<Self> {
|
||
let conn = Connection::open_in_memory()
|
||
.map_err(|e| DuckDbError::ConnectionError(e.to_string()))?;
|
||
|
||
// 设置 DuckDB 内存限制，防止无限占用内存
|
||
// 注意：memory_limit 只限制 buffer manager，部分数据结构（如向量、查询结果）
|
||
// 可能在 buffer manager 外部分配，因此实际内存可能略高于此限制
|
||
conn.execute(
|
||
&format!("PRAGMA memory_limit='{}'", Self::DEFAULT_MEMORY_LIMIT),
|
||
[],
|
||
)
|
||
.map_err(|e| DuckDbError::ConnectionError(format!("Failed to set memory_limit: {}", e)))?;
|
||
|
||
Ok(Self {
|
||
inner: Arc::new(Mutex::new(conn)),
|
||
})
|
||
}
|
||
|
||
/// 创建内存模式的数据库连接（带自定义内存限制）
|
||
///
|
||
/// # Arguments
|
||
/// * `memory_limit` - 内存限制字符串，如 "2GB", "512MB", "1GB"
|
||
pub fn open_in_memory_with_limit(memory_limit: &str) -> DuckDbResult<Self> {
|
||
let conn = Connection::open_in_memory()
|
||
.map_err(|e| DuckDbError::ConnectionError(e.to_string()))?;
|
||
|
||
// 设置自定义内存限制
|
||
conn.execute(
|
||
&format!("PRAGMA memory_limit='{}'", memory_limit),
|
||
[],
|
||
)
|
||
.map_err(|e| DuckDbError::ConnectionError(format!("Failed to set memory_limit: {}", e)))?;
|
||
|
||
Ok(Self {
|
||
inner: Arc::new(Mutex::new(conn)),
|
||
})
|
||
}
|
||
|
||
/// 执行需要独占连接的操作
|
||
///
|
||
/// 使用闭包模式确保锁的正确释放
|
||
pub fn with_connection<F, T>(&self, f: F) -> DuckDbResult<T>
|
||
where
|
||
F: FnOnce(&Connection) -> DuckDbResult<T>,
|
||
{
|
||
let conn = self.inner.lock();
|
||
f(&conn)
|
||
}
|
||
|
||
/// 执行需要可变连接的操作
|
||
pub fn with_connection_mut<F, T>(&self, f: F) -> DuckDbResult<T>
|
||
where
|
||
F: FnOnce(&mut Connection) -> DuckDbResult<T>,
|
||
{
|
||
let mut conn = self.inner.lock();
|
||
f(&mut conn)
|
||
}
|
||
|
||
/// 尝试获取连接（非阻塞）
|
||
pub fn try_with_connection<F, T>(&self, f: F) -> DuckDbResult<Option<T>>
|
||
where
|
||
F: FnOnce(&Connection) -> DuckDbResult<T>,
|
||
{
|
||
match self.inner.try_lock() {
|
||
Some(conn) => Ok(Some(f(&conn)?)),
|
||
None => Ok(None),
|
||
}
|
||
}
|
||
|
||
/// 获取 DuckDB 内存使用统计
|
||
///
|
||
/// 返回格式化的内存使用信息字符串，用于调试和监控
|
||
pub fn get_memory_stats(&self) -> DuckDbResult<String> {
|
||
self.with_connection(|c| {
|
||
// 查询内存使用情况
|
||
let mut stmt = c.prepare(
|
||
"SELECT name, size, reservation FROM duckdb_memory() ORDER BY size DESC LIMIT 10",
|
||
)?;
|
||
let mut rows = stmt.query([])?;
|
||
|
||
let mut result = String::from("DuckDB Memory Usage (Top 10):\n");
|
||
while let Some(row) = rows.next()? {
|
||
let name: String = row.get(0).unwrap_or_default();
|
||
let size: i64 = row.get(1).unwrap_or(0);
|
||
let reservation: i64 = row.get(2).unwrap_or(0);
|
||
result.push_str(&format!(
|
||
" {}: size={}, reservation={}\n",
|
||
name, size, reservation
|
||
));
|
||
}
|
||
Ok(result)
|
||
})
|
||
}
|
||
|
||
/// 执行事务
|
||
///
|
||
/// 自动处理事务的提交和回滚
|
||
pub fn transaction<F, T>(&self, f: F) -> DuckDbResult<T>
|
||
where
|
||
F: FnOnce(&Connection) -> DuckDbResult<T>,
|
||
{
|
||
let conn = self.inner.lock();
|
||
|
||
// 开始事务
|
||
conn.execute("BEGIN TRANSACTION", [])
|
||
.map_err(|e| DuckDbError::TransactionError(format!("failed to begin transaction: {}", e)))?;
|
||
|
||
// 执行操作
|
||
match f(&conn) {
|
||
Ok(result) => {
|
||
// 提交事务
|
||
conn.execute("COMMIT", [])
|
||
.map_err(|e| DuckDbError::TransactionError(format!("failed to commit transaction: {}", e)))?;
|
||
Ok(result)
|
||
}
|
||
Err(e) => {
|
||
// 回滚事务
|
||
let _ = conn.execute("ROLLBACK", []);
|
||
Err(e)
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
impl std::fmt::Debug for DuckDbConnection {
|
||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||
f.debug_struct("DuckDbConnection")
|
||
.field("inner", &"<Connection>")
|
||
.finish()
|
||
}
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use super::*;
|
||
|
||
#[test]
|
||
fn test_open_in_memory() {
|
||
let conn = DuckDbConnection::open_in_memory();
|
||
assert!(conn.is_ok());
|
||
}
|
||
|
||
#[test]
|
||
fn test_with_connection() {
|
||
let conn = DuckDbConnection::open_in_memory().unwrap();
|
||
let result = conn.with_connection(|c| {
|
||
c.execute("SELECT 1", [])?;
|
||
Ok(42)
|
||
});
|
||
assert_eq!(result.unwrap(), 42);
|
||
}
|
||
|
||
#[test]
|
||
fn test_transaction_commit() {
|
||
let conn = DuckDbConnection::open_in_memory().unwrap();
|
||
|
||
// 创建测试表
|
||
conn.with_connection(|c| {
|
||
c.execute("CREATE TABLE test (id INTEGER)", [])?;
|
||
Ok(())
|
||
})
|
||
.unwrap();
|
||
|
||
// 使用事务插入数据
|
||
conn.transaction(|c| {
|
||
c.execute("INSERT INTO test VALUES (1)", [])?;
|
||
Ok(())
|
||
})
|
||
.unwrap();
|
||
|
||
// 验证数据存在
|
||
let count: i32 = conn
|
||
.with_connection(|c| {
|
||
let mut stmt = c.prepare("SELECT COUNT(*) FROM test")?;
|
||
let mut rows = stmt.query([])?;
|
||
let row = rows.next()?.unwrap();
|
||
Ok(row.get(0)?)
|
||
})
|
||
.unwrap();
|
||
|
||
assert_eq!(count, 1);
|
||
}
|
||
|
||
#[test]
|
||
fn test_transaction_rollback() {
|
||
let conn = DuckDbConnection::open_in_memory().unwrap();
|
||
|
||
// 创建测试表
|
||
conn.with_connection(|c| {
|
||
c.execute("CREATE TABLE test (id INTEGER)", [])?;
|
||
Ok(())
|
||
})
|
||
.unwrap();
|
||
|
||
// 使用事务插入数据，但返回错误触发回滚
|
||
let result: DuckDbResult<()> = conn.transaction(|c| {
|
||
c.execute("INSERT INTO test VALUES (1)", [])?;
|
||
Err(DuckDbError::InternalError("故意触发回滚".to_string()))
|
||
});
|
||
|
||
assert!(result.is_err());
|
||
|
||
// 验证数据不存在（已回滚）
|
||
let count: i32 = conn
|
||
.with_connection(|c| {
|
||
let mut stmt = c.prepare("SELECT COUNT(*) FROM test")?;
|
||
let mut rows = stmt.query([])?;
|
||
let row = rows.next()?.unwrap();
|
||
Ok(row.get(0)?)
|
||
})
|
||
.unwrap();
|
||
|
||
assert_eq!(count, 0);
|
||
}
|
||
}
|