diff --git a/lib/crates/fabro-cli/src/commands/store/dump.rs b/lib/crates/fabro-cli/src/commands/store/dump.rs index 897c6da69..7c19b0ca9 100644 --- a/lib/crates/fabro-cli/src/commands/store/dump.rs +++ b/lib/crates/fabro-cli/src/commands/store/dump.rs @@ -2,8 +2,8 @@ use anyhow::{Context, Result}; use bytes::Bytes; #[cfg(test)] use fabro_store::{ArtifactStore, RunDatabase}; -use fabro_store::{RunProjection, StageId}; -use fabro_types::RunBlobId; +use fabro_store::{EventEnvelope, RunProjection, StageId}; +use fabro_types::{RunBlobId, RunId}; use fabro_workflow::run_dump::RunDump; use futures::future::BoxFuture; #[cfg(test)] @@ -23,7 +23,8 @@ pub(crate) async fn dump_command(args: &StoreDumpArgs, globals: &GlobalArgs) -> let run = lookup.resolve(&args.run)?; let run_id = run.run_id(); 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?; + let source = ServerDumpSource::new(lookup.client(), &run_id); + let file_count = export_run_from_source(&source, &state, &args.output).await?; if globals.json { print_json_pretty(&serde_json::json!({ "run_id": run_id, @@ -47,64 +48,13 @@ pub(crate) async fn export_run( output_dir: &Path, ) -> Result { let state = run_store.state().await?; - anyhow::ensure!(state.run.is_some(), "run has no data in the store"); - - 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_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, - ) + let run_id = state + .run + .as_ref() + .map(|run| run.run_id) + .context("run has no data in the store")?; + let source = LocalDumpSource::new(run_store, artifact_store, run_id); + export_run_from_source(&source, &state, output_dir).await } fn finalize_export( @@ -130,59 +80,172 @@ fn finalize_export( Ok(file_count) } -#[cfg(test)] -async fn export_run_to_dir( - run_store: &RunDatabase, - artifact_store: &ArtifactStore, - state: &RunProjection, - output_dir: &Path, -) -> Result { - let dump = RunDump::store_export(run_store, artifact_store, state).await?; - dump.write_to_dir(output_dir) +struct DumpArtifact { + stage_id: StageId, + relative_path: String, + data: Vec, } -async fn export_server_run_to_dir( - client: &ServerStoreClient, - run_id: &fabro_types::RunId, +trait DumpDataSource { + fn list_events(&self) -> BoxFuture<'_, Result>>; + + fn read_blob(&self, blob_id: RunBlobId) -> BoxFuture<'_, Result>>; + + fn list_artifacts(&self) -> BoxFuture<'_, Result>>; +} + +#[cfg(test)] +struct LocalDumpSource<'a> { + run_store: &'a RunDatabase, + artifact_store: &'a ArtifactStore, + run_id: RunId, +} + +#[cfg(test)] +impl<'a> LocalDumpSource<'a> { + fn new(run_store: &'a RunDatabase, artifact_store: &'a ArtifactStore, run_id: RunId) -> Self { + Self { + run_store, + artifact_store, + run_id, + } + } +} + +#[cfg(test)] +impl DumpDataSource for LocalDumpSource<'_> { + fn list_events(&self) -> BoxFuture<'_, Result>> { + Box::pin(async move { Ok(self.run_store.list_events().await?) }) + } + + fn read_blob(&self, blob_id: RunBlobId) -> BoxFuture<'_, Result>> { + Box::pin(async move { Ok(self.run_store.read_blob(&blob_id).await?) }) + } + + fn list_artifacts(&self) -> BoxFuture<'_, Result>> { + Box::pin(async move { + let mut artifacts = Vec::new(); + for asset in self.artifact_store.list_for_run(&self.run_id).await? { + let data = self + .artifact_store + .get(&self.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() + ) + })?; + artifacts.push(DumpArtifact { + stage_id: asset.node, + relative_path: asset.filename, + data: data.to_vec(), + }); + } + Ok(artifacts) + }) + } +} + +struct ServerDumpSource<'a> { + client: &'a ServerStoreClient, + run_id: &'a RunId, +} + +impl<'a> ServerDumpSource<'a> { + fn new(client: &'a ServerStoreClient, run_id: &'a RunId) -> Self { + Self { client, run_id } + } +} + +impl DumpDataSource for ServerDumpSource<'_> { + fn list_events(&self) -> BoxFuture<'_, Result>> { + Box::pin(async move { self.client.list_run_events(self.run_id, None, None).await }) + } + + fn read_blob(&self, blob_id: RunBlobId) -> BoxFuture<'_, Result>> { + Box::pin(async move { self.client.read_run_blob(self.run_id, &blob_id).await }) + } + + fn list_artifacts(&self) -> BoxFuture<'_, Result>> { + Box::pin(async move { + let mut artifacts = Vec::new(); + for artifact in self.client.list_run_artifacts(self.run_id).await? { + let stage_id: StageId = artifact.stage_id.parse().with_context(|| { + format!("server returned invalid stage id {:?}", artifact.stage_id) + })?; + let data = self + .client + .download_stage_artifact(self.run_id, &stage_id, &artifact.relative_path) + .await + .with_context(|| { + format!( + "failed to download artifact {} for stage {}", + artifact.relative_path, artifact.stage_id + ) + })?; + artifacts.push(DumpArtifact { + stage_id, + relative_path: artifact.relative_path, + data, + }); + } + Ok(artifacts) + }) + } +} + +async fn export_run_from_source( + source: &impl DumpDataSource, state: &RunProjection, output_dir: &Path, ) -> Result { - let events = client.list_run_events(run_id, None, None).await?; + 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 = write_run_dump(source, state, &staging_path).await?; + finalize_export( + output_dir, + output_state, + staging_dir, + &staging_path, + file_count, + ) +} + +async fn write_run_dump( + source: &impl DumpDataSource, + state: &RunProjection, + output_dir: &Path, +) -> Result { + let events = source.list_events().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?; + dump.hydrate_referenced_blobs_with_reader(|blob_id| source.read_blob(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)?; + for artifact in source.list_artifacts().await? { + dump.add_artifact_bytes(&artifact.stage_id, &artifact.relative_path, artifact.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, 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 0a668a659..db88197d3 100644 --- a/lib/crates/fabro-cli/tests/it/cmd/store_dump.rs +++ b/lib/crates/fabro-cli/tests/it/cmd/store_dump.rs @@ -114,6 +114,110 @@ fn store_dump_exports_large_command_output_backed_by_blob_refs() { ); } +#[test] +fn store_dump_exports_blob_refs_and_artifacts_together() { + let context = test_context!(); + let workspace_dir = context.temp_dir.join("mixed-export"); + fs::create_dir_all(&workspace_dir).unwrap(); + + fs::write( + workspace_dir.join("mixed-export.fabro"), + r#"digraph MixedExport { + graph [goal="Generate oversized command output and artifacts"] + rankdir=LR + + start [shape=Mdiamond, label="Start"] + exit [shape=Msquare, label="Exit"] + big [shape=parallelogram, label="Big", script="mkdir -p assets/shared && printf exported > assets/shared/report.txt && printf '%*s' 120000 '' | tr ' ' x"] + + start -> big -> exit +} +"#, + ) + .unwrap(); + fs::write( + workspace_dir.join("run.toml"), + r#"version = 1 +graph = "mixed-export.fabro" +goal = "Generate oversized command output and artifacts" + +[sandbox] +provider = "local" +preserve = true + +[sandbox.local] +worktree_mode = "never" + +[artifacts] +include = ["assets/**"] +"#, + ) + .unwrap(); + + let run_id = unique_run_id(); + let mut run_cmd = context.run_cmd(); + run_cmd.current_dir(&workspace_dir); + run_cmd.timeout(Duration::from_secs(30)); + run_cmd.args([ + "--run-id", + run_id.as_str(), + "--no-retro", + "--sandbox", + "local", + "run.toml", + ]); + 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-mixed"); + 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}" + ); + assert_eq!( + fs::read_to_string(output_dir.join("artifacts/nodes/big/visit-1/assets/shared/report.txt")) + .unwrap(), + "exported" + ); +} + #[test] fn store_dump_exports_completed_run_snapshot() { let context = test_context!(); diff --git a/lib/crates/fabro-workflow/src/run_dump.rs b/lib/crates/fabro-workflow/src/run_dump.rs index a4b07150b..ec8f573c2 100644 --- a/lib/crates/fabro-workflow/src/run_dump.rs +++ b/lib/crates/fabro-workflow/src/run_dump.rs @@ -4,7 +4,7 @@ use std::path::{Component, Path, PathBuf}; use anyhow::{Context, Result, bail}; use bytes::Bytes; -use fabro_store::{ArtifactStore, EventEnvelope, RunDatabase, RunProjection, StageId}; +use fabro_store::{EventEnvelope, RunProjection, StageId}; use fabro_types::{RunBlobId, parse_blob_ref, parse_legacy_blob_file_ref}; use futures::future::BoxFuture; @@ -120,40 +120,6 @@ 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], @@ -502,14 +468,6 @@ fn replace_blob_refs_in_value( 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?) }) -} - 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)?;