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 { // 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(()) }