Newer
Older
TelosDB / src-tauri / src / db.rs
use sea_orm::{ConnectOptions, Database, DatabaseConnection};
use std::time::Duration;

pub async fn init_db(db_path: &str, extension_path: &str) -> anyhow::Result<DatabaseConnection> {
    // SeaORM (sqlx) は URL 形式でパスを指定する
    let db_url = format!("sqlite:{}?mode=rwc", db_path);

    let mut opt = ConnectOptions::new(db_url);
    opt.max_connections(10)
        .min_connections(5)
        .connect_timeout(Duration::from_secs(8))
        .idle_timeout(Duration::from_secs(8))
        .max_lifetime(Duration::from_secs(8))
        .sqlx_logging(true);

    // 拡張機能のロードが必要なため、sqlx の初期化フックがあればいいが、
    // SeaORM/sqlx で動的にロードするのは少し工夫が必要。
    // ここではまず rusqlite か何かでスキーマ初期化と拡張ロードの確認をしてから SeaORM で繋ぐか、
    // sqlx の ConnectOptions にて初期化SQLを流す。

    // NOTE: sqlite-vec のロードは非常に特殊なため、
    // 生の rusqlite で拡張をロードしてテーブルを作った後、SeaORM で接続する。
    {
        use rusqlite::Connection;
        let conn = Connection::open(db_path)?;
        unsafe {
            conn.load_extension(extension_path, None)?;
        }
        conn.execute_batch(
            "PRAGMA journal_mode = WAL;
             CREATE TABLE IF NOT EXISTS items (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                content TEXT NOT NULL,
                path TEXT,
                created_at TEXT DEFAULT (datetime('now')),
                updated_at TEXT DEFAULT (datetime('now'))
            );
            CREATE VIRTUAL TABLE IF NOT EXISTS vec_items USING vec0(
                id INTEGER PRIMARY KEY,
                embedding FLOAT[3]
            );",
        )?;
    }

    let db = Database::connect(opt).await?;

    // SeaORM の接続後にも拡張をロードしておく必要がある場合がある(セッション毎)
    // sqlx-sqlite の場合、初期化 SQL でロードできるケースがあるが設定が複雑。
    // 今回は rusqlite での初期化を優先し、クエリ実行時に MATCH 句が動くか確認する。

    Ok(db)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[tokio::test]
    async fn test_init_db_basic() {
        // NOTE: We test with a dummy extension path and expect it to fail if it's not found,
        // or we test just the SeaORM part if we could separation.
        // For now, let's just ensure we can connect to an in-memory sqlite via SeaORM.
        let mut opt = ConnectOptions::new("sqlite::memory:");
        let db = Database::connect(opt).await;
        assert!(db.is_ok());
    }
}