unistore-sqlite 0.1.0

SQLite embedded database capability for UniStore
Documentation
//! SQLite 连接管理
//!
//! 职责:封装 rusqlite 连接,提供底层 SQL 执行能力

use crate::config::SqliteConfig;
use crate::error::SqliteError;
use crate::types::{Param, Row, Rows, SqlValue};
use parking_lot::Mutex;
use rusqlite::Connection as RawConnection;
use std::sync::Arc;

/// 连接状态
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ConnectionState {
    /// 已打开
    Open,
    /// 已关闭
    Closed,
}

/// SQLite 连接包装器
///
/// 线程安全的连接封装,提供执行和查询方法
pub struct Connection {
    /// 底层连接(Mutex 保护)
    inner: Arc<Mutex<Option<RawConnection>>>,
    /// 配置
    config: SqliteConfig,
    /// 状态
    state: Arc<Mutex<ConnectionState>>,
}

impl Connection {
    /// 打开数据库连接
    pub fn open(config: SqliteConfig) -> Result<Self, SqliteError> {
        let path = config.path_string();

        // 打开连接
        let conn = if config.read_only {
            RawConnection::open_with_flags(
                &path,
                rusqlite::OpenFlags::SQLITE_OPEN_READ_ONLY | rusqlite::OpenFlags::SQLITE_OPEN_URI,
            )
        } else if config.create_if_missing {
            RawConnection::open(&path)
        } else {
            RawConnection::open_with_flags(
                &path,
                rusqlite::OpenFlags::SQLITE_OPEN_READ_WRITE | rusqlite::OpenFlags::SQLITE_OPEN_URI,
            )
        }
        .map_err(|e| SqliteError::OpenFailed(e.to_string()))?;

        // 应用 PRAGMA 配置
        for pragma in config.to_pragmas() {
            conn.execute_batch(&pragma)
                .map_err(|e| SqliteError::OpenFailed(format!("PRAGMA failed: {}", e)))?;
        }

        Ok(Self {
            inner: Arc::new(Mutex::new(Some(conn))),
            config,
            state: Arc::new(Mutex::new(ConnectionState::Open)),
        })
    }

    /// 打开内存数据库
    pub fn open_in_memory() -> Result<Self, SqliteError> {
        Self::open(SqliteConfig::memory())
    }

    /// 获取连接状态
    pub fn state(&self) -> ConnectionState {
        *self.state.lock()
    }

    /// 判断连接是否打开
    pub fn is_open(&self) -> bool {
        self.state() == ConnectionState::Open
    }

    /// 获取配置引用
    pub fn config(&self) -> &SqliteConfig {
        &self.config
    }

    /// 执行 SQL(无返回结果)
    pub fn execute(&self, sql: &str, params: &[Param]) -> Result<usize, SqliteError> {
        let guard = self.inner.lock();
        let conn = guard.as_ref().ok_or(SqliteError::DatabaseClosed)?;

        let params_refs: Vec<&dyn rusqlite::ToSql> =
            params.iter().map(|p| p as &dyn rusqlite::ToSql).collect();

        conn.execute(sql, params_refs.as_slice())
            .map_err(SqliteError::from)
    }

    /// 执行多条 SQL 语句
    pub fn execute_batch(&self, sql: &str) -> Result<(), SqliteError> {
        let guard = self.inner.lock();
        let conn = guard.as_ref().ok_or(SqliteError::DatabaseClosed)?;

        conn.execute_batch(sql).map_err(SqliteError::from)
    }

    /// 查询单行
    pub fn query_row(&self, sql: &str, params: &[Param]) -> Result<Option<Row>, SqliteError> {
        let guard = self.inner.lock();
        let conn = guard.as_ref().ok_or(SqliteError::DatabaseClosed)?;

        let params_refs: Vec<&dyn rusqlite::ToSql> =
            params.iter().map(|p| p as &dyn rusqlite::ToSql).collect();

        let mut stmt = conn.prepare(sql).map_err(SqliteError::from)?;
        let columns: Vec<String> = stmt.column_names().iter().map(|s: &&str| s.to_string()).collect();

        let result = stmt.query_row(params_refs.as_slice(), |row: &rusqlite::Row| {
            let mut r = Row::new();
            for (i, col) in columns.iter().enumerate() {
                let value: SqlValue = row.get(i)?;
                r.push(col.to_string(), value);
            }
            Ok(r)
        });

        match result {
            Ok(row) => Ok(Some(row)),
            Err(rusqlite::Error::QueryReturnedNoRows) => Ok(None),
            Err(e) => Err(SqliteError::from(e)),
        }
    }

    /// 查询多行
    pub fn query(&self, sql: &str, params: &[Param]) -> Result<Rows, SqliteError> {
        let guard = self.inner.lock();
        let conn = guard.as_ref().ok_or(SqliteError::DatabaseClosed)?;

        let params_refs: Vec<&dyn rusqlite::ToSql> =
            params.iter().map(|p| p as &dyn rusqlite::ToSql).collect();

        let mut stmt = conn.prepare(sql).map_err(SqliteError::from)?;
        let columns: Vec<String> = stmt.column_names().iter().map(|s: &&str| s.to_string()).collect();

        let rows = stmt
            .query_map(params_refs.as_slice(), |row: &rusqlite::Row| {
                let mut r = Row::new();
                for (i, col) in columns.iter().enumerate() {
                    let value: SqlValue = row.get(i)?;
                    r.push(col.to_string(), value);
                }
                Ok(r)
            })
            .map_err(SqliteError::from)?;

        let mut result: Vec<Row> = Vec::new();
        for row in rows {
            result.push(row.map_err(SqliteError::from)?);
        }

        Ok(result)
    }

    /// 获取最后插入的行 ID
    pub fn last_insert_rowid(&self) -> Result<i64, SqliteError> {
        let guard = self.inner.lock();
        let conn = guard.as_ref().ok_or(SqliteError::DatabaseClosed)?;
        Ok(conn.last_insert_rowid())
    }

    /// 获取上次操作影响的行数
    pub fn changes(&self) -> Result<usize, SqliteError> {
        let guard = self.inner.lock();
        let conn = guard.as_ref().ok_or(SqliteError::DatabaseClosed)?;
        Ok(conn.changes() as usize)
    }

    /// 关闭连接
    pub fn close(&self) -> Result<(), SqliteError> {
        let mut guard = self.inner.lock();
        let mut state = self.state.lock();

        if let Some(conn) = guard.take() {
            // 尝试关闭,忽略错误(连接会自动清理)
            drop(conn);
        }

        *state = ConnectionState::Closed;
        Ok(())
    }

    /// 访问底层连接(高级用法)
    ///
    /// # Safety
    ///
    /// 调用者需要确保不会破坏连接状态
    pub fn with_raw<F, R>(&self, f: F) -> Result<R, SqliteError>
    where
        F: FnOnce(&RawConnection) -> Result<R, SqliteError>,
    {
        let guard = self.inner.lock();
        let conn = guard.as_ref().ok_or(SqliteError::DatabaseClosed)?;
        f(conn)
    }
}

impl Clone for Connection {
    fn clone(&self) -> Self {
        Self {
            inner: self.inner.clone(),
            config: self.config.clone(),
            state: self.state.clone(),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_open_memory() {
        let conn = Connection::open_in_memory().unwrap();
        assert!(conn.is_open());
    }

    #[test]
    fn test_execute_and_query() {
        let conn = Connection::open_in_memory().unwrap();

        // 创建表
        conn.execute_batch("CREATE TABLE test (id INTEGER PRIMARY KEY, name TEXT)")
            .unwrap();

        // 插入数据
        conn.execute("INSERT INTO test (name) VALUES (?)", &["Alice".into()])
            .unwrap();
        conn.execute("INSERT INTO test (name) VALUES (?)", &["Bob".into()])
            .unwrap();

        let id = conn.last_insert_rowid().unwrap();
        assert_eq!(id, 2);

        // 查询
        let rows = conn.query("SELECT * FROM test", &[]).unwrap();
        assert_eq!(rows.len(), 2);
        assert_eq!(rows[0].get_str("name"), Some("Alice"));
        assert_eq!(rows[1].get_str("name"), Some("Bob"));
    }

    #[test]
    fn test_query_row() {
        let conn = Connection::open_in_memory().unwrap();
        conn.execute_batch("CREATE TABLE test (id INTEGER PRIMARY KEY, value INTEGER)")
            .unwrap();
        conn.execute("INSERT INTO test (value) VALUES (?)", &[42i32.into()])
            .unwrap();

        let row = conn
            .query_row("SELECT * FROM test WHERE id = ?", &[1i32.into()])
            .unwrap();
        assert!(row.is_some());
        assert_eq!(row.unwrap().get_i64("value"), Some(42));

        let row = conn
            .query_row("SELECT * FROM test WHERE id = ?", &[999i32.into()])
            .unwrap();
        assert!(row.is_none());
    }

    #[test]
    fn test_close() {
        let conn = Connection::open_in_memory().unwrap();
        assert!(conn.is_open());

        conn.close().unwrap();
        assert!(!conn.is_open());

        // 关闭后操作应该失败
        let result = conn.execute("SELECT 1", &[]);
        assert!(result.is_err());
    }
}