From 54dbb2828947317e0c2b1f4e4015eb04c040efd5 Mon Sep 17 00:00:00 2001 From: Bryan Helmkamp Date: Tue, 7 Apr 2026 18:16:29 -0400 Subject: [PATCH] fix(store): make run export use server-backed blob reads Route `fabro store dump` through the server client for run state, events, blob hydration, and artifact downloads instead of reopening storage directly from the CLI process. This fixes blob-backed checkpoint exports and restores store-dump coverage under the in-memory test server. --- .../fabro-cli/src/commands/store/dump.rs | 118 ++++++-- .../fabro-cli/tests/it/cmd/store_dump.rs | 259 ++++++++++++------ lib/crates/fabro-workflow/src/run_dump.rs | 224 +++++++++------ 3 files changed, 414 insertions(+), 187 deletions(-) diff --git a/lib/crates/fabro-cli/src/commands/store/dump.rs b/lib/crates/fabro-cli/src/commands/store/dump.rs index a5f49bbfa..897c6da69 100644 --- a/lib/crates/fabro-cli/src/commands/store/dump.rs +++ b/lib/crates/fabro-cli/src/commands/store/dump.rs @@ -1,35 +1,29 @@ use anyhow::{Context, Result}; -use fabro_config::Storage; +use bytes::Bytes; #[cfg(test)] -use fabro_store::StageId; -use fabro_store::{ArtifactStore, RunDatabase, RunProjection}; +use fabro_store::{ArtifactStore, RunDatabase}; +use fabro_store::{RunProjection, StageId}; +use fabro_types::RunBlobId; use fabro_workflow::run_dump::RunDump; +use futures::future::BoxFuture; #[cfg(test)] use serde::de::DeserializeOwned; use std::io::ErrorKind; use std::path::Path; -use std::sync::Arc; use crate::args::{GlobalArgs, StoreDumpArgs}; -use crate::commands::store::rebuild::rebuild_run_store; +use crate::server_client::ServerStoreClient; use crate::server_runs::ServerRunLookup; use crate::shared::{absolute_or_current, print_json_pretty}; use crate::user_config::load_settings_with_storage_dir; -use object_store::{ObjectStore, local::LocalFileSystem}; pub(crate) async fn dump_command(args: &StoreDumpArgs, globals: &GlobalArgs) -> Result<()> { let cli_settings = load_settings_with_storage_dir(args.storage_dir.as_deref())?; let lookup = ServerRunLookup::connect(&cli_settings.storage_dir()).await?; let run = lookup.resolve(&args.run)?; let run_id = run.run_id(); - let events = lookup.client().list_run_events(&run_id, None, None).await?; - let run_store = rebuild_run_store(&run_id, &events).await?; - let artifact_object_store: Arc = Arc::new(LocalFileSystem::new_with_prefix( - Storage::new(cli_settings.storage_dir()).store_dir(), - )?); - let artifact_store = ArtifactStore::new(artifact_object_store, "artifacts"); - - let file_count = export_run(&run_store, &artifact_store, &args.output).await?; + let state = lookup.client().get_run_state(&run_id).await?; + let file_count = export_run_from_server(lookup.client(), &run_id, &state, &args.output).await?; if globals.json { print_json_pretty(&serde_json::json!({ "run_id": run_id, @@ -46,6 +40,7 @@ pub(crate) async fn dump_command(args: &StoreDumpArgs, globals: &GlobalArgs) -> Ok(()) } +#[cfg(test)] pub(crate) async fn export_run( run_store: &RunDatabase, artifact_store: &ArtifactStore, @@ -71,12 +66,59 @@ pub(crate) async fn export_run( let staging_path = staging_dir.path().to_path_buf(); let file_count = export_run_to_dir(run_store, artifact_store, &state, &staging_path).await?; + finalize_export( + output_dir, + output_state, + staging_dir, + &staging_path, + file_count, + ) +} +async fn export_run_from_server( + client: &ServerStoreClient, + run_id: &fabro_types::RunId, + state: &RunProjection, + output_dir: &Path, +) -> Result { + let output_state = inspect_output_dir(output_dir)?; + let staging_parent = output_parent_dir(output_dir); + std::fs::create_dir_all(staging_parent) + .with_context(|| format!("failed to create {}", staging_parent.display()))?; + + let staging_dir = tempfile::Builder::new() + .prefix(".fabro-store-dump-") + .tempdir_in(staging_parent) + .with_context(|| { + format!( + "failed to create staging dir in {}", + staging_parent.display() + ) + })?; + let staging_path = staging_dir.path().to_path_buf(); + + let file_count = export_server_run_to_dir(client, run_id, state, &staging_path).await?; + finalize_export( + output_dir, + output_state, + staging_dir, + &staging_path, + file_count, + ) +} + +fn finalize_export( + output_dir: &Path, + output_state: OutputDirState, + staging_dir: tempfile::TempDir, + staging_path: &Path, + file_count: usize, +) -> Result { if matches!(output_state, OutputDirState::ExistingEmpty) { std::fs::remove_dir(output_dir) .with_context(|| format!("failed to replace {}", output_dir.display()))?; } - std::fs::rename(&staging_path, output_dir).with_context(|| { + std::fs::rename(staging_path, output_dir).with_context(|| { format!( "failed to move staged export {} into {}", staging_path.display(), @@ -88,6 +130,7 @@ pub(crate) async fn export_run( Ok(file_count) } +#[cfg(test)] async fn export_run_to_dir( run_store: &RunDatabase, artifact_store: &ArtifactStore, @@ -98,6 +141,48 @@ async fn export_run_to_dir( dump.write_to_dir(output_dir) } +async fn export_server_run_to_dir( + client: &ServerStoreClient, + run_id: &fabro_types::RunId, + state: &RunProjection, + output_dir: &Path, +) -> Result { + let events = client.list_run_events(run_id, None, None).await?; + let mut dump = RunDump::from_store_state_and_events(state, &events)?; + + dump.hydrate_referenced_blobs_with_reader(|blob_id| { + read_blob_from_client(client, run_id, blob_id) + }) + .await?; + + for artifact in client.list_run_artifacts(run_id).await? { + let stage_id: StageId = artifact + .stage_id + .parse() + .with_context(|| format!("server returned invalid stage id {:?}", artifact.stage_id))?; + let data = client + .download_stage_artifact(run_id, &stage_id, &artifact.relative_path) + .await + .with_context(|| { + format!( + "failed to download artifact {} for stage {}", + artifact.relative_path, artifact.stage_id + ) + })?; + dump.add_artifact_bytes(&stage_id, &artifact.relative_path, data)?; + } + + dump.write_to_dir(output_dir) +} + +fn read_blob_from_client<'a>( + client: &'a ServerStoreClient, + run_id: &'a fabro_types::RunId, + blob_id: RunBlobId, +) -> BoxFuture<'a, Result>> { + Box::pin(async move { client.read_run_blob(run_id, &blob_id).await }) +} + #[derive(Debug, Clone, Copy, PartialEq, Eq)] enum OutputDirState { Missing, @@ -144,6 +229,7 @@ mod tests { use std::collections::HashMap; use std::path::PathBuf; + use std::sync::Arc; use std::time::Duration; use chrono::{DateTime, Utc}; @@ -154,7 +240,7 @@ mod tests { Settings, StageStatus, StartRecord, StatusReason, fixtures, }; use fabro_workflow::event::{Event, append_event}; - use object_store::memory::InMemory; + use object_store::{ObjectStore, memory::InMemory}; fn dt(rfc3339: &str) -> DateTime { DateTime::parse_from_rfc3339(rfc3339) diff --git a/lib/crates/fabro-cli/tests/it/cmd/store_dump.rs b/lib/crates/fabro-cli/tests/it/cmd/store_dump.rs index 44bc68deb..0a668a659 100644 --- a/lib/crates/fabro-cli/tests/it/cmd/store_dump.rs +++ b/lib/crates/fabro-cli/tests/it/cmd/store_dump.rs @@ -1,3 +1,9 @@ +use super::support::setup_completed_dry_run; +use insta::assert_snapshot; +use std::fs; +use std::time::Duration; + +use crate::support::unique_run_id; use fabro_test::{fabro_snapshot, test_context}; #[test] @@ -29,92 +35,167 @@ fn help() { "); } -// TODO: Re-enable once `store dump` funnels all SlateDB/artifact access through -// the server's live store handles instead of opening `storage/store` directly -// from the CLI process. -// -// #[test] -// fn store_dump_exports_completed_run_snapshot() { -// let context = test_context!(); -// let run = setup_completed_dry_run(&context); -// let output_dir = context.temp_dir.join("export"); -// -// let mut cmd = context.command(); -// cmd.args([ -// "store", -// "dump", -// "--output", -// output_dir.to_str().unwrap(), -// &run.run_id, -// ]); -// fabro_snapshot!(context.filters(), cmd, @" -// success: true -// exit_code: 0 -// ----- stdout ----- -// Exported 17 files for run [ULID] to [TEMP_DIR]/export -// ----- stderr ----- -// "); -// -// assert_snapshot!(dump_file_summary(&output_dir), @" -// checkpoint.json -// checkpoints/0012.json -// checkpoints/0016.json -// checkpoints/0020.json -// conclusion.json -// events.jsonl -// graph.fabro -// nodes/exit/visit-1/status.json -// nodes/report/visit-1/response.md -// nodes/report/visit-1/status.json -// nodes/run_tests/visit-1/response.md -// nodes/run_tests/visit-1/status.json -// nodes/start/visit-1/status.json -// run.json -// sandbox.json -// start.json -// status.json -// "); -// } -// -// #[test] -// fn store_dump_rejects_non_empty_output_dir() { -// let context = test_context!(); -// let run = setup_completed_dry_run(&context); -// let output_dir = context.temp_dir.join("nonempty"); -// std::fs::create_dir_all(&output_dir).unwrap(); -// std::fs::write(output_dir.join("file.txt"), "x").unwrap(); -// -// let mut cmd = context.command(); -// cmd.args([ -// "store", -// "dump", -// "--output", -// output_dir.to_str().unwrap(), -// &run.run_id, -// ]); -// fabro_snapshot!(context.filters(), cmd, @" -// success: false -// exit_code: 1 -// ----- stdout ----- -// ----- stderr ----- -// error: output path [TEMP_DIR]/nonempty already exists and is not an empty directory; remove it first or choose a different path -// "); -// } -// -// fn dump_file_summary(output_dir: &std::path::Path) -> String { -// let mut files: Vec = walkdir::WalkDir::new(output_dir) -// .into_iter() -// .filter_map(Result::ok) -// .filter(|entry| entry.file_type().is_file()) -// .map(|entry| { -// entry -// .path() -// .strip_prefix(output_dir) -// .unwrap() -// .to_string_lossy() -// .replace('\\', "/") -// }) -// .collect(); -// files.sort(); -// files.join("\n") + "\n" -// } +#[test] +fn store_dump_exports_large_command_output_backed_by_blob_refs() { + let context = test_context!(); + let workflow = context.temp_dir.join("large-output.fabro"); + fs::write( + &workflow, + r#"digraph LargeOutput { + graph [goal="Generate oversized command output"] + rankdir=LR + + start [shape=Mdiamond, label="Start"] + exit [shape=Msquare, label="Exit"] + big [shape=parallelogram, label="Big", script="printf '%*s' 120000 '' | tr ' ' x"] + + start -> big -> exit +} +"#, + ) + .unwrap(); + + let run_id = unique_run_id(); + let mut run_cmd = context.run_cmd(); + run_cmd.current_dir(&context.temp_dir); + run_cmd.timeout(Duration::from_secs(30)); + run_cmd.args([ + "--run-id", + run_id.as_str(), + "--no-retro", + "--sandbox", + "local", + ]); + run_cmd.arg(&workflow); + let run_output = run_cmd.output().expect("command should execute"); + assert!( + run_output.status.success(), + "workflow run failed\nstdout:\n{}\nstderr:\n{}", + String::from_utf8_lossy(&run_output.stdout), + String::from_utf8_lossy(&run_output.stderr) + ); + + let mut inspect_cmd = context.command(); + inspect_cmd.args(["inspect", "--json", &run_id]); + let inspect_output = inspect_cmd.output().expect("inspect should execute"); + assert!( + inspect_output.status.success(), + "inspect failed\nstdout:\n{}\nstderr:\n{}", + String::from_utf8_lossy(&inspect_output.stdout), + String::from_utf8_lossy(&inspect_output.stderr) + ); + let inspect_json = String::from_utf8(inspect_output.stdout).unwrap(); + assert!( + inspect_json.contains("blob://sha256/"), + "inspect output should contain blob refs to exercise hydration\n{inspect_json}" + ); + + let output_dir = context.temp_dir.join("export"); + let mut dump_cmd = context.command(); + dump_cmd.args([ + "store", + "dump", + "--output", + output_dir.to_str().unwrap(), + &run_id, + ]); + let dump_output = dump_cmd.output().expect("store dump should execute"); + assert!( + dump_output.status.success(), + "store dump failed\nstdout:\n{}\nstderr:\n{}", + String::from_utf8_lossy(&dump_output.stdout), + String::from_utf8_lossy(&dump_output.stderr) + ); + + let checkpoint = fs::read_to_string(output_dir.join("checkpoint.json")).unwrap(); + assert!( + !checkpoint.contains("blob://sha256/"), + "checkpoint export should hydrate blob refs\n{checkpoint}" + ); +} + +#[test] +fn store_dump_exports_completed_run_snapshot() { + let context = test_context!(); + let run = setup_completed_dry_run(&context); + let output_dir = context.temp_dir.join("export"); + + let mut cmd = context.command(); + cmd.args([ + "store", + "dump", + "--output", + output_dir.to_str().unwrap(), + &run.run_id, + ]); + fabro_snapshot!(context.filters(), cmd, @" + success: true + exit_code: 0 + ----- stdout ----- + Exported 17 files for run [ULID] to [TEMP_DIR]/export + ----- stderr ----- + "); + + assert_snapshot!(dump_file_summary(&output_dir), @" + checkpoint.json + checkpoints/0012.json + checkpoints/0016.json + checkpoints/0020.json + conclusion.json + events.jsonl + graph.fabro + nodes/exit/visit-1/status.json + nodes/report/visit-1/response.md + nodes/report/visit-1/status.json + nodes/run_tests/visit-1/response.md + nodes/run_tests/visit-1/status.json + nodes/start/visit-1/status.json + run.json + sandbox.json + start.json + status.json + "); +} + +#[test] +fn store_dump_rejects_non_empty_output_dir() { + let context = test_context!(); + let run = setup_completed_dry_run(&context); + let output_dir = context.temp_dir.join("nonempty"); + std::fs::create_dir_all(&output_dir).unwrap(); + std::fs::write(output_dir.join("file.txt"), "x").unwrap(); + + let mut cmd = context.command(); + cmd.args([ + "store", + "dump", + "--output", + output_dir.to_str().unwrap(), + &run.run_id, + ]); + fabro_snapshot!(context.filters(), cmd, @" + success: false + exit_code: 1 + ----- stdout ----- + ----- stderr ----- + error: output path [TEMP_DIR]/nonempty already exists and is not an empty directory; remove it first or choose a different path + "); +} + +fn dump_file_summary(output_dir: &std::path::Path) -> String { + let mut files: Vec = walkdir::WalkDir::new(output_dir) + .into_iter() + .filter_map(Result::ok) + .filter(|entry| entry.file_type().is_file()) + .map(|entry| { + entry + .path() + .strip_prefix(output_dir) + .unwrap() + .to_string_lossy() + .replace('\\', "/") + }) + .collect(); + files.sort(); + files.join("\n") + "\n" +} diff --git a/lib/crates/fabro-workflow/src/run_dump.rs b/lib/crates/fabro-workflow/src/run_dump.rs index 9c851ed9a..a4b07150b 100644 --- a/lib/crates/fabro-workflow/src/run_dump.rs +++ b/lib/crates/fabro-workflow/src/run_dump.rs @@ -3,7 +3,8 @@ use std::io::Write; use std::path::{Component, Path, PathBuf}; use anyhow::{Context, Result, bail}; -use fabro_store::{ArtifactStore, RunDatabase, RunProjection}; +use bytes::Bytes; +use fabro_store::{ArtifactStore, EventEnvelope, RunDatabase, RunProjection, StageId}; use fabro_types::{RunBlobId, parse_blob_ref, parse_legacy_blob_file_ref}; use futures::future::BoxFuture; @@ -119,15 +120,47 @@ impl RunDump { dump } + #[allow(dead_code)] pub async fn store_export( run_store: &RunDatabase, artifact_store: &ArtifactStore, state: &RunProjection, + ) -> Result { + let events = run_store.list_events().await?; + let mut dump = Self::from_store_state_and_events(state, &events)?; + + if let Some(run_record) = state.run.as_ref() { + for asset in artifact_store.list_for_run(&run_record.run_id).await? { + let data = artifact_store + .get(&run_record.run_id, &asset.node, &asset.filename) + .await? + .with_context(|| { + format!( + "asset {:?} for node {:?} visit {} is missing from the store", + asset.filename, + asset.node.node_id(), + asset.node.visit() + ) + })?; + dump.add_artifact_bytes(&asset.node, &asset.filename, data.to_vec())?; + } + } + + dump.hydrate_referenced_blobs_with_reader(|blob_id| { + read_blob_from_store(run_store, blob_id) + }) + .await?; + + Ok(dump) + } + + pub fn from_store_state_and_events( + state: &RunProjection, + events: &[EventEnvelope], ) -> Result { let mut entries = Vec::new(); - let run_record = state.run.as_ref(); - if let Some(record) = run_record { + if let Some(record) = state.run.as_ref() { push_json_entry(&mut entries, "run.json", record); } if let Some(record) = state.start.as_ref() { @@ -200,8 +233,8 @@ impl RunDump { } let mut events_jsonl = Vec::new(); - for event in run_store.list_events().await? { - serde_json::to_writer(&mut events_jsonl, &event)?; + for event in events { + serde_json::to_writer(&mut events_jsonl, event)?; events_jsonl.write_all(b"\n")?; } entries.push(RunDumpEntry::bytes("events.jsonl", events_jsonl)); @@ -214,36 +247,47 @@ impl RunDump { ); } - if let Some(run_record) = run_record { - for asset in artifact_store.list_for_run(&run_record.run_id).await? { - let node_id_segment = - validate_single_path_segment("node id", asset.node.node_id())?; - let filename_path = validate_relative_path("artifact filename", &asset.filename)?; - let data = artifact_store - .get(&run_record.run_id, &asset.node, &asset.filename) - .await? - .with_context(|| { - format!( - "asset {:?} for node {:?} visit {} is missing from the store", - asset.filename, - asset.node.node_id(), - asset.node.visit() - ) - })?; - entries.push(RunDumpEntry::bytes_path( - &PathBuf::from("artifacts") - .join("nodes") - .join(node_id_segment) - .join(format!("visit-{}", asset.node.visit())) - .join(filename_path), - data.to_vec(), - )); + Ok(Self { entries }) + } + + pub fn add_artifact_bytes( + &mut self, + stage_id: &StageId, + filename: &str, + data: Vec, + ) -> Result<()> { + let path = artifact_dump_path(stage_id, filename)?; + self.entries.push(RunDumpEntry::bytes_path(&path, data)); + Ok(()) + } + + pub async fn hydrate_referenced_blobs_with_reader<'a, F>( + &mut self, + mut read_blob: F, + ) -> Result<()> + where + F: FnMut(RunBlobId) -> BoxFuture<'a, Result>>, + { + let mut cache = HashMap::new(); + for entry in &mut self.entries { + if let RunDumpContents::Json(value) = &mut entry.contents { + let mut blob_ids = Vec::new(); + collect_blob_refs_in_value(value, &mut blob_ids); + for blob_id in blob_ids { + if cache.contains_key(&blob_id) { + continue; + } + let blob = read_blob(blob_id) + .await? + .with_context(|| format!("blob {blob_id:?} is missing from the store"))?; + let hydrated: serde_json::Value = serde_json::from_slice(&blob) + .with_context(|| format!("blob {blob_id:?} is not valid JSON"))?; + cache.insert(blob_id, hydrated); + } + replace_blob_refs_in_value(value, &cache)?; } } - - hydrate_referenced_blobs(&mut entries, run_store).await?; - - Ok(Self { entries }) + Ok(()) } pub fn entries(&self) -> &[RunDumpEntry] { @@ -403,61 +447,77 @@ fn validate_relative_path(kind: &str, value: &str) -> Result { Ok(normalized) } -async fn hydrate_referenced_blobs( - entries: &mut [RunDumpEntry], - run_store: &RunDatabase, -) -> Result<()> { - let mut cache = HashMap::new(); - for entry in entries { - if let RunDumpContents::Json(value) = &mut entry.contents { - hydrate_blob_refs_in_value(value, run_store, &mut cache).await?; +fn collect_blob_refs_in_value(value: &serde_json::Value, blob_ids: &mut Vec) { + match value { + serde_json::Value::String(current) => { + if let Some(blob_id) = + parse_blob_ref(current).or_else(|| parse_legacy_blob_file_ref(current)) + { + blob_ids.push(blob_id); + } } + serde_json::Value::Array(items) => { + for item in items { + collect_blob_refs_in_value(item, blob_ids); + } + } + serde_json::Value::Object(map) => { + for item in map.values() { + collect_blob_refs_in_value(item, blob_ids); + } + } + serde_json::Value::Null | serde_json::Value::Bool(_) | serde_json::Value::Number(_) => {} + } +} + +fn replace_blob_refs_in_value( + value: &mut serde_json::Value, + cache: &HashMap, +) -> Result<()> { + match value { + serde_json::Value::String(current) => { + let Some(blob_id) = + parse_blob_ref(current).or_else(|| parse_legacy_blob_file_ref(current)) + else { + return Ok(()); + }; + let hydrated = cache + .get(&blob_id) + .cloned() + .with_context(|| format!("blob {blob_id:?} is missing from the hydration cache"))?; + *value = hydrated; + } + serde_json::Value::Array(items) => { + for item in items { + replace_blob_refs_in_value(item, cache)?; + } + } + serde_json::Value::Object(map) => { + for item in map.values_mut() { + replace_blob_refs_in_value(item, cache)?; + } + } + serde_json::Value::Null | serde_json::Value::Bool(_) | serde_json::Value::Number(_) => {} } Ok(()) } -fn hydrate_blob_refs_in_value<'a>( - value: &'a mut serde_json::Value, - run_store: &'a RunDatabase, - cache: &'a mut HashMap, -) -> BoxFuture<'a, Result<()>> { - Box::pin(async move { - match value { - serde_json::Value::String(current) => { - let Some(blob_id) = - parse_blob_ref(current).or_else(|| parse_legacy_blob_file_ref(current)) - else { - return Ok(()); - }; - if let Some(cached) = cache.get(&blob_id).cloned() { - *value = cached; - return Ok(()); - } +#[allow(dead_code)] +fn read_blob_from_store( + run_store: &RunDatabase, + blob_id: RunBlobId, +) -> BoxFuture<'_, Result>> { + Box::pin(async move { Ok(run_store.read_blob(&blob_id).await?) }) +} - let blob = run_store - .read_blob(&blob_id) - .await? - .with_context(|| format!("blob {blob_id:?} is missing from the store"))?; - let hydrated: serde_json::Value = serde_json::from_slice(&blob) - .with_context(|| format!("blob {blob_id:?} is not valid JSON"))?; - cache.insert(blob_id, hydrated.clone()); - *value = hydrated; - } - serde_json::Value::Array(items) => { - for item in items { - hydrate_blob_refs_in_value(item, run_store, cache).await?; - } - } - serde_json::Value::Object(map) => { - for item in map.values_mut() { - hydrate_blob_refs_in_value(item, run_store, cache).await?; - } - } - serde_json::Value::Null | serde_json::Value::Bool(_) | serde_json::Value::Number(_) => { - } - } - Ok(()) - }) +fn artifact_dump_path(stage_id: &StageId, filename: &str) -> Result { + let node_id_segment = validate_single_path_segment("node id", stage_id.node_id())?; + let filename_path = validate_relative_path("artifact filename", filename)?; + Ok(PathBuf::from("artifacts") + .join("nodes") + .join(node_id_segment) + .join(format!("visit-{}", stage_id.visit())) + .join(filename_path)) } fn ensure_parent_dir(path: &Path) -> Result<()> {