GopherGate/src/database/mod.rs

use anyhow::Result;
use sqlx::sqlite::{SqlitePool, SqliteConnectOptions};
use std::str::FromStr;
use tracing::info;

use crate::config::DatabaseConfig;

pub type DbPool = SqlitePool;

pub async fn init(config: &DatabaseConfig) -> Result<DbPool> {
    // Ensure the database directory exists
    if let Some(parent) = config.path.parent() {
        if !parent.as_os_str().is_empty() {
            tokio::fs::create_dir_all(parent).await?;
        }
    }

    let database_path = config.path.to_string_lossy().to_string();
    info!("Connecting to database at {}", database_path);

    let options = SqliteConnectOptions::from_str(&format!("sqlite:{}", database_path))?
        .create_if_missing(true);

    let pool = SqlitePool::connect_with(options).await?;

    // Run migrations
    run_migrations(&pool).await?;
    info!("Database migrations completed");

    Ok(pool)
}

async fn run_migrations(pool: &DbPool) -> Result<()> {
    // Create clients table if it doesn't exist
    sqlx::query(
        r#"
        CREATE TABLE IF NOT EXISTS clients (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            client_id TEXT UNIQUE NOT NULL,
            name TEXT,
            description TEXT,
            created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
            updated_at DATETIME DEFAULT CURRENT_TIMESTAMP,
            is_active BOOLEAN DEFAULT TRUE,
            rate_limit_per_minute INTEGER DEFAULT 60,
            total_requests INTEGER DEFAULT 0,
            total_tokens INTEGER DEFAULT 0,
            total_cost REAL DEFAULT 0.0
        )
        "#,
    )
    .execute(pool)
    .await?;

    // Create llm_requests table if it doesn't exist
    sqlx::query(
        r#"
        CREATE TABLE IF NOT EXISTS llm_requests (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            timestamp DATETIME DEFAULT CURRENT_TIMESTAMP,
            client_id TEXT,
            provider TEXT,
            model TEXT,
            prompt_tokens INTEGER,
            completion_tokens INTEGER,
            total_tokens INTEGER,
            cost REAL,
            has_images BOOLEAN DEFAULT FALSE,
            status TEXT DEFAULT 'success',
            error_message TEXT,
            duration_ms INTEGER,
            request_body TEXT,
            response_body TEXT,
            FOREIGN KEY (client_id) REFERENCES clients(client_id) ON DELETE SET NULL
        )
        "#,
    )
    .execute(pool)
    .await?;

    // Create provider_configs table
    sqlx::query(
        r#"
        CREATE TABLE IF NOT EXISTS provider_configs (
            id TEXT PRIMARY KEY,
            display_name TEXT NOT NULL,
            enabled BOOLEAN DEFAULT TRUE,
            base_url TEXT,
            api_key TEXT,
            credit_balance REAL DEFAULT 0.0,
            low_credit_threshold REAL DEFAULT 5.0,
            updated_at DATETIME DEFAULT CURRENT_TIMESTAMP
        )
        "#
    )
    .execute(pool)
    .await?;

    // Create model_configs table
    sqlx::query(
        r#"
        CREATE TABLE IF NOT EXISTS model_configs (
            id TEXT PRIMARY KEY,
            provider_id TEXT NOT NULL,
            display_name TEXT,
            enabled BOOLEAN DEFAULT TRUE,
            prompt_cost_per_m REAL,
            completion_cost_per_m REAL,
            mapping TEXT,
            updated_at DATETIME DEFAULT CURRENT_TIMESTAMP,
            FOREIGN KEY (provider_id) REFERENCES provider_configs(id) ON DELETE CASCADE
        )
        "#
    )
    .execute(pool)
    .await?;

    // Create indices
    sqlx::query(
        "CREATE INDEX IF NOT EXISTS idx_clients_client_id ON clients(client_id)"
    )
    .execute(pool)
    .await?;

    sqlx::query(
        "CREATE INDEX IF NOT EXISTS idx_clients_created_at ON clients(created_at)"
    )
    .execute(pool)
    .await?;

    sqlx::query(
        "CREATE INDEX IF NOT EXISTS idx_llm_requests_timestamp ON llm_requests(timestamp)"
    )
    .execute(pool)
    .await?;

    sqlx::query(
        "CREATE INDEX IF NOT EXISTS idx_llm_requests_client_id ON llm_requests(client_id)"
    )
    .execute(pool)
    .await?;

    sqlx::query(
        "CREATE INDEX IF NOT EXISTS idx_llm_requests_provider ON llm_requests(provider)"
    )
    .execute(pool)
    .await?;

    sqlx::query(
        "CREATE INDEX IF NOT EXISTS idx_llm_requests_status ON llm_requests(status)"
    )
    .execute(pool)
    .await?;

    // Insert default client if none exists
    sqlx::query(
        r#"
        INSERT OR IGNORE INTO clients (client_id, name, description)
        VALUES ('default', 'Default Client', 'Default client for anonymous requests')
        "#,
    )
    .execute(pool)
    .await?;

    Ok(())
}

pub async fn test_connection(pool: &DbPool) -> Result<()> {
    sqlx::query("SELECT 1").execute(pool).await?;
    Ok(())
}