Init repo

src/database/mod.rs

@@ -0,0 +1,128 @@
+use anyhow::Result;
+use sqlx::SqlitePool;
+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() {
+        tokio::fs::create_dir_all(parent).await?;
+    }
+
+    let database_url = format!("sqlite:{}", config.path.display());
+    info!("Connecting to database at {}", database_url);
+
+    let pool = SqlitePool::connect(&database_url).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 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(())
+}