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', 'localtime')),
updated_at TEXT DEFAULT (datetime('now', 'localtime'))
);
CREATE TRIGGER IF NOT EXISTS update_items_updated_at
AFTER UPDATE ON items
FOR EACH ROW
BEGIN
UPDATE items SET updated_at = datetime('now', 'localtime') WHERE id = OLD.id;
END;
CREATE VIRTUAL TABLE IF NOT EXISTS vec_items USING vec0(
id INTEGER PRIMARY KEY,
embedding FLOAT[384]
);",
)?;
}
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());
}
}