zclaw_openfang/crates/zclaw-saas/src/db.rs

//! 数据库初始化与 Schema

use sqlx::SqlitePool;
use crate::error::SaasResult;

const SCHEMA_VERSION: i32 = 1;

const SCHEMA_SQL: &str = r#"
CREATE TABLE IF NOT EXISTS saas_schema_version (
    version INTEGER PRIMARY KEY
);

CREATE TABLE IF NOT EXISTS accounts (
    id TEXT PRIMARY KEY,
    username TEXT NOT NULL UNIQUE,
    email TEXT NOT NULL UNIQUE,
    password_hash TEXT NOT NULL,
    display_name TEXT NOT NULL DEFAULT '',
    avatar_url TEXT,
    role TEXT NOT NULL DEFAULT 'user',
    status TEXT NOT NULL DEFAULT 'active',
    totp_secret TEXT,
    totp_enabled INTEGER NOT NULL DEFAULT 0,
    last_login_at TEXT,
    created_at TEXT NOT NULL,
    updated_at TEXT NOT NULL
);
CREATE INDEX IF NOT EXISTS idx_accounts_email ON accounts(email);
CREATE INDEX IF NOT EXISTS idx_accounts_role ON accounts(role);

CREATE TABLE IF NOT EXISTS api_tokens (
    id TEXT PRIMARY KEY,
    account_id TEXT NOT NULL,
    name TEXT NOT NULL,
    token_hash TEXT NOT NULL,
    token_prefix TEXT NOT NULL,
    permissions TEXT NOT NULL DEFAULT '[]',
    last_used_at TEXT,
    expires_at TEXT,
    created_at TEXT NOT NULL,
    revoked_at TEXT,
    FOREIGN KEY (account_id) REFERENCES accounts(id) ON DELETE CASCADE
);
CREATE INDEX IF NOT EXISTS idx_api_tokens_account ON api_tokens(account_id);
CREATE INDEX IF NOT EXISTS idx_api_tokens_hash ON api_tokens(token_hash);

CREATE TABLE IF NOT EXISTS roles (
    id TEXT PRIMARY KEY,
    name TEXT NOT NULL,
    description TEXT,
    permissions TEXT NOT NULL DEFAULT '[]',
    is_system INTEGER NOT NULL DEFAULT 0,
    created_at TEXT NOT NULL,
    updated_at TEXT NOT NULL
);

CREATE TABLE IF NOT EXISTS permission_templates (
    id TEXT PRIMARY KEY,
    name TEXT NOT NULL,
    description TEXT,
    permissions TEXT NOT NULL DEFAULT '[]',
    created_at TEXT NOT NULL,
    updated_at TEXT NOT NULL
);

CREATE TABLE IF NOT EXISTS operation_logs (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    account_id TEXT,
    action TEXT NOT NULL,
    target_type TEXT,
    target_id TEXT,
    details TEXT,
    ip_address TEXT,
    created_at TEXT NOT NULL
);
CREATE INDEX IF NOT EXISTS idx_op_logs_account ON operation_logs(account_id);
CREATE INDEX IF NOT EXISTS idx_op_logs_action ON operation_logs(action);
CREATE INDEX IF NOT EXISTS idx_op_logs_time ON operation_logs(created_at);

CREATE TABLE IF NOT EXISTS providers (
    id TEXT PRIMARY KEY,
    name TEXT NOT NULL UNIQUE,
    display_name TEXT NOT NULL,
    api_key TEXT,
    base_url TEXT NOT NULL,
    api_protocol TEXT NOT NULL DEFAULT 'openai',
    enabled INTEGER NOT NULL DEFAULT 1,
    rate_limit_rpm INTEGER,
    rate_limit_tpm INTEGER,
    config_json TEXT DEFAULT '{}',
    created_at TEXT NOT NULL,
    updated_at TEXT NOT NULL
);

CREATE TABLE IF NOT EXISTS models (
    id TEXT PRIMARY KEY,
    provider_id TEXT NOT NULL,
    model_id TEXT NOT NULL,
    alias TEXT NOT NULL,
    context_window INTEGER NOT NULL DEFAULT 8192,
    max_output_tokens INTEGER NOT NULL DEFAULT 4096,
    supports_streaming INTEGER NOT NULL DEFAULT 1,
    supports_vision INTEGER NOT NULL DEFAULT 0,
    enabled INTEGER NOT NULL DEFAULT 1,
    pricing_input REAL DEFAULT 0,
    pricing_output REAL DEFAULT 0,
    created_at TEXT NOT NULL,
    updated_at TEXT NOT NULL,
    UNIQUE(provider_id, model_id),
    FOREIGN KEY (provider_id) REFERENCES providers(id) ON DELETE CASCADE
);
CREATE INDEX IF NOT EXISTS idx_models_provider ON models(provider_id);

CREATE TABLE IF NOT EXISTS account_api_keys (
    id TEXT PRIMARY KEY,
    account_id TEXT NOT NULL,
    provider_id TEXT NOT NULL,
    key_value TEXT NOT NULL,
    key_label TEXT,
    permissions TEXT NOT NULL DEFAULT '[]',
    enabled INTEGER NOT NULL DEFAULT 1,
    last_used_at TEXT,
    created_at TEXT NOT NULL,
    updated_at TEXT NOT NULL,
    revoked_at TEXT,
    FOREIGN KEY (account_id) REFERENCES accounts(id) ON DELETE CASCADE,
    FOREIGN KEY (provider_id) REFERENCES providers(id) ON DELETE CASCADE
);
CREATE INDEX IF NOT EXISTS idx_account_api_keys_account ON account_api_keys(account_id);

CREATE TABLE IF NOT EXISTS usage_records (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    account_id TEXT NOT NULL,
    provider_id TEXT NOT NULL,
    model_id TEXT NOT NULL,
    input_tokens INTEGER NOT NULL DEFAULT 0,
    output_tokens INTEGER NOT NULL DEFAULT 0,
    latency_ms INTEGER,
    status TEXT NOT NULL DEFAULT 'success',
    error_message TEXT,
    created_at TEXT NOT NULL
);
CREATE INDEX IF NOT EXISTS idx_usage_account ON usage_records(account_id);
CREATE INDEX IF NOT EXISTS idx_usage_time ON usage_records(created_at);

CREATE TABLE IF NOT EXISTS relay_tasks (
    id TEXT PRIMARY KEY,
    account_id TEXT NOT NULL,
    provider_id TEXT NOT NULL,
    model_id TEXT NOT NULL,
    request_hash TEXT NOT NULL,
    status TEXT NOT NULL DEFAULT 'queued',
    priority INTEGER NOT NULL DEFAULT 0,
    attempt_count INTEGER NOT NULL DEFAULT 0,
    max_attempts INTEGER NOT NULL DEFAULT 3,
    request_body TEXT NOT NULL,
    response_body TEXT,
    input_tokens INTEGER DEFAULT 0,
    output_tokens INTEGER DEFAULT 0,
    error_message TEXT,
    queued_at TEXT NOT NULL,
    started_at TEXT,
    completed_at TEXT,
    created_at TEXT NOT NULL
);
CREATE INDEX IF NOT EXISTS idx_relay_status ON relay_tasks(status);
CREATE INDEX IF NOT EXISTS idx_relay_account ON relay_tasks(account_id);
CREATE INDEX IF NOT EXISTS idx_relay_provider ON relay_tasks(provider_id);

CREATE TABLE IF NOT EXISTS config_items (
    id TEXT PRIMARY KEY,
    category TEXT NOT NULL,
    key_path TEXT NOT NULL,
    value_type TEXT NOT NULL,
    current_value TEXT,
    default_value TEXT,
    source TEXT NOT NULL DEFAULT 'local',
    description TEXT,
    requires_restart INTEGER NOT NULL DEFAULT 0,
    created_at TEXT NOT NULL,
    updated_at TEXT NOT NULL,
    UNIQUE(category, key_path)
);
CREATE INDEX IF NOT EXISTS idx_config_category ON config_items(category);

CREATE TABLE IF NOT EXISTS config_sync_log (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    account_id TEXT NOT NULL,
    client_fingerprint TEXT NOT NULL,
    action TEXT NOT NULL,
    config_keys TEXT NOT NULL,
    client_values TEXT,
    saas_values TEXT,
    resolution TEXT,
    created_at TEXT NOT NULL
);
CREATE INDEX IF NOT EXISTS idx_sync_account ON config_sync_log(account_id);
"#;

const SEED_ROLES: &str = r#"
INSERT OR IGNORE INTO roles (id, name, description, permissions, is_system, created_at, updated_at)
VALUES
    ('super_admin', '超级管理员', '拥有所有权限', '["admin:full"]', 1, datetime('now'), datetime('now')),
    ('admin', '管理员', '管理账号和配置', '["account:read","account:write","model:read","model:write","relay:use","relay:admin","config:read","config:write"]', 1, datetime('now'), datetime('now')),
    ('user', '普通用户', '基础使用权限', '["model:read","relay:use","config:read"]', 1, datetime('now'), datetime('now'));
"#;

/// 初始化数据库
pub async fn init_db(database_url: &str) -> SaasResult<SqlitePool> {
    if database_url.starts_with("sqlite:") {
        let path_part = database_url.strip_prefix("sqlite:").unwrap_or("");
        if path_part != ":memory:" {
            if let Some(parent) = std::path::Path::new(path_part).parent() {
                if !parent.as_os_str().is_empty() && !parent.exists() {
                    std::fs::create_dir_all(parent)?;
                }
            }
        }
    }

    let pool = SqlitePool::connect(database_url).await?;
    sqlx::query("PRAGMA journal_mode=WAL;")
        .execute(&pool)
        .await?;
    sqlx::query(SCHEMA_SQL).execute(&pool).await?;
    sqlx::query("INSERT OR IGNORE INTO saas_schema_version (version) VALUES (?1)")
        .bind(SCHEMA_VERSION)
        .execute(&pool)
        .await?;
    sqlx::query(SEED_ROLES).execute(&pool).await?;
    tracing::info!("Database initialized (schema v{})", SCHEMA_VERSION);
    Ok(pool)
}

/// 创建内存数据库 (测试用)
pub async fn init_memory_db() -> SaasResult<SqlitePool> {
    let pool = SqlitePool::connect("sqlite::memory:").await?;
    sqlx::query(SCHEMA_SQL).execute(&pool).await?;
    sqlx::query("INSERT OR IGNORE INTO saas_schema_version (version) VALUES (?1)")
        .bind(SCHEMA_VERSION)
        .execute(&pool)
        .await?;
    sqlx::query(SEED_ROLES).execute(&pool).await?;
    Ok(pool)
}

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

    #[tokio::test]
    async fn test_init_memory_db() {
        let pool = init_memory_db().await.unwrap();
        let roles: Vec<(String,)> = sqlx::query_as(
            "SELECT id FROM roles WHERE is_system = 1"
        )
        .fetch_all(&pool)
        .await
        .unwrap();
        assert_eq!(roles.len(), 3);
    }

    #[tokio::test]
    async fn test_schema_tables_exist() {
        let pool = init_memory_db().await.unwrap();
        let tables = [
            "accounts", "api_tokens", "roles", "permission_templates",
            "operation_logs", "providers", "models", "account_api_keys",
            "usage_records", "relay_tasks", "config_items", "config_sync_log",
        ];
        for table in tables {
            let count: (i64,) = sqlx::query_as(&format!(
                "SELECT COUNT(*) FROM sqlite_master WHERE type='table' AND name='{}'", table
            ))
            .fetch_one(&pool)
            .await
            .unwrap();
            assert_eq!(count.0, 1, "Table {} should exist", table);
        }
    }
}