From ff2aef45647a6ea00891074a9fe7383ceae8d537 Mon Sep 17 00:00:00 2001 From: Scott Werner Date: Thu, 27 Aug 2026 14:59:34 -0400 Subject: [PATCH] Simplify SQL run record store write path and test fixtures Share one bind helper across the runs insert/upsert/update statements, compute the next event sequence once per append, and decode stored sequence columns through a single helper. Check the run head before decoding events, rewrite the first-visit stage listing as a UNION ALL so each arm uses its partial index, and share the run_events insert SQL with the test seeder. Collapse the duplicated in-memory pool fixture, remove two tests that only asserted Arc sharing, and fold the fabro-db test row helpers. Co-Authored-By: Claude Fable 5 --- lib/apps/fabro-server/src/server/tests.rs | 9 - .../fabro-store/src/run_record_store.rs | 284 +++++++----------- lib/components/fabro-store/src/slate/mod.rs | 10 - .../fabro-store/src/slate/run_store.rs | 5 - .../fabro-store/src/test_support/mod.rs | 58 ++-- lib/foundation/fabro-db/tests/sqlite.rs | 153 +++++----- 6 files changed, 209 insertions(+), 310 deletions(-) diff --git a/lib/apps/fabro-server/src/server/tests.rs b/lib/apps/fabro-server/src/server/tests.rs index 10666634c..e2b22054f 100644 --- a/lib/apps/fabro-server/src/server/tests.rs +++ b/lib/apps/fabro-server/src/server/tests.rs @@ -210,15 +210,6 @@ methods = ["dev-token"] ) } -#[test] -fn run_record_store_is_database_owned() { - let state = test_app_state(); - assert!(Arc::ptr_eq( - &state.stores.run_records, - &state.stores.runs.run_record_store(), - )); -} - async fn body_json(body: Body) -> serde_json::Value { let bytes = to_bytes(body, usize::MAX).await.unwrap(); serde_json::from_slice(&bytes).unwrap() diff --git a/lib/components/fabro-store/src/run_record_store.rs b/lib/components/fabro-store/src/run_record_store.rs index 3068185b7..b3707a729 100644 --- a/lib/components/fabro-store/src/run_record_store.rs +++ b/lib/components/fabro-store/src/run_record_store.rs @@ -7,7 +7,8 @@ use fabro_types::{ BilledTokenCounts, EventEnvelope, Run, RunEvent, RunId, RunSize, RunStatusKind, RunTiming, SessionId, StageId, timing, }; -use sqlx::sqlite::{SqliteConnection, SqliteRow}; +use sqlx::query::Query; +use sqlx::sqlite::{SqliteArguments, SqliteConnection, SqliteRow}; use sqlx::{QueryBuilder, Row as _, Sqlite, SqlitePool}; use strum::VariantArray as _; @@ -95,6 +96,11 @@ WHERE id = ? AND source_last_seq = ? const SELECT_EVENT_COLUMNS: &str = "SELECT run_id, seq, event_name, node_id, stage_id, session_id, event_json FROM run_events"; +const INSERT_EVENT_SQL: &str = r" +INSERT INTO run_events (run_id, seq, event_name, node_id, stage_id, session_id, event_json) +VALUES (?, ?, ?, ?, ?, ?, ?) +"; + const SELECT_RUN_SUMMARIES_SQL: &str = r" SELECT runs.id, runs.summary_json, (SELECT COUNT(*) FROM runs AS child WHERE child.parent_id = runs.id) AS children_count @@ -363,12 +369,7 @@ impl RunRecordStore { entry: &CachedRunProjection, payload: &EventPayload, ) -> Result { - let next_seq = expected_last_seq - .checked_add(1) - .filter(|seq| *seq <= keys::MAX_EVENT_SEQ) - .ok_or(Error::EventSequenceExhausted { - max_seq: keys::MAX_EVENT_SEQ, - })?; + let next_seq = next_event_seq_after(expected_last_seq)?; let record = PreparedRunSummary::from_entry(entry); ensure_entry_identity(entry, &record, next_seq)?; ensure_prepared_head(&record, next_seq)?; @@ -388,12 +389,15 @@ impl RunRecordStore { .push(" WHERE run_id = ") .push_bind(run_id.to_string()) .push(" ORDER BY seq ASC"); - let rows = query.build().fetch_all(&mut *connection).await?; - let events = decode_event_rows(&rows, run_id)?; - let expected_last_seq = select_run_head(connection, run_id) + let expected_last_seq = select_run_head(&mut *connection, run_id) .await? .ok_or_else(|| Error::RunNotFound(run_id.to_string()))?; - let actual_last_seq = events.last().map(|event| event.seq); + let rows = query.build().fetch_all(&mut *connection).await?; + let actual_last_seq = rows + .last() + .map(|row| row.try_get::("seq")) + .transpose()? + .and_then(stored_seq); if actual_last_seq != Some(expected_last_seq) { return Err(Error::RunHeadMismatch { run_id: run_id.to_string(), @@ -401,7 +405,7 @@ impl RunRecordStore { actual_last_seq, }); } - Ok(events) + decode_event_rows(&rows, run_id) } pub(crate) async fn list_events_from_with_limit_on_connection( @@ -453,7 +457,7 @@ impl RunRecordStore { .push_bind(i64::from(seq)); let row = query.build().fetch_optional(&mut *connection).await?; row.as_ref() - .map(|row| decode_event_row(row, run_id)) + .map(|row| decode_event_row(row, run_id, &run_id.to_string())) .transpose() } @@ -464,23 +468,27 @@ impl RunRecordStore { start_seq: u32, limit: usize, ) -> Result> { + // Legacy rows for a first visit carry only `node_id`. Query them as a + // second `UNION ALL` arm instead of an `OR` so each arm can use its + // own partial index rather than scanning the run's primary key range. let mut query = QueryBuilder::::new(SELECT_EVENT_COLUMNS); query .push(" WHERE run_id = ") .push_bind(run_id.to_string()) .push(" AND seq >= ") - .push_bind(i64::from(start_seq)); + .push_bind(i64::from(start_seq)) + .push(" AND stage_id = ") + .push_bind(stage_id.to_string()); if stage_id.visit() == 1 { query - .push(" AND (stage_id = ") - .push_bind(stage_id.to_string()) - .push(" OR (stage_id IS NULL AND node_id = ") - .push_bind(stage_id.node_id().to_string()) - .push("))"); - } else { - query - .push(" AND stage_id = ") - .push_bind(stage_id.to_string()); + .push(" UNION ALL ") + .push(SELECT_EVENT_COLUMNS) + .push(" WHERE run_id = ") + .push_bind(run_id.to_string()) + .push(" AND seq >= ") + .push_bind(i64::from(start_seq)) + .push(" AND stage_id IS NULL AND node_id = ") + .push_bind(stage_id.node_id().to_string()); } query .push(" ORDER BY seq ASC LIMIT ") @@ -594,9 +602,6 @@ fn validate_event_for_record( payload: &EventPayload, seq: u32, ) -> Result { - if seq == 0 || seq > keys::MAX_EVENT_SEQ { - return Err(run_event_mismatch(&record.run.id, seq, "seq")); - } payload.validate(&record.run.id)?; let event = RunEvent::try_from(payload)?; if event.run_id != record.run.id { @@ -619,23 +624,17 @@ async fn insert_event_on_connection( payload: &EventPayload, envelope: &EventEnvelope, ) -> Result<()> { - ensure_prepared_head(record, envelope.seq)?; let event_json = serde_json::to_string(payload)?; - sqlx::query( - r" -INSERT INTO run_events (run_id, seq, event_name, node_id, stage_id, session_id, event_json) -VALUES (?, ?, ?, ?, ?, ?, ?) -", - ) - .bind(record.run.id.to_string()) - .bind(i64::from(envelope.seq)) - .bind(envelope.event.event_name()) - .bind(envelope.event.node_id.as_deref()) - .bind(envelope.event.stage_id.as_ref().map(ToString::to_string)) - .bind(envelope.event.session_id.as_deref()) - .bind(event_json) - .execute(connection) - .await?; + sqlx::query(INSERT_EVENT_SQL) + .bind(record.run.id.to_string()) + .bind(i64::from(envelope.seq)) + .bind(envelope.event.event_name()) + .bind(envelope.event.node_id.as_deref()) + .bind(envelope.event.stage_id.as_ref().map(ToString::to_string)) + .bind(envelope.event.session_id.as_deref()) + .bind(event_json) + .execute(connection) + .await?; Ok(()) } @@ -643,20 +642,39 @@ fn sql_limit(limit: usize) -> i64 { i64::try_from(limit.saturating_add(1)).unwrap_or(i64::MAX) } +fn next_event_seq_after(last_seq: u32) -> Result { + last_seq + .checked_add(1) + .filter(|seq| *seq <= keys::MAX_EVENT_SEQ) + .ok_or(Error::EventSequenceExhausted { + max_seq: keys::MAX_EVENT_SEQ, + }) +} + +/// Decodes a stored sequence column, rejecting anything outside the valid +/// `1..=MAX_EVENT_SEQ` range. +fn stored_seq(value: i64) -> Option { + u32::try_from(value) + .ok() + .filter(|seq| (1..=keys::MAX_EVENT_SEQ).contains(seq)) +} + fn decode_event_rows(rows: &[SqliteRow], run_id: &RunId) -> Result> { + let run_id_text = run_id.to_string(); rows.iter() - .map(|row| decode_event_row(row, run_id)) + .map(|row| decode_event_row(row, run_id, &run_id_text)) .collect() } -fn decode_event_row(row: &SqliteRow, expected_run_id: &RunId) -> Result { +fn decode_event_row( + row: &SqliteRow, + expected_run_id: &RunId, + expected_run_id_text: &str, +) -> Result { let stored_run_id: String = row.try_get("run_id")?; - let stored_seq: i64 = row.try_get("seq")?; - let seq = u32::try_from(stored_seq) - .ok() - .filter(|seq| (1..=keys::MAX_EVENT_SEQ).contains(seq)) - .ok_or_else(|| run_event_mismatch(expected_run_id, 0, "seq"))?; - if stored_run_id != expected_run_id.to_string() { + let raw_seq: i64 = row.try_get("seq")?; + let seq = stored_seq(raw_seq).ok_or_else(|| run_event_mismatch(expected_run_id, 0, "seq"))?; + if stored_run_id != expected_run_id_text { return Err(run_event_mismatch(expected_run_id, seq, "run_id")); } @@ -695,10 +713,7 @@ async fn select_run_head(connection: &mut SqliteConnection, run_id: &RunId) -> R .await?; stored .map(|value| { - u32::try_from(value) - .ok() - .filter(|seq| (1..=keys::MAX_EVENT_SEQ).contains(seq)) - .ok_or_else(|| run_event_mismatch(run_id, 0, "source_last_seq")) + stored_seq(value).ok_or_else(|| run_event_mismatch(run_id, 0, "source_last_seq")) }) .transpose() } @@ -741,7 +756,6 @@ async fn insert_run_on_connection( connection: &mut SqliteConnection, record: &PreparedRunSummary, ) -> Result<()> { - ensure_prepared_head(record, 1)?; write_insert_shaped_run(connection, record, INSERT_RUN_SQL).await } @@ -750,11 +764,23 @@ async fn write_insert_shaped_run( record: &PreparedRunSummary, sql: &'static str, ) -> Result<()> { + bind_run_columns(sqlx::query(sql).bind(record.run.id.to_string()), record)? + .execute(connection) + .await?; + Ok(()) +} + +/// Binds the `runs` columns shared by the insert, upsert, and update +/// statements, in the positional order those statements declare them +/// (`source_last_seq` through `summary_json`). +fn bind_run_columns<'q>( + query: Query<'q, Sqlite, SqliteArguments>, + record: &'q PreparedRunSummary, +) -> Result> { let run = &record.run; let diff = run.diff.unwrap_or_default(); let summary_json = serde_json::to_string(run)?; - sqlx::query(sql) - .bind(run.id.to_string()) + Ok(query .bind(i64::from(record.last_seq)) .bind(run.timestamps.created_at.timestamp_millis()) .bind( @@ -794,10 +820,7 @@ async fn write_insert_shaped_run( .bind(record.cache_read_tokens) .bind(record.cache_write_tokens) .bind(record.total_usd_micros) - .bind(summary_json) - .execute(connection) - .await?; - Ok(()) + .bind(summary_json)) } async fn update_run_on_connection( @@ -805,58 +828,8 @@ async fn update_run_on_connection( record: &PreparedRunSummary, expected_last_seq: u32, ) -> Result<()> { - let next_seq = expected_last_seq - .checked_add(1) - .filter(|seq| *seq <= keys::MAX_EVENT_SEQ) - .ok_or(Error::EventSequenceExhausted { - max_seq: keys::MAX_EVENT_SEQ, - })?; - ensure_prepared_head(record, next_seq)?; - let run = &record.run; - let diff = run.diff.unwrap_or_default(); - let summary_json = serde_json::to_string(run)?; - let result = sqlx::query(UPDATE_RUN_SQL) - .bind(i64::from(record.last_seq)) - .bind(run.timestamps.created_at.timestamp_millis()) - .bind( - run.timestamps - .started_at - .map(|value| value.timestamp_millis()), - ) - .bind( - run.timestamps - .last_event_at - .unwrap_or(run.timestamps.created_at) - .timestamp_millis(), - ) - .bind( - run.timestamps - .completed_at - .map(|value| value.timestamp_millis()), - ) - .bind(run.lifecycle.status.kind().to_string()) - .bind( - run.lifecycle - .archived_at - .map(|value| value.timestamp_millis()), - ) - .bind(run.parent_id.map(|value| value.to_string())) - .bind(&run.title) - .bind(&run.workflow.slug) - .bind(&record.workflow_name) - .bind(&record.repository_name) - .bind(run.automation.as_ref().map(|automation| &automation.id)) - .bind(diff.files_changed) - .bind(diff.additions) - .bind(diff.deletions) - .bind(record.input_tokens) - .bind(record.output_tokens) - .bind(record.reasoning_tokens) - .bind(record.cache_read_tokens) - .bind(record.cache_write_tokens) - .bind(record.total_usd_micros) - .bind(summary_json) + let result = bind_run_columns(sqlx::query(UPDATE_RUN_SQL), record)? .bind(run.id.to_string()) .bind(i64::from(expected_last_seq)) .execute(&mut *connection) @@ -1017,17 +990,17 @@ mod tests { use chrono::{DateTime, Utc}; use fabro_types::{ - AutomationRef, BilledTokenCounts, BlockedReason, Conclusion, DiffSummary, FailureReason, - Graph, PendingReason, RunDiff, RunId, RunProjection, RunSize, RunSpec, RunStatus, - RunStatusKind, RunTiming, SessionId, StageId, StageOutcome, SuccessReason, + AutomationRef, BilledTokenCounts, BlockedReason, Conclusion, DiffSummary, EventEnvelope, + FailureReason, Graph, PendingReason, RunDiff, RunId, RunProjection, RunSize, RunSpec, + RunStatus, RunStatusKind, RunTiming, SessionId, StageId, StageOutcome, SuccessReason, WorkflowSettings, test_support, }; use strum::VariantArray as _; use ulid::Ulid; use super::{ - RunRecordStore, RunSummaryListQuery, RunSummarySort, RunSummarySortDirection, - RunSummaryVisibility, decode_event_row, + INSERT_EVENT_SQL, RunRecordStore, RunSummaryListQuery, RunSummarySort, + RunSummarySortDirection, RunSummaryVisibility, decode_event_row, }; use crate::slate::CachedRunProjection; use crate::{Error, EventPayload, test_support as store_test_support}; @@ -1073,10 +1046,6 @@ mod tests { store_test_support::sqlite_run_record_store().await } - #[expect( - clippy::needless_pass_by_value, - reason = "test call sites construct one-off JSON properties" - )] fn sql_event_payload( run_id: &RunId, event: &str, @@ -1090,9 +1059,9 @@ mod tests { "ts": "2026-08-27T12:00:00Z", "run_id": run_id.to_string(), "event": event, - "properties": properties, }); let object = value.as_object_mut().unwrap(); + object.insert("properties".to_string(), properties); if let Some(node_id) = node_id { object.insert("node_id".to_string(), node_id.into()); } @@ -1105,6 +1074,10 @@ mod tests { EventPayload::new(value, run_id).unwrap() } + fn seqs(events: &[EventEnvelope]) -> Vec { + events.iter().map(|event| event.seq).collect() + } + fn created_payload(run_id: &RunId) -> EventPayload { sql_event_payload( run_id, @@ -1130,22 +1103,17 @@ mod tests { payload: &EventPayload, ) { let event = fabro_types::RunEvent::try_from(payload).unwrap(); - sqlx::query( - r" -INSERT INTO run_events (run_id, seq, event_name, node_id, stage_id, session_id, event_json) -VALUES (?, ?, ?, ?, ?, ?, ?) -", - ) - .bind(run_id.to_string()) - .bind(i64::from(seq)) - .bind(event.event_name()) - .bind(event.node_id) - .bind(event.stage_id.map(|stage_id| stage_id.to_string())) - .bind(event.session_id) - .bind(serde_json::to_string(payload).unwrap()) - .execute(&store.pool) - .await - .unwrap(); + sqlx::query(INSERT_EVENT_SQL) + .bind(run_id.to_string()) + .bind(i64::from(seq)) + .bind(event.event_name()) + .bind(event.node_id) + .bind(event.stage_id.map(|stage_id| stage_id.to_string())) + .bind(event.session_id) + .bind(serde_json::to_string(payload).unwrap()) + .execute(&store.pool) + .await + .unwrap(); } fn sample_status(kind: RunStatusKind) -> RunStatus { @@ -1473,17 +1441,12 @@ VALUES (?, ?, ?, ?, ?, ?, ?) let all = RunRecordStore::list_events_on_connection(&mut connection, &id) .await .unwrap(); - assert_eq!(all.iter().map(|event| event.seq).collect::>(), vec![ - 1, 2, 3, 4, 5, 6 - ]); + assert_eq!(seqs(&all), vec![1, 2, 3, 4, 5, 6]); let forward = RunRecordStore::list_events_from_with_limit_on_connection(&mut connection, &id, 2, 2) .await .unwrap(); - assert_eq!( - forward.iter().map(|event| event.seq).collect::>(), - vec![2, 3, 4] - ); + assert_eq!(seqs(&forward), vec![2, 3, 4]); let reverse = RunRecordStore::list_events_before_with_limit_on_connection( &mut connection, &id, @@ -1492,10 +1455,7 @@ VALUES (?, ?, ?, ?, ?, ?, ?) ) .await .unwrap(); - assert_eq!( - reverse.iter().map(|event| event.seq).collect::>(), - vec![4, 3, 2] - ); + assert_eq!(seqs(&reverse), vec![4, 3, 2]); let exact = RunRecordStore::get_event_on_connection(&mut connection, &id, 5) .await .unwrap() @@ -1515,13 +1475,7 @@ VALUES (?, ?, ?, ?, ?, ?, ?) ) .await .unwrap(); - assert_eq!( - visit_one_events - .iter() - .map(|event| event.seq) - .collect::>(), - vec![2, 4] - ); + assert_eq!(seqs(&visit_one_events), vec![2, 4]); let visit_two_events = RunRecordStore::list_events_for_stage_from_with_limit_on_connection( &mut connection, &id, @@ -1531,13 +1485,7 @@ VALUES (?, ?, ?, ?, ?, ?, ?) ) .await .unwrap(); - assert_eq!( - visit_two_events - .iter() - .map(|event| event.seq) - .collect::>(), - vec![3] - ); + assert_eq!(seqs(&visit_two_events), vec![3]); let session_events = RunRecordStore::list_events_for_session_from_with_limit_on_connection( &mut connection, &id, @@ -1547,13 +1495,7 @@ VALUES (?, ?, ?, ?, ?, ?, ?) ) .await .unwrap(); - assert_eq!( - session_events - .iter() - .map(|event| event.seq) - .collect::>(), - vec![5] - ); + assert_eq!(seqs(&session_events), vec![5]); drop(connection); sqlx::query("DELETE FROM run_events WHERE run_id = ? AND seq = 3") @@ -1666,7 +1608,7 @@ VALUES (?, ?, ?, ?, ?, ?, ?) .await .unwrap(); assert!(matches!( - decode_event_row(&row, &id).unwrap_err(), + decode_event_row(&row, &id, &id.to_string()).unwrap_err(), Error::RunEventMismatch { field: "run_id", .. diff --git a/lib/components/fabro-store/src/slate/mod.rs b/lib/components/fabro-store/src/slate/mod.rs index 972f5eab5..1ec399d5d 100644 --- a/lib/components/fabro-store/src/slate/mod.rs +++ b/lib/components/fabro-store/src/slate/mod.rs @@ -567,16 +567,6 @@ mod tests { (object_store, store) } - #[tokio::test] - async fn required_run_record_store_is_shared_with_run_handles() { - let (_object_store, store) = make_store(); - let records = store.run_record_store(); - let run = store.create_run(&test_run_id("run-1")).await.unwrap(); - - assert!(run.shares_run_record_store(&records)); - assert!(Arc::ptr_eq(&records, &store.clone().run_record_store())); - } - #[tokio::test] async fn retire_refresh_token_keyspace_clears_the_prefix_and_is_idempotent() { let (_object_store, store) = make_store(); diff --git a/lib/components/fabro-store/src/slate/run_store.rs b/lib/components/fabro-store/src/slate/run_store.rs index 6ab5fc1a4..226936ceb 100644 --- a/lib/components/fabro-store/src/slate/run_store.rs +++ b/lib/components/fabro-store/src/slate/run_store.rs @@ -121,11 +121,6 @@ impl RunDatabase { self.inner.run_id } - #[cfg(test)] - pub(crate) fn shares_run_record_store(&self, store: &Arc) -> bool { - Arc::ptr_eq(&self.inner.run_record_store, store) - } - pub fn subscribe(&self) -> broadcast::Receiver { self.inner.event_tx.subscribe() } diff --git a/lib/components/fabro-store/src/test_support/mod.rs b/lib/components/fabro-store/src/test_support/mod.rs index 3124f8ebe..0b9b05489 100644 --- a/lib/components/fabro-store/src/test_support/mod.rs +++ b/lib/components/fabro-store/src/test_support/mod.rs @@ -18,58 +18,48 @@ use crate::{BlobStore, Database, Result, RunRecordStore}; /// by other tests in the same process. Reopen-style tests that model one /// process-wide blob authority across several store handles should call this /// once and share the result through [`test_database_with_blobs`]. -/// -/// The pool connects lazily so synchronous fixture builders can remain -/// synchronous. Its single connection installs the production blob schema on -/// first use. #[must_use] pub fn test_blob_store() -> Arc { - let options = SqliteConnectOptions::new() - .filename(":memory:") - .foreign_keys(true); - let pool = SqlitePoolOptions::new() - .max_connections(1) - // A single in-memory test connection never needs reaping. Disabling - // both timers also keeps this lazy fixture constructible from sync - // tests, where SQLx has no Tokio runtime for maintenance tasks. - .max_lifetime(None) - .idle_timeout(None) - .after_connect(|connection, _metadata| { - Box::pin(async move { - sqlx::query(fabro_db::BLOBS_MIGRATION_SQL) - .execute(&mut *connection) - .await?; - Ok(()) - }) - }) - .connect_lazy_with(options); - Arc::new(BlobStore::new(pool)) + Arc::new(BlobStore::new(lazy_in_memory_pool(&[ + fabro_db::BLOBS_MIGRATION_SQL, + ]))) } /// Returns an isolated SQLite run-record store backed by its own in-memory /// database and the production `runs` and `run_events` schemas. #[must_use] pub fn test_run_record_store() -> Arc { + Arc::new(RunRecordStore::new(lazy_in_memory_pool(&[ + fabro_db::RUNS_MIGRATION_SQL, + fabro_db::RUN_EVENTS_MIGRATION_SQL, + ]))) +} + +/// Builds a single-connection in-memory SQLite pool that installs +/// `migrations` on first use. +/// +/// The pool connects lazily so synchronous fixture builders can remain +/// synchronous. +fn lazy_in_memory_pool(migrations: &'static [&'static str]) -> sqlx::SqlitePool { let options = SqliteConnectOptions::new() .filename(":memory:") .foreign_keys(true); - let pool = SqlitePoolOptions::new() + SqlitePoolOptions::new() .max_connections(1) + // A single in-memory test connection never needs reaping. Disabling + // both timers also keeps this lazy fixture constructible from sync + // tests, where SQLx has no Tokio runtime for maintenance tasks. .max_lifetime(None) .idle_timeout(None) - .after_connect(|connection, _metadata| { + .after_connect(move |connection, _metadata| { Box::pin(async move { - sqlx::raw_sql(fabro_db::RUNS_MIGRATION_SQL) - .execute(&mut *connection) - .await?; - sqlx::raw_sql(fabro_db::RUN_EVENTS_MIGRATION_SQL) - .execute(&mut *connection) - .await?; + for migration in migrations { + sqlx::raw_sql(*migration).execute(&mut *connection).await?; + } Ok(()) }) }) - .connect_lazy_with(options); - Arc::new(RunRecordStore::new(pool)) + .connect_lazy_with(options) } /// Returns the SQLite file backing [`test_blob_store_at`] for `store_dir`. diff --git a/lib/foundation/fabro-db/tests/sqlite.rs b/lib/foundation/fabro-db/tests/sqlite.rs index 07e6673b4..2a5c4a88c 100644 --- a/lib/foundation/fabro-db/tests/sqlite.rs +++ b/lib/foundation/fabro-db/tests/sqlite.rs @@ -554,58 +554,13 @@ async fn run_events_schema_has_final_shape_constraints_and_indexes() -> anyhow:: insert_run_with_id(database.pool(), "parent", None).await?; insert_run_with_id(database.pool(), "child", Some("parent")).await?; - insert_run_event( - database.pool(), - "parent", - 1, - "run.created", - None, - None, - None, - ) - .await?; + insert_run_event(database.pool(), "parent", 1, "run.created").await?; for invalid in [ - insert_run_event( - database.pool(), - "parent", - 1, - "run.created", - None, - None, - None, - ) - .await, - insert_run_event( - database.pool(), - "missing", - 1, - "run.created", - None, - None, - None, - ) - .await, - insert_run_event( - database.pool(), - "parent", - 0, - "run.created", - None, - None, - None, - ) - .await, - insert_run_event( - database.pool(), - "parent", - 1_000_000, - "run.created", - None, - None, - None, - ) - .await, + insert_run_event(database.pool(), "parent", 1, "run.created").await, + insert_run_event(database.pool(), "missing", 1, "run.created").await, + insert_run_event(database.pool(), "parent", 0, "run.created").await, + insert_run_event(database.pool(), "parent", 1_000_000, "run.created").await, ] { assert!(invalid.is_err()); } @@ -695,27 +650,33 @@ async fn run_events_schema_query_plans_use_candidate_indexes() -> anyhow::Result ); } - Ok(()) -} - -async fn insert_run_with_id( - pool: &fabro_db::DbPool, - id: &str, - parent_id: Option<&str>, -) -> Result<(), sqlx::Error> { - sqlx::query( - r" -INSERT INTO runs ( - id, source_last_seq, created_at_ms, last_event_at_ms, status, parent_id, title, - input_tokens, summary_json -) VALUES (?, 1, 0, 0, 'submitted', ?, 'title', 0, ?) -", + // The first-visit stage listing unions both shapes so each arm keeps its + // own partial index instead of scanning the run's primary key range. + let details = sqlx::query( + "EXPLAIN QUERY PLAN SELECT * FROM run_events WHERE run_id = ? AND seq >= ? AND stage_id = ? \ + UNION ALL SELECT * FROM run_events WHERE run_id = ? AND seq >= ? AND stage_id IS NULL AND node_id = ? \ + ORDER BY seq ASC LIMIT ?", ) - .bind(id) - .bind(parent_id) - .bind(format!(r#"{{"id":"{id}"}}"#)) - .execute(pool) - .await?; + .bind("run") + .bind(1_i64) + .bind("stage") + .bind("run") + .bind(1_i64) + .bind("node") + .bind(10_i64) + .fetch_all(database.pool()) + .await? + .into_iter() + .map(|row| row.get::("detail")) + .collect::>() + .join("; "); + for expected_index in ["run_events_by_stage", "run_events_by_legacy_node"] { + assert!( + details.contains(expected_index), + "expected {expected_index} in query plan: {details}" + ); + } + Ok(()) } @@ -724,22 +685,16 @@ async fn insert_run_event( run_id: &str, seq: i64, event_name: &str, - node_id: Option<&str>, - stage_id: Option<&str>, - session_id: Option<&str>, ) -> Result<(), sqlx::Error> { sqlx::query( r" -INSERT INTO run_events (run_id, seq, event_name, node_id, stage_id, session_id, event_json) -VALUES (?, ?, ?, ?, ?, ?, '{}') +INSERT INTO run_events (run_id, seq, event_name, event_json) +VALUES (?, ?, ?, '{}') ", ) .bind(run_id) .bind(seq) .bind(event_name) - .bind(node_id) - .bind(stage_id) - .bind(session_id) .execute(pool) .await?; Ok(()) @@ -750,17 +705,53 @@ async fn insert_minimal_run( status: &str, input_tokens: i64, summary_json: &str, +) -> Result<(), sqlx::Error> { + insert_run_row( + pool, + &format!("run-{status}-{input_tokens}"), + None, + status, + input_tokens, + summary_json, + ) + .await +} + +async fn insert_run_with_id( + pool: &fabro_db::DbPool, + id: &str, + parent_id: Option<&str>, +) -> Result<(), sqlx::Error> { + insert_run_row( + pool, + id, + parent_id, + "submitted", + 0, + &format!(r#"{{"id":"{id}"}}"#), + ) + .await +} + +async fn insert_run_row( + pool: &fabro_db::DbPool, + id: &str, + parent_id: Option<&str>, + status: &str, + input_tokens: i64, + summary_json: &str, ) -> Result<(), sqlx::Error> { sqlx::query( r" INSERT INTO runs ( - id, source_last_seq, created_at_ms, last_event_at_ms, status, title, + id, source_last_seq, created_at_ms, last_event_at_ms, status, parent_id, title, input_tokens, summary_json -) VALUES (?, 1, 0, 0, ?, 'title', ?, ?) +) VALUES (?, 1, 0, 0, ?, ?, 'title', ?, ?) ", ) - .bind(format!("run-{status}-{input_tokens}")) + .bind(id) .bind(status) + .bind(parent_id) .bind(input_tokens) .bind(summary_json) .execute(pool)