diff --git a/src/storage/mod.rs b/src/storage/mod.rs index 36f83d5..21035d6 100644 --- a/src/storage/mod.rs +++ b/src/storage/mod.rs @@ -3,3 +3,90 @@ pub mod session; pub mod message; pub use error::StorageError; + +use sqlx::{Pool, Sqlite, SqlitePool}; +use std::path::Path; + +pub struct Storage { + pool: Pool, +} + +impl Storage { + /// 打开或创建数据库 + pub async fn new(db_path: &Path) -> Result { + let database_url = format!("sqlite:{}?mode=rwc", db_path.display()); + let pool = SqlitePool::connect(&database_url).await?; + + let storage = Self { pool }; + storage.init_schema().await?; + Ok(storage) + } + + /// 初始化数据库 schema + async fn init_schema(&self) -> Result<(), StorageError> { + sqlx::query( + r#" + CREATE TABLE IF NOT EXISTS sessions ( + id TEXT PRIMARY KEY, + channel TEXT NOT NULL, + chat_id TEXT NOT NULL, + dialog_id TEXT NOT NULL, + title TEXT NOT NULL DEFAULT '新对话', + created_at INTEGER NOT NULL, + last_active_at INTEGER NOT NULL, + message_count INTEGER DEFAULT 0, + routing_info TEXT, + deleted_at INTEGER, + UNIQUE(channel, chat_id, dialog_id) + ) + "#, + ) + .execute(&self.pool) + .await?; + + sqlx::query( + r#" + CREATE INDEX IF NOT EXISTS idx_sessions_chat + ON sessions(channel, chat_id, deleted_at) + "#, + ) + .execute(&self.pool) + .await?; + + sqlx::query( + r#" + CREATE TABLE IF NOT EXISTS messages ( + id TEXT PRIMARY KEY, + session_id TEXT NOT NULL, + seq INTEGER NOT NULL, + role TEXT NOT NULL, + content TEXT NOT NULL, + media_refs TEXT, + tool_call_id TEXT, + tool_name TEXT, + tool_calls TEXT, + created_at INTEGER NOT NULL, + FOREIGN KEY (session_id) REFERENCES sessions(id) ON DELETE CASCADE + ) + "#, + ) + .execute(&self.pool) + .await?; + + sqlx::query( + r#" + CREATE INDEX IF NOT EXISTS idx_messages_session_seq + ON messages(session_id, seq) + "#, + ) + .execute(&self.pool) + .await?; + + Ok(()) + } + + /// 获取连接池引用(供内部 CRUD 使用) + pub(crate) fn pool(&self) -> &Pool { + &self.pool + } +}