olai-stack-state 0.0.2

//! The store's records, revisions, leases and schema guard.

use std::collections::BTreeSet;
use std::sync::Arc;

use olai_stack_state::{Desired, Draft, Error, SCHEMA_VERSION, SecretRef, Store, VarSet, VarValue};
use olai_stack_topology::{EnvManifest, MAX_SLOT, PlanCtx, Selection};

fn store() -> (tempfile::TempDir, Store) {
    let dir = tempfile::tempdir().unwrap();
    let store = Store::open(&dir.path().join("state.db")).unwrap();
    (dir, store)
}

fn draft(modules: &[&str]) -> Draft {
    Draft::new(EnvManifest::new(
        Selection::modules(modules.iter().copied()),
        PlanCtx::default(),
    ))
}

#[test]
fn drafts_commit_to_numbered_revisions() {
    let (_dir, store) = store();
    let ws = store.create_workspace("demo", true).unwrap();
    let env = store
        .create_environment(&ws.id, "dev", draft(&["unity-catalog"]))
        .unwrap();
    assert_eq!(env.head, None);
    assert_eq!(env.draft.manifest.context.env_name, "dev");

    let r1 = store.commit_draft(&env.id).unwrap();
    assert_eq!(r1.number(), 1);

    let mut next = draft(&["unity-catalog", "mlflow"]);
    next.vars
        .set("LOG_LEVEL", VarValue::Literal("debug".into()));
    store.update_draft(&env.id, &next).unwrap();
    let r2 = store.commit_draft(&env.id).unwrap();
    assert_eq!(r2.number(), 2);
    assert_ne!(r1.digest(), r2.digest());

    // Revision 1 is as it was committed; the draft and head moved on.
    assert_eq!(store.revision(&env.id, 1).unwrap(), r1);
    assert_eq!(r1.manifest().selection.modules.len(), 1);
    assert_eq!(store.environment(&env.id).unwrap().head, Some(2));
    assert_eq!(store.revisions(&env.id).unwrap().len(), 2);

    // The same content digests the same.
    store.update_draft(&env.id, &next).unwrap();
    assert_eq!(store.commit_draft(&env.id).unwrap().digest(), r2.digest());
}

#[test]
fn a_revision_cant_be_updated_even_bypassing_the_api() {
    let (dir, store) = store();
    let ws = store.create_workspace("demo", false).unwrap();
    let env = store.create_environment(&ws.id, "dev", draft(&[])).unwrap();
    store.commit_draft(&env.id).unwrap();

    let raw = rusqlite::Connection::open(dir.path().join("state.db")).unwrap();
    let err = raw
        .execute("UPDATE revision SET manifest = 'tampered'", [])
        .unwrap_err();
    assert!(err.to_string().contains("revisions are immutable"), "{err}");
}

#[test]
fn names_are_validated_and_unique() {
    let (_dir, store) = store();
    assert!(matches!(
        store.create_workspace("Demo", false),
        Err(Error::InvalidName(_))
    ));
    let ws = store.create_workspace("demo", false).unwrap();
    assert!(matches!(
        store.create_workspace("demo", false),
        Err(Error::Duplicate {
            kind: "workspace",
            ..
        })
    ));
    store.create_environment(&ws.id, "dev", draft(&[])).unwrap();
    assert!(matches!(
        store.create_environment(&ws.id, "dev", draft(&[])),
        Err(Error::Duplicate {
            kind: "environment",
            ..
        })
    ));
    // Another workspace may reuse the name.
    let other = store.create_workspace("other", false).unwrap();
    store
        .create_environment(&other.id, "dev", draft(&[]))
        .unwrap();
}

#[test]
fn vars_hold_literals_and_secret_references() {
    let (_dir, store) = store();
    let ws = store.create_workspace("demo", false).unwrap();
    let mut vars = VarSet::new();
    vars.set("REGION", VarValue::Literal("eu-west-1".into()));
    vars.set(
        "PG_PASSWORD",
        VarValue::Secret(SecretRef {
            provider: "keychain".into(),
            key: "demo/pg".into(),
        }),
    );
    store.set_workspace_vars(&ws.id, &vars).unwrap();
    assert_eq!(store.workspace("demo").unwrap().unwrap().vars, vars);
    // Debug never prints a value.
    let debug = format!("{vars:?}");
    assert!(!debug.contains("eu-west-1"), "{debug}");
}

#[test]
fn instances_lease_host_wide_slots() {
    let (_dir, store) = store();
    let mut instances = Vec::new();
    for ws_name in ["alpha", "beta"] {
        let ws = store.create_workspace(ws_name, false).unwrap();
        let env = store.create_environment(&ws.id, "dev", draft(&[])).unwrap();
        let rev = store.commit_draft(&env.id).unwrap();
        instances.push(store.create_instance(&env.id, rev.number()).unwrap());
    }
    // Both workspaces have a `dev`; they still get different slots.
    let a = store.lease_slot(&instances[0].id).unwrap();
    let b = store.lease_slot(&instances[1].id).unwrap();
    assert_eq!((a.slot, b.slot), (0, 1));
    // Leasing again returns the held slot.
    assert_eq!(store.lease_slot(&instances[0].id).unwrap(), a);
    assert_eq!(store.instance(&instances[1].id).unwrap().slot, Some(1));

    // A released slot is handed out again, lowest first.
    store.release_slot(&instances[0].id).unwrap();
    assert_eq!(store.instance(&instances[0].id).unwrap().slot, None);
    let env = instances[0].environment.clone();
    let third = store.create_instance(&env, 1).unwrap();
    assert_eq!(store.lease_slot(&third.id).unwrap().slot, 0);

    // Deleting an instance frees its slot.
    store.delete_instance(&instances[1].id).unwrap();
    assert!(store.leases().unwrap().iter().all(|l| l.slot != 1));

    store.set_desired(&third.id, Desired::Running).unwrap();
    assert_eq!(store.instance(&third.id).unwrap().desired, Desired::Running);
}

#[test]
fn an_instance_needs_an_existing_revision() {
    let (_dir, store) = store();
    let ws = store.create_workspace("demo", false).unwrap();
    let env = store.create_environment(&ws.id, "dev", draft(&[])).unwrap();
    assert!(matches!(
        store.create_instance(&env.id, 1),
        Err(Error::NotFound {
            kind: "revision",
            ..
        })
    ));
}

#[test]
fn concurrent_leases_never_share_a_slot_and_run_out_cleanly() {
    let (dir, store) = store();
    let ws = store.create_workspace("demo", false).unwrap();
    let env = store.create_environment(&ws.id, "dev", draft(&[])).unwrap();
    store.commit_draft(&env.id).unwrap();
    let total = usize::from(MAX_SLOT) + 1 + 8;
    let ids: Vec<_> = (0..total)
        .map(|_| store.create_instance(&env.id, 1).unwrap().id)
        .collect();
    let path = Arc::new(dir.path().join("state.db"));

    // Eight threads, each on its own connection, race for slots.
    let chunks: Vec<Vec<_>> = ids.chunks(total.div_ceil(8)).map(<[_]>::to_vec).collect();
    let handles: Vec<_> = chunks
        .into_iter()
        .map(|chunk| {
            let path = Arc::clone(&path);
            std::thread::spawn(move || {
                let store = Store::open(&path).unwrap();
                chunk
                    .iter()
                    .map(|id| store.lease_slot(id))
                    .collect::<Vec<_>>()
            })
        })
        .collect();
    let results: Vec<_> = handles
        .into_iter()
        .flat_map(|h| h.join().unwrap())
        .collect();
    let slots: Vec<u16> = results
        .iter()
        .filter_map(|r| r.as_ref().ok().map(|l| l.slot))
        .collect();
    let distinct: BTreeSet<u16> = slots.iter().copied().collect();
    assert_eq!(slots.len(), usize::from(MAX_SLOT) + 1);
    assert_eq!(distinct.len(), slots.len(), "a slot was leased twice");
    let exhausted = results
        .iter()
        .filter(|r| matches!(r, Err(Error::SlotsExhausted)))
        .count();
    assert_eq!(exhausted, 8);
}

#[test]
fn a_newer_schema_is_refused() {
    let (dir, store) = store();
    drop(store);
    let path = dir.path().join("state.db");
    let raw = rusqlite::Connection::open(&path).unwrap();
    raw.pragma_update(None, "user_version", SCHEMA_VERSION + 1)
        .unwrap();
    drop(raw);
    assert!(matches!(
        Store::open(&path),
        Err(Error::UnsupportedVersion { found, supported })
            if found == SCHEMA_VERSION + 1 && supported == SCHEMA_VERSION
    ));
}

#[test]
fn reopening_keeps_the_data() {
    let (dir, store) = store();
    store.create_workspace("demo", false).unwrap();
    drop(store);
    let store = Store::open(&dir.path().join("state.db")).unwrap();
    assert_eq!(store.workspaces().unwrap().len(), 1);
}

#[test]
fn trestle_home_overrides_the_data_dir() {
    // Only this test touches the variable in this binary.
    let dir = tempfile::tempdir().unwrap();
    // SAFETY: no other thread in this test binary reads or writes TRESTLE_HOME.
    unsafe { std::env::set_var("TRESTLE_HOME", dir.path()) };
    assert_eq!(olai_stack_state::trestle_home().unwrap(), dir.path());
    let store = Store::open_default().unwrap();
    assert_eq!(store.path(), dir.path().join("state.db"));
    unsafe { std::env::remove_var("TRESTLE_HOME") };
}