use anyhow::Result; use sqlx::sqlite::{SqliteConnectOptions, SqlitePool}; 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() && !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 users table for dashboard access sqlx::query( r#" CREATE TABLE IF NOT EXISTS users ( id INTEGER PRIMARY KEY AUTOINCREMENT, username TEXT UNIQUE NOT NULL, password_hash TEXT NOT NULL, display_name TEXT, role TEXT DEFAULT 'admin', must_change_password BOOLEAN DEFAULT FALSE, created_at DATETIME DEFAULT CURRENT_TIMESTAMP ) "#, ) .execute(pool) .await?; // Create client_tokens table for DB-based token auth sqlx::query( r#" CREATE TABLE IF NOT EXISTS client_tokens ( id INTEGER PRIMARY KEY AUTOINCREMENT, client_id TEXT NOT NULL, token TEXT NOT NULL UNIQUE, name TEXT DEFAULT 'default', is_active BOOLEAN DEFAULT TRUE, created_at DATETIME DEFAULT CURRENT_TIMESTAMP, last_used_at DATETIME, FOREIGN KEY (client_id) REFERENCES clients(client_id) ON DELETE CASCADE ) "#, ) .execute(pool) .await?; // Add must_change_password column if it doesn't exist (migration for existing DBs) let _ = sqlx::query("ALTER TABLE users ADD COLUMN must_change_password BOOLEAN DEFAULT FALSE") .execute(pool) .await; // Add display_name column if it doesn't exist (migration for existing DBs) let _ = sqlx::query("ALTER TABLE users ADD COLUMN display_name TEXT") .execute(pool) .await; // Add cache token columns if they don't exist (migration for existing DBs) let _ = sqlx::query("ALTER TABLE llm_requests ADD COLUMN cache_read_tokens INTEGER DEFAULT 0") .execute(pool) .await; let _ = sqlx::query("ALTER TABLE llm_requests ADD COLUMN cache_write_tokens INTEGER DEFAULT 0") .execute(pool) .await; // Insert default admin user if none exists (default password: admin) let user_count: (i64,) = sqlx::query_as("SELECT COUNT(*) FROM users").fetch_one(pool).await?; if user_count.0 == 0 { // 'admin' hashed with default cost (12) let default_admin_hash = bcrypt::hash("admin", 12).map_err(|e| anyhow::anyhow!("Failed to hash default password: {}", e))?; sqlx::query( "INSERT INTO users (username, password_hash, role, must_change_password) VALUES ('admin', ?, 'admin', TRUE)" ) .bind(default_admin_hash) .execute(pool) .await?; info!("Created default admin user with password 'admin' (must change on first login)"); } // 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?; sqlx::query("CREATE UNIQUE INDEX IF NOT EXISTS idx_client_tokens_token ON client_tokens(token)") .execute(pool) .await?; sqlx::query("CREATE INDEX IF NOT EXISTS idx_client_tokens_client_id ON client_tokens(client_id)") .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(()) }