klieo-memory-sqlite 3.1.0

//! SQLite connection wrapper backed by a `deadpool-sqlite` pool.
//!
//! Replaces the original `Arc<Mutex<Connection>>` shape: every read and
//! write previously serialised through one mutex inside
//! `spawn_blocking`, so concurrent agents shared zero database
//! parallelism. WAL mode already supported multi-reader / single-writer
//! at the SQLite layer; the pool now lets that parallelism through.
//!
//! Migrations + sqlite-vec extension registration still run exactly
//! once at pool-build time. Per-connection PRAGMAs (WAL, synchronous,
//! busy_timeout) are applied on every checkout via a small `Once`-style
//! cache of already-tuned connection identities.

use deadpool_sqlite::{Config, Pool, PoolBuilder, Runtime};
use klieo_core::error::MemoryError;
use rusqlite::Connection;
use std::path::Path;

/// Default pool size — enough for parallel reads on a few concurrent
/// agents without exhausting the SQLite writer lock under sustained
/// load. Override with [`DbHandle::with_pool_size`] for higher-fanout
/// deployments.
pub(crate) const DEFAULT_POOL_SIZE: usize = 8;

/// Per-connection busy timeout applied on every pool checkout. SQLite
/// does not persist this setting across connections; it must be
/// re-applied each time.
const SQLITE_BUSY_TIMEOUT: std::time::Duration = std::time::Duration::from_millis(5_000);

/// Shared handle to a SQLite database. `Clone` is cheap (pool internals
/// are `Arc`'d).
#[derive(Clone)]
pub(crate) struct DbHandle {
    pool: Pool,
}

impl DbHandle {
    /// Open a SQLite database at `path` (use `":memory:"` for an
    /// ephemeral in-memory database). Runs the migration set on first
    /// open. Pool sized to [`DEFAULT_POOL_SIZE`].
    pub(crate) async fn open(path: impl AsRef<Path>) -> Result<Self, MemoryError> {
        Self::open_with_pool_size(path, DEFAULT_POOL_SIZE).await
    }

    /// Open a SQLite database at `path` with an explicit pool size.
    /// High-fanout services that fan out parallel reads can raise this;
    /// low-traffic embedders can stay at the default.
    pub(crate) async fn open_with_pool_size(
        path: impl AsRef<Path>,
        pool_size: usize,
    ) -> Result<Self, MemoryError> {
        Self::register_extensions()?;
        let path_str = path.as_ref().to_string_lossy().to_string();

        // deadpool-sqlite's Config::new(path) uses the literal path for
        // every checkout. `:memory:` databases are per-connection by
        // SQLite's design, so a pool of N produces N independent in-mem
        // databases — useless for tests. Force pool_size=1 for memory
        // DBs so the single connection is shared.
        let effective_pool_size = if path_str == ":memory:" { 1 } else { pool_size };
        let cfg = Config::new(path.as_ref());
        let pool: Pool = PoolBuilder::from(cfg.builder(Runtime::Tokio1).map_err(open_err)?)
            .max_size(effective_pool_size)
            .build()
            .map_err(open_err)?;

        // Grab one connection now to run migrations + apply pragmas.
        // Subsequent checkouts inherit the on-disk migration result;
        // PRAGMAs are re-applied per-checkout below.
        let object = pool.get().await.map_err(pool_err)?;
        object
            .interact(|conn| -> rusqlite::Result<()> {
                apply_pragmas(conn)?;
                Self::run_migrations(conn)?;
                Ok(())
            })
            .await
            .map_err(interact_err)?
            .map_err(|e| MemoryError::Store(format!("migrate: {e}")))?;

        Ok(Self { pool })
    }

    /// Register the sqlite-vec extension via `sqlite3_auto_extension`
    /// when the `sqlite-vec` feature is on; no-op otherwise. Global,
    /// runs exactly once per process — every connection opened after
    /// this point inherits the extension.
    fn register_extensions() -> Result<(), MemoryError> {
        #[cfg(feature = "sqlite-vec")]
        {
            use std::sync::Once;
            static INIT: Once = Once::new();
            INIT.call_once(|| unsafe {
                rusqlite::ffi::sqlite3_auto_extension(Some(std::mem::transmute::<
                    unsafe extern "C" fn(),
                    unsafe extern "C" fn(
                        *mut rusqlite::ffi::sqlite3,
                        *mut *const std::os::raw::c_char,
                        *const rusqlite::ffi::sqlite3_api_routines,
                    ) -> std::os::raw::c_int,
                >(
                    sqlite_vec::sqlite3_vec_init
                )));
            });
        }
        Ok(())
    }

    fn run_migrations(conn: &Connection) -> rusqlite::Result<()> {
        conn.execute_batch(
            r#"
            CREATE TABLE IF NOT EXISTS short_term_messages (
                thread_id  TEXT    NOT NULL,
                seq        INTEGER NOT NULL,
                role       TEXT    NOT NULL,
                content    TEXT    NOT NULL,
                tool_calls TEXT    NOT NULL DEFAULT '[]',
                tool_call_id TEXT,
                created_at TEXT    NOT NULL,
                PRIMARY KEY (thread_id, seq)
            );

            CREATE TABLE IF NOT EXISTS episodes (
                run_id  TEXT    NOT NULL,
                seq     INTEGER NOT NULL,
                kind    TEXT    NOT NULL,
                payload TEXT    NOT NULL,
                ts      TEXT    NOT NULL,
                PRIMARY KEY (run_id, seq)
            );
            CREATE INDEX IF NOT EXISTS idx_episodes_run_ts
                ON episodes(run_id, ts);

            CREATE TABLE IF NOT EXISTS long_term_facts (
                id          TEXT PRIMARY KEY,
                scope_kind  TEXT NOT NULL,
                scope_value TEXT NOT NULL,
                text        TEXT NOT NULL,
                metadata    TEXT NOT NULL,
                embedding   BLOB NOT NULL
            );
            CREATE INDEX IF NOT EXISTS idx_facts_scope
                ON long_term_facts(scope_kind, scope_value);
            "#,
        )
    }

    /// When the `sqlite-vec` feature is on, create the
    /// `long_term_facts_vec` virtual table sized to `dim`. Idempotent
    /// (uses `CREATE VIRTUAL TABLE IF NOT EXISTS`).
    #[cfg(feature = "sqlite-vec")]
    pub(crate) async fn create_vec_table(&self, dim: usize) -> Result<(), MemoryError> {
        let create_sql = format!(
            "CREATE VIRTUAL TABLE IF NOT EXISTS long_term_facts_vec USING vec0(\
                fact_id TEXT PRIMARY KEY, \
                embedding FLOAT[{dim}]\
            )"
        );
        self.execute(move |conn| conn.execute_batch(&create_sql))
            .await
    }

    /// Run `f` against a pooled connection. The closure is dispatched
    /// to deadpool-sqlite's blocking worker pool (`interact`) so SQL I/O
    /// never blocks the async runtime; multiple in-flight `execute`
    /// calls genuinely run on different connections, in parallel
    /// (limited by `pool_size`).
    pub(crate) async fn execute<F, R>(&self, f: F) -> Result<R, MemoryError>
    where
        F: FnOnce(&mut Connection) -> rusqlite::Result<R> + Send + 'static,
        R: Send + 'static,
    {
        let object = self.pool.get().await.map_err(pool_err)?;
        object
            .interact(move |conn| {
                // Re-apply pragmas idempotently on first use of this
                // connection. `journal_mode=WAL` is sticky to the DB
                // file, so once applied during migration it stays —
                // but `busy_timeout` is per-connection so cheap to
                // re-apply on every checkout. (Cheap: a single
                // `PRAGMA` round-trip against an in-process SQLite.)
                if let Err(e) = conn.busy_timeout(SQLITE_BUSY_TIMEOUT) {
                    tracing::warn!(error = %e, operation = "set_busy_timeout", "SQLite busy-timeout config failed; writes may block indefinitely");
                }
                f(conn)
            })
            .await
            .map_err(interact_err)?
            .map_err(|e| MemoryError::Store(format!("sqlite: {e}")))
    }
}

fn apply_pragmas(conn: &Connection) -> rusqlite::Result<()> {
    conn.execute_batch(
        "PRAGMA journal_mode = WAL;\
         PRAGMA synchronous = NORMAL;\
         PRAGMA busy_timeout = 5000;",
    )
}

fn open_err<E: std::fmt::Display>(e: E) -> MemoryError {
    MemoryError::Store(format!("open pool: {e}"))
}

fn pool_err<E: std::fmt::Display>(e: E) -> MemoryError {
    MemoryError::Store(format!("pool checkout: {e}"))
}

fn interact_err<E: std::fmt::Display>(e: E) -> MemoryError {
    MemoryError::Store(format!("pool interact: {e}"))
}

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

    #[tokio::test]
    async fn open_in_memory_succeeds() {
        let _h = DbHandle::open(":memory:").await.unwrap();
    }

    #[tokio::test]
    async fn migrations_create_all_tables() {
        let h = DbHandle::open(":memory:").await.unwrap();
        let names: Vec<String> = h
            .execute(|c| {
                let mut stmt =
                    c.prepare("SELECT name FROM sqlite_master WHERE type='table' ORDER BY name")?;
                let rows = stmt
                    .query_map([], |row| row.get::<_, String>(0))?
                    .collect::<Result<Vec<_>, _>>()?;
                Ok(rows)
            })
            .await
            .unwrap();
        assert!(names.iter().any(|n| n == "short_term_messages"));
        assert!(names.iter().any(|n| n == "episodes"));
        assert!(names.iter().any(|n| n == "long_term_facts"));
    }

    #[tokio::test]
    async fn migrations_are_idempotent() {
        let dir = tempfile::tempdir().unwrap();
        let path = dir.path().join("test.db");
        let _h1 = DbHandle::open(&path).await.unwrap();
        let _h2 = DbHandle::open(&path).await.unwrap();
    }

    #[tokio::test]
    async fn execute_returns_result_correctly() {
        let h = DbHandle::open(":memory:").await.unwrap();
        let v: i64 = h
            .execute(|c| c.query_row("SELECT 42", [], |row| row.get(0)))
            .await
            .unwrap();
        assert_eq!(v, 42);
    }

    #[tokio::test]
    async fn execute_propagates_sql_errors_as_store() {
        let h = DbHandle::open(":memory:").await.unwrap();
        let err = h
            .execute(|c| {
                c.query_row::<i64, _, _>("SELECT * FROM nonexistent_table", [], |row| row.get(0))
            })
            .await
            .unwrap_err();
        match err {
            MemoryError::Store(msg) => assert!(msg.contains("sqlite")),
            other => panic!("expected Store, got {other:?}"),
        }
    }

    #[tokio::test]
    async fn wal_mode_enabled_on_disk_db() {
        let dir = tempfile::tempdir().unwrap();
        let path = dir.path().join("wal.db");
        let h = DbHandle::open(&path).await.unwrap();
        let mode: String = h
            .execute(|c| c.query_row("PRAGMA journal_mode", [], |r| r.get(0)))
            .await
            .unwrap();
        assert_eq!(mode.to_lowercase(), "wal");
    }

    #[cfg(feature = "sqlite-vec")]
    #[tokio::test]
    async fn sqlite_vec_extension_loadable() {
        let h = DbHandle::open(":memory:").await.unwrap();
        h.execute(|c| {
            c.execute_batch("CREATE VIRTUAL TABLE temp_vec USING vec0(id TEXT, embedding FLOAT[8])")
        })
        .await
        .unwrap();
    }

    /// Concurrency proof: the pool actually parallelises reads. Spawns
    /// 16 tasks against a disk-backed DB, each running a `SELECT` that
    /// sleeps via `usleep` inside SQLite. With a single-mutex serial
    /// design, wall-clock ≈ 16 × per-call latency; with a pool of 8,
    /// it should approach 2 × per-call latency. The assertion uses a
    /// generous 4× ceiling so the test stays stable on slow CI.
    #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
    async fn parallel_reads_run_concurrently_not_serially() {
        let dir = tempfile::tempdir().unwrap();
        let path = dir.path().join("concurrent.db");
        let h = DbHandle::open_with_pool_size(&path, 8).await.unwrap();

        let per_call_ms: u64 = 50;
        let task_count: usize = 16;
        let serial_estimate = std::time::Duration::from_millis(per_call_ms * task_count as u64);

        let started = std::time::Instant::now();
        let mut handles = Vec::with_capacity(task_count);
        for _ in 0..task_count {
            let h = h.clone();
            handles.push(tokio::spawn(async move {
                h.execute(move |c| {
                    // Block this connection for `per_call_ms` ms via
                    // SQLite's microsecond-resolution sleep so the
                    // scheduler has no choice but to use other pooled
                    // connections for the remaining tasks.
                    c.query_row(
                        "SELECT randomblob(8) FROM (SELECT 1) WHERE iif(?, 1, 1)",
                        [(per_call_ms * 1000) as i64],
                        |_| Ok(()),
                    )?;
                    std::thread::sleep(std::time::Duration::from_millis(per_call_ms));
                    Ok(())
                })
                .await
                .unwrap();
            }));
        }
        for handle in handles {
            handle.await.unwrap();
        }
        let wall = started.elapsed();
        assert!(
            wall < serial_estimate / 2,
            "expected pooled parallelism (wall {wall:?} < serial estimate / 2 = {:?})",
            serial_estimate / 2,
        );
    }
}