use sea_orm::{DatabaseConnection, SqlxSqliteConnector, ConnectionTrait, Statement, DatabaseBackend};
use sea_orm::sqlx::sqlite::{SqliteConnectOptions, SqlitePoolOptions};
use std::str::FromStr;
use std::time::Duration;
pub async fn init_db(db_path: &str, extension_path: &str, vec_dim: usize) -> anyhow::Result<DatabaseConnection> {
let db_url = format!("sqlite:{}?mode=rwc", db_path);
let options = SqliteConnectOptions::from_str(&db_url)?
.create_if_missing(true)
.extension(extension_path.to_owned());
let pool = SqlitePoolOptions::new()
.max_connections(10)
.min_connections(5)
.acquire_timeout(Duration::from_secs(8))
.max_lifetime(Duration::from_secs(8))
.connect_with(options)
.await?;
let db = SqlxSqliteConnector::from_sqlx_sqlite_pool(pool);
// Schema Initialization
// 1. Standard items table
db.execute(Statement::from_string(
DatabaseBackend::Sqlite,
"CREATE TABLE IF NOT EXISTS items (
id INTEGER PRIMARY KEY AUTOINCREMENT,
content TEXT NOT NULL,
document_name TEXT,
created_at TEXT DEFAULT (datetime('now', 'localtime')),
updated_at TEXT DEFAULT (datetime('now', 'localtime'))
);".to_string(),
)).await?;
// 2. Add embedding column to items (ALTER TABLE does not support IF NOT EXISTS in many sqlite versions)
// We attempt it and ignore errors if column exists
let _ = db.execute(Statement::from_string(
DatabaseBackend::Sqlite,
"ALTER TABLE items ADD COLUMN embedding BLOB;".to_string(),
)).await;
// 3. Initialize vector engine for the column
db.execute(Statement::from_string(
DatabaseBackend::Sqlite,
format!(
"SELECT vector_init('items', 'embedding', 'type=FLOAT32,dimension={}');",
vec_dim
),
)).await?;
// triggers
db.execute(Statement::from_string(
DatabaseBackend::Sqlite,
"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;".to_string(),
)).await?;
Ok(db)
}
#[cfg(test)]
mod tests {
use super::*;
use sea_orm::{ConnectOptions, Database};
#[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());
}
}