mirror of
https://github.com/fabro-sh/fabro.git
synced 2026-09-15 23:32:46 +00:00
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.
This commit is contained in:
Bryan Helmkamp
parent
668d7857e6
commit
54dbb28289
3 changed files with 414 additions and 187 deletions
|
|
@ -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<dyn ObjectStore> = 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<usize> {
|
||||
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<usize> {
|
||||
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<usize> {
|
||||
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<Option<Bytes>>> {
|
||||
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<Utc> {
|
||||
DateTime::parse_from_rfc3339(rfc3339)
|
||||
|
|
|
|||
|
|
@ -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<String> = 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<String> = 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"
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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<Self> {
|
||||
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<Self> {
|
||||
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<u8>,
|
||||
) -> 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<Option<Bytes>>>,
|
||||
{
|
||||
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<PathBuf> {
|
|||
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<RunBlobId>) {
|
||||
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<RunBlobId, serde_json::Value>,
|
||||
) -> 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<RunBlobId, serde_json::Value>,
|
||||
) -> 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<Option<Bytes>>> {
|
||||
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<PathBuf> {
|
||||
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<()> {
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue