khive_db/

migrations.rs

//! Schema migration system for the SQLite storage layer.
//!
//! Two APIs coexist:
//! - **Legacy per-service migrations** (`ServiceSchemaPlan` / `apply_schema_plan`):
//!   used by pack-scoped schemas.
//! - **Versioned migrations** (`MIGRATIONS` / `run_migrations`): the forward-only
//!   migration pipeline for the core tables.

use rusqlite::Connection;

use crate::error::SqliteError;

// =============================================================================
// Legacy per-service migration API (preserved for backward compatibility)
// =============================================================================

/// A single legacy migration step within a `ServiceSchemaPlan`.
pub struct Migration {
    /// Unique identifier for this migration.
    pub id: &'static str,
    /// SQL to apply (forward direction).
    pub up_sql: &'static str,
    /// SQL to revert (optional).
    pub down_sql: Option<&'static str>,
    /// Optional predicate: returns true if migration was already applied
    /// through a mechanism other than the migration tracker.
    pub is_already_applied: Option<fn(&Connection) -> bool>,
}

/// A pack-scoped schema plan containing migrations for SQLite and Postgres.
pub struct ServiceSchemaPlan {
    /// Service name used as a key in the `_schema_versions` tracking table.
    pub service: &'static str,
    /// SQLite-specific migration steps, applied in order.
    pub sqlite: &'static [Migration],
    /// Postgres-specific migration steps (reserved for future use).
    pub postgres: &'static [Migration],
}

const SCHEMA_VERSION_TABLE: &str = include_str!("../sql/schema-version-table.sql");

/// Apply a pack-scoped schema plan, tracking each migration in `_schema_versions`.
pub fn apply_schema_plan(conn: &Connection, plan: &ServiceSchemaPlan) -> Result<(), SqliteError> {
    conn.execute_batch(SCHEMA_VERSION_TABLE)?;

    for migration in plan.sqlite {
        // Check if custom predicate says it's already applied
        if let Some(check) = migration.is_already_applied {
            if check(conn) {
                continue;
            }
        }

        // Check if tracked as applied
        let already: bool = conn.query_row(
            "SELECT COUNT(*) > 0 FROM _schema_versions WHERE service = ?1 AND migration_id = ?2",
            rusqlite::params![plan.service, migration.id],
            |row| row.get(0),
        )?;

        if already {
            continue;
        }

        let tx =
            rusqlite::Transaction::new_unchecked(conn, rusqlite::TransactionBehavior::Immediate)?;
        tx.execute_batch(migration.up_sql)?;

        tx.execute(
            "INSERT INTO _schema_versions (service, migration_id, applied_at) VALUES (?1, ?2, ?3)",
            rusqlite::params![
                plan.service,
                migration.id,
                chrono::Utc::now().timestamp_micros(),
            ],
        )?;
        tx.commit()?;
    }

    Ok(())
}

// =============================================================================
// Versioned migration system
// =============================================================================

/// A single forward-only schema migration.
///
/// Migrations are applied in order from the current DB version to the target
/// version. Each migration runs in its own transaction; a failure rolls back
/// that migration and leaves the DB at the prior version.
pub struct VersionedMigration {
    /// Monotonically increasing version number, starting at 1.
    pub version: u32,
    /// Short human-readable name for the migration (used in the audit table).
    pub name: &'static str,
    /// SQL to apply this migration. May contain multiple statements separated
    /// by semicolons; `execute_batch` runs them all.
    pub up: &'static str,
}

// V1: complete schema, loaded from sql/schema.sql.
// Fresh-start repo (v0.2.8) — all schema in one migration, no incremental versions.
const V1_UP: &str = include_str!("../sql/schema.sql");

const V2_UP: &str = include_str!("../sql/002-narrow-fts-sections-update-trigger.sql");

const V3_UP: &str = include_str!("../sql/003-backfill-domain-mirror-atoms.sql");

const V4_UP: &str = include_str!("../sql/004-fts-consolidation.sql");

const V5_UP: &str = include_str!("../sql/005-unique-comm-external-id.sql");

const V6_UP: &str = include_str!("../sql/006-brain-retune-driver.sql");

const V7_UP: &str = include_str!("../sql/007-notes-seq.sql");

const V8_UP: &str = include_str!("../sql/008-notes-seq-repair.sql");

const V9_UP: &str = include_str!("../sql/009-entities-name-ci-index.sql");

const V10_UP: &str = include_str!("../sql/010-entities-content-ref.sql");

const V11_UP: &str = include_str!("../sql/011-ann-write-log.sql");

const V12_UP: &str = include_str!("../sql/012-ann-write-log-model-seq-index.sql");

const V13_UP: &str = include_str!("../sql/013-list-cursor-sequences.sql");

const V14_UP: &str = include_str!("../sql/014-graph-edges-id-unique.sql");

const V15_UP: &str = include_str!("../sql/015-serve-ledger-attribution.sql");

const V16_UP: &str = include_str!("../sql/016-gtd-dependency-cycle-guards.sql");

/// DDL for the `ann_write_log` delta table.
///
/// Shared between migration V11 and the belt-and-suspenders creation in
/// `StorageBackend::vectors_for_namespace` (same pattern as
/// [`EMBEDDING_MODELS_DDL`]): every database that hosts `vec_*` tables must
/// also have the write log, or vector writes would fail on databases opened
/// without `run_migrations()`. The `.sql` file is `IF NOT EXISTS`-idempotent.
pub const ANN_WRITE_LOG_DDL: &str = V11_UP;

/// DDL for the `ann_write_log` model/kind/field-leading index (ADR-118 §"Cost
/// bound"), shared between migration V12 and the belt-and-suspenders creation
/// in `StorageBackend::vectors_for_namespace` for the same reason as
/// [`ANN_WRITE_LOG_DDL`].
pub const ANN_WRITE_LOG_MODEL_SEQ_INDEX_DDL: &str = V12_UP;

/// DDL for the `_embedding_models` registry table.
///
/// Shared between the V1 schema and the belt-and-suspenders creation in
/// `StorageBackend::vectors_for_namespace`. Both sites reference this constant so
/// the schema cannot silently diverge if the registry evolves.
pub const EMBEDDING_MODELS_DDL: &str = include_str!("../sql/embedding-models-ddl.sql");

/// All versioned migrations in ascending order, applied by `run_migrations`.
pub const MIGRATIONS: &[VersionedMigration] = &[
    VersionedMigration {
        version: 1,
        name: "initial_schema",
        up: V1_UP,
    },
    VersionedMigration {
        version: 2,
        name: "narrow_fts_sections_update_trigger",
        up: V2_UP,
    },
    VersionedMigration {
        version: 3,
        name: "backfill_domain_mirror_atoms",
        up: V3_UP,
    },
    VersionedMigration {
        version: 4,
        name: "fts_consolidation",
        up: V4_UP,
    },
    VersionedMigration {
        version: 5,
        name: "unique_comm_message_external_id",
        up: V5_UP,
    },
    VersionedMigration {
        version: 6,
        name: "brain_retune_driver",
        up: V6_UP,
    },
    VersionedMigration {
        version: 7,
        name: "notes_seq",
        up: V7_UP,
    },
    VersionedMigration {
        version: 8,
        name: "notes_seq_repair",
        up: V8_UP,
    },
    VersionedMigration {
        version: 9,
        name: "entities_name_ci_index",
        up: V9_UP,
    },
    VersionedMigration {
        version: 10,
        name: "entities_content_ref",
        up: V10_UP,
    },
    VersionedMigration {
        version: 11,
        name: "ann_write_log",
        up: V11_UP,
    },
    VersionedMigration {
        version: 12,
        name: "ann_write_log_model_seq_index",
        up: V12_UP,
    },
    VersionedMigration {
        version: 13,
        name: "list_cursor_sequences",
        up: V13_UP,
    },
    VersionedMigration {
        version: 14,
        name: "graph_edges_id_unique",
        up: V14_UP,
    },
    VersionedMigration {
        version: 15,
        name: "serve_ledger_attribution",
        up: V15_UP,
    },
    VersionedMigration {
        version: 16,
        name: "gtd_dependency_cycle_guards",
        up: V16_UP,
    },
];

const MIGRATION_TRACKING_TABLE: &str = include_str!("../sql/schema-migrations-table.sql");

/// Apply all unapplied migrations in order. Idempotent; each migration runs in its own transaction.
/// Errors on non-contiguous version array or failed migration.
/// Read the applied schema version from an open connection **without** running
/// migrations. Returns 0 when the `_schema_migrations` ledger is absent (an
/// un-migrated or empty database); any other failure (BUSY, IO) propagates —
/// collapsing it to 0 would misreport a live database as un-migrated. Never
/// writes.
pub fn read_schema_version(conn: &Connection) -> Result<u32, SqliteError> {
    match conn.query_row(
        "SELECT COALESCE(MAX(version), 0) FROM _schema_migrations",
        [],
        |row| row.get(0),
    ) {
        Ok(version) => Ok(version),
        Err(rusqlite::Error::SqliteFailure(_, Some(ref msg)))
            if msg.contains("no such table: _schema_migrations") =>
        {
            Ok(0)
        }
        Err(e) => Err(e.into()),
    }
}

/// Open `path` read-only and report its applied schema version without creating
/// or migrating the file. The caller must ensure `path` exists — opening a
/// missing file read-only errors rather than creating it. This is the path used
/// by schema-inspection commands that must not mutate the database.
pub fn inspect_schema_version(path: &std::path::Path) -> Result<u32, SqliteError> {
    let conn = Connection::open_with_flags(
        path,
        rusqlite::OpenFlags::SQLITE_OPEN_READ_ONLY | rusqlite::OpenFlags::SQLITE_OPEN_NO_MUTEX,
    )?;
    read_schema_version(&conn)
}

#[cfg(test)]
pub(crate) mod test_sync {
    use std::sync::atomic::AtomicU32;
    use std::sync::{Arc, Barrier, Mutex};

    /// When set, `run_migrations_locked` parks after its initial (stale)
    /// ledger read until every racing thread has arrived — forcing the
    /// contended interleaving the concurrent-boot test asserts on.
    pub(crate) static STALE_READ_BARRIER: Mutex<Option<Arc<Barrier>>> = Mutex::new(None);
    /// Counts entries into the under-lock sibling fast-forward branch.
    pub(crate) static LOCKED_FAST_FORWARDS: AtomicU32 = AtomicU32::new(0);
    /// Set by the SQLite busy handler installed on participating connections:
    /// `true` means SQLite itself reported a blocked lock acquisition to the
    /// loser — actual contention, not merely an intended attempt.
    pub(crate) static BUSY_OBSERVED: std::sync::atomic::AtomicBool =
        std::sync::atomic::AtomicBool::new(false);

    /// Busy handler for participating test connections: records that SQLite
    /// observed a busy acquisition, then keeps retrying.
    pub(crate) fn record_busy(_count: i32) -> bool {
        BUSY_OBSERVED.store(true, std::sync::atomic::Ordering::SeqCst);
        std::thread::sleep(std::time::Duration::from_millis(1));
        true
    }

    /// Set by the winner immediately before committing its first migration
    /// transaction — i.e. before the write lock is first released.
    pub(crate) static WINNER_COMMITTED: std::sync::atomic::AtomicBool =
        std::sync::atomic::AtomicBool::new(false);
    /// Recorded by the loser when its first `BEGIN IMMEDIATE` returns: whether
    /// the winner had already committed at that moment. `true` is direct
    /// evidence the loser's lock acquisition blocked across the winner's held
    /// write lock rather than the two calls serializing by scheduler accident.
    pub(crate) static LOSER_SAW_WINNER_COMMIT: std::sync::atomic::AtomicBool =
        std::sync::atomic::AtomicBool::new(false);

    std::thread_local! {
        /// Opt-in flag: only threads that set this participate in the barrier,
        /// so unrelated tests migrating in parallel are never parked.
        pub(crate) static PARTICIPATE: std::cell::Cell<bool> =
            const { std::cell::Cell::new(false) };
        /// Whether this thread has already instrumented its first BEGIN.
        pub(crate) static FIRST_BEGIN_DONE: std::cell::Cell<bool> =
            const { std::cell::Cell::new(false) };
    }
}

pub fn run_migrations(conn: &mut Connection) -> Result<u32, SqliteError> {
    // Concurrent boots (multiple processes migrating the same file) contend on
    // the write lock below; a short hot-path busy_timeout cannot wait out a
    // sibling's migration. Raise-only to a 5s floor — never reduce a caller
    // whose configured timeout is already longer — and restore after.
    let prior_busy_ms: i64 = conn.query_row("PRAGMA busy_timeout", [], |row| row.get(0))?;
    let raised = prior_busy_ms < 5_000;
    if raised {
        conn.busy_timeout(std::time::Duration::from_secs(5))?;
    }
    let result = run_migrations_locked(conn);
    if raised {
        let _ = conn.busy_timeout(std::time::Duration::from_millis(prior_busy_ms.max(0) as u64));
    }
    result
}

fn run_migrations_locked(conn: &mut Connection) -> Result<u32, SqliteError> {
    conn.execute_batch(MIGRATION_TRACKING_TABLE)?;

    let current_version: u32 = read_schema_version(conn)?;

    // Deterministic-contention hook: parks every caller after the stale ledger
    // read (no lock held) until all racing test threads have observed it, so
    // they are then released to compete for the IMMEDIATE write lock below.
    #[cfg(test)]
    if test_sync::PARTICIPATE.with(|p| p.get()) {
        // Replaces the busy_timeout raised by `run_migrations` on this test
        // connection: records SQLite-observed contention, then keeps retrying.
        conn.busy_handler(Some(test_sync::record_busy))?;
        let barrier = test_sync::STALE_READ_BARRIER.lock().unwrap().clone();
        if let Some(barrier) = barrier {
            barrier.wait();
        }
    }

    // A database whose recorded version is ahead of the latest known migration
    // predates the consolidated V1 baseline (ADR-015) — e.g. it still carries the
    // pre-consolidation V2..V22 ledger — or was written by a newer build. Either
    // way the baseline schema would be silently skipped, leaving the process on a
    // stale schema. Fail loudly instead of corrupting silently.
    let latest_version = MIGRATIONS.last().map(|m| m.version).unwrap_or(0);
    if current_version > latest_version {
        return Err(SqliteError::InvalidData(format!(
            "database schema version {current_version} is ahead of the latest known migration \
             {latest_version}. This database predates the consolidated baseline (ADR-015) or was \
             written by a newer build. Recreate it from the current schema; in-place downgrade is \
             not supported."
        )));
    }

    let mut applied_version = current_version;
    // Floor advanced when a sibling's work is observed under the write lock,
    // so a losing process skips the remaining already-applied migrations
    // without opening a transaction for each.
    let mut skip_through = current_version;

    for migration in MIGRATIONS {
        if migration.version <= skip_through {
            applied_version = applied_version.max(migration.version);
            continue;
        }

        // IMMEDIATE: take the write lock up front so concurrent boots serialize
        // here instead of failing mid-migration when a DEFERRED transaction
        // upgrades to a write.
        #[cfg(test)]
        let instrumented_first_begin = test_sync::PARTICIPATE.with(|p| p.get())
            && !test_sync::FIRST_BEGIN_DONE.with(|f| f.get());
        #[cfg(test)]
        if instrumented_first_begin {
            test_sync::FIRST_BEGIN_DONE.with(|f| f.set(true));
        }
        let tx = conn
            .transaction_with_behavior(rusqlite::TransactionBehavior::Immediate)
            .map_err(|e| SqliteError::Migration {
                version: migration.version,
                error: e.to_string(),
            })?;

        // Re-check under the write lock: a sibling process may have applied
        // this migration (and possibly later ones) while we waited. Running
        // its DDL again would fail; fast-forward past everything it applied.
        let sibling_version: u32 = tx
            .query_row(
                "SELECT COALESCE(MAX(version), 0) FROM _schema_migrations",
                [],
                |row| row.get(0),
            )
            .map_err(|e| SqliteError::Migration {
                version: migration.version,
                error: e.to_string(),
            })?;
        #[cfg(test)]
        if instrumented_first_begin {
            use std::sync::atomic::Ordering::SeqCst;
            if sibling_version == 0 {
                // Winner: hold the write lock until SQLite has reported a
                // busy acquisition to the loser (its busy handler fired) —
                // proof the loser's BEGIN is actually blocked on this held
                // lock, not merely intended. Bounded so a regression fails
                // the assertion instead of hanging the test.
                let deadline = std::time::Instant::now() + std::time::Duration::from_secs(5);
                while !test_sync::BUSY_OBSERVED.load(SeqCst) && std::time::Instant::now() < deadline
                {
                    std::thread::yield_now();
                }
            } else {
                // Loser: our first BEGIN just returned. Record whether the
                // winner had already committed — true means we blocked across
                // its held lock.
                test_sync::LOSER_SAW_WINNER_COMMIT
                    .store(test_sync::WINNER_COMMITTED.load(SeqCst), SeqCst);
            }
        }

        // The ahead-of-latest guard above ran on a pre-lock read; a newer
        // build may have committed a version past ours while we waited for
        // the write lock. Accepting it (clamped) would return Ok on a schema
        // this binary does not understand — reject it the same way.
        if sibling_version > latest_version {
            return Err(SqliteError::InvalidData(format!(
                "database schema version {sibling_version} is ahead of the latest known \
                 migration {latest_version} (committed by a concurrent process while this \
                 one waited for the migration write lock). This build cannot run against \
                 the newer schema; upgrade the binary or recreate the database."
            )));
        }

        if sibling_version >= migration.version {
            #[cfg(test)]
            test_sync::LOCKED_FAST_FORWARDS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
            skip_through = sibling_version.min(latest_version);
            applied_version = applied_version.max(migration.version);
            continue;
        }

        tx.execute_batch(migration.up)
            .map_err(|e| SqliteError::Migration {
                version: migration.version,
                error: e.to_string(),
            })?;

        let now = chrono::Utc::now().timestamp_micros();
        tx.execute(
            "INSERT INTO _schema_migrations (version, name, applied_at) VALUES (?1, ?2, ?3) \
             ON CONFLICT(version) DO NOTHING",
            rusqlite::params![migration.version, migration.name, now],
        )
        .map_err(|e| SqliteError::Migration {
            version: migration.version,
            error: e.to_string(),
        })?;

        #[cfg(test)]
        if instrumented_first_begin {
            test_sync::WINNER_COMMITTED.store(true, std::sync::atomic::Ordering::SeqCst);
        }

        tx.commit().map_err(|e| SqliteError::Migration {
            version: migration.version,
            error: e.to_string(),
        })?;

        applied_version = migration.version;
    }

    Ok(applied_version)
}

#[derive(Debug)]
pub struct EmbeddingModelRegistryRecord {
    /// Vector engine name (e.g. `"paraphrase"`).
    pub engine_name: String,
    /// Model identifier (e.g. `"all-minilm-l6-v2"`).
    pub model_id: String,
    /// Canonical deduplication key combining engine and model.
    pub key_version: String,
    /// Embedding dimensionality.
    pub dimensions: u32,
    /// Lifecycle status (`"active"` or `"superseded"`).
    pub status: String,
    /// Epoch timestamp when the model was activated.
    pub activated_at: Option<i64>,
    /// Epoch timestamp when the model was superseded.
    pub superseded_at: Option<i64>,
}

/// Query the `_embedding_models` registry.
///
/// Opens the database at `db` (defaults to `~/.khive/khive.db`) and
/// returns all registry rows, optionally filtered by `engine_name`.
/// Returns an empty vec if the database or table does not exist.
pub fn query_embedding_models(
    db: Option<&std::path::Path>,
    engine_filter: Option<&str>,
) -> Result<Vec<EmbeddingModelRegistryRecord>, SqliteError> {
    let path = db.map(std::path::Path::to_path_buf).unwrap_or_else(|| {
        std::env::var("HOME")
            .map(std::path::PathBuf::from)
            .unwrap_or_else(|_| std::path::PathBuf::from("."))
            .join(".khive/khive.db")
    });
    if !path.exists() {
        return Ok(Vec::new());
    }
    let conn = Connection::open(path)?;
    query_embedding_models_conn(&conn, engine_filter)
}

/// Query `_embedding_models` from an existing connection (testable without a file).
///
/// Returns an empty vec if the table does not exist.
pub(crate) fn query_embedding_models_conn(
    conn: &Connection,
    engine_filter: Option<&str>,
) -> Result<Vec<EmbeddingModelRegistryRecord>, SqliteError> {
    let exists: bool = conn.query_row(
        "SELECT COUNT(*) > 0 FROM sqlite_master \
         WHERE type='table' AND name='_embedding_models'",
        [],
        |row| row.get(0),
    )?;
    if !exists {
        return Ok(Vec::new());
    }

    let sql = if engine_filter.is_some() {
        "SELECT engine_name, model_id, key_version, dim, status, activated_at, superseded_at \
         FROM _embedding_models WHERE engine_name = ?1 \
         ORDER BY engine_name, activated_at IS NULL, activated_at"
    } else {
        "SELECT engine_name, model_id, key_version, dim, status, activated_at, superseded_at \
         FROM _embedding_models \
         ORDER BY engine_name, activated_at IS NULL, activated_at"
    };
    let mut stmt = conn.prepare(sql)?;
    let map_row = |row: &rusqlite::Row<'_>| {
        let dim_raw: i64 = row.get(3)?;
        let dimensions = u32::try_from(dim_raw).map_err(|_| {
            rusqlite::Error::FromSqlConversionFailure(
                3,
                rusqlite::types::Type::Integer,
                Box::new(std::io::Error::other(format!(
                    "_embedding_models.dim value {dim_raw} is outside the valid u32 range [0, {}]",
                    u32::MAX,
                ))),
            )
        })?;
        Ok(EmbeddingModelRegistryRecord {
            engine_name: row.get(0)?,
            model_id: row.get(1)?,
            key_version: row.get(2)?,
            dimensions,
            status: row.get(4)?,
            activated_at: row.get(5)?,
            superseded_at: row.get(6)?,
        })
    };

    if let Some(engine) = engine_filter {
        stmt.query_map([engine], map_row)?
            .collect::<Result<Vec<_>, _>>()
            .map_err(Into::into)
    } else {
        stmt.query_map([], map_row)?
            .collect::<Result<Vec<_>, _>>()
            .map_err(Into::into)
    }
}

// =============================================================================
// Tests
// =============================================================================

#[cfg(test)]
#[path = "migrations_tests.rs"]
mod tests;