128 lines
3.4 KiB
Rust
128 lines
3.4 KiB
Rust
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(())
|
|
} |