use async_trait::async_trait;
use chrono::{DateTime, Utc};
use sqlx::{SqlitePool, Row};
use std::path::Path;
use crate::core::error::StorageError;
use crate::core::session::Session;
use crate::storage::{Storage, Message, StorageStatistics};
pub struct SqliteStorage {
pool: SqlitePool,
}
impl SqliteStorage {
pub async fn new<P: AsRef<Path>>(database_path: P) -> Result<Self, StorageError> {
let database_url = format!("sqlite:{}", database_path.as_ref().display());
let pool = SqlitePool::connect(&database_url).await
.map_err(|e| StorageError::Connection(e.to_string()))?;
let storage = Self { pool };
storage.run_migrations().await?;
Ok(storage)
}
async fn run_migrations(&self) -> Result<(), StorageError> {
sqlx::query(
r#"
CREATE TABLE IF NOT EXISTS sessions (
id TEXT PRIMARY KEY,
title TEXT NOT NULL,
created_at TEXT NOT NULL,
updated_at TEXT NOT NULL,
model TEXT NOT NULL,
cost REAL NOT NULL DEFAULT 0.0,
input_tokens INTEGER NOT NULL DEFAULT 0,
output_tokens INTEGER NOT NULL DEFAULT 0,
cache_creation_tokens INTEGER NOT NULL DEFAULT 0,
cache_read_tokens INTEGER NOT NULL DEFAULT 0,
metadata TEXT
)
"#
)
.execute(&self.pool)
.await
.map_err(|e| StorageError::Migration(e.to_string()))?;
sqlx::query(
r#"
CREATE TABLE IF NOT EXISTS messages (
id TEXT PRIMARY KEY,
session_id TEXT NOT NULL,
role TEXT NOT NULL,
content TEXT NOT NULL,
timestamp TEXT NOT NULL,
metadata TEXT,
FOREIGN KEY (session_id) REFERENCES sessions (id) ON DELETE CASCADE
)
"#
)
.execute(&self.pool)
.await
.map_err(|e| StorageError::Migration(e.to_string()))?;
sqlx::query("CREATE INDEX IF NOT EXISTS idx_messages_session_id ON messages (session_id)")
.execute(&self.pool)
.await
.map_err(|e| StorageError::Migration(e.to_string()))?;
sqlx::query("CREATE INDEX IF NOT EXISTS idx_messages_timestamp ON messages (timestamp)")
.execute(&self.pool)
.await
.map_err(|e| StorageError::Migration(e.to_string()))?;
Ok(())
}
}
#[async_trait]
impl Storage for SqliteStorage {
async fn create_session(&self, session: &Session) -> Result<String, StorageError> {
sqlx::query(
r#"
INSERT INTO sessions (
id, title, created_at, updated_at, model, cost,
input_tokens, output_tokens, cache_creation_tokens, cache_read_tokens, metadata
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
"#
)
.bind(&session.id)
.bind(&session.title)
.bind(session.created_at.to_rfc3339())
.bind(session.updated_at.to_rfc3339())
.bind(&session.model)
.bind(session.cost)
.bind(session.token_usage.input_tokens as i64)
.bind(session.token_usage.output_tokens as i64)
.bind(session.token_usage.cache_creation_tokens as i64)
.bind(session.token_usage.cache_read_tokens as i64)
.bind(serde_json::to_string(&session.metadata).unwrap_or_default())
.execute(&self.pool)
.await
.map_err(|e| StorageError::Database(e))?;
Ok(session.id.clone())
}
async fn get_session(&self, id: &str) -> Result<Option<Session>, StorageError> {
let row = sqlx::query(
r#"
SELECT id, title, created_at, updated_at, model, cost,
input_tokens, output_tokens, cache_creation_tokens, cache_read_tokens, metadata
FROM sessions WHERE id = ?
"#
)
.bind(id)
.fetch_optional(&self.pool)
.await
.map_err(|e| StorageError::Database(e))?;
if let Some(row) = row {
let created_at: String = row.get("created_at");
let updated_at: String = row.get("updated_at");
let metadata: String = row.get("metadata");
let session = Session {
id: row.get("id"),
title: row.get("title"),
created_at: DateTime::parse_from_rfc3339(&created_at)
.map_err(|e| StorageError::Corruption(e.to_string()))?
.with_timezone(&Utc),
updated_at: DateTime::parse_from_rfc3339(&updated_at)
.map_err(|e| StorageError::Corruption(e.to_string()))?
.with_timezone(&Utc),
model: row.get("model"),
cost: row.get("cost"),
token_usage: crate::core::session::TokenUsage {
input_tokens: row.get::<i64, _>("input_tokens") as u32,
output_tokens: row.get::<i64, _>("output_tokens") as u32,
cache_creation_tokens: row.get::<i64, _>("cache_creation_tokens") as u32,
cache_read_tokens: row.get::<i64, _>("cache_read_tokens") as u32,
},
metadata: serde_json::from_str(&metadata).unwrap_or(serde_json::Value::Null),
};
Ok(Some(session))
} else {
Ok(None)
}
}
async fn update_session(&self, session: &Session) -> Result<(), StorageError> {
sqlx::query(
r#"
UPDATE sessions SET
title = ?, updated_at = ?, model = ?, cost = ?,
input_tokens = ?, output_tokens = ?, cache_creation_tokens = ?, cache_read_tokens = ?,
metadata = ?
WHERE id = ?
"#
)
.bind(&session.title)
.bind(session.updated_at.to_rfc3339())
.bind(&session.model)
.bind(session.cost)
.bind(session.token_usage.input_tokens as i64)
.bind(session.token_usage.output_tokens as i64)
.bind(session.token_usage.cache_creation_tokens as i64)
.bind(session.token_usage.cache_read_tokens as i64)
.bind(serde_json::to_string(&session.metadata).unwrap_or_default())
.bind(&session.id)
.execute(&self.pool)
.await
.map_err(|e| StorageError::Database(e))?;
Ok(())
}
async fn delete_session(&self, id: &str) -> Result<(), StorageError> {
sqlx::query("DELETE FROM sessions WHERE id = ?")
.bind(id)
.execute(&self.pool)
.await
.map_err(|e| StorageError::Database(e))?;
Ok(())
}
async fn list_sessions(&self) -> Result<Vec<Session>, StorageError> {
let rows = sqlx::query(
r#"
SELECT id, title, created_at, updated_at, model, cost,
input_tokens, output_tokens, cache_creation_tokens, cache_read_tokens, metadata
FROM sessions ORDER BY updated_at DESC
"#
)
.fetch_all(&self.pool)
.await
.map_err(|e| StorageError::Database(e))?;
let mut sessions = Vec::new();
for row in rows {
let created_at: String = row.get("created_at");
let updated_at: String = row.get("updated_at");
let metadata: String = row.get("metadata");
let session = Session {
id: row.get("id"),
title: row.get("title"),
created_at: DateTime::parse_from_rfc3339(&created_at)
.map_err(|e| StorageError::Corruption(e.to_string()))?
.with_timezone(&Utc),
updated_at: DateTime::parse_from_rfc3339(&updated_at)
.map_err(|e| StorageError::Corruption(e.to_string()))?
.with_timezone(&Utc),
model: row.get("model"),
cost: row.get("cost"),
token_usage: crate::core::session::TokenUsage {
input_tokens: row.get::<i64, _>("input_tokens") as u32,
output_tokens: row.get::<i64, _>("output_tokens") as u32,
cache_creation_tokens: row.get::<i64, _>("cache_creation_tokens") as u32,
cache_read_tokens: row.get::<i64, _>("cache_read_tokens") as u32,
},
metadata: serde_json::from_str(&metadata).unwrap_or(serde_json::Value::Null),
};
sessions.push(session);
}
Ok(sessions)
}
async fn save_message(&self, message: &Message) -> Result<String, StorageError> {
let content_json = serde_json::to_string(&message.content)
.map_err(|e| StorageError::Serialization(e.to_string()))?;
sqlx::query(
r#"
INSERT INTO messages (id, session_id, role, content, timestamp, metadata)
VALUES (?, ?, ?, ?, ?, ?)
"#
)
.bind(&message.id)
.bind(&message.session_id)
.bind(message.role.to_string())
.bind(content_json)
.bind(message.timestamp.to_rfc3339())
.bind(serde_json::to_string(&message.metadata).unwrap_or_default())
.execute(&self.pool)
.await
.map_err(|e| StorageError::Database(e))?;
Ok(message.id.clone())
}
async fn get_messages(&self, session_id: &str) -> Result<Vec<Message>, StorageError> {
let rows = sqlx::query(
"SELECT id, session_id, role, content, timestamp, metadata FROM messages WHERE session_id = ? ORDER BY timestamp ASC"
)
.bind(session_id)
.fetch_all(&self.pool)
.await
.map_err(|e| StorageError::Database(e))?;
let mut messages = Vec::new();
for row in rows {
let role_str: String = row.get("role");
let content_str: String = row.get("content");
let timestamp_str: String = row.get("timestamp");
let metadata_str: String = row.get("metadata");
let message = Message {
id: row.get("id"),
session_id: row.get("session_id"),
role: role_str.parse().map_err(|e| StorageError::Corruption(format!("Invalid role: {}", e)))?,
content: serde_json::from_str(&content_str)
.map_err(|e| StorageError::Corruption(format!("Invalid content: {}", e)))?,
timestamp: DateTime::parse_from_rfc3339(×tamp_str)
.map_err(|e| StorageError::Corruption(e.to_string()))?
.with_timezone(&Utc),
metadata: serde_json::from_str(&metadata_str).unwrap_or(serde_json::Value::Null),
};
messages.push(message);
}
Ok(messages)
}
async fn delete_message(&self, id: &str) -> Result<(), StorageError> {
sqlx::query("DELETE FROM messages WHERE id = ?")
.bind(id)
.execute(&self.pool)
.await
.map_err(|e| StorageError::Database(e))?;
Ok(())
}
async fn search_messages(&self, query: &str) -> Result<Vec<Message>, StorageError> {
let rows = sqlx::query(
r#"
SELECT id, session_id, role, content, timestamp, metadata
FROM messages
WHERE content LIKE ?
ORDER BY timestamp DESC
LIMIT 100
"#
)
.bind(format!("%{}%", query))
.fetch_all(&self.pool)
.await
.map_err(|e| StorageError::Database(e))?;
let mut messages = Vec::new();
for row in rows {
let role_str: String = row.get("role");
let content_str: String = row.get("content");
let timestamp_str: String = row.get("timestamp");
let metadata_str: String = row.get("metadata");
let message = Message {
id: row.get("id"),
session_id: row.get("session_id"),
role: role_str.parse().map_err(|e| StorageError::Corruption(format!("Invalid role: {}", e)))?,
content: serde_json::from_str(&content_str)
.map_err(|e| StorageError::Corruption(format!("Invalid content: {}", e)))?,
timestamp: DateTime::parse_from_rfc3339(×tamp_str)
.map_err(|e| StorageError::Corruption(e.to_string()))?
.with_timezone(&Utc),
metadata: serde_json::from_str(&metadata_str).unwrap_or(serde_json::Value::Null),
};
messages.push(message);
}
Ok(messages)
}
async fn get_statistics(&self) -> Result<StorageStatistics, StorageError> {
let session_stats = sqlx::query(
"SELECT COUNT(*) as count, MIN(created_at) as oldest, MAX(updated_at) as newest FROM sessions"
)
.fetch_one(&self.pool)
.await
.map_err(|e| StorageError::Database(e))?;
let message_count: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM messages")
.fetch_one(&self.pool)
.await
.map_err(|e| StorageError::Database(e))?;
let total_sessions: i64 = session_stats.get("count");
let oldest_str: Option<String> = session_stats.get("oldest");
let newest_str: Option<String> = session_stats.get("newest");
let oldest_session = oldest_str
.and_then(|s| DateTime::parse_from_rfc3339(&s).ok())
.map(|dt| dt.with_timezone(&Utc));
let newest_session = newest_str
.and_then(|s| DateTime::parse_from_rfc3339(&s).ok())
.map(|dt| dt.with_timezone(&Utc));
let avg_messages_per_session = if total_sessions > 0 {
message_count as f64 / total_sessions as f64
} else {
0.0
};
Ok(StorageStatistics {
total_sessions: total_sessions as u64,
total_messages: message_count as u64,
storage_size: 0, oldest_session,
newest_session,
avg_messages_per_session,
})
}
async fn cleanup(&self, older_than: DateTime<Utc>) -> Result<u64, StorageError> {
let result = sqlx::query("DELETE FROM sessions WHERE updated_at < ?")
.bind(older_than.to_rfc3339())
.execute(&self.pool)
.await
.map_err(|e| StorageError::Database(e))?;
Ok(result.rows_affected())
}
}