qiming/qiming-rcoder/crates/duckdb_manager/src/schema.rs

//! DuckDB Schema 定义和初始化
//!
//! 包含数据库表结构定义和初始化逻辑

use crate::connection::DuckDbConnection;
use crate::error::{DuckDbError, DuckDbResult};

/// 容器表 DDL
const CREATE_CONTAINERS_TABLE: &str = r#"
CREATE TABLE IF NOT EXISTS containers (
    container_id VARCHAR PRIMARY KEY,
    container_name VARCHAR NOT NULL,
    container_ip VARCHAR NOT NULL,
    internal_port INTEGER NOT NULL,
    external_port INTEGER NOT NULL,
    service_type VARCHAR NOT NULL,
    status VARCHAR NOT NULL,
    service_url VARCHAR NOT NULL,
    created_at TIMESTAMP NOT NULL,
    last_activity TIMESTAMP NOT NULL
)
"#;

/// 项目表 DDL
const CREATE_PROJECTS_TABLE: &str = r#"
CREATE TABLE IF NOT EXISTS projects (
    project_id VARCHAR PRIMARY KEY,
    session_id VARCHAR,
    service_type VARCHAR NOT NULL,
    container_id VARCHAR NOT NULL,
    user_id VARCHAR,
    pod_id VARCHAR,
    agent_status_code INTEGER,
    agent_status_name VARCHAR,
    request_id VARCHAR,
    model_provider_json VARCHAR,
    created_at TIMESTAMP NOT NULL,
    last_activity TIMESTAMP NOT NULL,
    session_created_at TIMESTAMP,
    session_last_activity TIMESTAMP
)
"#;

/// 容器表索引
const CREATE_CONTAINERS_INDEXES: &[&str] = &[
    "CREATE INDEX IF NOT EXISTS idx_containers_service_type ON containers(service_type)",
    "CREATE INDEX IF NOT EXISTS idx_containers_status ON containers(status)",
    "CREATE INDEX IF NOT EXISTS idx_containers_last_activity ON containers(last_activity)",
];

/// 项目表索引
const CREATE_PROJECTS_INDEXES: &[&str] = &[
    "CREATE INDEX IF NOT EXISTS idx_projects_session_id ON projects(session_id)",
    "CREATE INDEX IF NOT EXISTS idx_projects_container_id ON projects(container_id)",
    "CREATE INDEX IF NOT EXISTS idx_projects_user_id ON projects(user_id)",
    "CREATE INDEX IF NOT EXISTS idx_projects_pod_id ON projects(pod_id)",
    "CREATE INDEX IF NOT EXISTS idx_projects_service_type ON projects(service_type)",
    "CREATE INDEX IF NOT EXISTS idx_projects_last_activity ON projects(last_activity)",
];

/// Schema 初始化器
pub struct SchemaInitializer;

impl SchemaInitializer {
    /// 初始化数据库 Schema
    ///
    /// 创建所有必需的表和索引
    pub fn initialize(conn: &DuckDbConnection) -> DuckDbResult<()> {
        conn.with_connection(|c| {
            // 创建表
            c.execute(CREATE_CONTAINERS_TABLE, []).map_err(|e| {
                DuckDbError::InitializationError(format!("failed to create containers table: {}", e))
            })?;

            c.execute(CREATE_PROJECTS_TABLE, []).map_err(|e| {
                DuckDbError::InitializationError(format!("failed to create projects table: {}", e))
            })?;

            // 创建容器表索引
            for sql in CREATE_CONTAINERS_INDEXES {
                c.execute(sql, []).map_err(|e| {
                    DuckDbError::InitializationError(format!("failed to create containers table index: {}", e))
                })?;
            }

            // 创建项目表索引
            for sql in CREATE_PROJECTS_INDEXES {
                c.execute(sql, []).map_err(|e| {
                    DuckDbError::InitializationError(format!("failed to create projects table index: {}", e))
                })?;
            }

            tracing::info!("DuckDB Schema initialized");
            Ok(())
        })
    }

    /// 重置数据库（仅用于测试）
    #[cfg(test)]
    pub fn reset(conn: &DuckDbConnection) -> DuckDbResult<()> {
        conn.with_connection(|c| {
            c.execute("DROP TABLE IF EXISTS projects", [])?;
            c.execute("DROP TABLE IF EXISTS containers", [])?;
            Ok(())
        })?;

        Self::initialize(conn)
    }

    /// 验证 Schema 是否正确初始化
    pub fn verify(conn: &DuckDbConnection) -> DuckDbResult<bool> {
        conn.with_connection(|c| {
            // 检查 containers 表
            let containers_exists: i32 = {
                let mut stmt = c.prepare(
                    "SELECT COUNT(*) FROM information_schema.tables WHERE table_name = 'containers'"
                )?;
                let mut rows = stmt.query([])?;
                let row = rows
                    .next()?
                    .ok_or_else(|| DuckDbError::InternalError("unable to query table info".to_string()))?;
                row.get(0)?
            };

            // 检查 projects 表
            let projects_exists: i32 = {
                let mut stmt = c.prepare(
                    "SELECT COUNT(*) FROM information_schema.tables WHERE table_name = 'projects'",
                )?;
                let mut rows = stmt.query([])?;
                let row = rows
                    .next()?
                    .ok_or_else(|| DuckDbError::InternalError("unable to query table info".to_string()))?;
                row.get(0)?
            };

            Ok(containers_exists > 0 && projects_exists > 0)
        })
    }
}

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

    #[test]
    fn test_initialize_schema() {
        let conn = DuckDbConnection::open_in_memory().unwrap();
        let result = SchemaInitializer::initialize(&conn);
        assert!(result.is_ok());
    }

    #[test]
    fn test_verify_schema() {
        let conn = DuckDbConnection::open_in_memory().unwrap();
        SchemaInitializer::initialize(&conn).unwrap();

        let verified = SchemaInitializer::verify(&conn).unwrap();
        assert!(verified);
    }

    #[test]
    fn test_reset_schema() {
        let conn = DuckDbConnection::open_in_memory().unwrap();
        SchemaInitializer::initialize(&conn).unwrap();

        // 插入一些数据
        conn.with_connection(|c| {
            c.execute(
                "INSERT INTO containers (container_id, container_name, container_ip, internal_port, external_port, service_type, status, service_url, created_at, last_activity) VALUES ('c1', 'name1', '127.0.0.1', 8080, 8080, 'rcoder', 'running', 'http://localhost', NOW(), NOW())",
                [],
            )?;
            Ok(())
        }).unwrap();

        // 重置
        SchemaInitializer::reset(&conn).unwrap();

        // 验证数据已清空
        let count: i32 = conn
            .with_connection(|c| {
                let mut stmt = c.prepare("SELECT COUNT(*) FROM containers")?;
                let mut rows = stmt.query([])?;
                let row = rows.next()?.unwrap();
                Ok(row.get(0)?)
            })
            .unwrap();

        assert_eq!(count, 0);
    }

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

        // 多次初始化应该是幂等的
        SchemaInitializer::initialize(&conn).unwrap();
        SchemaInitializer::initialize(&conn).unwrap();
        SchemaInitializer::initialize(&conn).unwrap();

        let verified = SchemaInitializer::verify(&conn).unwrap();
        assert!(verified);
    }
}