Newer
Older
TelosDB / src-backend / src / db.rs
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());
    }
}