From 169e31a415fc8c0da4c88cf9e2a07460d01111bf Mon Sep 17 00:00:00 2001 From: Reiase Date: Tue, 8 Sep 2026 22:35:27 +0800 Subject: [PATCH 1/8] fix(pchronicle): allow non-loopback serve and lazy Directory discovery Stop recursive object-store listing that hit max_files on large prefixes, classify Directory vs Dataset from shallow markers, and let catalog serve bind/listen before discovery finishes. Co-authored-by: Cursor --- crates/persisting-pchronicle-cli/README.md | 4 +- .../persisting-pchronicle-cli/src/control.rs | 4 - .../src/gateway_ingest.rs | 4 - crates/persisting-pchronicle-cli/src/lib.rs | 25 +- .../persisting-pchronicle-cli/src/onboard.rs | 3 +- .../src/server/catalog.rs | 13 +- .../src/server/mod.rs | 42 +- .../src/server/tests.rs | 23 +- crates/persisting-pchronicle-cli/src/tests.rs | 7 +- .../src/store/catalog/discovery.rs | 386 +++++++++++------- .../src/store/opendal_store.rs | 34 ++ .../rfcs/0013-pchronicle-warehouse-catalog.md | 13 +- docs/src/en/rfcs/0015-chronicle-manifest.md | 17 +- .../rfcs/0013-pchronicle-warehouse-catalog.md | 13 +- docs/src/zh/rfcs/0015-chronicle-manifest.md | 10 +- 15 files changed, 366 insertions(+), 232 deletions(-) diff --git a/crates/persisting-pchronicle-cli/README.md b/crates/persisting-pchronicle-cli/README.md index 603d96a4..f57d64f0 100644 --- a/crates/persisting-pchronicle-cli/README.md +++ b/crates/persisting-pchronicle-cli/README.md @@ -3,7 +3,7 @@ **Standalone `pchronicle` CLI for onboarding, browsing, querying, importing, exporting, and serving trajectory Datasets.** -Owns the `pchronicle` binary, loopback-only Warehouse HTTP, the write-capable +Owns the `pchronicle` binary, Warehouse HTTP, the write-capable `--control` plane used by pPilot and pVisor, optional Gateway ingest/forwarding flags, and the embed of staged `pchronicle-web` assets at build time. @@ -17,7 +17,7 @@ Current commands include `onboard`, `dataset` (pin/unpin/list/show/set/rename), `list`/`ls`, `stats`, bounded read-only `query`, built-in `stats` reports, assisted `agent` sessions, Source-local `find`, create/append/replace `import`, destructive `drop`, complete-trajectory `export`, directory `sync`, `echo`, and -loopback-only `serve`. Import and export support ATIF, OpenAI Messages, ACTF, +`serve`. Import and export support ATIF, OpenAI Messages, ACTF, Storyline JSON, and record-level Compact JSONL. `sync --from SOURCE --to WAREHOUSE --convert OUTPUT` polls a local source directory, atomically mirrors supported JSON files into a local Warehouse Dataset byte-for-byte, and rebuilds diff --git a/crates/persisting-pchronicle-cli/src/control.rs b/crates/persisting-pchronicle-cli/src/control.rs index d82995e3..7967bd8f 100644 --- a/crates/persisting-pchronicle-cli/src/control.rs +++ b/crates/persisting-pchronicle-cli/src/control.rs @@ -32,10 +32,6 @@ pub(super) struct PreparedControl { impl PreparedControl { pub(super) async fn bind(storage: &str, listen: SocketAddr) -> Result { - anyhow::ensure!( - listen.ip().is_loopback(), - "pChronicle control may only bind to a loopback address" - ); let control = Arc::new( RunControlStore::open(storage) .await diff --git a/crates/persisting-pchronicle-cli/src/gateway_ingest.rs b/crates/persisting-pchronicle-cli/src/gateway_ingest.rs index a5dc9eb1..76e725fb 100644 --- a/crates/persisting-pchronicle-cli/src/gateway_ingest.rs +++ b/crates/persisting-pchronicle-cli/src/gateway_ingest.rs @@ -67,10 +67,6 @@ impl PreparedIngestGateway { split: Option, manifest_write_mode: ObjectStoreManifestWriteMode, ) -> Result { - anyhow::ensure!( - listen.ip().is_loopback(), - "pChronicle ingest Gateway may only bind to a loopback address" - ); let listener = tokio::net::TcpListener::bind(listen) .await .with_context(|| format!("bind pChronicle ingest Gateway to {listen}"))?; diff --git a/crates/persisting-pchronicle-cli/src/lib.rs b/crates/persisting-pchronicle-cli/src/lib.rs index b3751837..b34eeef6 100644 --- a/crates/persisting-pchronicle-cli/src/lib.rs +++ b/crates/persisting-pchronicle-cli/src/lib.rs @@ -1124,13 +1124,9 @@ fn parse_gateway_bind(value: &str) -> std::result::Result { if value.eq_ignore_ascii_case("auto") { return Ok(SocketAddr::from(([127, 0, 0, 1], 0))); } - let address = value + value .parse::() - .map_err(|error| format!("invalid Gateway address '{value}': {error}"))?; - if !address.ip().is_loopback() { - return Err("the embedded Gateway is loopback-only; use 127.0.0.1:PORT or 'auto'".into()); - } - Ok(address) + .map_err(|error| format!("invalid Gateway address '{value}': {error}")) } #[derive(Debug, Args)] @@ -1677,14 +1673,9 @@ fn local_dataset_path(uri: &str) -> Result> { } fn parse_gateway_listener(value: &str, label: &str) -> Result { - let addr = value + value .parse::() - .with_context(|| format!("parse {label} address '{value}'"))?; - anyhow::ensure!( - addr.ip().is_loopback(), - "pChronicle embedded {label} may only bind to a loopback address" - ); - Ok(addr) + .with_context(|| format!("parse {label} address '{value}'")) } async fn prepare_gateway( @@ -2278,10 +2269,6 @@ async fn run_serve( projections.converge_before_readiness().await?; let warehouse = match warehouse_listen(&args) { Some(listen) => { - anyhow::ensure!( - listen.ip().is_loopback(), - "pChronicle Warehouse may only bind to a loopback address" - ); let listener = tokio::net::TcpListener::bind(listen) .await .with_context(|| format!("bind pChronicle Warehouse to {listen}"))?; @@ -2611,10 +2598,6 @@ fn control_storage_uri(config: &server::ChronicleServerConfig) -> Result<&str> { } async fn run_echo(args: EchoArgs, stderr: &mut dyn Write) -> Result<()> { - anyhow::ensure!( - args.listen.ip().is_loopback(), - "pChronicle Echo may only bind to a loopback address" - ); let listener = tokio::net::TcpListener::bind(args.listen) .await .with_context(|| format!("bind pChronicle Echo to {}", args.listen))?; diff --git a/crates/persisting-pchronicle-cli/src/onboard.rs b/crates/persisting-pchronicle-cli/src/onboard.rs index 1b634efa..3e9435a2 100644 --- a/crates/persisting-pchronicle-cli/src/onboard.rs +++ b/crates/persisting-pchronicle-cli/src/onboard.rs @@ -551,7 +551,8 @@ fn render_serve(renderer: &mut WalkthroughRenderer<'_>) -> Result<()> { pchronicle serve --listen 127.0.0.1:8080 --open evals=../data/atif ``` -服务只允许 loopback 地址，因为这个本地表面不提供认证；Dataset API 和 Web UI 都是只读的。 +默认示例仍使用 loopback；`--listen` 也可绑定非 loopback 地址。无认证时不要把 +只读 Warehouse 暴露到不可信网络。Dataset API 和 Web UI 都是只读的。 Runs 页面检索使用与 `find --match` 相同的 FTS/JSONB 语义，命中的轨迹会展示上下文预览； 可以先用 CLI `find` 定位，再在 Web 中继续钻取。 diff --git a/crates/persisting-pchronicle-cli/src/server/catalog.rs b/crates/persisting-pchronicle-cli/src/server/catalog.rs index aa5ef19e..d389bf61 100644 --- a/crates/persisting-pchronicle-cli/src/server/catalog.rs +++ b/crates/persisting-pchronicle-cli/src/server/catalog.rs @@ -690,13 +690,9 @@ pub(crate) fn parse_catalog_pin_target(input: &str) -> Result { let host = url .host_str() .ok_or_else(|| anyhow!("catalog pin URL must include a host"))?; - let address: std::net::IpAddr = host + let _: std::net::IpAddr = host .parse() - .with_context(|| format!("catalog pin host '{host}' must be a loopback IP"))?; - anyhow::ensure!( - address.is_loopback(), - "catalog pin host must be a loopback address" - ); + .with_context(|| format!("catalog pin host '{host}' must be an IP address"))?; let port = url .port() .ok_or_else(|| anyhow!("catalog pin URL must include a port"))?; @@ -1228,9 +1224,10 @@ dataset = "prod" } #[test] - fn catalog_pin_target_must_be_loopback_with_port() { + fn catalog_pin_target_accepts_any_ip_with_port() { assert!(parse_catalog_pin_target("catalog://127.0.0.1:8081").is_ok()); - assert!(parse_catalog_pin_target("catalog://8.8.8.8:8081").is_err()); + assert!(parse_catalog_pin_target("catalog://8.8.8.8:8081").is_ok()); + assert!(parse_catalog_pin_target("catalog://10.12.111.136:8000").is_ok()); assert!(parse_catalog_pin_target("catalog://127.0.0.1").is_err()); assert!(parse_catalog_pin_target("s3://bucket/prod").is_err()); } diff --git a/crates/persisting-pchronicle-cli/src/server/mod.rs b/crates/persisting-pchronicle-cli/src/server/mod.rs index b5d083ff..e60b67a4 100644 --- a/crates/persisting-pchronicle-cli/src/server/mod.rs +++ b/crates/persisting-pchronicle-cli/src/server/mod.rs @@ -1,4 +1,4 @@ -//! Local, loopback-only pChronicle browser. +//! Local pChronicle browser Warehouse. mod acceleration; mod asset; @@ -202,6 +202,9 @@ impl PreparedWarehouse { /// Mount every library from `catalog.toml` into the Warehouse process. /// Directory ticket routes remain available when users exist; the data /// plane serves in-process mounts instead of spawning query workers. + /// + /// Discovery runs in the background so `serve --listen` can accept + /// connections before large object prefixes finish classifying. pub(crate) async fn prepare_catalog(acl: catalog::CatalogAcl) -> anyhow::Result { acl.apply_backend_env(); let mounts = acl.mounts()?; @@ -213,7 +216,28 @@ impl PreparedWarehouse { let mut state = app_state(config); state.catalog_acl = Some(Arc::new(acl)); let warehouse = Self { state }; - warehouse.install_initial_runtime().await?; + let background = warehouse.state.clone(); + tokio::spawn(async move { + match build_catalog_runtime(&background.config).await { + Ok(runtime) => { + let snapshot_id = runtime.snapshot.snapshot_id().to_string(); + *background.catalog.write().await = Some(runtime); + *background.trajectory_cache.write().await = None; + tracing::info!( + target: "pchronicle.serve", + snapshot_id = %snapshot_id, + "catalog discovery ready" + ); + } + Err(error) => { + tracing::error!( + target: "pchronicle.serve", + error = %error, + "catalog discovery failed" + ); + } + } + }); Ok(warehouse) } @@ -328,10 +352,6 @@ pub async fn serve_warehouse( config: ChronicleServerConfig, addr: SocketAddr, ) -> anyhow::Result<()> { - anyhow::ensure!( - addr.ip().is_loopback(), - "pChronicle Warehouse may only bind to a loopback address" - ); let listener = tokio::net::TcpListener::bind(addr).await?; serve_warehouse_with_listener(config, listener).await } @@ -354,10 +374,7 @@ pub async fn serve_warehouse_with_listener_and_shutdown( let addr = listener .local_addr() .context("read Warehouse listen address")?; - anyhow::ensure!( - addr.ip().is_loopback(), - "pChronicle Warehouse may only bind to a loopback address" - ); + let _ = addr; axum::serve(listener, warehouse_router(config)) .with_graceful_shutdown(shutdown) .await @@ -372,10 +389,7 @@ pub(crate) async fn serve_prepared_warehouse_with_listener_and_shutdown( let addr = listener .local_addr() .context("read Warehouse listen address")?; - anyhow::ensure!( - addr.ip().is_loopback(), - "pChronicle Warehouse may only bind to a loopback address" - ); + let _ = addr; axum::serve(listener, warehouse.router()) .with_graceful_shutdown(shutdown) .await diff --git a/crates/persisting-pchronicle-cli/src/server/tests.rs b/crates/persisting-pchronicle-cli/src/server/tests.rs index 8355bb53..f16ae642 100644 --- a/crates/persisting-pchronicle-cli/src/server/tests.rs +++ b/crates/persisting-pchronicle-cli/src/server/tests.rs @@ -635,18 +635,25 @@ fn write_gateway_fixture_with_status( } #[tokio::test] -async fn warehouse_rejects_non_loopback_bind() { +async fn warehouse_binds_non_loopback() { let config = ChronicleServerConfig::mounted(vec![ DatasetMount::default("/tmp/none").expect("test Dataset mount must be valid"), ]) .expect("test server config must be valid"); - let error = serve_warehouse( - config, - SocketAddr::new(std::net::IpAddr::from([0, 0, 0, 0]), 0), - ) - .await - .unwrap_err(); - assert!(error.to_string().contains("loopback")); + let listener = tokio::net::TcpListener::bind("0.0.0.0:0") + .await + .expect("bind non-loopback warehouse"); + let addr = listener.local_addr().expect("local addr"); + assert!(!addr.ip().is_loopback()); + let (stop_tx, stop_rx) = tokio::sync::oneshot::channel::<()>(); + let serve = tokio::spawn(async move { + serve_warehouse_with_listener_and_shutdown(config, listener, async move { + let _ = stop_rx.await; + }) + .await + }); + stop_tx.send(()).expect("stop warehouse"); + serve.await.expect("join").expect("serve warehouse"); } #[test] diff --git a/crates/persisting-pchronicle-cli/src/tests.rs b/crates/persisting-pchronicle-cli/src/tests.rs index b865ab89..925c14d5 100644 --- a/crates/persisting-pchronicle-cli/src/tests.rs +++ b/crates/persisting-pchronicle-cli/src/tests.rs @@ -5029,11 +5029,10 @@ fn gateway_dataset_uri_is_auto_mounted_and_deduplicated() -> Result<()> { } #[test] -fn embedded_gateway_rejects_public_listeners() { - let error = parse_gateway_listener("0.0.0.0:8787", "Gateway").unwrap_err(); - assert!(error.to_string().contains("loopback")); +fn embedded_gateway_accepts_public_listeners() { + assert!(parse_gateway_listener("0.0.0.0:8787", "Gateway").is_ok()); assert!(parse_gateway_listener("127.0.0.1:0", "Gateway").is_ok()); - assert!(parse_gateway_bind("0.0.0.0:0").is_err()); + assert!(parse_gateway_bind("0.0.0.0:0").is_ok()); assert_eq!( parse_gateway_bind("auto").unwrap(), "127.0.0.1:0".parse::().unwrap() diff --git a/crates/persisting-pchronicle/src/store/catalog/discovery.rs b/crates/persisting-pchronicle/src/store/catalog/discovery.rs index f7391e56..46cdedde 100644 --- a/crates/persisting-pchronicle/src/store/catalog/discovery.rs +++ b/crates/persisting-pchronicle/src/store/catalog/discovery.rs @@ -586,79 +586,62 @@ async fn discover_local_candidates( }]); } + // Plain directory: only inspect immediate children for Dataset markers. + // Loose files are not registered as sources (lazy Directory navigation). let mut candidates = Vec::new(); - let mut pending = vec![root.to_path_buf()]; - let mut visited = 0usize; - while let Some(directory) = pending.pop() { - let mut entries = fs::read_dir(&directory) - .with_context(|| format!("read Dataset directory {}", directory.display()))? - .collect::>>()?; - entries.sort_by_key(|entry| entry.path()); - for entry in entries { - visited = visited.saturating_add(1); - anyhow::ensure!( - visited <= options.max_entries, - "Dataset traversal exceeds max_entries limit of {}", - options.max_entries - ); - let file_type = entry.file_type()?; - if file_type.is_symlink() { - continue; - } - let path = entry.path(); - if file_type.is_dir() { - if let Some(manifest) = try_load_manifest(&path) { - let nested = collect_manifest_subtree(root, &path, &manifest, options).await?; - candidates.extend(nested); - } else if path.join("CURRENT").is_file() { - let metadata = fs::metadata(path.join("CURRENT"))?; - candidates.push(Candidate::Storyline { - file: relative_catalog_path(root, &path, true)?, - uri: canonical_local_uri(&path)?, - size_bytes: Some(metadata.len()), - last_modified: modified_string(&metadata), - }); - } else if path.join("_manifest.json").is_file() - && path.file_name().is_some_and(|name| name == "events.lance") - { - let metadata = fs::metadata(path.join("_manifest.json"))?; - candidates.push(Candidate::Events { - file: relative_catalog_path(root, &path, true)?, - uri: canonical_local_uri(&path)?, - size_bytes: Some(metadata.len()), - last_modified: modified_string(&metadata), - }); - } else if is_lance_directory(&path) { - if is_compact_jsonl_directory(&path).await? { - let metadata = fs::metadata(&path)?; - candidates.push(Candidate::Compact { - file: relative_catalog_path(root, &path, true)?, - uri: canonical_local_uri(&path)?, - size_bytes: Some(metadata.len()), - last_modified: modified_string(&metadata), - }); - } - // Derived Lance datasets are sidecars of a canonical Run, - // not trajectory sources. Never descend into their internal - // metadata and register it as an outer file source. - } else { - pending.push(path); - } - } else if file_type.is_file() && is_json_candidate(&path) { - let metadata = entry.metadata()?; - candidates.push(Candidate::LocalFile { - file: relative_catalog_path(root, &path, false)?, - root: root.to_path_buf(), - path, - size_bytes: metadata.len(), - last_modified: modified_string(&metadata), - }); - } - anyhow::ensure!( - candidates.len() <= options.max_files, - "Dataset manifest exceeds max_files limit of {}", - options.max_files - ); + let mut entries = fs::read_dir(root) + .with_context(|| format!("read Dataset directory {}", root.display()))? + .collect::>>()?; + entries.sort_by_key(|entry| entry.path()); + for entry in entries { + let file_type = entry.file_type()?; + if file_type.is_symlink() || !file_type.is_dir() { + continue; + } + anyhow::ensure!( + candidates.len() < options.max_files, + "Dataset manifest exceeds max_files limit of {}", + options.max_files + ); + let path = entry.path(); + if let Some(manifest) = try_load_manifest(&path) { + let nested = collect_manifest_subtree(root, &path, &manifest, options).await?; + candidates.extend(nested); + } else if path.join("CURRENT").is_file() { + let metadata = fs::metadata(path.join("CURRENT"))?; + candidates.push(Candidate::Storyline { + file: relative_catalog_path(root, &path, true)?, + uri: canonical_local_uri(&path)?, + size_bytes: Some(metadata.len()), + last_modified: modified_string(&metadata), + }); + } else if path.join("_manifest.json").is_file() + && path.file_name().is_some_and(|name| name == "events.lance") + { + let metadata = fs::metadata(path.join("_manifest.json"))?; + candidates.push(Candidate::Events { + file: relative_catalog_path(root, &path, true)?, + uri: canonical_local_uri(&path)?, + size_bytes: Some(metadata.len()), + last_modified: modified_string(&metadata), + }); + } else if path.join("events.lance/_manifest.json").is_file() { + let events = path.join("events.lance"); + let metadata = fs::metadata(events.join("_manifest.json"))?; + candidates.push(Candidate::Events { + file: relative_catalog_path(root, &events, true)?, + uri: canonical_local_uri(&events)?, + size_bytes: Some(metadata.len()), + last_modified: modified_string(&metadata), + }); + } else if is_lance_directory(&path) && is_compact_jsonl_directory(&path).await? { + let metadata = fs::metadata(&path)?; + candidates.push(Candidate::Compact { + file: relative_catalog_path(root, &path, true)?, + uri: canonical_local_uri(&path)?, + size_bytes: Some(metadata.len()), + last_modified: modified_string(&metadata), + }); } } candidates.sort_by(|left, right| left.source_stub().file.cmp(&right.source_stub().file)); @@ -754,97 +737,206 @@ async fn discover_object_candidates( uri: &str, options: LocalQueryManifestOptions, ) -> Result> { + anyhow::ensure!(options.max_files > 0, "catalog max_files must be positive"); let store = OpendalStore::from_uri(uri).await?; - let mut metas = Vec::new(); - for entry in store - .list("") - .await - .with_context(|| format!("list Dataset object prefix {uri}"))? - { + + // Prefer a Dataset root (chronicle.manifest / CURRENT / events) over a flat + // recursive object walk. Plain prefixes navigate one directory level only. + match probe_object_prefix(&store, uri, "", ".").await? { + Some(ObjectProbe::Source(candidate)) => return Ok(vec![candidate]), + Some(ObjectProbe::Branch) => { + return collect_object_branch_children(&store, uri, "", options).await; + } + None => {} + } + + let mut candidates = Vec::new(); + for child in object_child_names(&store, "").await? { anyhow::ensure!( - metas.len() < options.max_entries, - "Dataset traversal exceeds max_entries limit of {}", - options.max_entries + candidates.len() < options.max_files, + "Dataset manifest exceeds max_files limit of {}", + options.max_files ); - metas.push(RemoteObjectMeta::from(entry)); + match probe_object_prefix(&store, uri, &child, root_source_path(&child)).await? { + Some(ObjectProbe::Source(candidate)) => candidates.push(candidate), + Some(ObjectProbe::Branch) => { + let nested = + collect_object_branch_children(&store, uri, &child, options).await?; + candidates.extend(nested); + } + None => {} + } } - metas.sort_by(|left, right| left.location.cmp(&right.location)); - - let root_is_events = uri.trim_end_matches('/').ends_with("events.lance"); - let mut storyline_roots = BTreeMap::::new(); - let mut event_roots = BTreeMap::::new(); - let mut relative_metas = Vec::with_capacity(metas.len()); - for meta in metas { - let relative = meta.location.clone(); - if relative == "CURRENT" || relative.ends_with("/CURRENT") { - storyline_roots.insert(parent_relative_path(&relative, "CURRENT"), meta.clone()); + + candidates.sort_by(|left, right| left.source_stub().file.cmp(&right.source_stub().file)); + Ok(candidates) +} + +enum ObjectProbe { + Source(Candidate), + Branch, +} + +async fn object_child_names(store: &OpendalStore, relative: &str) -> Result> { + let prefix = if relative.is_empty() { + String::new() + } else { + format!("{}/", relative.trim_end_matches('/')) + }; + let entries = store + .list_shallow(&prefix) + .await + .with_context(|| format!("list object prefix '{prefix}'"))?; + let mut child_names = BTreeSet::new(); + for entry in entries { + let path = entry.path.trim_start_matches(&prefix).trim_matches('/'); + if path.is_empty() { + continue; } - if (relative == "_manifest.json" && root_is_events) - || relative.ends_with("/events.lance/_manifest.json") - { - event_roots.insert( - parent_relative_path(&relative, "_manifest.json"), - meta.clone(), - ); + let child = path.split('/').next().unwrap_or(path); + if child.is_empty() { + continue; } - relative_metas.push((relative, meta)); + // Loose files at this level are ignored (Directory navigation only). + if entry.mode == opendal::EntryMode::FILE && !path.contains('/') { + continue; + } + child_names.insert(child.to_string()); } + Ok(child_names) +} +async fn collect_object_branch_children( + store: &OpendalStore, + root_uri: &str, + relative: &str, + options: LocalQueryManifestOptions, +) -> Result> { let mut candidates = Vec::new(); - for (relative, meta) in &storyline_roots { - candidates.push(Candidate::Storyline { - file: root_source_path(relative), - uri: child_uri(uri, relative), - size_bytes: Some(meta.size), - last_modified: Some(meta.last_modified.clone()), - }); + let mut stack = vec![relative.to_string()]; + while let Some(current) = stack.pop() { + for child in object_child_names(store, ¤t).await? { + anyhow::ensure!( + candidates.len() < options.max_files, + "Dataset manifest exceeds max_files limit of {}", + options.max_files + ); + let child_relative = if current.is_empty() { + child.clone() + } else { + format!("{}/{}", current.trim_end_matches('/'), child) + }; + match probe_object_prefix( + store, + root_uri, + &child_relative, + root_source_path(&child_relative), + ) + .await? + { + Some(ObjectProbe::Source(candidate)) => candidates.push(candidate), + Some(ObjectProbe::Branch) => stack.push(child_relative), + None => {} + } + } } - for (relative, meta) in &event_roots { - if is_nested_in_any(relative, storyline_roots.keys()) { - continue; + candidates.sort_by(|left, right| left.source_stub().file.cmp(&right.source_stub().file)); + Ok(candidates) +} + +async fn probe_object_prefix( + store: &OpendalStore, + root_uri: &str, + relative: &str, + source_file: impl Into, +) -> Result> { + let source_file = source_file.into(); + let prefix = if relative.is_empty() { + String::new() + } else { + format!("{}/", relative.trim_end_matches('/')) + }; + let join = |name: &str| { + if prefix.is_empty() { + name.to_string() + } else { + format!("{prefix}{name}") + } + }; + + if let Some(entry) = store + .stat_file(&join(crate::store::CHRONICLE_MANIFEST_FILE)) + .await? + { + let bytes = store + .read(&entry.path) + .await? + .map(|(bytes, _)| bytes) + .unwrap_or_default(); + let text = std::str::from_utf8(&bytes).context("chronicle.manifest must be UTF-8")?; + let manifest: crate::store::ChronicleManifest = + toml::from_str(text).context("parse chronicle.manifest")?; + manifest.validate()?; + match manifest.kind { + ManifestKind::Leaf => { + anyhow::ensure!( + manifest.is_compact_jsonl_leaf(), + "chronicle.manifest leaf format {:?} is not supported for discovery yet", + manifest.format + ); + let meta = RemoteObjectMeta::from(entry); + return Ok(Some(ObjectProbe::Source(Candidate::Compact { + file: source_file, + uri: child_uri(root_uri, relative), + size_bytes: Some(meta.size), + last_modified: Some(meta.last_modified), + }))); + } + ManifestKind::Branch => return Ok(Some(ObjectProbe::Branch)), } - candidates.push(Candidate::Events { - file: root_source_path(relative), - uri: child_uri(uri, relative), + } + + if let Some(entry) = store.stat_file(&join("CURRENT")).await? { + let meta = RemoteObjectMeta::from(entry); + return Ok(Some(ObjectProbe::Source(Candidate::Storyline { + file: source_file, + uri: child_uri(root_uri, relative), size_bytes: Some(meta.size), - last_modified: Some(meta.last_modified.clone()), - }); + last_modified: Some(meta.last_modified), + }))); } - let composite_roots = storyline_roots - .keys() - .chain(event_roots.keys()) - .cloned() - .collect::>(); - for (relative, meta) in relative_metas { - if is_nested_in_any(&relative, composite_roots.iter()) - || path_is_inside_lance_directory(&relative) - { - continue; - } - let candidate_path = if relative.is_empty() { - Path::new(uri) + let events_manifest = if relative.is_empty() { + "_manifest.json".to_string() + } else if relative.trim_end_matches('/').ends_with("events.lance") { + join("_manifest.json") + } else { + join("events.lance/_manifest.json") + }; + if let Some(entry) = store.stat_file(&events_manifest).await? { + let meta = RemoteObjectMeta::from(entry); + let events_relative = if relative.is_empty() { + if root_uri.trim_end_matches('/').ends_with("events.lance") { + String::new() + } else { + "events.lance".to_string() + } + } else if relative.trim_end_matches('/').ends_with("events.lance") { + relative.to_string() } else { - Path::new(&relative) + format!("{}/events.lance", relative.trim_end_matches('/')) }; - if is_json_candidate(candidate_path) { - let file = if relative.is_empty() { - uri.rsplit('/').next().unwrap_or("dataset.json").to_string() + return Ok(Some(ObjectProbe::Source(Candidate::Events { + file: if events_relative.is_empty() { + ".".into() } else { - relative - }; - candidates.push(Candidate::RemoteFile { - file, - store: store.clone(), - meta, - }); - } + events_relative.clone() + }, + uri: child_uri(root_uri, &events_relative), + size_bytes: Some(meta.size), + last_modified: Some(meta.last_modified), + }))); } - anyhow::ensure!( - candidates.len() <= options.max_files, - "Dataset manifest exceeds max_files limit of {}", - options.max_files - ); - candidates.sort_by(|left, right| left.source_stub().file.cmp(&right.source_stub().file)); - Ok(candidates) + + Ok(None) } diff --git a/crates/persisting-pchronicle/src/store/opendal_store.rs b/crates/persisting-pchronicle/src/store/opendal_store.rs index 10a3cbd8..bd15b216 100644 --- a/crates/persisting-pchronicle/src/store/opendal_store.rs +++ b/crates/persisting-pchronicle/src/store/opendal_store.rs @@ -36,6 +36,13 @@ pub(crate) struct Entry { pub(crate) metadata: Metadata, } +#[derive(Clone, Debug)] +pub(crate) struct ShallowEntry { + pub(crate) path: String, + pub(crate) mode: EntryMode, + pub(crate) metadata: Metadata, +} + static SHARED_MEMORY: OnceLock>> = OnceLock::new(); static SHARED_LOCKS: OnceLock>>>> = OnceLock::new(); @@ -148,6 +155,33 @@ impl Store { Ok(entries) } + /// Non-recursive listing of the immediate children under `prefix`. + /// Returns both files and directories so callers can navigate lazily. + pub(crate) async fn list_shallow(&self, prefix: &str) -> Result> { + let mut lister = self.operator.lister_with(prefix).recursive(false).await?; + let mut entries = Vec::new(); + while let Some(entry) = lister.try_next().await? { + entries.push(ShallowEntry { + path: entry.path().to_string(), + mode: entry.metadata().mode(), + metadata: entry.metadata().clone(), + }); + } + Ok(entries) + } + + pub(crate) async fn stat_file(&self, path: &str) -> Result> { + match self.operator.stat(path).await { + Ok(metadata) if metadata.mode() == EntryMode::FILE => Ok(Some(Entry { + path: path.to_string(), + metadata, + })), + Ok(_) => Ok(None), + Err(error) if error.kind() == ErrorKind::NotFound => Ok(None), + Err(error) => Err(error.into()), + } + } + pub(crate) async fn exists(&self) -> Result { Ok(self .operator diff --git a/docs/src/en/rfcs/0013-pchronicle-warehouse-catalog.md b/docs/src/en/rfcs/0013-pchronicle-warehouse-catalog.md index 9b17d5d8..b989b33e 100644 --- a/docs/src/en/rfcs/0013-pchronicle-warehouse-catalog.md +++ b/docs/src/en/rfcs/0013-pchronicle-warehouse-catalog.md @@ -17,7 +17,8 @@ Dataset 身份始终是 path（本机路径或 `s3://` / `az://` / `gs://` URI CLI 标志、配置文件和 HTTP 路径为兼容性仍使用 `catalog` 一词（`--catalog-config`、`catalog.toml`、`catalog://`、`/api/v1/catalog/datasets`）。产品与 RFC 口径称 Directory。 -规范实现挂在现有 `pchronicle serve --catalog-config` 上，不引入独立 `catalog serve` 进程，也不把 listener 从 loopback 打开。 +规范实现挂在现有 `pchronicle serve --catalog-config` 上，不引入独立 `catalog serve` 进程。 +Listener 默认可为 loopback；也允许绑定非环回地址，但部署方 MUST 自行保证网络边界。 - **Serve 挂载**：`pchronicle serve --catalog-config FILE` MUST 把 `catalog.toml` 中的 **全部** `[datasets.*]` 挂进 Warehouse（与位置参数挂载等价）。本机 Web / 无用户钥的数据面请求在 @@ -56,14 +57,14 @@ pchronicle query @team/prod 'SELECT 1' - 让 `@name/library` 解析为一条 path（换票后的 `uri`）；引擎随后只打开该 path。 - 换票后 CLI 自己访问存储；后端密钥只出现在票和 worker stdin 中，不写入用户 `config.toml`。 - Web 用用户钥换授权范围，查询只看到该用户的 mounts。 -- 保持 Warehouse 为 loopback-only 本地检查面，而不是公网多租户服务。 +- 允许 Warehouse 绑定任意 listen 地址；默认示例仍用 loopback。Catalog 头不是公网认证边界，不可信网络上的暴露由部署方负责。 ### 非目标 - STS、临时凭证轮换、或把用户钥映射成短时 AWS session。 - 热加载 `catalog.toml`；改配置 MUST 重启 serve。 - 在运行中的 Warehouse 上提供 HTTP 签发接口。 -- 把 listener bind 到非环回地址，或提供独立 `catalog serve` 二进制。 +- 提供独立 `catalog serve` 二进制。 - 在已运行的 Tokio runtime 上 `fork(2)`（未定义行为）。 - 把后端对象存储密钥写入本机 dataset pin 配置。 - 改变 Snapshot 协议、SQL schema 或 Gateway/Control 协议。 @@ -89,11 +90,11 @@ Directory 挂在现有 Warehouse listener 上。未传 `--catalog-config` 时， ```text 浏览器 / CLI - → loopback Warehouse + → Warehouse listener ├─ GET /health ├─ GET /api/v1/catalog/datasets[/{name}] 父进程：鉴权 + 目录/票 ├─ 静态 UI - └─ 其余 /api/* 父进程鉴权后 spawn worker + └─ 其余 /api/* 父进程内挂载 / 或 spawn worker → pchronicle serve --catalog-query-worker stdin: mounts + HTTP 请求 stdout: status / content-type / body @@ -102,7 +103,7 @@ Directory 挂在现有 Warehouse listener 上。未传 `--catalog-config` 时， 约束： -1. Listener MUST 为 loopback。本 RFC 不把 catalog 头当作公网认证边界。 +1. Listener MAY 绑定非 loopback 地址。本 RFC 不把 catalog 头当作公网认证边界；部署方 MUST 在不可信网络上自行加边界。 2. 父进程 MUST NOT 打开 `catalog.toml` 中的 libraries。父进程使用空 mount 的 front-only Warehouse。 3. Worker MUST 由 `Command` 启动新进程，MUST NOT `fork(2)` 已运行的 Tokio runtime。 4. Worker MUST NOT 监听端口、MUST NOT 读取 `catalog.toml`、MUST NOT 读取用户钥。它只消费 stdin 中过滤后的 mounts 和原始请求。 diff --git a/docs/src/en/rfcs/0015-chronicle-manifest.md b/docs/src/en/rfcs/0015-chronicle-manifest.md index 600e4b08..8e835cb6 100644 --- a/docs/src/en/rfcs/0015-chronicle-manifest.md +++ b/docs/src/en/rfcs/0015-chronicle-manifest.md @@ -45,7 +45,12 @@ Goals: - Persist aggregate stats used by explorer tree / dataset summaries. - Support nested Dataset trees by **automatically scanning** child directories for `chronicle.manifest`. -- Keep missing or stale manifests compatible with existing heuristic discovery. +- Treat a prefix without `chronicle.manifest` as a **Directory**: inspect only + **immediate** child directories for Dataset markers, and do not register loose + files as Sources. +- When the sidecar is missing, still classify Datasets via `CURRENT` / events / + compact-jsonl markers, but MUST NOT recursively list an entire object-store + prefix just to classify. Non-goals (v1): @@ -84,10 +89,12 @@ Parents MUST NOT require an explicit children list. Discovery MUST: 3. If `kind = "branch"`, scan **immediate** child directories only; for each child that contains `chronicle.manifest`, treat that child as a nested Dataset node and continue according to that child's kind. -4. If the current directory has no `chronicle.manifest`, keep the existing - heuristic discovery, but when a subdirectory contains - `chronicle.manifest`, prefer that node and MUST NOT open Lance solely to - classify it. +4. If the current directory has no `chronicle.manifest`, treat it as a + **Directory**: inspect **immediate** child directories only; classify each + child via Dataset markers (`chronicle.manifest`, `CURRENT`, + `events.lance/_manifest.json`, compact-jsonl Lance). Loose files MUST NOT + be registered as Sources. Discovery MUST NOT recursively list an entire + object-store prefix tree to classify. Symlinks MUST be ignored. Existing `max_entries` / `max_files` limits still apply to traversal. diff --git a/docs/src/zh/rfcs/0013-pchronicle-warehouse-catalog.md b/docs/src/zh/rfcs/0013-pchronicle-warehouse-catalog.md index e46ea5f5..0bbdef4a 100644 --- a/docs/src/zh/rfcs/0013-pchronicle-warehouse-catalog.md +++ b/docs/src/zh/rfcs/0013-pchronicle-warehouse-catalog.md @@ -17,7 +17,8 @@ Dataset 身份始终是 path（本机路径或 `s3://` / `az://` / `gs://` URI CLI 标志、配置文件和 HTTP 路径为兼容性仍使用 `catalog` 一词（`--catalog-config`、`catalog.toml`、`catalog://`、`/api/v1/catalog/datasets`）。产品与 RFC 口径称 Directory。 -规范实现挂在现有 `pchronicle serve --catalog-config` 上，不引入独立 `catalog serve` 进程，也不把 listener 从 loopback 打开。 +规范实现挂在现有 `pchronicle serve --catalog-config` 上，不引入独立 `catalog serve` 进程。 +Listener 默认可为 loopback；也允许绑定非环回地址，但部署方 MUST 自行保证网络边界。 - **Serve 挂载**：`pchronicle serve --catalog-config FILE` MUST 把 `catalog.toml` 中的 **全部** `[datasets.*]` 挂进 Warehouse（与位置参数挂载等价）。本机 Web / 无用户钥的数据面请求在 @@ -56,14 +57,14 @@ pchronicle query @team/prod 'SELECT 1' - 让 `@name/library` 解析为一条 path（换票后的 `uri`）；引擎随后只打开该 path。 - 换票后 CLI 自己访问存储；后端密钥只出现在票和 worker stdin 中，不写入用户 `config.toml`。 - Web 用用户钥换授权范围，查询只看到该用户的 mounts。 -- 保持 Warehouse 为 loopback-only 本地检查面，而不是公网多租户服务。 +- 允许 Warehouse 绑定任意 listen 地址；默认示例仍用 loopback。Catalog 头不是公网认证边界，不可信网络上的暴露由部署方负责。 ### 非目标 - STS、临时凭证轮换、或把用户钥映射成短时 AWS session。 - 热加载 `catalog.toml`；改配置 MUST 重启 serve。 - 在运行中的 Warehouse 上提供 HTTP 签发接口。 -- 把 listener bind 到非环回地址，或提供独立 `catalog serve` 二进制。 +- 提供独立 `catalog serve` 二进制。 - 在已运行的 Tokio runtime 上 `fork(2)`（未定义行为）。 - 把后端对象存储密钥写入本机 dataset pin 配置。 - 改变 Snapshot 协议、SQL schema 或 Gateway/Control 协议。 @@ -89,11 +90,11 @@ Directory 挂在现有 Warehouse listener 上。未传 `--catalog-config` 时， ```text 浏览器 / CLI - → loopback Warehouse + → Warehouse listener ├─ GET /health ├─ GET /api/v1/catalog/datasets[/{name}] 父进程：鉴权 + 目录/票 ├─ 静态 UI - └─ 其余 /api/* 父进程鉴权后 spawn worker + └─ 其余 /api/* 父进程内挂载 / 或 spawn worker → pchronicle serve --catalog-query-worker stdin: mounts + HTTP 请求 stdout: status / content-type / body @@ -102,7 +103,7 @@ Directory 挂在现有 Warehouse listener 上。未传 `--catalog-config` 时， 约束： -1. Listener MUST 为 loopback。本 RFC 不把 catalog 头当作公网认证边界。 +1. Listener MAY 绑定非 loopback 地址。本 RFC 不把 catalog 头当作公网认证边界；部署方 MUST 在不可信网络上自行加边界。 2. 父进程 MUST NOT 打开 `catalog.toml` 中的 datasets。父进程使用空 mount 的 front-only Warehouse。 3. Worker MUST 由 `Command` 启动新进程，MUST NOT `fork(2)` 已运行的 Tokio runtime。 4. Worker MUST NOT 监听端口、MUST NOT 读取 `catalog.toml`、MUST NOT 读取用户钥。它只消费 stdin 中过滤后的 mounts 和原始请求。 diff --git a/docs/src/zh/rfcs/0015-chronicle-manifest.md b/docs/src/zh/rfcs/0015-chronicle-manifest.md index a66250de..a220077e 100644 --- a/docs/src/zh/rfcs/0015-chronicle-manifest.md +++ b/docs/src/zh/rfcs/0015-chronicle-manifest.md @@ -40,7 +40,10 @@ pChronicle 已对其它布局使用应用控制文件（Storyline 的 `CURRENT` - 让 discovery 通过读小 TOML 文件即可分类 Dataset 节点； - 持久化 explorer tree / dataset 摘要所需的聚合统计； - 通过**自动扫描**子目录中的 `chronicle.manifest` 支持嵌套 Dataset 树； -- 在 sidecar 缺失或过期时，仍兼容现有启发式发现。 +- 无 manifest 的普通目录按 **Directory** 处理：只检查**一层**子目录是否为 + Dataset（manifest / `CURRENT` / events），不把松散文件登记为 Source； +- 在 sidecar 缺失时，仍可用 `CURRENT` / events / compact-jsonl 标记做 Dataset + 分类，但 MUST NOT 为分类而全量递归列举对象存储前缀。 非目标（v1）： @@ -71,7 +74,10 @@ pChronicle 已对其它布局使用应用控制文件（Storyline 的 `CURRENT` 1. 若当前目录存在 `chronicle.manifest`，则解析它； 2. 若 `kind = "leaf"`，将该目录视为对应 `format` 的一个 source 候选，且 MUST NOT 再递归其内部寻找其它 source； 3. 若 `kind = "branch"`，只扫描**一层**子目录；对每个含有 `chronicle.manifest` 的子目录，按该子节点的 kind 继续处理； -4. 若当前目录没有 `chronicle.manifest`，保留现有启发式发现，但当子目录含有 `chronicle.manifest` 时，优先采用该节点，且 MUST NOT 仅为分类而打开 Lance。 +4. 若当前目录没有 `chronicle.manifest`，则视为 **Directory**：只检查**一层** + 子目录；对每个子目录用 Dataset 标记（`chronicle.manifest`、`CURRENT`、 + `events.lance/_manifest.json`、compact-jsonl Lance）分类。松散文件 MUST NOT + 登记为 Source。MUST NOT 为发现而递归列举整棵对象前缀树。 MUST 忽略符号链接。现有 `max_entries` / `max_files` 遍历上限仍然适用。 From da485413a57b4501ba707c90e346aa7c28314320 Mon Sep 17 00:00:00 2001 From: Reiase Date: Tue, 8 Sep 2026 22:53:50 +0800 Subject: [PATCH 2/8] refactor: simplify code formatting in discovery and generic modules Consolidated multiple lines of code into single lines for improved clarity and readability in the `discovery.rs` and `generic.rs` files. This change enhances the overall code structure without altering functionality. --- crates/persisting-pchronicle/src/store/catalog/discovery.rs | 3 +-- crates/persisting-replay/src/adapter/generic.rs | 5 +---- 2 files changed, 2 insertions(+), 6 deletions(-) diff --git a/crates/persisting-pchronicle/src/store/catalog/discovery.rs b/crates/persisting-pchronicle/src/store/catalog/discovery.rs index 46cdedde..dd0fef2c 100644 --- a/crates/persisting-pchronicle/src/store/catalog/discovery.rs +++ b/crates/persisting-pchronicle/src/store/catalog/discovery.rs @@ -760,8 +760,7 @@ async fn discover_object_candidates( match probe_object_prefix(&store, uri, &child, root_source_path(&child)).await? { Some(ObjectProbe::Source(candidate)) => candidates.push(candidate), Some(ObjectProbe::Branch) => { - let nested = - collect_object_branch_children(&store, uri, &child, options).await?; + let nested = collect_object_branch_children(&store, uri, &child, options).await?; candidates.extend(nested); } None => {} diff --git a/crates/persisting-replay/src/adapter/generic.rs b/crates/persisting-replay/src/adapter/generic.rs index 88baba1d..5f84e5de 100644 --- a/crates/persisting-replay/src/adapter/generic.rs +++ b/crates/persisting-replay/src/adapter/generic.rs @@ -1150,10 +1150,7 @@ fn continue_native_cli( log_path: log_path.clone(), }) .map_err(|error| ReplayError::new(ReplayErrorKind::Continuation, error.message))?; - let bridge_result = codex_bridge.take().map(|bridge| { - let result = bridge.finish(); - result - }); + let bridge_result = codex_bridge.take().map(|bridge| bridge.finish()); let bridge_error = bridge_result.and_then(|result| result.err()); if !output.status.success() { let process_error = ReplayError::classify_continuation( From 835c595ffbab2542a5acec04f4909ff830d63f6e Mon Sep 17 00:00:00 2001 From: Reiase Date: Wed, 9 Sep 2026 00:16:30 +0800 Subject: [PATCH 3/8] feat(catalog): introduce Directory candidate type and enhance discovery logic Added a new `Directory` variant to the `Candidate` enum to represent navigational directories in the catalog. Updated the discovery logic to classify immediate child directories and dataset sources separately, ensuring that loose files are not registered as sources. Enhanced the sorting and counting of sources in the catalog to accommodate the new directory type. Updated documentation to reflect these changes. --- .../persisting-pchronicle-cli/src/exchange.rs | 57 +++-- crates/persisting-pchronicle-cli/src/lib.rs | 35 ++- crates/persisting-pchronicle-cli/src/sync.rs | 42 +--- .../src/store/catalog/discovery.rs | 210 +++++++++++++----- .../src/store/catalog/mod.rs | 41 +++- .../src/store/catalog/provider.rs | 1 + .../src/store/catalog/tests.rs | 46 +++- docs/src/en/rfcs/0015-chronicle-manifest.md | 13 +- docs/src/zh/rfcs/0015-chronicle-manifest.md | 8 +- 9 files changed, 321 insertions(+), 132 deletions(-) diff --git a/crates/persisting-pchronicle-cli/src/exchange.rs b/crates/persisting-pchronicle-cli/src/exchange.rs index af153ad2..928e030c 100644 --- a/crates/persisting-pchronicle-cli/src/exchange.rs +++ b/crates/persisting-pchronicle-cli/src/exchange.rs @@ -1362,11 +1362,38 @@ fn collect_import_candidates(input: &Path) -> Result<(bool, Vec Result> { + let mut pending = vec![root.to_path_buf()]; + let mut files = Vec::new(); while let Some(directory) = pending.pop() { let mut entries = std::fs::read_dir(&directory) - .with_context(|| format!("read import directory {}", directory.display()))? + .with_context(|| format!("read directory {}", directory.display()))? .collect::>>()?; entries.sort_by_key(std::fs::DirEntry::path); for entry in entries { @@ -1377,30 +1404,16 @@ fn collect_import_candidates(input: &Path) -> Result<(bool, Vec bool { +fn is_visible_json_file(path: &Path) -> bool { path.extension() .and_then(|extension| extension.to_str()) .is_some_and(|extension| { diff --git a/crates/persisting-pchronicle-cli/src/lib.rs b/crates/persisting-pchronicle-cli/src/lib.rs index b34eeef6..6930e4a8 100644 --- a/crates/persisting-pchronicle-cli/src/lib.rs +++ b/crates/persisting-pchronicle-cli/src/lib.rs @@ -2673,11 +2673,17 @@ async fn run_list( let dataset = snapshot .dataset(DEFAULT_DATASET_NAME) .context("default Dataset missing from Snapshot")?; + let mut sources: Vec = dataset.sources.iter().map(source_response).collect(); + sources.sort_by(|left, right| { + directory_list_sort_key(left.kind) + .cmp(&directory_list_sort_key(right.kind)) + .then_with(|| left.source_path.cmp(&right.source_path)) + }); let response = ListResponse { dataset_uri, snapshot_id: snapshot.snapshot_id().to_string(), created_at: snapshot.created_at().to_string(), - sources: dataset.sources.iter().map(source_response).collect(), + sources, }; let output_format = match args.format { @@ -2694,12 +2700,18 @@ async fn run_list( } OutputFormat::Auto => unreachable!("auto output format was resolved"), } + let queryable = response + .sources + .iter() + .filter(|source| source.kind != CatalogSourceKind::Directory) + .count(); writeln!( stderr, - "snapshot_id={} dataset_uri={} sources={} ready={} errors={}", + "snapshot_id={} dataset_uri={} sources={} directories={} ready={} errors={}", response.snapshot_id, response.dataset_uri, - response.sources.len(), + queryable, + dataset.directory_count(), dataset.ready_source_count(), dataset.error_source_count(), ) @@ -2707,6 +2719,13 @@ async fn run_list( Ok(()) } +fn directory_list_sort_key(kind: CatalogSourceKind) -> u8 { + match kind { + CatalogSourceKind::Directory => 0, + CatalogSourceKind::Store | CatalogSourceKind::File => 1, + } +} + fn write_catalog_pin_dataset_list( listing: CatalogPinDatasetList, format: OutputFormat, @@ -2754,8 +2773,16 @@ fn write_catalog_pin_dataset_list( } fn source_response(source: &DiscoveredSource) -> SourceResponse { + let source_path = if source.kind == CatalogSourceKind::Directory + && !source.file.ends_with('/') + && source.file != "." + { + format!("{}/", source.file) + } else { + source.file.clone() + }; SourceResponse { - source_path: source.file.clone(), + source_path, format: source.format.clone(), kind: source.kind, snapshot_ref: source.snapshot_ref(), diff --git a/crates/persisting-pchronicle-cli/src/sync.rs b/crates/persisting-pchronicle-cli/src/sync.rs index a945372f..d2e94792 100644 --- a/crates/persisting-pchronicle-cli/src/sync.rs +++ b/crates/persisting-pchronicle-cli/src/sync.rs @@ -174,28 +174,17 @@ fn prepare_target(path: &Path, name: &str) -> Result { } fn scan_files(root: &Path) -> Result> { - let mut pending = vec![root.to_path_buf()]; let mut files = BTreeMap::new(); - while let Some(directory) = pending.pop() { - for entry in fs::read_dir(&directory) - .with_context(|| format!("read sync directory {}", directory.display()))? - { - let entry = entry?; - let file_type = entry.file_type()?; - let path = entry.path(); - if file_type.is_dir() { - pending.push(path); - } else if file_type.is_file() && is_sync_candidate(&path) { - let metadata = entry.metadata()?; - files.insert( - path.strip_prefix(root)?.to_path_buf(), - FileStamp { - size: metadata.len(), - modified: metadata.modified().ok(), - }, - ); - } - } + for path in crate::exchange::collect_visible_json_files(root)? { + let metadata = fs::metadata(&path) + .with_context(|| format!("stat sync file {}", path.display()))?; + files.insert( + path.strip_prefix(root)?.to_path_buf(), + FileStamp { + size: metadata.len(), + modified: metadata.modified().ok(), + }, + ); } Ok(files) } @@ -212,17 +201,6 @@ fn changed_paths( .collect() } -fn is_sync_candidate(path: &Path) -> bool { - path.extension() - .and_then(|extension| extension.to_str()) - .is_some_and(|extension| { - matches!( - extension.to_ascii_lowercase().as_str(), - "json" | "jsonl" | "ndjson" - ) - }) -} - #[cfg(test)] mod tests { use super::*; diff --git a/crates/persisting-pchronicle/src/store/catalog/discovery.rs b/crates/persisting-pchronicle/src/store/catalog/discovery.rs index dd0fef2c..01f35c3a 100644 --- a/crates/persisting-pchronicle/src/store/catalog/discovery.rs +++ b/crates/persisting-pchronicle/src/store/catalog/discovery.rs @@ -34,6 +34,9 @@ pub(super) enum Candidate { store: OpendalStore, meta: RemoteObjectMeta, }, + Directory { + file: String, + }, } impl Candidate { @@ -104,6 +107,14 @@ impl Candidate { Some(meta.last_modified.clone()), Some(remote_source_revision(meta)), ), + Self::Directory { file } => ( + file.clone(), + None, + CatalogSourceKind::Directory, + None, + None, + None, + ), }; DiscoveredSource { file, @@ -251,6 +262,9 @@ pub(super) async fn freeze_candidate( )), )) } + Candidate::Directory { file } => Err(anyhow::anyhow!( + "directory entry '{file}' is not a queryable Source" + )), } } @@ -586,63 +600,99 @@ async fn discover_local_candidates( }]); } - // Plain directory: only inspect immediate children for Dataset markers. - // Loose files are not registered as sources (lazy Directory navigation). + // Directory: inspect immediate children only. Dataset markers become + // Sources; unlabeled child dirs become navigational Directory entries. + // Loose JSON is accepted only as a flat Dataset when the mount root has no + // child directories. let mut candidates = Vec::new(); + let mut root_json = Vec::new(); + let mut has_child_dirs = false; let mut entries = fs::read_dir(root) .with_context(|| format!("read Dataset directory {}", root.display()))? .collect::>>()?; entries.sort_by_key(|entry| entry.path()); for entry in entries { let file_type = entry.file_type()?; - if file_type.is_symlink() || !file_type.is_dir() { + if file_type.is_symlink() { continue; } - anyhow::ensure!( - candidates.len() < options.max_files, - "Dataset manifest exceeds max_files limit of {}", - options.max_files - ); let path = entry.path(); - if let Some(manifest) = try_load_manifest(&path) { - let nested = collect_manifest_subtree(root, &path, &manifest, options).await?; - candidates.extend(nested); - } else if path.join("CURRENT").is_file() { - let metadata = fs::metadata(path.join("CURRENT"))?; - candidates.push(Candidate::Storyline { - file: relative_catalog_path(root, &path, true)?, - uri: canonical_local_uri(&path)?, - size_bytes: Some(metadata.len()), - last_modified: modified_string(&metadata), - }); - } else if path.join("_manifest.json").is_file() - && path.file_name().is_some_and(|name| name == "events.lance") - { - let metadata = fs::metadata(path.join("_manifest.json"))?; - candidates.push(Candidate::Events { - file: relative_catalog_path(root, &path, true)?, - uri: canonical_local_uri(&path)?, - size_bytes: Some(metadata.len()), - last_modified: modified_string(&metadata), - }); - } else if path.join("events.lance/_manifest.json").is_file() { - let events = path.join("events.lance"); - let metadata = fs::metadata(events.join("_manifest.json"))?; - candidates.push(Candidate::Events { - file: relative_catalog_path(root, &events, true)?, - uri: canonical_local_uri(&events)?, - size_bytes: Some(metadata.len()), - last_modified: modified_string(&metadata), - }); - } else if is_lance_directory(&path) && is_compact_jsonl_directory(&path).await? { - let metadata = fs::metadata(&path)?; - candidates.push(Candidate::Compact { - file: relative_catalog_path(root, &path, true)?, - uri: canonical_local_uri(&path)?, - size_bytes: Some(metadata.len()), + if file_type.is_dir() { + has_child_dirs = true; + anyhow::ensure!( + candidates.len() < options.max_files, + "Dataset manifest exceeds max_files limit of {}", + options.max_files + ); + if let Some(manifest) = try_load_manifest(&path) { + let nested = collect_manifest_subtree(root, &path, &manifest, options).await?; + candidates.extend(nested); + } else if path.join("CURRENT").is_file() { + let metadata = fs::metadata(path.join("CURRENT"))?; + candidates.push(Candidate::Storyline { + file: relative_catalog_path(root, &path, true)?, + uri: canonical_local_uri(&path)?, + size_bytes: Some(metadata.len()), + last_modified: modified_string(&metadata), + }); + } else if path.join("_manifest.json").is_file() + && path.file_name().is_some_and(|name| name == "events.lance") + { + let metadata = fs::metadata(path.join("_manifest.json"))?; + candidates.push(Candidate::Events { + file: relative_catalog_path(root, &path, true)?, + uri: canonical_local_uri(&path)?, + size_bytes: Some(metadata.len()), + last_modified: modified_string(&metadata), + }); + } else if path.join("events.lance/_manifest.json").is_file() { + let events = path.join("events.lance"); + let metadata = fs::metadata(events.join("_manifest.json"))?; + candidates.push(Candidate::Events { + file: relative_catalog_path(root, &events, true)?, + uri: canonical_local_uri(&events)?, + size_bytes: Some(metadata.len()), + last_modified: modified_string(&metadata), + }); + } else if is_lance_directory(&path) { + if is_compact_jsonl_directory(&path).await? { + let metadata = fs::metadata(&path)?; + candidates.push(Candidate::Compact { + file: relative_catalog_path(root, &path, true)?, + uri: canonical_local_uri(&path)?, + size_bytes: Some(metadata.len()), + last_modified: modified_string(&metadata), + }); + } + // Unknown Lance sidecars are not navigational Directory entries. + } else { + candidates.push(Candidate::Directory { + file: relative_catalog_path(root, &path, true)?, + }); + } + } else if file_type.is_file() && is_json_candidate(&path) { + let metadata = entry.metadata()?; + root_json.push(Candidate::LocalFile { + file: relative_catalog_path(root, &path, false)?, + root: root.to_path_buf(), + path, + size_bytes: metadata.len(), last_modified: modified_string(&metadata), }); } + anyhow::ensure!( + candidates.len() <= options.max_files, + "Dataset manifest exceeds max_files limit of {}", + options.max_files + ); + } + if !has_child_dirs && candidates.is_empty() { + anyhow::ensure!( + root_json.len() <= options.max_files, + "Dataset manifest exceeds max_files limit of {}", + options.max_files + ); + candidates = root_json; } candidates.sort_by(|left, right| left.source_stub().file.cmp(&right.source_stub().file)); Ok(candidates) @@ -740,8 +790,6 @@ async fn discover_object_candidates( anyhow::ensure!(options.max_files > 0, "catalog max_files must be positive"); let store = OpendalStore::from_uri(uri).await?; - // Prefer a Dataset root (chronicle.manifest / CURRENT / events) over a flat - // recursive object walk. Plain prefixes navigate one directory level only. match probe_object_prefix(&store, uri, "", ".").await? { Some(ObjectProbe::Source(candidate)) => return Ok(vec![candidate]), Some(ObjectProbe::Branch) => { @@ -750,8 +798,11 @@ async fn discover_object_candidates( None => {} } + // Directory: one shallow level only — never recursive list(""). let mut candidates = Vec::new(); - for child in object_child_names(&store, "").await? { + let (child_dirs, files) = object_shallow_children(&store, "").await?; + let has_child_dirs = !child_dirs.is_empty(); + for child in child_dirs { anyhow::ensure!( candidates.len() < options.max_files, "Dataset manifest exceeds max_files limit of {}", @@ -763,10 +814,41 @@ async fn discover_object_candidates( let nested = collect_object_branch_children(&store, uri, &child, options).await?; candidates.extend(nested); } - None => {} + None => { + if child.ends_with(".lance") { + continue; + } + candidates.push(Candidate::Directory { + file: root_source_path(&child), + }); + } } } + if !has_child_dirs && candidates.is_empty() { + let mut root_json = Vec::new(); + for (name, meta) in files { + if is_json_candidate(Path::new(&name)) { + root_json.push(Candidate::RemoteFile { + file: name, + store: store.clone(), + meta, + }); + } + } + anyhow::ensure!( + root_json.len() <= options.max_files, + "Dataset manifest exceeds max_files limit of {}", + options.max_files + ); + candidates = root_json; + } + + anyhow::ensure!( + candidates.len() <= options.max_files, + "Dataset manifest exceeds max_files limit of {}", + options.max_files + ); candidates.sort_by(|left, right| left.source_stub().file.cmp(&right.source_stub().file)); Ok(candidates) } @@ -776,7 +858,10 @@ enum ObjectProbe { Branch, } -async fn object_child_names(store: &OpendalStore, relative: &str) -> Result> { +async fn object_shallow_children( + store: &OpendalStore, + relative: &str, +) -> Result<(BTreeSet, Vec<(String, RemoteObjectMeta)>)> { let prefix = if relative.is_empty() { String::new() } else { @@ -786,7 +871,8 @@ async fn object_child_names(store: &OpendalStore, relative: &str) -> Result Result Result> { + let (dirs, _) = object_shallow_children(store, relative).await?; + Ok(dirs) } async fn collect_object_branch_children( diff --git a/crates/persisting-pchronicle/src/store/catalog/mod.rs b/crates/persisting-pchronicle/src/store/catalog/mod.rs index b9baa999..ada57a6e 100644 --- a/crates/persisting-pchronicle/src/store/catalog/mod.rs +++ b/crates/persisting-pchronicle/src/store/catalog/mod.rs @@ -17,7 +17,7 @@ use provider::*; use source::*; use discovery::{ - bind_canonical_storyline_projections, discover_candidates, freeze_candidate, + Candidate, bind_canonical_storyline_projections, discover_candidates, freeze_candidate, normalize_event_storylines, }; @@ -118,6 +118,8 @@ pub enum CatalogErrorPolicy { pub enum CatalogSourceKind { Store, File, + /// Navigational Directory child under a non-Dataset mount. Not queryable. + Directory, } #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)] @@ -228,12 +230,28 @@ impl CatalogDataset { pub fn ready_source_count(&self) -> usize { self.sources .iter() - .filter(|source| source.status == CatalogSourceStatus::Ready) + .filter(|source| { + source.status == CatalogSourceStatus::Ready + && source.kind != CatalogSourceKind::Directory + }) + .count() + } + + pub fn directory_count(&self) -> usize { + self.sources + .iter() + .filter(|source| source.kind == CatalogSourceKind::Directory) .count() } pub fn error_source_count(&self) -> usize { - self.sources.len().saturating_sub(self.ready_source_count()) + self.sources + .iter() + .filter(|source| { + source.status == CatalogSourceStatus::Error + && source.kind != CatalogSourceKind::Directory + }) + .count() } } @@ -351,6 +369,10 @@ impl DatasetCatalogSnapshot { let mut source_rows = Vec::with_capacity(candidates.len()); let mut prepared_sources = Vec::with_capacity(candidates.len()); for candidate in candidates { + if matches!(candidate, Candidate::Directory { .. }) { + source_rows.push(candidate.source_stub()); + continue; + } let stub = candidate.source_stub(); match freeze_candidate(&mount, candidate, temporary_files.clone(), options).await { Ok((source, lazy_source)) => { @@ -368,7 +390,11 @@ impl DatasetCatalogSnapshot { } } bind_canonical_storyline_projections(&mut source_rows, &mut prepared_sources)?; - source_rows.sort_by(|left, right| left.file.cmp(&right.file)); + source_rows.sort_by(|left, right| { + directory_sort_key(left.kind) + .cmp(&directory_sort_key(right.kind)) + .then_with(|| left.file.cmp(&right.file)) + }); prepared_sources.sort_by(|left, right| left.file().cmp(right.file())); datasets.push(CatalogDataset { mount: mount.clone(), @@ -807,6 +833,13 @@ impl DatasetCatalogSnapshot { } } +fn directory_sort_key(kind: CatalogSourceKind) -> u8 { + match kind { + CatalogSourceKind::Directory => 0, + CatalogSourceKind::Store | CatalogSourceKind::File => 1, + } +} + fn validate_catalog_options(options: CatalogSnapshotOptions) -> Result<()> { anyhow::ensure!( options.manifest.max_files > 0, diff --git a/crates/persisting-pchronicle/src/store/catalog/provider.rs b/crates/persisting-pchronicle/src/store/catalog/provider.rs index 05aa1ae0..5838f15e 100644 --- a/crates/persisting-pchronicle/src/store/catalog/provider.rs +++ b/crates/persisting-pchronicle/src/store/catalog/provider.rs @@ -595,6 +595,7 @@ pub(super) fn sources_table_provider( |source| match source.kind { CatalogSourceKind::Store => "store", CatalogSourceKind::File => "file", + CatalogSourceKind::Directory => "directory", }, ))), Arc::new(StringArray::from( diff --git a/crates/persisting-pchronicle/src/store/catalog/tests.rs b/crates/persisting-pchronicle/src/store/catalog/tests.rs index 462759e8..5c4f971e 100644 --- a/crates/persisting-pchronicle/src/store/catalog/tests.rs +++ b/crates/persisting-pchronicle/src/store/catalog/tests.rs @@ -475,6 +475,33 @@ async fn empty_dataset_still_exposes_the_stable_catalog_tables() -> Result<()> { Ok(()) } +#[tokio::test] +async fn directory_lists_child_dirs_and_dataset_sources_separately() -> Result<()> { + let temp = tempfile::tempdir()?; + let plain = temp.path().join("plain"); + fs::create_dir_all(&plain)?; + fs::write(plain.join("notes.txt"), "skip")?; + let story = temp.path().join("story"); + let store = StorylineLanceStore::open(&story).await?; + store + .replace_storyline(&storyline("session-story", "run-story")) + .await?; + let snapshot = DatasetCatalogSnapshot::discover( + vec![DatasetMount::default(temp.path().to_string_lossy())?], + Some(DEFAULT_DATASET_NAME.into()), + CatalogSnapshotOptions::default(), + ) + .await?; + let dataset = &snapshot.datasets()[0]; + assert_eq!(dataset.directory_count(), 1); + assert_eq!(dataset.ready_source_count(), 1); + assert_eq!(dataset.sources[0].kind, CatalogSourceKind::Directory); + assert_eq!(dataset.sources[0].file, "plain"); + assert_eq!(dataset.sources[1].kind, CatalogSourceKind::Store); + assert_eq!(dataset.sources[1].file, "story"); + Ok(()) +} + #[tokio::test] async fn catalog_prunes_file_sources_before_lazy_resolution() -> Result<()> { let temp = tempfile::tempdir()?; @@ -930,9 +957,10 @@ async fn canonical_event_source_exposes_and_loads_each_storyline_independently() panic!("initial catalog projection build unexpectedly reported nonempty output") }; + let mount_root = storage.join("agent"); let snapshot = Arc::new( DatasetCatalogSnapshot::discover( - vec![DatasetMount::default(storage.to_string_lossy())?], + vec![DatasetMount::default(mount_root.to_string_lossy())?], Some(DEFAULT_DATASET_NAME.into()), CatalogSnapshotOptions::default(), ) @@ -941,7 +969,7 @@ async fn canonical_event_source_exposes_and_loads_each_storyline_independently() assert_eq!(snapshot.datasets()[0].sources.len(), 1); assert_eq!( snapshot.datasets()[0].sources[0].file, - "agent/run-1/events.lance" + "run-1/events.lance" ); assert_eq!( snapshot.datasets()[0].sources[0].projection_status, @@ -977,7 +1005,7 @@ async fn canonical_event_source_exposes_and_loads_each_storyline_independently() .await?; let live_key = CatalogStorylineKey { dataset: DEFAULT_DATASET_NAME.into(), - file: "agent/run-1/events.lance".into(), + file: "run-1/events.lance".into(), document_id: "root".into(), session_id: "root".into(), }; @@ -999,7 +1027,7 @@ async fn canonical_event_source_exposes_and_loads_each_storyline_independently() let event_count = engine .query_jsonl( "SELECT COUNT(*) AS rows FROM dataset.events \ - WHERE _file_ = 'agent/run-1/events.lance' AND seq = 0", + WHERE _file_ = 'run-1/events.lance' AND seq = 0", ) .await?; assert_eq!(event_count.trim(), r#"{"rows":2}"#); @@ -1037,7 +1065,7 @@ async fn canonical_event_source_exposes_and_loads_each_storyline_independently() let stale_snapshot = Arc::new( DatasetCatalogSnapshot::discover( - vec![DatasetMount::default(storage.to_string_lossy())?], + vec![DatasetMount::default(mount_root.to_string_lossy())?], Some(DEFAULT_DATASET_NAME.into()), CatalogSnapshotOptions::default(), ) @@ -1059,7 +1087,7 @@ async fn canonical_event_source_exposes_and_loads_each_storyline_independently() stale_snapshot .load_events(&CatalogStorylineKey { dataset: DEFAULT_DATASET_NAME.into(), - file: "agent/run-1/events.lance".into(), + file: "run-1/events.lance".into(), document_id: "root".into(), session_id: "root".into(), }) @@ -1081,7 +1109,7 @@ async fn canonical_event_source_exposes_and_loads_each_storyline_independently() let limited_snapshot = Arc::new( DatasetCatalogSnapshot::discover( - vec![DatasetMount::default(storage.to_string_lossy())?], + vec![DatasetMount::default(mount_root.to_string_lossy())?], Some(DEFAULT_DATASET_NAME.into()), CatalogSnapshotOptions { max_event_fallback_rows: 1, @@ -1193,14 +1221,14 @@ async fn multiple_fresh_projections_choose_one_without_hiding_canonical_events() } let snapshot = DatasetCatalogSnapshot::discover( - vec![DatasetMount::default(storage.to_string_lossy())?], + vec![DatasetMount::default(storage.join("agent").to_string_lossy())?], Some(DEFAULT_DATASET_NAME.into()), CatalogSnapshotOptions::default(), ) .await?; assert_eq!(snapshot.datasets()[0].sources.len(), 1); let source = &snapshot.datasets()[0].sources[0]; - assert_eq!(source.file, "agent/run-1/events.lance"); + assert_eq!(source.file, "run-1/events.lance"); assert_eq!( source.projection_status, Some(CatalogProjectionStatus::Fresh) diff --git a/docs/src/en/rfcs/0015-chronicle-manifest.md b/docs/src/en/rfcs/0015-chronicle-manifest.md index 8e835cb6..498a5a72 100644 --- a/docs/src/en/rfcs/0015-chronicle-manifest.md +++ b/docs/src/en/rfcs/0015-chronicle-manifest.md @@ -90,11 +90,14 @@ Parents MUST NOT require an explicit children list. Discovery MUST: child that contains `chronicle.manifest`, treat that child as a nested Dataset node and continue according to that child's kind. 4. If the current directory has no `chronicle.manifest`, treat it as a - **Directory**: inspect **immediate** child directories only; classify each - child via Dataset markers (`chronicle.manifest`, `CURRENT`, - `events.lance/_manifest.json`, compact-jsonl Lance). Loose files MUST NOT - be registered as Sources. Discovery MUST NOT recursively list an entire - object-store prefix tree to classify. + **Directory**: inspect **immediate** child directories only. Children with + Dataset markers (`chronicle.manifest`, `CURRENT`, + `events.lance/_manifest.json`, compact-jsonl Lance) become queryable + Sources; other children become navigational entries (`kind = directory`, + visible to `ls`, not queryable). Loose files MUST NOT be registered as + Sources. Discovery MUST NOT recursively list an entire object-store prefix + tree to classify. **`import` / `sync` use a separate recursive JSON scan** + and are not bound by this Directory shallow rule. Symlinks MUST be ignored. Existing `max_entries` / `max_files` limits still apply to traversal. diff --git a/docs/src/zh/rfcs/0015-chronicle-manifest.md b/docs/src/zh/rfcs/0015-chronicle-manifest.md index a220077e..313aad2b 100644 --- a/docs/src/zh/rfcs/0015-chronicle-manifest.md +++ b/docs/src/zh/rfcs/0015-chronicle-manifest.md @@ -75,9 +75,11 @@ pChronicle 已对其它布局使用应用控制文件（Storyline 的 `CURRENT` 2. 若 `kind = "leaf"`，将该目录视为对应 `format` 的一个 source 候选，且 MUST NOT 再递归其内部寻找其它 source； 3. 若 `kind = "branch"`，只扫描**一层**子目录；对每个含有 `chronicle.manifest` 的子目录，按该子节点的 kind 继续处理； 4. 若当前目录没有 `chronicle.manifest`，则视为 **Directory**：只检查**一层** - 子目录；对每个子目录用 Dataset 标记（`chronicle.manifest`、`CURRENT`、 - `events.lance/_manifest.json`、compact-jsonl Lance）分类。松散文件 MUST NOT - 登记为 Source。MUST NOT 为发现而递归列举整棵对象前缀树。 + 子目录。子目录若含 Dataset 标记（`chronicle.manifest`、`CURRENT`、 + `events.lance/_manifest.json`、compact-jsonl Lance）则登记为可查询 Source； + 否则登记为导航项（`kind = directory`，`ls` 可见，不可 query）。松散文件 + MUST NOT 登记为 Source。MUST NOT 为发现而递归列举整棵对象前缀树。 + **`import` / `sync` 使用独立递归 JSON 扫描**，不受本条 Directory 浅层约束。 MUST 忽略符号链接。现有 `max_entries` / `max_files` 遍历上限仍然适用。 From d96b7fb69247de4320bc24677de6cfa10f876ed1 Mon Sep 17 00:00:00 2001 From: Reiase Date: Wed, 9 Sep 2026 01:10:21 +0800 Subject: [PATCH 4/8] feat(import): enhance object store import functionality and replace behavior Added support for writing and reading relative bytes in the DatasetLocation, enabling recursive listing of importable JSON objects. Improved the import process to clear existing prefixes in object stores before writing new data, ensuring a clean slate for imports. Updated documentation to clarify the behavior of the replace mode for object-store datasets, emphasizing that it clears the destination prefix before writing. Added tests to verify the new import behavior and ensure correct handling of existing data. --- .../persisting-pchronicle-cli/src/exchange.rs | 164 +++++++++++++--- crates/persisting-pchronicle-cli/src/lib.rs | 2 +- crates/persisting-pchronicle-cli/src/sync.rs | 168 ++++++++++++----- crates/persisting-pchronicle-cli/src/tests.rs | 107 +++++++++++ .../src/formats/actf/mod.rs | 26 ++- .../src/store/catalog/discovery.rs | 34 +++- .../src/store/catalog/mod.rs | 25 --- .../src/store/location.rs | 175 ++++++++++++++++++ docs/src/en/pchronicle/guides/exchange.md | 3 +- docs/src/zh/pchronicle/guides/exchange.md | 2 +- docs/src/zh/pchronicle/reference/cli.md | 2 +- 11 files changed, 604 insertions(+), 104 deletions(-) diff --git a/crates/persisting-pchronicle-cli/src/exchange.rs b/crates/persisting-pchronicle-cli/src/exchange.rs index 928e030c..1c0315ff 100644 --- a/crates/persisting-pchronicle-cli/src/exchange.rs +++ b/crates/persisting-pchronicle-cli/src/exchange.rs @@ -104,10 +104,6 @@ async fn prepare_import_destination( }) }; } - anyhow::ensure!( - !parsed.is_object_store(), - "replace mode for an existing object-store Dataset is unsupported; use a new URI" - ); let existing = parsed.into_existing()?; ensure_import_source_outside_destination(args, &existing)?; confirm_destructive_dataset( @@ -290,9 +286,14 @@ pub(super) async fn run_import( ) .await; } - let input_path = (!args.stream).then(|| Path::new(&args.from)); - let (directory_input, candidates) = if let Some(input_path) = input_path { - collect_import_candidates(input_path)? + let (directory_input, candidates) = if args.stream { + (false, Vec::new()) + } else if let Some(location) = &from_location { + if location.is_object_store() { + collect_object_store_import_candidates(location, stderr).await? + } else { + collect_import_candidates(Path::new(&args.from))? + } } else { (false, Vec::new()) }; @@ -348,10 +349,28 @@ pub(super) async fn run_import( ) } else if destination.is_object_store() { if destination.exists().await? { - return Err(cli_boundary_error( - BoundaryCode::Conflict, - "import output already exists", - )); + if replace_existing { + writeln!( + stderr, + "import to={} status=replacing", + destination.as_str() + ) + .context("write pChronicle import replace progress")?; + destination + .remove_all() + .await + .with_context(|| { + format!( + "remove existing object-store Dataset {}", + destination.as_str() + ) + })?; + } else { + return Err(cli_boundary_error( + BoundaryCode::Conflict, + "import output already exists", + )); + } } let store = StorylineLanceStore::open_uri(destination.as_str()) .await @@ -422,9 +441,11 @@ pub(super) async fn run_import( "processing", None, )?; - let file = std::fs::File::open(&candidate.path) - .with_context(|| format!("open {label}"))?; - let input = read_bounded(file, max_input_bytes, &label)?; + let input = read_import_candidate_bytes( + candidate, + max_input_bytes, + &label, + )?; if let Some(source) = stage_preserved_import_source( args.format, Some(&candidate.path), @@ -633,9 +654,9 @@ async fn run_compact_jsonl_import( /// directory. Keeping the orchestration here avoids a second decoder or /// Dataset publication protocol in the sync command. pub(crate) async fn sync_snapshot( - source: &Path, - warehouse: &Path, - storyline: &Path, + source: &str, + warehouse: &str, + storyline: &str, input_format: ExchangeFormat, columns: &[String], ) -> Result<()> { @@ -644,8 +665,8 @@ pub(crate) async fn sync_snapshot( let mut stderr = std::io::sink(); return run_compact_jsonl_import( ImportArgs { - from: source.to_string_lossy().into_owned(), - output: Some(storyline.to_string_lossy().into_owned()), + from: source.to_owned(), + output: Some(storyline.to_owned()), format: ExchangeFormat::CompactJsonl, output_format: Some(ImportOutputFormat::CompactJsonl), mode: ImportMode::Replace, @@ -655,7 +676,7 @@ pub(crate) async fn sync_snapshot( max_input_bytes: Some(256 * 1024 * 1024), columns: columns.to_vec(), }, - storyline.to_string_lossy().as_ref(), + storyline, &mut stdout, &mut stderr, ) @@ -668,8 +689,8 @@ pub(crate) async fn sync_snapshot( let mut stdin = std::io::empty(); run_import( ImportArgs { - from: source.to_string_lossy().into_owned(), - output: Some(warehouse.to_string_lossy().into_owned()), + from: source.to_owned(), + output: Some(warehouse.to_owned()), format: input_format, output_format: Some(ImportOutputFormat::Preserve), mode: ImportMode::Replace, @@ -689,8 +710,8 @@ pub(crate) async fn sync_snapshot( .context("sync source into Warehouse")?; run_import( ImportArgs { - from: source.to_string_lossy().into_owned(), - output: Some(storyline.to_string_lossy().into_owned()), + from: source.to_owned(), + output: Some(storyline.to_owned()), format: input_format, output_format: Some(ImportOutputFormat::Storyline), mode: ImportMode::Replace, @@ -1301,6 +1322,9 @@ struct ImportFileCandidate { path: PathBuf, relative_path: PathBuf, output_relative_path: Option, + /// Object-store imports preload file bytes so the sync decode loop can + /// stay synchronous. Local imports leave this empty and open `path`. + content: Option>, } #[derive(Debug)] @@ -1354,6 +1378,7 @@ fn collect_import_candidates(input: &Path) -> Result<(bool, Vec Result<(bool, Vec Result> { if file_type.is_dir() { pending.push(path); } else if file_type.is_file() && is_visible_json_file(&path) { + let relative = path + .strip_prefix(root) + .unwrap_or(path.as_path()) + .to_string_lossy() + .replace('\\', "/"); + if relative.split('/').any(|part| part == "_meta") { + continue; + } files.push(path); } } @@ -1424,6 +1458,68 @@ fn is_visible_json_file(path: &Path) -> bool { }) } +async fn collect_object_store_import_candidates( + location: &DatasetLocation, + stderr: &mut dyn Write, +) -> Result<(bool, Vec)> { + writeln!( + stderr, + "import from={} status=discovering", + location.as_str() + ) + .context("write pChronicle import discovery progress")?; + let keys = location + .list_importable_json_objects(persisting_pchronicle::storage::DEFAULT_MAX_LOCAL_QUERY_FILES) + .await + .with_context(|| format!("discover importable objects under {}", location.as_str()))?; + if keys.is_empty() { + return Err(cli_boundary_error( + BoundaryCode::InvalidRequest, + "import object prefix contains no .json, .jsonl, or .ndjson files", + )); + } + writeln!( + stderr, + "import from={} status=discovered files={}", + location.as_str(), + keys.len() + ) + .context("write pChronicle import discovery progress")?; + + let mut candidates = Vec::with_capacity(keys.len()); + for key in keys { + let relative_path = PathBuf::from(&key); + write_import_progress(stderr, &key, "fetching", None)?; + let content = location + .read_relative_bytes(&key) + .await + .with_context(|| format!("read import object {key} under {}", location.as_str()))?; + candidates.push(ImportFileCandidate { + path: relative_path.clone(), + output_relative_path: Some(relative_path.clone()), + relative_path, + content: Some(content), + }); + } + Ok((true, candidates)) +} + +fn read_import_candidate_bytes( + candidate: &ImportFileCandidate, + max_input_bytes: usize, + label: &str, +) -> Result> { + if let Some(content) = &candidate.content { + anyhow::ensure!( + content.len() <= max_input_bytes, + "{label} exceeds max_input_bytes limit of {max_input_bytes}" + ); + return Ok(content.clone()); + } + let file = std::fs::File::open(&candidate.path).with_context(|| format!("open {label}"))?; + read_bounded(file, max_input_bytes, label) +} + fn scope_import_source_error(error: anyhow::Error, source_path: &Path) -> anyhow::Error { if let Some(boundary) = error.downcast_ref::() { return cli_boundary_error( @@ -1557,9 +1653,11 @@ impl<'a> StorylineImportIterator<'a> { "processing", None, )?; - let file = std::fs::File::open(&candidate.path) - .with_context(|| format!("open {label}"))?; - let input = read_bounded(file, self.max_input_bytes, &label)?; + let input = read_import_candidate_bytes( + candidate, + self.max_input_bytes, + &label, + )?; decode_import_source( self.requested_format, ImportOutputFormat::Storyline, @@ -1748,7 +1846,17 @@ fn decode_import_source( code, import_input_issue_message(&issue, decode_relative_path), ) - })?; + }); + let storylines = match storylines { + Ok(storylines) => storylines, + Err(error) if allow_skip => { + return Ok(DecodeImportOutcome::Skipped { + path: diagnostic_path, + reason: error.to_string(), + }); + } + Err(error) => return Err(error), + }; unknown_field_warnings .observe_storylines(&storylines) .map_err(|issue| { diff --git a/crates/persisting-pchronicle-cli/src/lib.rs b/crates/persisting-pchronicle-cli/src/lib.rs index 6930e4a8..71ebd85f 100644 --- a/crates/persisting-pchronicle-cli/src/lib.rs +++ b/crates/persisting-pchronicle-cli/src/lib.rs @@ -1600,7 +1600,7 @@ pub async fn run_with_stdio( .await } Command::Export(args) => run_export(args, config, stdout, &mut diagnostics).await, - Command::Sync(args) => sync::run(args, &mut diagnostics).await, + Command::Sync(args) => sync::run(args, config, &mut diagnostics).await, Command::Echo(args) => run_echo(args, &mut diagnostics).await, Command::Dev(DevArgs { command: DevCommand::Echo(args), diff --git a/crates/persisting-pchronicle-cli/src/sync.rs b/crates/persisting-pchronicle-cli/src/sync.rs index d2e94792..3b3f0ae6 100644 --- a/crates/persisting-pchronicle-cli/src/sync.rs +++ b/crates/persisting-pchronicle-cli/src/sync.rs @@ -5,20 +5,21 @@ use std::fs; use std::time::{Duration, SystemTime}; use clap::Args; +use persisting_pchronicle::storage::DatasetLocation; #[derive(Debug, Args)] pub(crate) struct SyncArgs { - /// Source directory to mirror. - #[arg(long, value_name = "DIRECTORY")] - pub(crate) from: PathBuf, + /// Source Dataset path, URI, or pin (for example `@origin/agentcompass`). + #[arg(long, value_name = "DATASET")] + pub(crate) from: String, - /// Local Warehouse Dataset receiving source files; unused for compact-jsonl. - #[arg(long = "to", alias = "warehouse", value_name = "DIRECTORY")] - pub(crate) to: PathBuf, + /// Warehouse Dataset receiving source files; unused for compact-jsonl. + #[arg(long = "to", alias = "warehouse", value_name = "DATASET")] + pub(crate) to: String, - /// Local Storyline or compact JSONL Lance Dataset receiving each snapshot. - #[arg(long = "convert", alias = "storyline", value_name = "DIRECTORY")] - pub(crate) convert: PathBuf, + /// Storyline or compact JSONL Lance Dataset receiving each snapshot. + #[arg(long = "convert", alias = "storyline", value_name = "DATASET")] + pub(crate) convert: String, /// Input format. compact-jsonl requires a tree of .jsonl files. #[arg(long = "input-format", value_enum, default_value_t = ExchangeFormat::Auto)] @@ -28,7 +29,7 @@ pub(crate) struct SyncArgs { #[arg(long = "column", value_name = "NAME=JSON_PATH", action = clap::ArgAction::Append)] pub(crate) columns: Vec, - /// Polling and update interval. Supports ms, s, m, and h. + /// Polling and update interval. Supports ms, s, and h. #[arg(long = "interval", value_name = "DURATION", value_parser = super::parse_duration_seconds, default_value = "1s")] pub(crate) interval_seconds: u64, @@ -43,29 +44,44 @@ struct FileStamp { modified: Option, } -pub(crate) async fn run(args: SyncArgs, stderr: &mut dyn Write) -> Result<()> { - let source = fs::canonicalize(&args.from) - .with_context(|| format!("canonicalize sync source {}", args.from.display()))?; - anyhow::ensure!(source.is_dir(), "sync source must be a directory"); - let warehouse = prepare_target(&args.to, "Warehouse")?; - let storyline = prepare_target(&args.convert, "conversion")?; - anyhow::ensure!(warehouse != storyline, "sync targets must be different"); +pub(crate) async fn run( + args: SyncArgs, + settings_override: Option<&Path>, + stderr: &mut dyn Write, +) -> Result<()> { + let source_uri = expand_dataset_reference(&args.from, settings_override, true) + .with_context(|| format!("resolve sync source '{}'", args.from))?; + let warehouse_uri = expand_dataset_reference(&args.to, settings_override, false) + .with_context(|| format!("resolve sync Warehouse '{}'", args.to))?; + let convert_uri = expand_dataset_reference(&args.convert, settings_override, false) + .with_context(|| format!("resolve sync convert '{}'", args.convert))?; + + let warehouse_uri = prepare_destination(&warehouse_uri, "Warehouse")?; + let convert_uri = prepare_destination(&convert_uri, "conversion")?; anyhow::ensure!( - !warehouse.starts_with(&source) && !storyline.starts_with(&source), - "sync targets must be outside the source directory" + warehouse_uri != convert_uri, + "sync targets must be different" ); + ensure_targets_outside_source(&source_uri, &warehouse_uri, &convert_uri)?; + + writeln!( + stderr, + "sync from={} to={} convert={}", + source_uri, warehouse_uri, convert_uri + ) + .context("write sync resolved targets")?; let interval = Duration::from_secs(args.interval_seconds.max(1)); if args.once { - let initial = scan_files(&source)?; + let initial = scan_source(&source_uri).await?; anyhow::ensure!( !initial.is_empty(), "sync source contains no supported JSON files" ); super::exchange::sync_snapshot( - &source, - &warehouse, - &storyline, + &source_uri, + &warehouse_uri, + &convert_uri, args.input_format, &args.columns, ) @@ -76,13 +92,13 @@ pub(crate) async fn run(args: SyncArgs, stderr: &mut dyn Write) -> Result<()> { } let (changes_tx, mut changes_rx) = tokio::sync::mpsc::channel::(1024); - let watcher_source = source.clone(); + let watcher_source = source_uri.clone(); let watcher = tokio::spawn(async move { let mut previous = BTreeMap::new(); loop { // ponytail: dependency-free polling; use an OS watcher when tree size or latency // makes recursive scans measurable. - let current = scan_files(&watcher_source)?; + let current = scan_source(&watcher_source).await?; for path in changed_paths(&previous, ¤t) { if changes_tx.send(path).await.is_err() { return Ok::<(), anyhow::Error>(()); @@ -106,9 +122,9 @@ pub(crate) async fn run(args: SyncArgs, stderr: &mut dyn Write) -> Result<()> { } match super::exchange::sync_snapshot( - &source, - &warehouse, - &storyline, + &source_uri, + &warehouse_uri, + &convert_uri, args.input_format, &args.columns, ) @@ -159,7 +175,12 @@ pub(crate) async fn run(args: SyncArgs, stderr: &mut dyn Write) -> Result<()> { } } -fn prepare_target(path: &Path, name: &str) -> Result { +fn prepare_destination(uri: &str, name: &str) -> Result { + anyhow::ensure!(!uri.is_empty(), "sync {name} target must not be empty"); + let location = DatasetLocation::parse(uri)?; + let Some(path) = location.local_path() else { + return Ok(location.as_str().to_owned()); + }; anyhow::ensure!( !path.as_os_str().is_empty(), "sync {name} target must not be empty" @@ -170,25 +191,75 @@ fn prepare_target(path: &Path, name: &str) -> Result { let filename = path .file_name() .with_context(|| format!("sync {name} target must name a directory"))?; - Ok(parent.join(filename)) + Ok(parent.join(filename).to_string_lossy().into_owned()) } -fn scan_files(root: &Path) -> Result> { +fn ensure_targets_outside_source(source: &str, warehouse: &str, convert: &str) -> Result<()> { + let source = DatasetLocation::parse(source)?; + let warehouse = DatasetLocation::parse(warehouse)?; + let convert = DatasetLocation::parse(convert)?; + let Some(source_path) = source.local_path() else { + return Ok(()); + }; + if let Some(warehouse_path) = warehouse.local_path() { + anyhow::ensure!( + !warehouse_path.starts_with(source_path), + "sync Warehouse target must be outside the source directory" + ); + } + if let Some(convert_path) = convert.local_path() { + anyhow::ensure!( + !convert_path.starts_with(source_path), + "sync conversion target must be outside the source directory" + ); + } + Ok(()) +} + +async fn scan_source(uri: &str) -> Result> { + let location = DatasetLocation::parse(uri)?; + if let Some(root) = location.local_path() { + anyhow::ensure!(root.is_dir(), "sync source must be a directory"); + let mut files = BTreeMap::new(); + for path in crate::exchange::collect_visible_json_files(root)? { + let metadata = + fs::metadata(&path).with_context(|| format!("stat sync file {}", path.display()))?; + files.insert( + path.strip_prefix(root)?.to_path_buf(), + FileStamp { + size: metadata.len(), + modified: metadata.modified().ok(), + }, + ); + } + return Ok(files); + } + + let stamps = location + .list_importable_json_object_stamps( + persisting_pchronicle::storage::DEFAULT_MAX_LOCAL_QUERY_FILES, + ) + .await + .with_context(|| format!("list sync source objects under {uri}"))?; let mut files = BTreeMap::new(); - for path in crate::exchange::collect_visible_json_files(root)? { - let metadata = fs::metadata(&path) - .with_context(|| format!("stat sync file {}", path.display()))?; + for (key, size, modified) in stamps { files.insert( - path.strip_prefix(root)?.to_path_buf(), + PathBuf::from(key), FileStamp { - size: metadata.len(), - modified: metadata.modified().ok(), + size, + modified: modified.and_then(parse_rfc3339_system_time), }, ); } Ok(files) } +fn parse_rfc3339_system_time(value: String) -> Option { + chrono::DateTime::parse_from_rfc3339(&value) + .ok() + .map(|value| SystemTime::UNIX_EPOCH + Duration::from_secs(value.timestamp().max(0) as u64)) +} + fn changed_paths( previous: &BTreeMap, current: &BTreeMap, @@ -228,36 +299,45 @@ mod tests { } #[tokio::test] - async fn once_mirrors_files_and_builds_storyline() -> Result<()> { + async fn sync_once_rebuilds_warehouse_and_storyline() -> Result<()> { let temporary = tempfile::tempdir()?; let source = temporary.path().join("source"); - let warehouse = temporary.path().join("warehouse"); - let storyline = temporary.path().join("storyline"); fs::create_dir_all(&source)?; fs::copy( Path::new(env!("CARGO_MANIFEST_DIR")).join("assets/onboard/support-ticket.json"), source.join("support-ticket.json"), )?; let source_bytes = fs::read(source.join("support-ticket.json"))?; + let mut stderr = Vec::new(); run( SyncArgs { - from: source, - to: warehouse, - convert: storyline.clone(), + from: source.to_string_lossy().into_owned(), + to: temporary + .path() + .join("warehouse") + .to_string_lossy() + .into_owned(), + convert: temporary + .path() + .join("storyline") + .to_string_lossy() + .into_owned(), input_format: ExchangeFormat::Auto, columns: Vec::new(), interval_seconds: 1, once: true, }, + None, &mut stderr, ) .await?; + assert_eq!( fs::read(temporary.path().join("warehouse/support-ticket.json"))?, source_bytes ); - assert!(storyline.join("CURRENT").is_file()); + assert!(temporary.path().join("storyline/CURRENT").is_file()); Ok(()) } } diff --git a/crates/persisting-pchronicle-cli/src/tests.rs b/crates/persisting-pchronicle-cli/src/tests.rs index 925c14d5..5b8df1d4 100644 --- a/crates/persisting-pchronicle-cli/src/tests.rs +++ b/crates/persisting-pchronicle-cli/src/tests.rs @@ -2755,6 +2755,60 @@ async fn import_storyline_output_writes_one_root_lance_store() -> Result<()> { Ok(()) } +#[tokio::test] +async fn object_store_replace_clears_existing_prefix_before_import() -> Result<()> { + let source = format!( + "shared-memory://pchronicle-object-replace-src-{}/corpus", + uuid::Uuid::new_v4().simple() + ); + let output = format!( + "shared-memory://pchronicle-object-replace-dst-{}/dataset", + uuid::Uuid::new_v4().simple() + ); + let input = DatasetLocation::parse(&source)?; + input + .write_relative_bytes( + "run.json", + &serde_json::to_vec(&atif_identity_document("document-new", "session-new"))?, + ) + .await?; + + // Seed an existing destination so replace must clear it. + let existing = DatasetLocation::parse(&output)?; + existing + .write_relative_bytes(".dataset-marker", b"old") + .await?; + assert!(existing.exists().await?); + + let cli = Cli::try_parse_from([ + "pchronicle", + "import", + "--from", + &source, + "--to", + &output, + "--output-format", + "storyline", + "--mode", + "replace", + "--yes", + ])?; + let mut stdout = Vec::new(); + let mut stderr = Vec::new(); + run(cli, false, &mut stdout, &mut stderr).await?; + let stderr = String::from_utf8(stderr)?; + assert!(stderr.contains("status=replacing")); + + let store = StorylineLanceStore::open_uri(&output).await?; + let ids = store + .document_ids_snapshot() + .await? + .context("replaced storyline snapshot")? + .1; + assert!(ids.iter().any(|id| id == "document-new")); + Ok(()) +} + #[tokio::test] async fn import_object_store_output_requires_storyline_format() -> Result<()> { let temp = tempfile::tempdir()?; @@ -2947,6 +3001,59 @@ async fn canonical_event_import_auto_detects_and_is_create_only() -> Result<()> Ok(()) } +#[tokio::test] +async fn object_store_directory_import_recurses_json_files() -> Result<()> { + let source = format!( + "shared-memory://pchronicle-object-import-{}/corpus", + uuid::Uuid::new_v4().simple() + ); + let location = DatasetLocation::parse(&source)?; + location + .write_relative_bytes( + "nested/run-a.json", + &serde_json::to_vec(&atif_identity_document("document-a", "session-a"))?, + ) + .await?; + location + .write_relative_bytes( + "nested/deeper/run-b.jsonl", + &serde_json::to_vec(&atif_identity_document("document-b", "session-b"))?, + ) + .await?; + // Lance interiors must be ignored even when they contain .json names. + location + .write_relative_bytes("keep/events.lance/_manifest.json", b"{\"not\":\"importable\"}") + .await?; + + let output = tempfile::tempdir()?; + let cli = Cli::try_parse_from([ + "pchronicle", + "import", + "--from", + &source, + "--to", + output.path().to_str().unwrap(), + "--output-format", + "storyline", + ])?; + let mut stdout = Vec::new(); + let mut stderr = Vec::new(); + run(cli, false, &mut stdout, &mut stderr).await?; + let stderr = String::from_utf8(stderr)?; + assert!(stderr.contains("status=discovering")); + assert!(stderr.contains("status=discovered files=2")); + + let store = StorylineLanceStore::open(output.path()).await?; + let ids = store + .document_ids_snapshot() + .await? + .context("imported storyline snapshot")? + .1; + assert!(ids.iter().any(|id| id == "document-a")); + assert!(ids.iter().any(|id| id == "document-b")); + Ok(()) +} + #[tokio::test] async fn canonical_event_import_supports_object_store_uris() -> Result<()> { let temp = tempfile::tempdir()?; diff --git a/crates/persisting-pchronicle/src/formats/actf/mod.rs b/crates/persisting-pchronicle/src/formats/actf/mod.rs index 4231d5e9..c77c6e4e 100644 --- a/crates/persisting-pchronicle/src/formats/actf/mod.rs +++ b/crates/persisting-pchronicle/src/formats/actf/mod.rs @@ -149,8 +149,9 @@ fn decode_json( let mut value: Value = serde_json::from_str(&input).map_err(|error| InputIssue::invalid(error.to_string()))?; let envelope = take_unknown_fields_envelope(&mut value)?; - let document: ActfDocument = + let mut document: ActfDocument = serde_json::from_value(value).map_err(|error| InputIssue::invalid(error.to_string()))?; + normalize_solved_at(&mut document.solved_at); document.validate()?; let mut stories = actf_to_storylines(&document).map_err(|error| InputIssue::invalid(error.to_string()))?; @@ -423,11 +424,24 @@ where Ok(Option::::deserialize(deserializer)?.unwrap_or_default()) } +/// Corpus exporters sometimes emit unix timestamps or booleans for `solved_at`. +/// Coerce scalars into the documented string-or-null shape before validate. +fn normalize_solved_at(value: &mut Value) { + match value { + Value::Null | Value::String(_) => {} + Value::Number(number) => *value = Value::String(number.to_string()), + Value::Bool(false) => *value = Value::Null, + Value::Bool(true) => *value = Value::String("true".into()), + _ => {} + } +} + impl ActfDocument { #[cfg(any(test, feature = "lance-store"))] pub fn from_json_str(input: &str) -> InputResult { - let document: Self = + let mut document: Self = serde_json::from_str(input).map_err(|error| InputIssue::invalid(error.to_string()))?; + normalize_solved_at(&mut document.solved_at); document.validate()?; Ok(document) } @@ -636,6 +650,14 @@ mod tests { .unwrap() } + #[test] + fn accepts_numeric_solved_at_by_coercing_to_string() { + let mut value = serde_json::to_value(fixture()).unwrap(); + value["solved_at"] = json!(1_714_000_000); + let document = ActfDocument::from_json_str(&value.to_string()).unwrap(); + assert_eq!(document.solved_at, json!("1714000000")); + } + #[test] fn accepts_name_arguments_tool_without_type_or_id() { let mut value = serde_json::to_value(fixture()).unwrap(); diff --git a/crates/persisting-pchronicle/src/store/catalog/discovery.rs b/crates/persisting-pchronicle/src/store/catalog/discovery.rs index 01f35c3a..7d077796 100644 --- a/crates/persisting-pchronicle/src/store/catalog/discovery.rs +++ b/crates/persisting-pchronicle/src/store/catalog/discovery.rs @@ -654,6 +654,16 @@ async fn discover_local_candidates( size_bytes: Some(metadata.len()), last_modified: modified_string(&metadata), }); + let storyline = path.join("storyline"); + if storyline.join("CURRENT").is_file() { + let metadata = fs::metadata(storyline.join("CURRENT"))?; + candidates.push(Candidate::Storyline { + file: relative_catalog_path(root, &storyline, true)?, + uri: canonical_local_uri(&storyline)?, + size_bytes: Some(metadata.len()), + last_modified: modified_string(&metadata), + }); + } } else if is_lance_directory(&path) { if is_compact_jsonl_directory(&path).await? { let metadata = fs::metadata(&path)?; @@ -809,7 +819,29 @@ async fn discover_object_candidates( options.max_files ); match probe_object_prefix(&store, uri, &child, root_source_path(&child)).await? { - Some(ObjectProbe::Source(candidate)) => candidates.push(candidate), + Some(ObjectProbe::Source(candidate)) => { + let maybe_storyline = match &candidate { + Candidate::Events { file, .. } if file.ends_with("/events.lance") => { + let parent = file.trim_end_matches("/events.lance"); + let storyline_rel = format!("{parent}/storyline"); + probe_object_prefix( + &store, + uri, + &storyline_rel, + root_source_path(&storyline_rel), + ) + .await? + } + Candidate::Events { file, .. } if file == "events.lance" => { + probe_object_prefix(&store, uri, "storyline", "storyline").await? + } + _ => None, + }; + candidates.push(candidate); + if let Some(ObjectProbe::Source(storyline)) = maybe_storyline { + candidates.push(storyline); + } + } Some(ObjectProbe::Branch) => { let nested = collect_object_branch_children(&store, uri, &child, options).await?; candidates.extend(nested); diff --git a/crates/persisting-pchronicle/src/store/catalog/mod.rs b/crates/persisting-pchronicle/src/store/catalog/mod.rs index ada57a6e..865a73c3 100644 --- a/crates/persisting-pchronicle/src/store/catalog/mod.rs +++ b/crates/persisting-pchronicle/src/store/catalog/mod.rs @@ -900,18 +900,6 @@ fn is_lance_directory(path: &Path) -> bool { || path.join("_versions").is_dir() } -fn path_is_inside_lance_directory(path: &str) -> bool { - Path::new(path) - .components() - .any(|component| match component { - std::path::Component::Normal(name) => Path::new(name) - .extension() - .and_then(|extension| extension.to_str()) - .is_some_and(|extension| extension.eq_ignore_ascii_case("lance")), - _ => false, - }) -} - fn relative_catalog_path(root: &Path, path: &Path, allow_root: bool) -> Result { let relative = path .strip_prefix(root) @@ -978,13 +966,6 @@ fn remote_source_revision(meta: &RemoteObjectMeta) -> CatalogSourceRevision { } } -fn parent_relative_path(path: &str, leaf: &str) -> String { - path.strip_suffix(leaf) - .unwrap_or(path) - .trim_end_matches('/') - .to_string() -} - fn root_source_path(relative: &str) -> String { if relative.is_empty() { ".".into() @@ -1001,12 +982,6 @@ fn child_uri(root: &str, relative: &str) -> String { } } -fn is_nested_in_any<'a>(path: &str, roots: impl Iterator) -> bool { - roots - .into_iter() - .any(|root| root.is_empty() || path == root || path.starts_with(&format!("{root}/"))) -} - fn catalog_snapshot_id(datasets: &[CatalogDataset]) -> String { let mut hasher = blake3::Hasher::new(); for dataset in datasets { diff --git a/crates/persisting-pchronicle/src/store/location.rs b/crates/persisting-pchronicle/src/store/location.rs index 200bdcc5..7645cd3f 100644 --- a/crates/persisting-pchronicle/src/store/location.rs +++ b/crates/persisting-pchronicle/src/store/location.rs @@ -167,6 +167,127 @@ impl DatasetLocation { store.exists().await } + /// Write `bytes` at a relative object key (or local path under this Dataset). + pub async fn write_relative_bytes(&self, relative: &str, bytes: &[u8]) -> Result<()> { + let relative = relative.trim_start_matches('/'); + anyhow::ensure!(!relative.is_empty(), "relative object path must not be empty"); + anyhow::ensure!( + !relative.split('/').any(|part| part == ".."), + "relative object path must not contain '..'" + ); + if let Some(root) = &self.local_path { + let path = root.join(relative); + if let Some(parent) = path.parent() { + std::fs::create_dir_all(parent) + .with_context(|| format!("create {}", parent.display()))?; + } + return put_local_bytes(&path, bytes, true); + } + let store = OpendalStore::from_uri(&self.uri).await?; + store + .write_overwrite(relative, bytes.to_vec()) + .await + .with_context(|| format!("write object {} under {}", relative, self.uri)) + } + + /// Read bytes at a relative object key (or local path under this Dataset). + pub async fn read_relative_bytes(&self, relative: &str) -> Result> { + let relative = relative.trim_start_matches('/'); + anyhow::ensure!(!relative.is_empty(), "relative object path must not be empty"); + if let Some(root) = &self.local_path { + let path = root.join(relative); + return std::fs::read(&path) + .with_context(|| format!("read {}", path.display())); + } + let store = OpendalStore::from_uri(&self.uri).await?; + let Some((bytes, _)) = store.read(relative).await? else { + return Err(anyhow!("object not found: {relative} under {}", self.uri)); + }; + Ok(bytes) + } + + /// Recursively list importable `.json` / `.jsonl` / `.ndjson` object keys. + /// Skips Lance table interiors (any path segment ending in `.lance`). + pub async fn list_importable_json_objects(&self, max_files: usize) -> Result> { + Ok(self + .list_importable_json_object_stamps(max_files) + .await? + .into_iter() + .map(|(key, _, _)| key) + .collect()) + } + + /// Like [`Self::list_importable_json_objects`], but also returns size and + /// last-modified metadata for change detection (`sync`). + pub async fn list_importable_json_object_stamps( + &self, + max_files: usize, + ) -> Result)>> { + anyhow::ensure!(max_files > 0, "import max_files must be positive"); + if let Some(root) = &self.local_path { + let paths = list_local_importable_json_files(root)?; + anyhow::ensure!( + paths.len() <= max_files, + "import input exceeds max_files limit of {max_files}" + ); + let mut stamps = Vec::with_capacity(paths.len()); + for path in paths { + let relative = path + .strip_prefix(root) + .context("derive Dataset-relative import source path")? + .to_string_lossy() + .replace('\\', "/"); + let metadata = std::fs::metadata(&path) + .with_context(|| format!("stat importable file {}", path.display()))?; + stamps.push(( + relative, + metadata.len(), + metadata.modified().ok().and_then(|modified| { + modified + .duration_since(std::time::UNIX_EPOCH) + .ok() + .map(|duration| { + chrono::DateTime::::from_timestamp( + duration.as_secs() as i64, + duration.subsec_nanos(), + ) + .map(|value| value.to_rfc3339()) + }) + .flatten() + }), + )); + } + return Ok(stamps); + } + + let store = OpendalStore::from_uri(&self.uri).await?; + let entries = store + .list("") + .await + .with_context(|| format!("list importable objects under {}", self.uri))?; + let mut stamps = Vec::new(); + for entry in entries { + let key = entry.path.trim_matches('/').to_string(); + if key.is_empty() || !is_importable_json_object_key(&key) { + continue; + } + anyhow::ensure!( + stamps.len() < max_files, + "import input exceeds max_files limit of {max_files}" + ); + stamps.push(( + key, + entry.metadata.content_length(), + entry + .metadata + .last_modified() + .map(|value| value.to_string()), + )); + } + stamps.sort_by(|left, right| left.0.cmp(&right.0)); + Ok(stamps) + } + pub async fn put_bytes(&self, bytes: &[u8], overwrite: bool) -> Result<()> { if let Some(path) = &self.local_path { return put_local_bytes(path, bytes, overwrite); @@ -208,6 +329,60 @@ impl DatasetLocation { } } +fn is_importable_json_object_key(key: &str) -> bool { + if key.split('/').any(|part| part == "_meta" || part.ends_with(".lance")) { + return false; + } + Path::new(key) + .extension() + .and_then(|extension| extension.to_str()) + .is_some_and(|extension| { + matches!( + extension.to_ascii_lowercase().as_str(), + "json" | "jsonl" | "ndjson" + ) + }) +} + +fn list_local_importable_json_files(root: &Path) -> Result> { + let mut pending = vec![root.to_path_buf()]; + let mut files = Vec::new(); + while let Some(directory) = pending.pop() { + let mut entries = std::fs::read_dir(&directory) + .with_context(|| format!("read directory {}", directory.display()))? + .collect::>>()?; + entries.sort_by_key(std::fs::DirEntry::path); + for entry in entries { + let file_type = entry.file_type()?; + if file_type.is_symlink() { + continue; + } + let path = entry.path(); + if file_type.is_dir() { + if path + .file_name() + .and_then(|name| name.to_str()) + .is_some_and(|name| name.ends_with(".lance")) + { + continue; + } + pending.push(path); + } else if file_type.is_file() { + let relative = path + .strip_prefix(root) + .unwrap_or(path.as_path()) + .to_string_lossy() + .replace('\\', "/"); + if is_importable_json_object_key(&relative) { + files.push(path); + } + } + } + } + files.sort(); + Ok(files) +} + fn validate_object_store_bucket(scheme: &str, bucket: &str) -> Result<()> { if matches!(scheme, "memory" | "shared-memory") { return Ok(()); diff --git a/docs/src/en/pchronicle/guides/exchange.md b/docs/src/en/pchronicle/guides/exchange.md index 3a622e84..13b28abc 100644 --- a/docs/src/en/pchronicle/guides/exchange.md +++ b/docs/src/en/pchronicle/guides/exchange.md @@ -28,7 +28,8 @@ for an existing Storyline Dataset; duplicate `document_id` values receive a `#N` suffix by default, or can be skipped with `--on-duplicate skip`. Use `--mode replace` to stage the complete import and atomically replace an existing local Dataset after confirmation; replacement requires interactive confirmation -or `--yes`. Existing object-store Datasets cannot currently be replaced in place. +or `--yes`. Object-store Dataset replace clears the destination prefix before writing +(not atomic; an interrupted replace may leave the target empty). Regular files can be auto-detected. A directory recursively imports `.json`, `.jsonl`, and `.ndjson` files while preserving their relative paths in the default output. When `--input-format` is diff --git a/docs/src/zh/pchronicle/guides/exchange.md b/docs/src/zh/pchronicle/guides/exchange.md index 3fa1ba7b..4ae26f87 100644 --- a/docs/src/zh/pchronicle/guides/exchange.md +++ b/docs/src/zh/pchronicle/guides/exchange.md @@ -22,7 +22,7 @@ pchronicle import --from input.json \ 默认 `--mode create` 会拒绝已有目标。`--mode append` 用于已有 Storyline Dataset；重复 `document_id` 默认增加 `#N` 后缀，也可用 `--on-duplicate skip` 跳过。`--mode replace` 会先 把完整导入写入临时路径，确认后以 rename 事务替换已有的本地 Dataset，最后才删除旧数据；要求 -交互确认或 `--yes`。已有对象存储 Dataset 当前不支持原地 replace。普通文件可以自动识别。目录输入会递归扫描 +交互确认或传入 `--yes`。对象存储 Dataset 的 replace 会先清空目标前缀再写入（非原子；中断可能导致目标暂时为空）。普通文件可以自动识别。目录输入会递归扫描 `.json`、`.jsonl` 与 `.ndjson` 文件；默认输出会保留其相对 路径。未指定 `--input-format` 时按文件分别探测类型；无法识别为运行数据格式的 JSON 会跳过并警告： diff --git a/docs/src/zh/pchronicle/reference/cli.md b/docs/src/zh/pchronicle/reference/cli.md index 9972ed5f..d9b1228c 100644 --- a/docs/src/zh/pchronicle/reference/cli.md +++ b/docs/src/zh/pchronicle/reference/cli.md @@ -312,7 +312,7 @@ Codex 和 Claude Code session 是 decode-only 输入格式。Canonical Event Sto Storyline Dataset。默认 `create` 模式要求目标不存在。`append` 要求目标是已有 Storyline Dataset； 重复 `document_id` 默认增加 `#N` 后缀，也可用 `--on-duplicate skip` 跳过。`replace` 会先将完整导入 写入临时路径，再将旧本地 Dataset rename 到备份路径、将新 Dataset rename 到正式路径，确认新路径 -发布后才删除备份；因此必须交互确认或传入 `--yes`。已有对象存储 Dataset 当前不支持原地 replace。 +发布后才删除备份；因此必须交互确认或传入 `--yes`。对象存储 Dataset 的 replace 会先清空目标前缀再写入（非原子）。 Compact JSONL 是记录存储，不会转换或推断轨迹语义。指定 `--input-format compact-jsonl` 或 `--output-format compact-jsonl` 均会选择该格式。输入必须是本地 From 4c361175d0e452c9b9f61a458346d230ae1309cb Mon Sep 17 00:00:00 2001 From: Reiase Date: Wed, 9 Sep 2026 01:13:49 +0800 Subject: [PATCH 5/8] refactor: improve code formatting and readability in location and CLI modules Enhanced the formatting of code in the `location.rs`, `exchange.rs`, `sync.rs`, and `tests.rs` files for better clarity. This includes consolidating multiple lines into single lines where appropriate and ensuring consistent indentation. No functional changes were made, focusing solely on code aesthetics and maintainability. --- .../persisting-pchronicle-cli/src/exchange.rs | 29 +++++++------------ crates/persisting-pchronicle-cli/src/sync.rs | 4 +-- crates/persisting-pchronicle-cli/src/tests.rs | 5 +++- .../src/store/catalog/tests.rs | 9 +++--- .../src/store/location.rs | 18 ++++++++---- 5 files changed, 33 insertions(+), 32 deletions(-) diff --git a/crates/persisting-pchronicle-cli/src/exchange.rs b/crates/persisting-pchronicle-cli/src/exchange.rs index 1c0315ff..70981143 100644 --- a/crates/persisting-pchronicle-cli/src/exchange.rs +++ b/crates/persisting-pchronicle-cli/src/exchange.rs @@ -356,15 +356,12 @@ pub(super) async fn run_import( destination.as_str() ) .context("write pChronicle import replace progress")?; - destination - .remove_all() - .await - .with_context(|| { - format!( - "remove existing object-store Dataset {}", - destination.as_str() - ) - })?; + destination.remove_all().await.with_context(|| { + format!( + "remove existing object-store Dataset {}", + destination.as_str() + ) + })?; } else { return Err(cli_boundary_error( BoundaryCode::Conflict, @@ -441,11 +438,8 @@ pub(super) async fn run_import( "processing", None, )?; - let input = read_import_candidate_bytes( - candidate, - max_input_bytes, - &label, - )?; + let input = + read_import_candidate_bytes(candidate, max_input_bytes, &label)?; if let Some(source) = stage_preserved_import_source( args.format, Some(&candidate.path), @@ -1653,11 +1647,8 @@ impl<'a> StorylineImportIterator<'a> { "processing", None, )?; - let input = read_import_candidate_bytes( - candidate, - self.max_input_bytes, - &label, - )?; + let input = + read_import_candidate_bytes(candidate, self.max_input_bytes, &label)?; decode_import_source( self.requested_format, ImportOutputFormat::Storyline, diff --git a/crates/persisting-pchronicle-cli/src/sync.rs b/crates/persisting-pchronicle-cli/src/sync.rs index 3b3f0ae6..26af8c3f 100644 --- a/crates/persisting-pchronicle-cli/src/sync.rs +++ b/crates/persisting-pchronicle-cli/src/sync.rs @@ -222,8 +222,8 @@ async fn scan_source(uri: &str) -> Result> { anyhow::ensure!(root.is_dir(), "sync source must be a directory"); let mut files = BTreeMap::new(); for path in crate::exchange::collect_visible_json_files(root)? { - let metadata = - fs::metadata(&path).with_context(|| format!("stat sync file {}", path.display()))?; + let metadata = fs::metadata(&path) + .with_context(|| format!("stat sync file {}", path.display()))?; files.insert( path.strip_prefix(root)?.to_path_buf(), FileStamp { diff --git a/crates/persisting-pchronicle-cli/src/tests.rs b/crates/persisting-pchronicle-cli/src/tests.rs index 5b8df1d4..5d1760c0 100644 --- a/crates/persisting-pchronicle-cli/src/tests.rs +++ b/crates/persisting-pchronicle-cli/src/tests.rs @@ -3022,7 +3022,10 @@ async fn object_store_directory_import_recurses_json_files() -> Result<()> { .await?; // Lance interiors must be ignored even when they contain .json names. location - .write_relative_bytes("keep/events.lance/_manifest.json", b"{\"not\":\"importable\"}") + .write_relative_bytes( + "keep/events.lance/_manifest.json", + b"{\"not\":\"importable\"}", + ) .await?; let output = tempfile::tempdir()?; diff --git a/crates/persisting-pchronicle/src/store/catalog/tests.rs b/crates/persisting-pchronicle/src/store/catalog/tests.rs index 5c4f971e..d72320dd 100644 --- a/crates/persisting-pchronicle/src/store/catalog/tests.rs +++ b/crates/persisting-pchronicle/src/store/catalog/tests.rs @@ -967,10 +967,7 @@ async fn canonical_event_source_exposes_and_loads_each_storyline_independently() .await?, ); assert_eq!(snapshot.datasets()[0].sources.len(), 1); - assert_eq!( - snapshot.datasets()[0].sources[0].file, - "run-1/events.lance" - ); + assert_eq!(snapshot.datasets()[0].sources[0].file, "run-1/events.lance"); assert_eq!( snapshot.datasets()[0].sources[0].projection_status, Some(CatalogProjectionStatus::Fresh) @@ -1221,7 +1218,9 @@ async fn multiple_fresh_projections_choose_one_without_hiding_canonical_events() } let snapshot = DatasetCatalogSnapshot::discover( - vec![DatasetMount::default(storage.join("agent").to_string_lossy())?], + vec![DatasetMount::default( + storage.join("agent").to_string_lossy(), + )?], Some(DEFAULT_DATASET_NAME.into()), CatalogSnapshotOptions::default(), ) diff --git a/crates/persisting-pchronicle/src/store/location.rs b/crates/persisting-pchronicle/src/store/location.rs index 7645cd3f..fcf46a67 100644 --- a/crates/persisting-pchronicle/src/store/location.rs +++ b/crates/persisting-pchronicle/src/store/location.rs @@ -170,7 +170,10 @@ impl DatasetLocation { /// Write `bytes` at a relative object key (or local path under this Dataset). pub async fn write_relative_bytes(&self, relative: &str, bytes: &[u8]) -> Result<()> { let relative = relative.trim_start_matches('/'); - anyhow::ensure!(!relative.is_empty(), "relative object path must not be empty"); + anyhow::ensure!( + !relative.is_empty(), + "relative object path must not be empty" + ); anyhow::ensure!( !relative.split('/').any(|part| part == ".."), "relative object path must not contain '..'" @@ -193,11 +196,13 @@ impl DatasetLocation { /// Read bytes at a relative object key (or local path under this Dataset). pub async fn read_relative_bytes(&self, relative: &str) -> Result> { let relative = relative.trim_start_matches('/'); - anyhow::ensure!(!relative.is_empty(), "relative object path must not be empty"); + anyhow::ensure!( + !relative.is_empty(), + "relative object path must not be empty" + ); if let Some(root) = &self.local_path { let path = root.join(relative); - return std::fs::read(&path) - .with_context(|| format!("read {}", path.display())); + return std::fs::read(&path).with_context(|| format!("read {}", path.display())); } let store = OpendalStore::from_uri(&self.uri).await?; let Some((bytes, _)) = store.read(relative).await? else { @@ -330,7 +335,10 @@ impl DatasetLocation { } fn is_importable_json_object_key(key: &str) -> bool { - if key.split('/').any(|part| part == "_meta" || part.ends_with(".lance")) { + if key + .split('/') + .any(|part| part == "_meta" || part.ends_with(".lance")) + { return false; } Path::new(key) From dd7ee81e6959a9dd027e2e7f90e04bfbcc2b9a49 Mon Sep 17 00:00:00 2001 From: Reiase Date: Thu, 10 Sep 2026 04:27:38 +0800 Subject: [PATCH 6/8] feat(indexing): introduce index build progress tracking and enhance storyline manifest handling Added a new module for tracking index build progress, allowing for better visibility during long-running operations. Enhanced the `ChronicleManifest` to support a new storyline format and added methods for writing and loading storyline manifests. Updated the `commit_pending_content` function to conditionally build indexes based on the new options. Improved the discovery logic to classify storyline datasets correctly and ensure proper handling of their metadata. This update aims to improve the user experience during data imports and management. --- .../persisting-pchronicle-cli/src/exchange.rs | 1933 ++++++++++++++--- crates/persisting-pchronicle-cli/src/lib.rs | 45 +- crates/persisting-pchronicle-cli/src/main.rs | 2 + .../persisting-pchronicle-cli/src/onboard.rs | 14 +- .../src/server/explorer.rs | 261 ++- .../src/server/mod.rs | 457 +++- .../src/server/request_log.rs | 25 +- .../src/server/tests.rs | 4 +- crates/persisting-pchronicle-cli/src/sync.rs | 37 + crates/persisting-pchronicle-cli/src/tests.rs | 74 +- .../src/search/storyline.rs | 15 +- crates/persisting-pchronicle/src/storage.rs | 13 +- .../src/store/catalog/discovery.rs | 81 +- .../src/store/chronicle_manifest.rs | 110 + .../src/store/index_build_progress.rs | 54 + .../src/store/location.rs | 539 ++++- crates/persisting-pchronicle/src/store/mod.rs | 19 +- .../src/store/object_store_io_gate.rs | 232 ++ .../src/store/opendal_store.rs | 68 +- .../src/store/storyline/content.rs | 44 +- .../src/store/storyline/datafusion.rs | 26 +- .../src/store/storyline/mod.rs | 422 +++- .../src/store/storyline/mutation.rs | 20 +- .../src/store/storyline/writer_control.rs | 305 ++- docs/src/en/pchronicle/guides/exchange.md | 16 +- .../src/en/pchronicle/reference/cases-self.md | 10 +- docs/src/en/pchronicle/reference/cli.md | 84 +- docs/src/en/rfcs/0015-chronicle-manifest.md | 2 +- docs/src/zh/pchronicle/guides/exchange.md | 16 +- .../src/zh/pchronicle/reference/cases-self.md | 10 +- docs/src/zh/pchronicle/reference/cli.md | 182 +- docs/src/zh/rfcs/0015-chronicle-manifest.md | 5 +- 32 files changed, 4364 insertions(+), 761 deletions(-) create mode 100644 crates/persisting-pchronicle/src/store/index_build_progress.rs create mode 100644 crates/persisting-pchronicle/src/store/object_store_io_gate.rs diff --git a/crates/persisting-pchronicle-cli/src/exchange.rs b/crates/persisting-pchronicle-cli/src/exchange.rs index 70981143..826a3679 100644 --- a/crates/persisting-pchronicle-cli/src/exchange.rs +++ b/crates/persisting-pchronicle-cli/src/exchange.rs @@ -58,7 +58,7 @@ async fn prepare_import_destination( ) -> Result { let parsed = DatasetLocation::parse(output_arg)?; let exists = parsed.exists().await?; - match args.mode { + match args.mode()? { ImportMode::Create => { if parsed.is_object_store() { anyhow::ensure!(!exists, "import output already exists"); @@ -194,6 +194,7 @@ pub(super) async fn run_import( mut args: ImportArgs, settings_override: Option<&Path>, stdin_is_terminal: bool, + stderr_is_terminal: bool, stdin: &mut dyn Read, stdout: &mut dyn Write, stderr: &mut dyn Write, @@ -216,23 +217,22 @@ pub(super) async fn run_import( "stdin import requires an explicit --input-format" ); } + let mode = args.mode()?; anyhow::ensure!( - args.mode == ImportMode::Append || args.on_duplicate.is_none(), - "--on-duplicate is only valid with --mode append" + mode == ImportMode::Append || args.on_duplicate.is_none(), + "--on-duplicate is only valid with --append" ); anyhow::ensure!( - args.mode == ImportMode::Replace || !args.yes, - "--yes is only valid with --mode replace" + mode == ImportMode::Replace || !args.yes, + "--yes is only valid with --replace" ); anyhow::ensure!( - !(args.stream && args.mode == ImportMode::Replace && !args.yes), + !(args.stream && mode == ImportMode::Replace && !args.yes), "stdin replace import requires --yes because stdin carries the import data" ); if args.from != "-" { args.from = expand_dataset_reference(&args.from, settings_override, true)?; } - writeln!(stderr, "import from={} status=started", args.from) - .context("write pChronicle import progress")?; let from_location = (!args.stream) .then(|| DatasetLocation::parse(&args.from)) .transpose()?; @@ -263,7 +263,7 @@ pub(super) async fn run_import( && args.output_format != Some(ImportOutputFormat::Storyline) { anyhow::ensure!( - args.mode == ImportMode::Append && args.output_format.is_none(), + mode == ImportMode::Append && args.output_format.is_none(), "object-store import requires --output-format storyline" ); } @@ -273,8 +273,8 @@ pub(super) async fn run_import( let replace_existing = prepared.replace_existing; if let Some(snapshot) = canonical { anyhow::ensure!( - args.mode != ImportMode::Append, - "canonical event import does not support --mode append" + mode != ImportMode::Append, + "canonical event import does not support --append" ); return run_canonical_event_import( args, @@ -286,30 +286,43 @@ pub(super) async fn run_import( ) .await; } + let mut progress = ImportProgress::new(stderr_is_terminal); + let object_store_from = from_location + .as_ref() + .filter(|location| location.is_object_store() && !args.stream) + .cloned(); let (directory_input, candidates) = if args.stream { + progress.set_discovered(1, 0)?; (false, Vec::new()) - } else if let Some(location) = &from_location { - if location.is_object_store() { - collect_object_store_import_candidates(location, stderr).await? - } else { - collect_import_candidates(Path::new(&args.from))? - } + } else if object_store_from.is_some() { + // Object-store Sources are discovered inside the Storyline pipeline so + // listing overlaps read/parse/write instead of buffering the full tree. + (true, Vec::new()) + } else if from_location.is_some() { + let (directory_input, candidates) = collect_import_candidates(Path::new(&args.from))?; + let discovered_bytes = candidates.iter().try_fold(0u64, |total, candidate| { + total + .checked_add(candidate.size_hint) + .context("import discovered byte count overflow") + })?; + progress.set_discovered(candidates.len() as u64, discovered_bytes)?; + (directory_input, candidates) } else { (false, Vec::new()) }; anyhow::ensure!( - args.mode != ImportMode::Append || args.output_format != Some(ImportOutputFormat::Preserve), + mode != ImportMode::Append || args.output_format != Some(ImportOutputFormat::Preserve), "append import requires --output-format storyline (or omit it)" ); let output_format = args .output_format - .unwrap_or(if args.mode == ImportMode::Append { + .unwrap_or(if mode == ImportMode::Append { ImportOutputFormat::Storyline } else { ImportOutputFormat::Preserve }); let duplicate_policy = args.on_duplicate.unwrap_or(DuplicateIdPolicy::Suffix); - let (dataset_uri, imported_sources, unknown_field_warnings, skipped_warnings) = if args.mode + let (dataset_uri, imported_sources, unknown_field_warnings, skipped_warnings) = if mode == ImportMode::Append { let store = StorylineLanceStore::open_uri(destination.as_str()) @@ -329,8 +342,9 @@ pub(super) async fn run_import( &store, &args, stdin, - stderr, + &mut progress, &candidates, + object_store_from.clone(), StorylineImportOptions { max_input_bytes, directory_input, @@ -347,21 +361,33 @@ pub(super) async fn run_import( unknown_field_warnings, skipped_warnings, ) - } else if destination.is_object_store() { + } else if destination.is_object_store() || output_format == ImportOutputFormat::Storyline { + // Storyline imports commit in place so progressive CURRENT + + // chronicle.manifest updates are visible to a live catalog mount. + // Remote object-store targets stage locally first: Lance index builds + // on S3 are extremely slow, so we write+index on disk then upload. if destination.exists().await? { if replace_existing { - writeln!( - stderr, - "import to={} status=replacing", - destination.as_str() - ) - .context("write pChronicle import replace progress")?; - destination.remove_all().await.with_context(|| { - format!( - "remove existing object-store Dataset {}", - destination.as_str() - ) - })?; + destination + .remove_all_with_progress(|deleted, total, path| { + progress.note_deleted(deleted, total, path) + }) + .await + .with_context(|| { + format!("remove existing Dataset {}", destination.as_str()) + })?; + progress.finish()?; + // Delete progress reuses the paint lines but must not wipe discovery + // totals collected before replace (local candidates only). + progress.reset_import_counters(); + if object_store_from.is_none() { + let discovered_bytes = candidates.iter().try_fold(0u64, |total, candidate| { + total + .checked_add(candidate.size_hint) + .context("import discovered byte count overflow") + })?; + progress.set_discovered(candidates.len() as u64, discovered_bytes)?; + } } else { return Err(cli_boundary_error( BoundaryCode::Conflict, @@ -369,19 +395,50 @@ pub(super) async fn run_import( )); } } - let store = StorylineLanceStore::open_uri(destination.as_str()) - .await - .context("create squashed Storyline Lance Dataset")?; let (imported_sources, unknown_field_warnings, skipped_warnings) = - squash_storyline_into_store( - &store, - &args, - stdin, - stderr, - &candidates, - StorylineImportOptions::create(max_input_bytes, directory_input), - ) - .await?; + if destination.is_object_store() { + progress.set_phase(ImportPhase::Writing, "local staging (indexes on disk)")?; + let staging = tempfile::Builder::new() + .prefix("pchronicle-storyline-stage-") + .tempdir() + .context("create local Storyline staging directory")?; + let store = StorylineLanceStore::open(staging.path()) + .await + .context("open local Storyline staging Dataset")?; + let result = squash_storyline_into_store( + &store, + &args, + stdin, + &mut progress, + &candidates, + object_store_from.clone(), + StorylineImportOptions::create(max_input_bytes, directory_input), + ) + .await?; + upload_local_storyline_dataset(staging.path(), &destination, &mut progress) + .await + .with_context(|| { + format!( + "upload staged Storyline Dataset to {}", + destination.as_str() + ) + })?; + result + } else { + let store = StorylineLanceStore::open_uri(destination.as_str()) + .await + .context("create squashed Storyline Lance Dataset")?; + squash_storyline_into_store( + &store, + &args, + stdin, + &mut progress, + &candidates, + object_store_from.clone(), + StorylineImportOptions::create(max_input_bytes, directory_input), + ) + .await? + }; ( destination.as_str().to_string(), imported_sources, @@ -407,8 +464,9 @@ pub(super) async fn run_import( let mut imported_sources = Vec::new(); let mut skipped_warnings = Vec::new(); if args.stream { - write_import_progress(stderr, "stdin", "processing", None)?; + progress.set_phase(ImportPhase::Reading, "stdin")?; let input = read_bounded(stdin, max_input_bytes, "stdin")?; + progress.set_phase(ImportPhase::Parsing, "stdin")?; if let Some(source) = stage_preserved_import_source( args.format, None, @@ -419,27 +477,20 @@ pub(super) async fn run_import( &mut unknown_field_warnings, &mut skipped_warnings, )? { - write_import_progress( - stderr, - &source.source_path, - "completed", - Some((&source.format, source.trajectories, source.input_bytes)), - )?; + progress.set_phase(ImportPhase::Writing, &source.source_path)?; + progress.note_imported(source.input_bytes as u64)?; imported_sources.push(source); } else { - write_import_progress(stderr, "stdin", "skipped", None)?; + progress.note_imported(input.len() as u64)?; } } else { for candidate in &candidates { - let label = format!("import source {}", candidate.relative_path.display()); - write_import_progress( - stderr, - &candidate.relative_path.to_string_lossy(), - "processing", - None, - )?; + let name = candidate.relative_path.to_string_lossy(); + let label = format!("import source {name}"); + progress.set_phase(ImportPhase::Reading, &name)?; let input = - read_import_candidate_bytes(candidate, max_input_bytes, &label)?; + load_import_candidate_bytes(candidate, max_input_bytes, &label).await?; + progress.set_phase(ImportPhase::Parsing, &name)?; if let Some(source) = stage_preserved_import_source( args.format, Some(&candidate.path), @@ -450,38 +501,18 @@ pub(super) async fn run_import( &mut unknown_field_warnings, &mut skipped_warnings, )? { - write_import_progress( - stderr, - &source.source_path, - "completed", - Some((&source.format, source.trajectories, source.input_bytes)), - )?; + progress.set_phase(ImportPhase::Writing, &source.source_path)?; + progress.note_imported(source.input_bytes as u64)?; imported_sources.push(source); } else { - write_import_progress( - stderr, - &candidate.relative_path.to_string_lossy(), - "skipped", - None, - )?; + progress.note_imported(input.len() as u64)?; } } } (imported_sources, unknown_field_warnings, skipped_warnings) } ImportOutputFormat::Storyline => { - let store = StorylineLanceStore::open(staging.path()) - .await - .context("create squashed Storyline Lance Dataset")?; - squash_storyline_into_store( - &store, - &args, - stdin, - stderr, - &candidates, - StorylineImportOptions::create(max_input_bytes, directory_input), - ) - .await? + unreachable!("storyline import commits in place above") } ImportOutputFormat::CompactJsonl => unreachable!("compact import handled above"), }; @@ -495,7 +526,7 @@ pub(super) async fn run_import( let staging_path = staging.keep(); let mut cleanup = StagingPathGuard::new(staging_path.clone()); - publish_staged_dataset(&staging_path, &output, replace_existing)?; + publish_staged_dataset(&staging_path, &output, replace_existing, Some(&mut progress)).await?; cleanup.disarm(); ( output.to_string_lossy().into_owned(), @@ -536,9 +567,9 @@ pub(super) async fn run_import( serde_json::to_writer_pretty(&mut *stdout, &response) .context("encode pChronicle import JSON")?; writeln!(stdout).context("write pChronicle import JSON")?; + progress.finish()?; if let (Some(source_path), Some(format)) = (&response.source_path, &response.format) { - writeln!( - stderr, + progress.notice(&format!( "dataset_uri={} source={} format={} output_format={} trajectories={} input_bytes={}", response.dataset_uri, source_path, @@ -548,11 +579,9 @@ pub(super) async fn run_import( response .input_bytes .expect("JSON imports always report input bytes"), - ) - .context("write pChronicle import metadata")?; + ))?; } else { - writeln!( - stderr, + progress.notice(&format!( "dataset_uri={} sources={} output_format={} trajectories={} input_bytes={}", response.dataset_uri, response.sources, @@ -561,15 +590,15 @@ pub(super) async fn run_import( response .input_bytes .expect("JSON imports always report input bytes"), - ) - .context("write pChronicle import metadata")?; + ))?; } for line in skipped_warnings { - writeln!(stderr, "{line}").context("write pChronicle skipped-source warning")?; + progress.notice(&line)?; } for line in unknown_field_warnings.warning_lines() { - writeln!(stderr, "{line}").context("write pChronicle unknown-field warning")?; + progress.notice(&line)?; } + progress.flush_log(stderr)?; Ok(()) } @@ -580,7 +609,7 @@ async fn run_compact_jsonl_import( stderr: &mut dyn Write, ) -> Result<()> { anyhow::ensure!( - args.mode != ImportMode::Append, + args.mode()? != ImportMode::Append, "compact JSONL append is not supported; use sync or replace" ); anyhow::ensure!( @@ -593,7 +622,7 @@ async fn run_compact_jsonl_import( !output_arg.starts_with("s3://") && !output_arg.starts_with("oss://"), "compact JSONL currently requires local paths" ); - if args.mode == ImportMode::Create { + if args.mode()? == ImportMode::Create { anyhow::ensure!(!output.exists(), "import output already exists"); } let columns = args @@ -626,7 +655,7 @@ async fn run_compact_jsonl_import( std::fs::File::open(staging.path())?.sync_all()?; let staging_path = staging.keep(); let mut cleanup = StagingPathGuard::new(staging_path.clone()); - publish_staged_dataset(&staging_path, output, output.exists())?; + publish_staged_dataset(&staging_path, output, output.exists(), None).await?; cleanup.disarm(); serde_json::to_writer_pretty( &mut *stdout, @@ -663,11 +692,14 @@ pub(crate) async fn sync_snapshot( output: Some(storyline.to_owned()), format: ExchangeFormat::CompactJsonl, output_format: Some(ImportOutputFormat::CompactJsonl), - mode: ImportMode::Replace, + replace: true, + append: false, + mode: None, on_duplicate: None, yes: true, stream: false, max_input_bytes: Some(256 * 1024 * 1024), + commit_every: None, columns: columns.to_vec(), }, storyline, @@ -687,15 +719,19 @@ pub(crate) async fn sync_snapshot( output: Some(warehouse.to_owned()), format: input_format, output_format: Some(ImportOutputFormat::Preserve), - mode: ImportMode::Replace, + replace: true, + append: false, + mode: None, on_duplicate: None, yes: true, stream: false, max_input_bytes: Some(256 * 1024 * 1024), + commit_every: None, columns: Vec::new(), }, None, false, + false, &mut stdin, &mut stdout, &mut stderr, @@ -708,15 +744,19 @@ pub(crate) async fn sync_snapshot( output: Some(storyline.to_owned()), format: input_format, output_format: Some(ImportOutputFormat::Storyline), - mode: ImportMode::Replace, + replace: true, + append: false, + mode: None, on_duplicate: None, yes: true, stream: false, max_input_bytes: Some(256 * 1024 * 1024), + commit_every: None, columns: Vec::new(), }, None, false, + false, &mut stdin, &mut stdout, &mut stderr, @@ -748,12 +788,168 @@ impl StorylineImportOptions { } } +/// How many Sources the reader may prefetch ahead of parse/write. +/// Bounded so large object-store imports do not buffer unbounded memory. +const IMPORT_READ_AHEAD: usize = 3; +/// Pipeline channel capacity for object-store discover/read events. Listing +/// emits Discovered first; this buffer only absorbs Loaded messages while a +/// commit is in flight. +const IMPORT_PIPELINE_CHANNEL: usize = 16; + +struct PipelineLoadedSource { + candidate: ImportFileCandidate, + bytes: Vec, +} + +enum PipelineMsg { + Scanning(String), + Discovered { path: String, bytes: u64 }, + Loaded(PipelineLoadedSource), +} + +fn spawn_candidates_load_producer( + candidates: Vec, + max_input_bytes: usize, + reading_ahead: Arc>, +) -> ( + tokio::sync::mpsc::Receiver>, + tokio::task::JoinHandle<()>, +) { + let (tx, rx) = tokio::sync::mpsc::channel::>(IMPORT_READ_AHEAD); + let producer = tokio::spawn(async move { + for candidate in candidates { + let name = candidate.relative_path.to_string_lossy().into_owned(); + if let Ok(mut guard) = reading_ahead.lock() { + *guard = name.clone(); + } + let label = format!("import source {name}"); + let loaded = match load_import_candidate_bytes(&candidate, max_input_bytes, &label).await + { + Ok(bytes) => Ok(PipelineMsg::Loaded(PipelineLoadedSource { candidate, bytes })), + Err(error) => Err(error), + }; + if tx.send(loaded).await.is_err() { + return; + } + } + if let Ok(mut guard) = reading_ahead.lock() { + guard.clear(); + } + }); + (rx, producer) +} + +fn spawn_object_store_discover_load_producer( + location: DatasetLocation, + max_input_bytes: usize, + reading_ahead: Arc>, +) -> ( + tokio::sync::mpsc::Receiver>, + tokio::task::JoinHandle<()>, +) { + let (tx, rx) = tokio::sync::mpsc::channel::>(IMPORT_PIPELINE_CHANNEL); + let producer = tokio::spawn(async move { + let remote_root = location.as_str().to_owned(); + // List completely before any Load so discovery totals keep moving even + // when a later commit/index stalls the consumer. + let pending_files = Arc::new(std::sync::Mutex::new(Vec::<(String, u64)>::new())); + let list_result = location + .for_each_importable_json_object_event( + persisting_pchronicle::storage::DEFAULT_MAX_LOCAL_QUERY_FILES, + |event| { + let tx = tx.clone(); + let pending_files = Arc::clone(&pending_files); + async move { + match event { + persisting_pchronicle::storage::ImportableObjectEvent::Scanning { + prefix, + } => { + let _ = tx.send(Ok(PipelineMsg::Scanning(prefix))).await; + Ok(()) + } + persisting_pchronicle::storage::ImportableObjectEvent::File { + key, + size, + .. + } => { + if tx + .send(Ok(PipelineMsg::Discovered { + path: key.clone(), + bytes: size, + })) + .await + .is_err() + { + return Ok(()); + } + if let Ok(mut guard) = pending_files.lock() { + guard.push((key, size)); + } + Ok(()) + } + } + } + }, + ) + .await; + if let Err(error) = list_result { + let _ = tx.send(Err(error)).await; + if let Ok(mut guard) = reading_ahead.lock() { + guard.clear(); + } + return; + } + let files = match pending_files.lock() { + Ok(mut guard) => std::mem::take(&mut *guard), + Err(_) => Vec::new(), + }; + for (key, size) in files { + if let Ok(mut guard) = reading_ahead.lock() { + *guard = key.clone(); + } + let relative_path = PathBuf::from(&key); + let candidate = ImportFileCandidate { + path: relative_path.clone(), + output_relative_path: Some(relative_path.clone()), + relative_path, + content: None, + remote_root: Some(remote_root.clone()), + size_hint: size, + }; + let label = format!("import source {key}"); + match load_import_candidate_bytes(&candidate, max_input_bytes, &label).await { + Ok(bytes) => { + if tx + .send(Ok(PipelineMsg::Loaded(PipelineLoadedSource { + candidate, + bytes, + }))) + .await + .is_err() + { + break; + } + } + Err(error) => { + let _ = tx.send(Err(error)).await; + break; + } + } + } + if let Ok(mut guard) = reading_ahead.lock() { + guard.clear(); + } + }); + (rx, producer) +} + async fn squash_storyline_into_store( store: &StorylineLanceStore, args: &ImportArgs, stdin: &mut dyn Read, - stderr: &mut dyn Write, + progress: &mut ImportProgress, candidates: &[ImportFileCandidate], + object_store_from: Option, options: StorylineImportOptions, ) -> Result<( Vec, @@ -768,61 +964,769 @@ async fn squash_storyline_into_store( allow_empty, append_generation, } = options; - let mut import = if args.stream { - StorylineImportIterator::stdin( + if args.stream { + return squash_storyline_stdin_into_store( + store, args.format, max_input_bytes, stdin, - stderr, - seen_document_ids, - duplicate_policy, - ) - } else { - StorylineImportIterator::files( - args.format, - max_input_bytes, - candidates, - stderr, + progress, seen_document_ids, duplicate_policy, + allow_empty, + directory_input, + append_generation, + commit_batch_schedule(args), ) + .await; + } + let source = match object_store_from { + Some(location) => ObjectStoreImportSource::Location(location), + None => ObjectStoreImportSource::Candidates(candidates.to_vec()), }; - let report_storylines = match import.next() { - Some(first) => match append_generation.as_deref() { - Some(generation) => { - store - .append_storyline_stream(std::iter::once(first).chain(&mut import), generation) - .await? - .storylines + squash_storyline_files_pipeline( + store, + args.format, + max_input_bytes, + progress, + source, + seen_document_ids, + duplicate_policy, + allow_empty, + directory_input, + append_generation, + commit_batch_schedule(args), + ) + .await +} + +const DEFAULT_COMMIT_BATCH_START: usize = 64; +const DEFAULT_COMMIT_BATCH_MAX: usize = 4096; + +#[derive(Debug, Clone)] +struct CommitBatchSchedule { + next: usize, + max: usize, + fixed: bool, +} + +impl CommitBatchSchedule { + fn adaptive() -> Self { + Self { + next: DEFAULT_COMMIT_BATCH_START, + max: DEFAULT_COMMIT_BATCH_MAX, + fixed: false, + } + } + + fn fixed(n: usize) -> Self { + let n = n.max(1); + Self { + next: n, + max: n, + fixed: true, + } + } + + fn current(&self) -> usize { + self.next + } + + fn after_commit(&mut self) { + if self.fixed { + return; + } + self.next = self.next.saturating_mul(2).min(self.max); + } +} + +fn commit_batch_schedule(args: &ImportArgs) -> CommitBatchSchedule { + match args.commit_every { + Some(n) => CommitBatchSchedule::fixed(n), + None => CommitBatchSchedule::adaptive(), + } +} + +enum ObjectStoreImportSource { + Candidates(Vec), + Location(DatasetLocation), +} + +#[allow(clippy::too_many_arguments)] +async fn squash_storyline_files_pipeline( + store: &StorylineLanceStore, + requested_format: ExchangeFormat, + max_input_bytes: usize, + progress: &mut ImportProgress, + source: ObjectStoreImportSource, + mut seen_document_ids: HashSet, + duplicate_policy: DuplicateIdPolicy, + allow_empty: bool, + directory_input: bool, + mut append_generation: Option, + mut commit_schedule: CommitBatchSchedule, +) -> Result<( + Vec, + persisting_pchronicle::model::UnknownFieldImportWarnings, + Vec, +)> { + let reading_ahead = Arc::new(std::sync::Mutex::new(String::new())); + let (mut rx, producer) = match source { + ObjectStoreImportSource::Candidates(candidates) => { + spawn_candidates_load_producer(candidates, max_input_bytes, Arc::clone(&reading_ahead)) + } + ObjectStoreImportSource::Location(location) => { + progress.set_phase(ImportPhase::Discovering, location.as_str())?; + spawn_object_store_discover_load_producer( + location, + max_input_bytes, + Arc::clone(&reading_ahead), + ) + } + }; + + let mut unknown_field_warnings = + persisting_pchronicle::model::UnknownFieldImportWarnings::default(); + let mut skipped_warnings = Vec::new(); + let mut imported_sources: Vec = Vec::new(); + let mut batch = Vec::with_capacity(commit_schedule.current()); + let mut committed_storylines = 0u64; + let mut skipped_commit_storylines = 0usize; + let mut saw_any = false; + let mut current_storylines = Vec::new().into_iter(); + let mut producer_done = false; + let mut discovered_any = false; + + loop { + if let Some(mut storyline) = current_storylines.next() { + saw_any = true; + if let Some(warning) = + apply_duplicate_document_policy(&mut storyline, &mut seen_document_ids, duplicate_policy) + { + if warning.contains("skipped") { + skipped_warnings.push(warning); + continue; + } + skipped_warnings.push(warning); + } + let metadata = imported_sources + .last_mut() + .expect("decoded Storyline has source metadata"); + metadata.trajectories = metadata + .trajectories + .checked_add(1) + .context("import trajectory count overflow")?; + batch.push(storyline); + if batch.len() >= commit_schedule.current() { + match commit_or_skip_storyline_import_batch( + store, + progress, + std::mem::take(&mut batch), + &mut append_generation, + committed_storylines, + &mut commit_schedule, + ) + .await? + { + StorylineBatchCommit::Committed(total) => { + committed_storylines = total; + } + StorylineBatchCommit::Skipped { batch_len, warning } => { + skipped_commit_storylines = skipped_commit_storylines + .saturating_add(batch_len as usize); + skipped_warnings.push(warning); + retract_imported_trajectories( + &mut imported_sources, + batch_len as usize, + ); + } + } + batch.reserve(commit_schedule.current()); + } + continue; + } + + if producer_done { + break; + } + + // Surface producer read activity while waiting for the next Source. + let msg = loop { + if let Ok(guard) = reading_ahead.lock() { + progress.set_reading_ahead(guard.as_str())?; + } + tokio::select! { + item = rx.recv() => break item, + _ = tokio::time::sleep(std::time::Duration::from_millis(100)) => {} + } + }; + match msg { + Some(Ok(PipelineMsg::Scanning(prefix))) => { + progress.note_scanning(&prefix)?; + } + Some(Ok(PipelineMsg::Discovered { path, bytes })) => { + discovered_any = true; + progress.note_discovered(&path, bytes)?; + } + Some(Ok(PipelineMsg::Loaded(loaded))) => { + let name = loaded.candidate.relative_path.to_string_lossy().into_owned(); + if let Ok(guard) = reading_ahead.lock() { + progress.set_reading_ahead(guard.as_str())?; + } + progress.set_phase(ImportPhase::Parsing, &name)?; + match decode_import_source( + requested_format, + ImportOutputFormat::Storyline, + Some(&loaded.candidate.path), + Some(&loaded.candidate.relative_path), + loaded.candidate.output_relative_path.as_deref(), + &loaded.bytes, + &mut unknown_field_warnings, + )? { + DecodeImportOutcome::Imported(decoded) => { + progress.set_phase( + ImportPhase::Writing, + &decoded.diagnostic_path.to_string_lossy(), + )?; + progress.note_imported(decoded.metadata.input_bytes as u64)?; + let mut metadata = decoded.metadata; + metadata.trajectories = 0; + imported_sources.push(metadata); + current_storylines = decoded.storylines.into_iter(); + } + DecodeImportOutcome::Skipped { path, reason } => { + progress.note_imported(0)?; + skipped_warnings.push(skipped_import_warning(&path, &reason)); + } + } + } + Some(Err(error)) => { + producer.abort(); + return Err(error); + } + None => { + producer_done = true; + progress.clear_reading_ahead()?; + } + } + } + + if !batch.is_empty() { + match commit_or_skip_storyline_import_batch( + store, + progress, + std::mem::take(&mut batch), + &mut append_generation, + committed_storylines, + &mut commit_schedule, + ) + .await? + { + StorylineBatchCommit::Committed(total) => { + committed_storylines = total; + } + StorylineBatchCommit::Skipped { batch_len, warning } => { + skipped_commit_storylines = + skipped_commit_storylines.saturating_add(batch_len as usize); + skipped_warnings.push(warning); + retract_imported_trajectories(&mut imported_sources, batch_len as usize); + } + } + } + progress.clear_reading_ahead()?; + + match producer.await { + Ok(()) => {} + Err(error) if error.is_cancelled() => {} + Err(error) => return Err(anyhow!("import reader task failed: {error}")), + } + + if imported_sources.is_empty() { + if allow_empty && !saw_any { + return Ok((imported_sources, unknown_field_warnings, skipped_warnings)); + } + if !discovered_any { + return Err(cli_boundary_error( + BoundaryCode::InvalidRequest, + "import object prefix contains no .json, .jsonl, or .ndjson files", + )); + } + return Err(empty_auto_directory_import_error(directory_input)); + } + // Drop Sources that lost every trajectory to skipped commits so empty + // placeholders do not inflate the import summary. + if skipped_commit_storylines > 0 { + imported_sources.retain(|source| source.trajectories > 0); + } + if imported_sources.is_empty() { + return Err(anyhow!( + "storyline import committed no trajectories after skipping failed batches" + )); + } + anyhow::ensure!( + store.current_table_paths().await?.is_some(), + "squashed Storyline Lance Dataset has no committed snapshot" + ); + let imported_trajectories = imported_sources.iter().try_fold(0usize, |total, source| { + total + .checked_add(source.trajectories) + .context("import trajectory count overflow") + })?; + anyhow::ensure!( + committed_storylines as usize == imported_trajectories, + "squashed Storyline import report does not match decoded trajectory count" + ); + finalize_storyline_import_indexes(store, progress).await?; + Ok((imported_sources, unknown_field_warnings, skipped_warnings)) +} + +#[allow(clippy::too_many_arguments)] +async fn squash_storyline_stdin_into_store( + store: &StorylineLanceStore, + requested_format: ExchangeFormat, + max_input_bytes: usize, + stdin: &mut dyn Read, + progress: &mut ImportProgress, + seen_document_ids: HashSet, + duplicate_policy: DuplicateIdPolicy, + allow_empty: bool, + directory_input: bool, + append_generation: Option, + commit_schedule: CommitBatchSchedule, +) -> Result<( + Vec, + persisting_pchronicle::model::UnknownFieldImportWarnings, + Vec, +)> { + let import = StorylineImportIterator::stdin( + requested_format, + max_input_bytes, + stdin, + progress, + seen_document_ids, + duplicate_policy, + ); + drain_storyline_import_batches( + store, + import, + append_generation, + commit_schedule, + allow_empty, + directory_input, + ) + .await +} + +fn apply_duplicate_document_policy( + storyline: &mut StorylineDocument, + seen_document_ids: &mut HashSet, + duplicate_policy: DuplicateIdPolicy, +) -> Option { + let original = storyline.document_id().to_string(); + match duplicate_policy { + DuplicateIdPolicy::Suffix => { + uniquify_storyline_document_id(storyline, seen_document_ids).map( + |(original, renamed)| { + format!("warning: duplicate document_id '{original}' renamed to '{renamed}'") + }, + ) + } + DuplicateIdPolicy::Skip => { + if !seen_document_ids.insert(original.clone()) { + Some(format!( + "warning: duplicate document_id '{original}' skipped" + )) + } else { + None + } + } + } +} + +async fn drain_storyline_import_batches( + store: &StorylineLanceStore, + mut import: StorylineImportIterator<'_>, + mut append_generation: Option, + mut commit_schedule: CommitBatchSchedule, + allow_empty: bool, + directory_input: bool, +) -> Result<( + Vec, + persisting_pchronicle::model::UnknownFieldImportWarnings, + Vec, +)> { + let mut batch = Vec::with_capacity(commit_schedule.current()); + let mut committed_storylines = 0u64; + let mut skipped_commit_storylines = 0usize; + let mut commit_skip_warnings = Vec::new(); + let mut saw_any = false; + + loop { + match import.next_document().await { + Some(item) => { + saw_any = true; + batch.push(item?); + if batch.len() < commit_schedule.current() { + continue; + } + match commit_or_skip_storyline_import_batch( + store, + import.progress, + std::mem::take(&mut batch), + &mut append_generation, + committed_storylines, + &mut commit_schedule, + ) + .await? + { + StorylineBatchCommit::Committed(total) => { + committed_storylines = total; + } + StorylineBatchCommit::Skipped { batch_len, warning } => { + skipped_commit_storylines = skipped_commit_storylines + .saturating_add(batch_len as usize); + commit_skip_warnings.push(warning); + } + } + batch.reserve(commit_schedule.current()); + } + None if batch.is_empty() => break, + None => { + match commit_or_skip_storyline_import_batch( + store, + import.progress, + std::mem::take(&mut batch), + &mut append_generation, + committed_storylines, + &mut commit_schedule, + ) + .await? + { + StorylineBatchCommit::Committed(total) => { + committed_storylines = total; + } + StorylineBatchCommit::Skipped { batch_len, warning } => { + skipped_commit_storylines = skipped_commit_storylines + .saturating_add(batch_len as usize); + commit_skip_warnings.push(warning); + } + } + break; + } + } + } + + let (mut imported_sources, unknown_field_warnings, mut skipped_warnings) = + import.into_result_parts(); + skipped_warnings.extend(commit_skip_warnings); + retract_imported_trajectories(&mut imported_sources, skipped_commit_storylines); + if skipped_commit_storylines > 0 { + imported_sources.retain(|source| source.trajectories > 0); + } + if imported_sources.is_empty() { + if allow_empty && !saw_any { + return Ok((imported_sources, unknown_field_warnings, skipped_warnings)); + } + if skipped_commit_storylines > 0 { + return Err(anyhow!( + "storyline import committed no trajectories after skipping failed batches" + )); + } + return Err(empty_auto_directory_import_error(directory_input)); + } + anyhow::ensure!( + store.current_table_paths().await?.is_some(), + "squashed Storyline Lance Dataset has no committed snapshot" + ); + let imported_trajectories = imported_sources.iter().try_fold(0usize, |total, source| { + total + .checked_add(source.trajectories) + .context("import trajectory count overflow") + })?; + anyhow::ensure!( + committed_storylines as usize == imported_trajectories, + "squashed Storyline import report does not match decoded trajectory count" + ); + finalize_storyline_import_indexes(store, import.progress).await?; + Ok((imported_sources, unknown_field_warnings, skipped_warnings)) +} + +enum StorylineBatchCommit { + Committed(u64), + Skipped { batch_len: u64, warning: String }, +} + +fn is_skippable_storyline_commit_error(error: &anyhow::Error) -> bool { + let text = format!("{error:#}").to_ascii_lowercase(); + text.contains("timeout") + || text.contains("timed out") + || text.contains("error sending request") + || text.contains("conditionnotmatch") + || text.contains("preconditionfailed") + || text.contains("precondition failed") + || text.contains("throttle") + || text.contains("slow down") + || text.contains("503") + || text.contains("429") + || text.contains("connection reset") + || text.contains("broken pipe") + || text.contains("lanceerror(io)") + || text.contains("generic s3 error") + || text.contains("client error (connect)") +} + +fn retract_imported_trajectories(sources: &mut [ImportedSource], mut count: usize) { + for source in sources.iter_mut().rev() { + if count == 0 { + break; + } + let take = source.trajectories.min(count); + source.trajectories -= take; + count -= take; + } +} + +async fn refresh_append_generation_after_skip( + store: &StorylineLanceStore, + append_generation: &mut Option, +) { + match store.current_table_paths().await { + Ok(Some(paths)) => { + *append_generation = Some(paths.generation); + } + Ok(None) => {} + Err(error) => { + tracing::warn!( + root = %store.root_uri(), + error = %error, + "failed to refresh Storyline generation after skipped commit batch" + ); + } + } +} + +async fn commit_or_skip_storyline_import_batch( + store: &StorylineLanceStore, + progress: &mut ImportProgress, + batch: Vec, + append_generation: &mut Option, + committed_storylines: u64, + commit_schedule: &mut CommitBatchSchedule, +) -> Result { + let batch_len = batch.len() as u64; + let sample_ids = batch + .iter() + .take(8) + .map(|storyline| storyline.document_id().to_string()) + .collect::>(); + match commit_storyline_import_batch( + store, + progress, + batch, + append_generation, + committed_storylines, + ) + .await + { + Ok(total) => { + commit_schedule.after_commit(); + Ok(StorylineBatchCommit::Committed(total)) + } + Err(error) if is_skippable_storyline_commit_error(&error) => { + tracing::warn!( + committed_before = committed_storylines, + batch_len, + root = %store.root_uri(), + sample_document_ids = ?sample_ids, + error = %format!("{error:#}"), + "skipping storyline commit batch after transient storage failure; continuing import" + ); + refresh_append_generation_after_skip(store, append_generation).await; + if !commit_schedule.fixed { + commit_schedule.next = DEFAULT_COMMIT_BATCH_START; + } + let warning = format!( + "warning: skipped storyline commit batch of {batch_len} trajectories (committed_before={committed_storylines}, sample_document_ids={sample_ids:?}): {error:#}" + ); + let _ = progress.notice(&warning); + Ok(StorylineBatchCommit::Skipped { batch_len, warning }) + } + Err(error) => Err(error), + } +} + +async fn finalize_storyline_import_indexes( + store: &StorylineLanceStore, + progress: &mut ImportProgress, +) -> Result<()> { + progress.set_phase(ImportPhase::Writing, "optimize indices (final)")?; + let _index_progress = progress.attach_index_progress(); + store + .maintain(&persisting_pchronicle::storage::LanceMaintenanceOptions { + compact: false, + optimize_indices: true, + vacuum_older_than: None, + ..Default::default() + }) + .await + .context("finalize Storyline indexes after progressive import")?; + progress.set_phase(ImportPhase::Writing, "optimize indices done")?; + Ok(()) +} + +fn collect_local_relative_files(root: &Path) -> Result> { + fn walk(root: &Path, dir: &Path, out: &mut Vec) -> Result<()> { + for entry in std::fs::read_dir(dir) + .with_context(|| format!("read staging directory {}", dir.display()))? + { + let entry = entry?; + let path = entry.path(); + if path.is_dir() { + walk(root, &path, out)?; + continue; } - None => { - store - .replace_storyline_stream(std::iter::once(first).chain(&mut import)) - .await? - .storylines + let relative = path + .strip_prefix(root) + .with_context(|| format!("strip staging root from {}", path.display()))? + .to_string_lossy() + .replace('\\', "/"); + if !relative.is_empty() { + out.push(relative); } - }, - None if allow_empty => 0, - None => return Err(empty_auto_directory_import_error(directory_input)), - }; - let (imported_sources, unknown_field_warnings, skipped_warnings) = import.into_result_parts(); - if imported_sources.is_empty() { - return Err(empty_auto_directory_import_error(directory_input)); + } + Ok(()) } + let mut files = Vec::new(); + walk(root, root, &mut files)?; + files.sort(); + Ok(files) +} + +fn is_deferred_storyline_publish_key(relative: &str) -> bool { + matches!( + relative, + "CURRENT" | "chronicle.manifest" | ".storyline-write.lock" + ) || relative.ends_with("/CURRENT") + || relative.ends_with("/chronicle.manifest") +} + +async fn upload_local_storyline_dataset( + local_root: &Path, + destination: &DatasetLocation, + progress: &mut ImportProgress, +) -> Result<()> { + let files = collect_local_relative_files(local_root)?; anyhow::ensure!( - store.current_table_paths().await?.is_some(), - "squashed Storyline Lance Dataset has no committed snapshot" + files.iter().any(|path| path == "CURRENT"), + "staged Storyline Dataset is missing CURRENT" ); - let imported_trajectories = imported_sources.iter().try_fold(0usize, |total, source| { - total - .checked_add(source.trajectories) - .context("import trajectory count overflow") - })?; + let (deferred, eager): (Vec<_>, Vec<_>) = files + .into_iter() + .partition(|path| is_deferred_storyline_publish_key(path)); + let total = eager.len().saturating_add(deferred.len()) as u64; + let mut uploaded = 0u64; + for relative in eager.into_iter().chain(deferred) { + if relative == ".storyline-write.lock" { + continue; + } + uploaded = uploaded.saturating_add(1); + progress.set_phase( + ImportPhase::Writing, + &format!( + "upload {uploaded}/{total} {}", + truncate_middle(&relative, 56) + ), + )?; + let bytes = tokio::fs::read(local_root.join(&relative)) + .await + .with_context(|| format!("read staged file {relative}"))?; + destination + .write_relative_bytes(&relative, &bytes) + .await + .with_context(|| format!("upload staged file {relative}"))?; + } + progress.set_phase(ImportPhase::Writing, "upload complete")?; + Ok(()) +} + +async fn commit_storyline_import_batch( + store: &StorylineLanceStore, + progress: &mut ImportProgress, + batch: Vec, + append_generation: &mut Option, + committed_storylines: u64, +) -> Result { + anyhow::ensure!(!batch.is_empty(), "storyline import commit batch is empty"); + let batch_len = batch.len() as u64; + progress.set_phase( + ImportPhase::Writing, + &format!("commit {batch_len} trajectories"), + )?; + let report = match append_generation.as_deref() { + Some(generation) => { + tracing::info!( + committed_before = committed_storylines, + batch_len, + expected_generation = generation, + root = %store.root_uri(), + "storyline progressive append commit starting" + ); + store + .append_storyline_stream_with_options( + batch.into_iter().map(Ok), + generation, + persisting_pchronicle::storage::StorylineStreamOptions::defer_index_optimize(), + ) + .await + .with_context(|| { + format!( + "storyline progressive append commit failed (committed_before={committed_storylines}, batch={batch_len}, expected_generation={generation}, root={})", + store.root_uri() + ) + })? + } + None => { + tracing::info!( + batch_len, + root = %store.root_uri(), + "storyline progressive replace commit starting" + ); + store + .replace_storyline_stream_with_options( + batch.into_iter().map(Ok), + persisting_pchronicle::storage::StorylineStreamOptions::defer_index_optimize(), + ) + .await + .with_context(|| { + format!( + "storyline progressive replace commit failed (batch={batch_len}, root={})", + store.root_uri() + ) + })? + } + }; anyhow::ensure!( - report_storylines == imported_trajectories, - "squashed Storyline import report does not match decoded trajectory count" + report.storylines as u64 == batch_len, + "storyline import batch report does not match batch size" ); - Ok((imported_sources, unknown_field_warnings, skipped_warnings)) + let paths = store + .current_table_paths() + .await? + .context("storyline import batch produced no committed snapshot")?; + let total = committed_storylines + .checked_add(batch_len) + .context("import trajectory count overflow")?; + persisting_pchronicle::storage::write_storyline_manifest_at_uri( + store.root_uri(), + &paths.generation, + total, + 0, + ) + .await + .context("write progressive chronicle.manifest after storyline commit")?; + *append_generation = Some(paths.generation.clone()); + progress.note_committed(total)?; + Ok(total) } async fn run_canonical_event_import( @@ -890,7 +1794,7 @@ async fn run_canonical_event_import( }; if let Some((staging_path, output)) = staged_path { let mut cleanup = StagingPathGuard::new(staging_path.clone()); - publish_staged_dataset(&staging_path, &output, true)?; + publish_staged_dataset(&staging_path, &output, true, None).await?; cleanup.disarm(); } let response = ImportResponse { @@ -1311,14 +2215,18 @@ fn local_file_snapshot_ref(path: &Path) -> String { format!("local:{}", hash.finalize().to_hex()) } -#[derive(Debug)] +#[derive(Debug, Clone)] struct ImportFileCandidate { path: PathBuf, relative_path: PathBuf, output_relative_path: Option, - /// Object-store imports preload file bytes so the sync decode loop can - /// stay synchronous. Local imports leave this empty and open `path`. + /// Prefetched bytes (tests / rare callers). Normal imports leave this empty + /// and read local paths or object-store keys on demand. content: Option>, + /// Object-store Dataset root URI; when set, bytes are fetched lazily. + remote_root: Option, + /// Size from discovery (`stat` / object metadata) for progress totals. + size_hint: u64, } #[derive(Debug)] @@ -1329,26 +2237,532 @@ struct ImportedSource { input_bytes: usize, } -fn write_import_progress( - stderr: &mut dyn Write, - source: &str, - status: &str, - details: Option<(&DocumentFormat, usize, usize)>, -) -> Result<()> { - if let Some((format, trajectories, input_bytes)) = details { - writeln!( - stderr, - "import source={} status={} format={} trajectories={} input_bytes={}", - source, - status, - format.as_str(), - trajectories, - input_bytes, - )?; +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum ImportPhase { + Discovering, + Deleting, + Reading, + Parsing, + Writing, +} + +impl ImportPhase { + fn as_str(self) -> &'static str { + match self { + Self::Discovering => "discovering", + Self::Deleting => "deleting", + Self::Reading => "reading", + Self::Parsing => "parsing", + Self::Writing => "writing", + } + } +} + +/// Dense import progress: TTY paints three in-place lines; redirected stderr gets +/// one summary line per completed source (buffered, flushed at the end). +struct ImportProgress { + tty: bool, + discovered_files: u64, + discovered_bytes: u64, + imported_files: u64, + imported_bytes: u64, + /// Storyline trajectories successfully committed so far. + committed: u64, + /// Replace/drop delete progress (separate from discovery totals). + deleted_files: u64, + delete_total: u64, + phase: ImportPhase, + file: String, + /// Producer side of the read→parse pipeline (empty when idle). + reading_ahead: String, + painted: bool, + log_lines: Vec, + last_paint: Option, + /// Shared with index-build callbacks so Lance work updates line 3 in place. + surface: Arc>, +} + +#[derive(Debug, Clone)] +struct ImportProgressSurface { + tty: bool, + painted: bool, + deleting: bool, + line1: String, + line2: String, + reading_ahead: String, + phase: String, + file: String, +} + +impl ImportProgressSurface { + fn paint_activity(&mut self, activity: &str) -> Result<()> { + if !self.tty { + return Ok(()); + } + let file = if activity.is_empty() { + if self.file.is_empty() { + "-".to_owned() + } else { + truncate_middle(&self.file, 72) + } + } else { + truncate_middle(activity, 96) + }; + let line3 = if !self.reading_ahead.is_empty() && !activity.is_empty() { + format!( + "[reading] {} | [writing] {file}", + truncate_middle(&self.reading_ahead, 40), + ) + } else if !self.reading_ahead.is_empty() && self.phase == "reading" { + format!( + "[reading] {}", + truncate_middle(&self.reading_ahead, 96) + ) + } else if !self.reading_ahead.is_empty() { + format!( + "[reading] {} | [{}] {file}", + truncate_middle(&self.reading_ahead, 40), + self.phase, + ) + } else { + format!("[{}] {file}", if activity.is_empty() { self.phase.as_str() } else { "writing" }) + }; + + let mut err = std::io::stderr(); + if self.painted { + write!(err, "\x1b[2A").context("move import progress cursor")?; + } + if self.deleting { + write!(err, "\r\x1b[2K{}\n\r\x1b[2K\n\r\x1b[2K{line3}", self.line1) + .context("paint delete progress")?; + } else { + write!( + err, + "\r\x1b[2K{}\n\r\x1b[2K{}\n\r\x1b[2K{line3}", + self.line1, self.line2 + ) + .context("paint import progress")?; + } + err.flush().context("flush import progress")?; + self.painted = true; + Ok(()) + } +} + +impl ImportProgress { + fn new(tty: bool) -> Self { + Self { + tty, + discovered_files: 0, + discovered_bytes: 0, + imported_files: 0, + imported_bytes: 0, + committed: 0, + deleted_files: 0, + delete_total: 0, + phase: ImportPhase::Discovering, + file: String::new(), + reading_ahead: String::new(), + painted: false, + log_lines: Vec::new(), + last_paint: None, + surface: Arc::new(std::sync::Mutex::new(ImportProgressSurface { + tty, + painted: false, + deleting: false, + line1: String::new(), + line2: String::new(), + reading_ahead: String::new(), + phase: ImportPhase::Discovering.as_str().to_owned(), + file: String::new(), + })), + } + } + + fn attach_index_progress(&self) -> persisting_pchronicle::storage::IndexBuildProgressGuard { + let surface = Arc::clone(&self.surface); + persisting_pchronicle::storage::install_index_build_progress(Arc::new(move |message| { + if let Ok(mut surface) = surface.lock() { + let _ = surface.paint_activity(message); + } + })) + } + + fn reset_import_counters(&mut self) { + self.imported_files = 0; + self.imported_bytes = 0; + self.committed = 0; + self.deleted_files = 0; + self.delete_total = 0; + self.reading_ahead.clear(); + self.file.clear(); + } + + fn set_discovered(&mut self, files: u64, bytes: u64) -> Result<()> { + self.discovered_files = files; + self.discovered_bytes = bytes; + self.phase = ImportPhase::Discovering; + self.file.clear(); + self.paint(false) + } + + fn note_discovered(&mut self, file: &str, bytes: u64) -> Result<()> { + self.discovered_files = self.discovered_files.saturating_add(1); + self.discovered_bytes = self.discovered_bytes.saturating_add(bytes); + self.phase = ImportPhase::Discovering; + self.file = file.to_owned(); + // Throttle TTY paints during large listings so discovery stays responsive. + let should_paint = !self.tty + || self + .last_paint + .map(|at| at.elapsed() >= std::time::Duration::from_millis(100)) + .unwrap_or(true) + || self.discovered_files == 1 + || self.discovered_files % 64 == 0; + if should_paint { + self.paint(true)?; + } + Ok(()) + } + + fn note_scanning(&mut self, prefix: &str) -> Result<()> { + self.phase = ImportPhase::Discovering; + self.file = if prefix.is_empty() { + "/".to_owned() + } else { + format!("{prefix}/") + }; + let should_paint = !self.tty + || self + .last_paint + .map(|at| at.elapsed() >= std::time::Duration::from_millis(100)) + .unwrap_or(true); + if should_paint { + self.paint(true)?; + } + Ok(()) + } + + fn note_deleted(&mut self, deleted: u64, total: u64, path: &str) -> Result<()> { + self.deleted_files = deleted; + self.delete_total = total; + self.phase = ImportPhase::Deleting; + self.file = path.to_owned(); + if deleted == total { + // Always emit a final summary line for non-TTY logs. + return self.paint(false); + } + let should_paint = !self.tty + || path.is_empty() + || self + .last_paint + .map(|at| at.elapsed() >= std::time::Duration::from_millis(100)) + .unwrap_or(true) + || deleted == 1 + || deleted % 64 == 0; + if should_paint { + self.paint(true)?; + } + Ok(()) + } + + fn set_phase(&mut self, phase: ImportPhase, file: &str) -> Result<()> { + self.phase = phase; + self.file = file.to_owned(); + self.paint(true) + } + + fn set_reading_ahead(&mut self, file: &str) -> Result<()> { + self.reading_ahead = file.to_owned(); + let should_paint = !self.tty + || self + .last_paint + .map(|at| at.elapsed() >= std::time::Duration::from_millis(100)) + .unwrap_or(true); + if should_paint { + self.paint(true)?; + } + Ok(()) + } + + fn clear_reading_ahead(&mut self) -> Result<()> { + if self.reading_ahead.is_empty() { + return Ok(()); + } + self.reading_ahead.clear(); + self.paint(true) + } + + fn note_imported(&mut self, bytes: u64) -> Result<()> { + self.imported_files = self.imported_files.saturating_add(1); + self.imported_bytes = self.imported_bytes.saturating_add(bytes); + self.paint(false) + } + + fn note_committed(&mut self, committed: u64) -> Result<()> { + self.committed = committed; + self.phase = ImportPhase::Writing; + self.file = format!("commit trajectories={committed}"); + self.paint(false) + } + + fn finish(&mut self) -> Result<()> { + if let Ok(surface) = self.surface.lock() { + self.painted = surface.painted; + } + if self.tty && self.painted { + let mut err = std::io::stderr(); + writeln!(err).context("finish import progress")?; + err.flush().context("flush import progress")?; + self.painted = false; + if let Ok(mut surface) = self.surface.lock() { + surface.painted = false; + } + } + Ok(()) + } + + fn notice(&mut self, message: &str) -> Result<()> { + self.finish()?; + if self.tty { + let mut err = std::io::stderr(); + writeln!(err, "{message}").context("write import notice")?; + err.flush().context("flush import notice")?; + } else { + self.log_lines.push(message.to_owned()); + } + Ok(()) + } + + fn flush_log(self, out: &mut dyn Write) -> Result<()> { + for line in self.log_lines { + writeln!(out, "{line}").context("flush import progress log")?; + } + Ok(()) + } + + fn paint(&mut self, phase_only: bool) -> Result<()> { + if let Ok(surface) = self.surface.lock() { + self.painted = surface.painted; + } + let deleting = self.phase == ImportPhase::Deleting; + let line1 = if deleting { + format!( + "deleted:total = {}/{}", + self.deleted_files, self.delete_total + ) + } else { + format!( + "imported:discovered = {}/{}", + self.imported_files, self.discovered_files + ) + }; + let line2 = if deleting { + String::new() + } else { + format!( + "committed = {} ; size = {}:{}", + self.committed, + format_byte_count(self.imported_bytes), + format_byte_count(self.discovered_bytes) + ) + }; + let file = if self.file.is_empty() { + "-".to_owned() + } else { + truncate_middle(&self.file, 72) + }; + let line3 = if !self.reading_ahead.is_empty() + && matches!( + self.phase, + ImportPhase::Parsing | ImportPhase::Writing + ) + { + format!( + "[reading] {} | [{}] {file}", + truncate_middle(&self.reading_ahead, 48), + self.phase.as_str(), + ) + } else if !self.reading_ahead.is_empty() && self.phase == ImportPhase::Reading { + format!( + "[reading] {}", + truncate_middle(&self.reading_ahead, 96) + ) + } else { + format!("[{}] {file}", self.phase.as_str()) + }; + + if let Ok(mut surface) = self.surface.lock() { + surface.tty = self.tty; + surface.deleting = deleting; + surface.line1 = line1.clone(); + surface.line2 = line2.clone(); + surface.reading_ahead = self.reading_ahead.clone(); + surface.phase = self.phase.as_str().to_owned(); + surface.file = self.file.clone(); + surface.painted = self.painted; + } + + if self.tty { + let mut err = std::io::stderr(); + if self.painted { + write!(err, "\x1b[2A").context("move import progress cursor")?; + } + if deleting { + write!(err, "\r\x1b[2K{line1}\n\r\x1b[2K\n\r\x1b[2K{line3}") + .context("paint delete progress")?; + } else { + write!(err, "\r\x1b[2K{line1}\n\r\x1b[2K{line2}\n\r\x1b[2K{line3}") + .context("paint import progress")?; + } + err.flush().context("flush import progress")?; + self.painted = true; + if let Ok(mut surface) = self.surface.lock() { + surface.painted = true; + } + self.last_paint = Some(std::time::Instant::now()); + return Ok(()); + } + + if phase_only { + return Ok(()); + } + if deleting { + self.log_lines + .push(format!("{line1}; {line3}")); + } else { + self.log_lines + .push(format!("{line1}; {line2}; {line3}")); + } + Ok(()) + } +} + +fn format_byte_count(bytes: u64) -> String { + const KIB: f64 = 1024.0; + const MIB: f64 = 1024.0 * 1024.0; + const GIB: f64 = 1024.0 * 1024.0 * 1024.0; + let value = bytes as f64; + if value >= GIB { + format!("{:.1}GiB", value / GIB) + } else if value >= MIB { + format!("{:.1}MiB", value / MIB) + } else if value >= KIB { + format!("{:.1}KiB", value / KIB) } else { - writeln!(stderr, "import source={} status={}", source, status)?; + format!("{bytes}B") + } +} + +fn truncate_middle(value: &str, max_chars: usize) -> String { + let chars: Vec = value.chars().collect(); + if chars.len() <= max_chars { + return value.to_owned(); + } + if max_chars <= 3 { + return chars.into_iter().take(max_chars).collect(); + } + let head = (max_chars - 1) / 2; + let tail = max_chars - 1 - head; + let mut out: String = chars.iter().take(head).collect(); + out.push('…'); + out.extend(chars.iter().skip(chars.len() - tail)); + out +} + +#[cfg(test)] +mod import_progress_tests { + use super::*; + + #[test] + fn commit_batch_schedule_grows_to_cap() { + let mut schedule = CommitBatchSchedule::adaptive(); + assert_eq!(schedule.current(), 64); + schedule.after_commit(); + assert_eq!(schedule.current(), 128); + schedule.after_commit(); + assert_eq!(schedule.current(), 256); + schedule.after_commit(); + assert_eq!(schedule.current(), 512); + schedule.after_commit(); + assert_eq!(schedule.current(), 1024); + schedule.after_commit(); + assert_eq!(schedule.current(), 2048); + schedule.after_commit(); + assert_eq!(schedule.current(), 4096); + schedule.after_commit(); + assert_eq!(schedule.current(), 4096); + } + + #[test] + fn skippable_commit_errors_cover_s3_timeouts_and_preconditions() { + assert!(is_skippable_storyline_commit_error(&anyhow!( + "LanceError(IO): Generic S3 error: operation timed out" + ))); + assert!(is_skippable_storyline_commit_error(&anyhow!( + "ConditionNotMatch (persistent) PreconditionFailed" + ))); + assert!(!is_skippable_storyline_commit_error(&anyhow!( + "duplicate document_id policy rejected payload" + ))); + } + + #[test] + fn retract_imported_trajectories_from_tail_sources() { + let mut sources = vec![ + ImportedSource { + source_path: "a.json".into(), + format: DocumentFormat::Atif, + trajectories: 3, + input_bytes: 10, + }, + ImportedSource { + source_path: "b.json".into(), + format: DocumentFormat::Atif, + trajectories: 2, + input_bytes: 10, + }, + ]; + retract_imported_trajectories(&mut sources, 3); + assert_eq!(sources[0].trajectories, 2); + assert_eq!(sources[1].trajectories, 0); + } + + #[test] + fn commit_batch_schedule_fixed_stays_put() { + let mut schedule = CommitBatchSchedule::fixed(50); + assert_eq!(schedule.current(), 50); + schedule.after_commit(); + assert_eq!(schedule.current(), 50); + } + + #[test] + fn format_byte_count_uses_binary_units() { + assert_eq!(format_byte_count(512), "512B"); + assert_eq!(format_byte_count(1536), "1.5KiB"); + assert_eq!(format_byte_count(2 * 1024 * 1024), "2.0MiB"); + } + + #[test] + fn non_tty_progress_emits_dense_completed_lines() { + let mut progress = ImportProgress::new(false); + progress.set_discovered(2, 300).unwrap(); + progress.set_phase(ImportPhase::Reading, "a/long.json").unwrap(); + progress.set_phase(ImportPhase::Parsing, "a/long.json").unwrap(); + progress.note_imported(100).unwrap(); + progress.set_phase(ImportPhase::Writing, "b.json").unwrap(); + progress.note_imported(200).unwrap(); + progress.note_committed(3).unwrap(); + let mut out = Vec::new(); + progress.flush_log(&mut out).unwrap(); + let text = String::from_utf8(out).unwrap(); + assert!(text.contains("imported:discovered = 1/2"), "{text}"); + assert!(text.contains("imported:discovered = 2/2"), "{text}"); + assert!(text.contains("committed = 3"), "{text}"); + assert!(text.contains("size ="), "{text}"); + assert!(text.contains("[writing] commit trajectories=3") || text.contains("[writing] b.json") || text.contains("[parsing] a/long.json"), "{text}"); + assert!(!text.contains("status=fetching"), "{text}"); } - Ok(()) } fn collect_import_candidates(input: &Path) -> Result<(bool, Vec)> { @@ -1366,6 +2780,7 @@ fn collect_import_candidates(input: &Path) -> Result<(bool, Vec Result<(bool, Vec Result<(bool, Vec bool { }) } -async fn collect_object_store_import_candidates( - location: &DatasetLocation, - stderr: &mut dyn Write, -) -> Result<(bool, Vec)> { - writeln!( - stderr, - "import from={} status=discovering", - location.as_str() - ) - .context("write pChronicle import discovery progress")?; - let keys = location - .list_importable_json_objects(persisting_pchronicle::storage::DEFAULT_MAX_LOCAL_QUERY_FILES) - .await - .with_context(|| format!("discover importable objects under {}", location.as_str()))?; - if keys.is_empty() { - return Err(cli_boundary_error( - BoundaryCode::InvalidRequest, - "import object prefix contains no .json, .jsonl, or .ndjson files", - )); - } - writeln!( - stderr, - "import from={} status=discovered files={}", - location.as_str(), - keys.len() - ) - .context("write pChronicle import discovery progress")?; - - let mut candidates = Vec::with_capacity(keys.len()); - for key in keys { - let relative_path = PathBuf::from(&key); - write_import_progress(stderr, &key, "fetching", None)?; - let content = location - .read_relative_bytes(&key) - .await - .with_context(|| format!("read import object {key} under {}", location.as_str()))?; - candidates.push(ImportFileCandidate { - path: relative_path.clone(), - output_relative_path: Some(relative_path.clone()), - relative_path, - content: Some(content), - }); - } - Ok((true, candidates)) -} - -fn read_import_candidate_bytes( +async fn load_import_candidate_bytes( candidate: &ImportFileCandidate, max_input_bytes: usize, label: &str, @@ -1510,6 +2886,19 @@ fn read_import_candidate_bytes( ); return Ok(content.clone()); } + if let Some(remote_root) = &candidate.remote_root { + let key = candidate.relative_path.to_string_lossy().replace('\\', "/"); + let location = DatasetLocation::parse(remote_root)?; + let bytes = location + .read_relative_bytes(&key) + .await + .with_context(|| format!("read import object {key} under {remote_root}"))?; + anyhow::ensure!( + bytes.len() <= max_input_bytes, + "{label} exceeds max_input_bytes limit of {max_input_bytes}" + ); + return Ok(bytes); + } let file = std::fs::File::open(&candidate.path).with_context(|| format!("open {label}"))?; read_bounded(file, max_input_bytes, label) } @@ -1542,16 +2931,12 @@ enum ImportFormatResolution { enum StorylineImportInputs<'a> { Stdin(Option<&'a mut dyn Read>), - Files { - candidates: &'a [ImportFileCandidate], - next: usize, - }, } struct StorylineImportIterator<'a> { requested_format: ExchangeFormat, max_input_bytes: usize, - progress: &'a mut dyn Write, + progress: &'a mut ImportProgress, inputs: StorylineImportInputs<'a>, current: std::vec::IntoIter, imported_sources: Vec, @@ -1567,7 +2952,7 @@ impl<'a> StorylineImportIterator<'a> { requested_format: ExchangeFormat, max_input_bytes: usize, stdin: &'a mut dyn Read, - progress: &'a mut dyn Write, + progress: &'a mut ImportProgress, seen_document_ids: HashSet, duplicate_policy: DuplicateIdPolicy, ) -> Self { @@ -1587,42 +2972,16 @@ impl<'a> StorylineImportIterator<'a> { } } - fn files( - requested_format: ExchangeFormat, - max_input_bytes: usize, - candidates: &'a [ImportFileCandidate], - progress: &'a mut dyn Write, - seen_document_ids: HashSet, - duplicate_policy: DuplicateIdPolicy, - ) -> Self { - Self { - requested_format, - max_input_bytes, - progress, - inputs: StorylineImportInputs::Files { - candidates, - next: 0, - }, - current: Vec::new().into_iter(), - imported_sources: Vec::new(), - unknown_field_warnings: - persisting_pchronicle::model::UnknownFieldImportWarnings::default(), - skipped_warnings: Vec::new(), - seen_document_ids, - duplicate_policy, - failed: false, - } - } - - fn decode_next_source(&mut self) -> Result> { + async fn decode_next_source(&mut self) -> Result> { loop { let outcome = match &mut self.inputs { StorylineImportInputs::Stdin(stdin) => { let Some(stdin) = stdin.take() else { return Ok(None); }; - write_import_progress(self.progress, "stdin", "processing", None)?; + self.progress.set_phase(ImportPhase::Reading, "stdin")?; let input = read_bounded(stdin, self.max_input_bytes, "stdin")?; + self.progress.set_phase(ImportPhase::Parsing, "stdin")?; decode_import_source( self.requested_format, ImportOutputFormat::Storyline, @@ -1633,49 +2992,19 @@ impl<'a> StorylineImportIterator<'a> { &mut self.unknown_field_warnings, )? } - StorylineImportInputs::Files { candidates, next } => { - let Some(candidate) = candidates.get(*next) else { - return Ok(None); - }; - *next = next - .checked_add(1) - .context("import Source index overflow")?; - let label = format!("import source {}", candidate.relative_path.display()); - write_import_progress( - self.progress, - &candidate.relative_path.to_string_lossy(), - "processing", - None, - )?; - let input = - read_import_candidate_bytes(candidate, self.max_input_bytes, &label)?; - decode_import_source( - self.requested_format, - ImportOutputFormat::Storyline, - Some(&candidate.path), - Some(&candidate.relative_path), - candidate.output_relative_path.as_deref(), - &input, - &mut self.unknown_field_warnings, - )? - } }; match outcome { DecodeImportOutcome::Imported(decoded) => { - write_import_progress( - self.progress, + self.progress.set_phase( + ImportPhase::Writing, &decoded.diagnostic_path.to_string_lossy(), - "completed", - Some(( - &decoded.metadata.format, - decoded.metadata.trajectories, - decoded.metadata.input_bytes, - )), )?; + self.progress + .note_imported(decoded.metadata.input_bytes as u64)?; return Ok(Some(decoded)); } DecodeImportOutcome::Skipped { path, reason } => { - write_import_progress(self.progress, &path.to_string_lossy(), "skipped", None)?; + self.progress.note_imported(0)?; self.skipped_warnings .push(skipped_import_warning(&path, &reason)); } @@ -1696,12 +3025,8 @@ impl<'a> StorylineImportIterator<'a> { self.skipped_warnings, ) } -} -impl Iterator for StorylineImportIterator<'_> { - type Item = Result; - - fn next(&mut self) -> Option { + async fn next_document(&mut self) -> Option> { loop { if let Some(mut storyline) = self.current.next() { let original = storyline.document_id().to_string(); @@ -1738,7 +3063,7 @@ impl Iterator for StorylineImportIterator<'_> { if self.failed { return None; } - match self.decode_next_source() { + match self.decode_next_source().await { Ok(Some(decoded)) => { let mut metadata = decoded.metadata; metadata.trajectories = 0; @@ -2147,7 +3472,12 @@ impl Drop for StagingPathGuard { } } -fn publish_staged_dataset(staging: &Path, output: &Path, replace_existing: bool) -> Result<()> { +async fn publish_staged_dataset( + staging: &Path, + output: &Path, + replace_existing: bool, + progress: Option<&mut ImportProgress>, +) -> Result<()> { let parent = output .parent() .context("Dataset output must have a parent directory")?; @@ -2183,8 +3513,25 @@ fn publish_staged_dataset(staging: &Path, output: &Path, replace_existing: bool) backup.display() ) })?; - std::fs::remove_dir_all(&backup) - .with_context(|| format!("delete replaced Dataset backup {}", backup.display()))?; + let backup_location = DatasetLocation::parse( + backup + .to_str() + .context("replaced Dataset backup path is not valid UTF-8")?, + )?; + if let Some(progress) = progress { + backup_location + .remove_all_with_progress(|deleted, total, path| { + progress.note_deleted(deleted, total, path) + }) + .await + .with_context(|| format!("delete replaced Dataset backup {}", backup.display()))?; + progress.finish()?; + } else { + backup_location + .remove_all() + .await + .with_context(|| format!("delete replaced Dataset backup {}", backup.display()))?; + } sync_dataset_parent(parent)?; Ok(()) } diff --git a/crates/persisting-pchronicle-cli/src/lib.rs b/crates/persisting-pchronicle-cli/src/lib.rs index 71ebd85f..3c6d83e1 100644 --- a/crates/persisting-pchronicle-cli/src/lib.rs +++ b/crates/persisting-pchronicle-cli/src/lib.rs @@ -751,15 +751,23 @@ struct ImportArgs { #[arg(short = 'o', long = "output-format", value_enum)] output_format: Option, - /// Destination behavior: create a new Dataset, append, or replace. - #[arg(long, value_enum, default_value_t = ImportMode::Create)] - mode: ImportMode, + /// Replace an existing destination Dataset (after confirmation unless --yes). + #[arg(long, conflicts_with = "append")] + replace: bool, + + /// Append trajectories into an existing Storyline Dataset. + #[arg(long, conflicts_with = "replace")] + append: bool, + + /// Deprecated alias for --replace/--append/--create. Prefer --replace or --append. + #[arg(long, value_enum, hide = true)] + mode: Option, /// How append handles an existing document ID. #[arg(long, value_enum, value_name = "suffix|skip")] on_duplicate: Option, - /// Skip the destructive confirmation required by --mode replace. + /// Skip the destructive confirmation required by --replace. #[arg(short = 'y', long)] yes: bool, @@ -771,6 +779,12 @@ struct ImportArgs { #[arg(long, value_parser = parse_byte_size, default_value = "256MiB")] max_input_bytes: Option, + /// Fixed Storyline commit batch size. When omitted, batch size grows + /// 64 → 128 → … → 4096 (then stays at 4096) so early progress stays fine + /// while later commits amortize CURRENT / Lance overhead. + #[arg(long, value_name = "N")] + commit_every: Option, + /// Compact JSONL mapping. id/timestamp override $.id/$.timestamp; missing or invalid id values /// use source_filename#line_number; other names add JSONB columns. /// Example: --column id=$.event.id --column model=$.payload.model. @@ -778,6 +792,24 @@ struct ImportArgs { columns: Vec, } +impl ImportArgs { + fn mode(&self) -> Result { + match (self.replace, self.append, self.mode) { + (true, true, _) => Err(anyhow!("--replace and --append cannot be combined")), + (true, false, Some(ImportMode::Append)) => Err(anyhow!( + "--replace conflicts with --mode append; omit --mode" + )), + (false, true, Some(ImportMode::Replace)) => Err(anyhow!( + "--append conflicts with --mode replace; omit --mode" + )), + (true, false, _) => Ok(ImportMode::Replace), + (false, true, _) => Ok(ImportMode::Append), + (false, false, Some(mode)) => Ok(mode), + (false, false, None) => Ok(ImportMode::Create), + } + } +} + #[derive(Debug, Args)] struct DropArgs { /// Dataset path, URI, or dataset pin to permanently delete. @@ -1508,17 +1540,19 @@ pub async fn run_with_stdin( stdout: &mut dyn Write, stderr: &mut dyn Write, ) -> Result<()> { - run_with_stdio(cli, false, stdout_is_terminal, stdin, stdout, stderr).await + run_with_stdio(cli, false, stdout_is_terminal, false, stdin, stdout, stderr).await } pub async fn run_with_stdio( cli: Cli, stdin_is_terminal: bool, stdout_is_terminal: bool, + stderr_is_terminal: bool, stdin: &mut dyn Read, stdout: &mut dyn Write, stderr: &mut dyn Write, ) -> Result<()> { + server::request_log::init_cli_tracing(cli.log_level); let config = cli.config.as_deref(); let mut diagnostics = DiagnosticWriter::new(cli.log_level, stderr); match cli.command { @@ -1582,6 +1616,7 @@ pub async fn run_with_stdio( args, config, stdin_is_terminal, + stderr_is_terminal, stdin, stdout, &mut diagnostics, diff --git a/crates/persisting-pchronicle-cli/src/main.rs b/crates/persisting-pchronicle-cli/src/main.rs index 413ce6b6..cc21a484 100644 --- a/crates/persisting-pchronicle-cli/src/main.rs +++ b/crates/persisting-pchronicle-cli/src/main.rs @@ -42,6 +42,7 @@ fn main() -> ExitCode { async fn async_main(cli: Cli, debug_errors: bool) -> ExitCode { let stdin_is_terminal = io::stdin().is_terminal(); let stdout_is_terminal = io::stdout().is_terminal(); + let stderr_is_terminal = io::stderr().is_terminal(); // Do not hold StdoutLock/StderrLock for the process lifetime. `pchronicle // serve` logs from Tokio worker threads via tracing; on macOS those writes // take the stdout lock, so a process-wide lock deadlocks the runtime. @@ -53,6 +54,7 @@ async fn async_main(cli: Cli, debug_errors: bool) -> ExitCode { cli, stdin_is_terminal, stdout_is_terminal, + stderr_is_terminal, &mut stdin, &mut stdout, &mut stderr, diff --git a/crates/persisting-pchronicle-cli/src/onboard.rs b/crates/persisting-pchronicle-cli/src/onboard.rs index 3e9435a2..fe488e7b 100644 --- a/crates/persisting-pchronicle-cli/src/onboard.rs +++ b/crates/persisting-pchronicle-cli/src/onboard.rs @@ -7,7 +7,7 @@ use clap::{Args, Subcommand}; use super::{ AnalysisOptions, DatasetArgs, DatasetCommand, ErrorMode, ExchangeFormat, ExportArgs, - ExportFormat, FindArgs, ImportArgs, ImportMode, ImportOutputFormat, ListArgs, OutputFormat, + ExportFormat, FindArgs, ImportArgs, ImportOutputFormat, ListArgs, OutputFormat, QueryArgs, QueryOutputFormat, StatsReport, StatusArgs, run_dataset, run_export, run_find, run_import, run_list, run_query, run_stats_report, run_status, }; @@ -803,15 +803,19 @@ async fn capture_exchange(demo: &DemoWorkspace) -> Result { ), format: ExchangeFormat::Atif, output_format: Some(ImportOutputFormat::Preserve), - mode: ImportMode::Create, + replace: false, + append: false, + mode: None, on_duplicate: None, yes: false, stream: false, max_input_bytes: None, + commit_every: None, columns: Vec::new(), }, Some(&settings), false, + false, &mut empty_stdin, &mut import_stdout, &mut import_stderr, @@ -831,15 +835,19 @@ async fn capture_exchange(demo: &DemoWorkspace) -> Result { output: Some(storyline_output.to_string_lossy().into_owned()), format: ExchangeFormat::Atif, output_format: Some(ImportOutputFormat::Storyline), - mode: ImportMode::Create, + replace: false, + append: false, + mode: None, on_duplicate: None, yes: false, stream: false, max_input_bytes: None, + commit_every: None, columns: Vec::new(), }, Some(&settings), false, + false, &mut std::io::empty(), &mut storyline_stdout, &mut storyline_stderr, diff --git a/crates/persisting-pchronicle-cli/src/server/explorer.rs b/crates/persisting-pchronicle-cli/src/server/explorer.rs index 43450e5b..3d3496e0 100644 --- a/crates/persisting-pchronicle-cli/src/server/explorer.rs +++ b/crates/persisting-pchronicle-cli/src/server/explorer.rs @@ -1,7 +1,9 @@ use std::collections::{BTreeMap, BTreeSet}; use persisting_pchronicle::model::EventRecord; -use persisting_pchronicle::storage::CatalogEventProvenance; +use persisting_pchronicle::storage::{ + CatalogDataset, CatalogEventProvenance, CatalogSourceKind, DiscoveredSource, ShallowNavEntry, +}; use serde::{Deserialize, Serialize}; use serde_json::Value; @@ -67,6 +69,18 @@ pub(crate) fn catalog_tree( dataset: Option<&str>, prefix: &str, max_children: usize, +) -> CatalogTree { + catalog_tree_with_mounts(summaries, &[], dataset, prefix, max_children) +} + +/// Build an explorer tree from run summaries, then fold in catalog mounts / +/// sources so Dataset and Directory nodes remain navigable before any runs exist. +pub(crate) fn catalog_tree_with_mounts( + summaries: &[RunSummary], + datasets: &[CatalogDataset], + dataset: Option<&str>, + prefix: &str, + max_children: usize, ) -> CatalogTree { let prefix = prefix.trim().trim_matches('/'); let scoped = summaries @@ -88,9 +102,23 @@ pub(crate) fn catalog_tree( }) .sum(); let children = if dataset.is_none() { - fold_tree_children(dataset_children(&scoped), max_children, prefix) + fold_tree_children( + merge_dataset_children(dataset_children(&scoped), datasets), + max_children, + prefix, + ) } else { - fold_tree_children(file_children(&scoped, prefix), max_children, prefix) + let dataset_name = dataset.expect("dataset scope is some"); + let sources = datasets + .iter() + .find(|row| row.mount.name == dataset_name) + .map(|row| row.sources.as_slice()) + .unwrap_or(&[]); + fold_tree_children( + merge_file_children(file_children(&scoped, prefix), sources, prefix), + max_children, + prefix, + ) }; CatalogTree { dataset: dataset.map(str::to_string), @@ -102,6 +130,180 @@ pub(crate) fn catalog_tree( } } +/// Append one-level object/local children when catalog sources do not yet +/// expose the next path segment (typical while a Directory is still importing). +pub(crate) fn append_shallow_nav_children( + tree: &mut CatalogTree, + prefix: &str, + entries: &[ShallowNavEntry], + max_children: usize, +) { + if entries.is_empty() { + return; + } + let prefix = prefix.trim().trim_matches('/'); + let mut children = std::mem::take(&mut tree.children); + let existing: BTreeSet<_> = children.iter().map(|child| child.name.clone()).collect(); + for entry in entries { + if existing.contains(&entry.name) { + continue; + } + let path = if prefix.is_empty() { + entry.name.clone() + } else { + format!("{prefix}/{}", entry.name) + }; + children.push(CatalogTreeChild { + name: entry.name.clone(), + kind: if entry.is_dir { + "dir".into() + } else { + "file".into() + }, + data_type: entry + .dataset_kind + .clone() + .unwrap_or_else(|| { + if entry.is_dir { + "directory".into() + } else { + "other".into() + } + }), + path, + run_count: 0, + failed_count: 0, + total_tokens: None, + entries: Vec::new(), + }); + } + tree.children = fold_tree_children(children, max_children, prefix); +} + +fn merge_dataset_children( + mut children: Vec, + datasets: &[CatalogDataset], +) -> Vec { + let existing: BTreeSet<_> = children.iter().map(|child| child.name.clone()).collect(); + for dataset in datasets { + if existing.contains(&dataset.mount.name) { + continue; + } + children.push(CatalogTreeChild { + name: dataset.mount.name.clone(), + kind: "dataset".into(), + data_type: "unknown".into(), + path: dataset.mount.name.clone(), + run_count: 0, + failed_count: 0, + total_tokens: None, + entries: Vec::new(), + }); + } + children +} + +fn merge_file_children( + mut children: Vec, + sources: &[DiscoveredSource], + prefix: &str, +) -> Vec { + let mut groups = BTreeMap::::new(); + for child in &children { + let entry = groups.entry(child.name.clone()).or_insert(ChildAcc { + run_count: 0, + failed_count: 0, + has_deeper: child.kind == "dir", + data_types: BTreeSet::new(), + }); + entry.run_count = entry.run_count.max(child.run_count); + entry.failed_count = entry.failed_count.max(child.failed_count); + entry.has_deeper |= child.kind == "dir"; + if !child.data_type.is_empty() { + entry.data_types.insert(child.data_type.clone()); + } + } + for source in sources { + if source.file == "." { + continue; + } + let rest = if prefix.is_empty() { + source.file.as_str() + } else if source.file == prefix { + // Standing on this source: Directory stays navigable via shallow + // listing; leaf Stores/Files show no further path children here. + continue; + } else { + match source.file.strip_prefix(&format!("{prefix}/")) { + Some(rest) => rest, + None => continue, + } + }; + if rest.is_empty() { + continue; + } + let (name, has_deeper) = match rest.split_once('/') { + Some((name, _)) => (name, true), + None => ( + rest, + source.kind == CatalogSourceKind::Directory + || source.file.contains('/'), + ), + }; + // A Directory leaf under this prefix is always a folder to open. + let has_deeper = has_deeper || source.kind == CatalogSourceKind::Directory; + if name.is_empty() { + continue; + } + let entry = groups.entry(name.to_string()).or_insert(ChildAcc { + run_count: 0, + failed_count: 0, + has_deeper: false, + data_types: BTreeSet::new(), + }); + if let Some(count) = source.record_count { + let weight = usize::try_from(count).unwrap_or(usize::MAX); + entry.run_count = entry.run_count.max(weight); + } + if let Some(count) = source.failed_count { + let weight = usize::try_from(count).unwrap_or(usize::MAX); + entry.failed_count = entry.failed_count.max(weight); + } + entry.has_deeper |= has_deeper; + entry.data_types.insert(match source.kind { + CatalogSourceKind::Directory => "directory".into(), + CatalogSourceKind::Store | CatalogSourceKind::File => { + data_type(source.format.as_deref()).into() + } + }); + } + if groups.is_empty() { + return children; + } + children = groups + .into_iter() + .map(|(name, acc)| CatalogTreeChild { + path: if prefix.is_empty() { + name.clone() + } else { + format!("{prefix}/{name}") + }, + kind: if acc.has_deeper { + "dir".into() + } else { + "file".into() + }, + name, + run_count: acc.run_count, + failed_count: acc.failed_count, + data_type: combined_data_type(&acc.data_types), + total_tokens: None, + entries: Vec::new(), + }) + .collect(); + children +} + fn is_failed_status(status: &str) -> bool { matches!(status, "failed" | "error") } @@ -1506,6 +1708,59 @@ mod tests { ); } + #[test] + fn empty_runs_still_list_catalog_mounts_and_directories() { + use persisting_pchronicle::storage::{ + CatalogDataset, CatalogSourceKind, CatalogSourceStatus, DatasetMount, DiscoveredSource, + }; + + let mounts = vec![ + CatalogDataset { + mount: DatasetMount::new("default", "/tmp/default").unwrap(), + sources: Vec::new(), + }, + CatalogDataset { + mount: DatasetMount::new("prod", "s3://prod").unwrap(), + sources: vec![DiscoveredSource { + file: "infra".into(), + format: None, + kind: CatalogSourceKind::Directory, + revision: None, + projection_status: None, + projection_generation: None, + projection_candidates: 0, + size_bytes: None, + last_modified: None, + status: CatalogSourceStatus::Ready, + error: None, + record_count: None, + failed_count: None, + }], + }, + ]; + + let root = catalog_tree_with_mounts(&[], &mounts, None, "", 16); + assert_eq!(root.run_count, 0); + let names: Vec<_> = root + .children + .iter() + .map(|child| (child.name.as_str(), child.kind.as_str(), child.run_count)) + .collect(); + assert_eq!( + names, + vec![("default", "dataset", 0), ("prod", "dataset", 0)] + ); + + let prod = catalog_tree_with_mounts(&[], &mounts, Some("prod"), "", 16); + assert_eq!( + prod.children + .iter() + .map(|child| (child.name.as_str(), child.kind.as_str(), child.data_type.as_str())) + .collect::>(), + vec![("infra", "dir", "directory")] + ); + } + #[test] fn dataset_tree_groups_the_next_file_segment() { let tree = catalog_tree( diff --git a/crates/persisting-pchronicle-cli/src/server/mod.rs b/crates/persisting-pchronicle-cli/src/server/mod.rs index e60b67a4..cd34b86c 100644 --- a/crates/persisting-pchronicle-cli/src/server/mod.rs +++ b/crates/persisting-pchronicle-cli/src/server/mod.rs @@ -741,6 +741,155 @@ async fn try_compact_jsonl_runs_page( })) } +/// Directory mounts only expose immediate children in the catalog. Nested +/// Storyline leaves reached via explorer navigation are therefore absent from +/// SQL acceleration. When the client asks for an exact `file=` that is a +/// Storyline store under the mount, list document IDs directly from CURRENT. +async fn try_on_demand_storyline_runs_page( + state: &AppState, + query: &explorer::ExplorerRunsQuery, + request_id: &RequestId, +) -> Result, ApiError> { + if query + .q + .as_deref() + .is_some_and(|value| !value.trim().is_empty()) + { + return Ok(None); + } + let Some(file) = query + .file + .as_deref() + .map(str::trim) + .filter(|value| !value.is_empty()) + else { + return Ok(None); + }; + let Some(dataset_name) = query + .dataset + .as_deref() + .map(str::trim) + .filter(|value| !value.is_empty() && *value != "all") + else { + return Ok(None); + }; + + let runtime = current_catalog(state, request_id).await?; + let Some(dataset) = runtime.snapshot.dataset(dataset_name) else { + return Ok(None); + }; + // Prefer catalog-backed sources; only fall through for nested Directory paths. + if dataset.sources.iter().any(|source| { + source.kind != persisting_pchronicle::storage::CatalogSourceKind::Directory + && (source.file == file || source.file.starts_with(&format!("{file}/"))) + }) { + return Ok(None); + } + let under_directory = dataset.sources.iter().any(|source| { + source.kind == persisting_pchronicle::storage::CatalogSourceKind::Directory + && (file == source.file || file.starts_with(&format!("{}/", source.file))) + }); + if !under_directory && !dataset.sources.is_empty() { + return Ok(None); + } + + let location = persisting_pchronicle::storage::DatasetLocation::parse(&dataset.mount.uri) + .map_err(|error| fail(request_id, "explorer_runs", error))?; + let kind = location + .probe_nav_dataset_kind(file) + .await + .map_err(|error| fail(request_id, "explorer_runs", error))?; + if kind != Some("storyline") { + return Ok(None); + } + let uri = format!( + "{}/{}", + dataset.mount.uri.trim_end_matches('/'), + file.trim_matches('/') + ); + let store = persisting_pchronicle::storage::StorylineLanceStore::open_uri(&uri) + .await + .map_err(|error| fail(request_id, "explorer_runs", error))?; + let Some((_generation, ids)) = store + .document_ids_snapshot() + .await + .map_err(|error| fail(request_id, "explorer_runs", error))? + else { + let limit = query.limit.unwrap_or(50).clamp(1, 200); + return Ok(Some(explorer::RunExplorerPage { + snapshot: explorer::PageSnapshot { + offset: 0, + next_offset: 0, + total: 0, + has_more: false, + limit, + }, + records: Vec::new(), + path_index: Vec::new(), + search: explorer::RunSearchStatus::default(), + })); + }; + + let offset = query.offset.unwrap_or(0); + let limit = query.limit.unwrap_or(50).clamp(1, 200); + let total = ids.len(); + let page_ids = ids.into_iter().skip(offset).take(limit).collect::>(); + let page_records = page_ids + .into_iter() + .map(|document_id| { + let path = explorer::explorer_run_path( + dataset_name, + file, + &document_id, + &document_id, + None, + None, + ); + explorer::RunExplorerItem { + model: None, + search_preview: None, + run: RunSummary { + dataset: dataset_name.to_string(), + file: file.to_string(), + document_id: document_id.clone(), + run_id: None, + agent_id: "storyline".into(), + model_name: None, + session_id: document_id, + root_session_id: None, + path, + row_count: 1, + duplicate_event_ids: 0, + status: "completed".into(), + format: Some("storyline-lance".into()), + explorer_weight: None, + }, + } + }) + .collect::>(); + let next_offset = offset.saturating_add(page_records.len()); + let path_index = page_records + .iter() + .map(|item| item.run.clone()) + .collect::>(); + Ok(Some(explorer::RunExplorerPage { + snapshot: explorer::PageSnapshot { + offset, + next_offset, + total, + has_more: next_offset < total, + limit, + }, + records: page_records, + path_index, + search: explorer::RunSearchStatus { + fts_available: false, + mode: "none", + tokenizer: None, + }, + })) +} + async fn explorer_runs( State(state): State, request_id: RequestId, @@ -751,6 +900,9 @@ async fn explorer_runs( if let Some(page) = try_compact_jsonl_runs_page(&state, &query, &request_id).await? { return Ok(Json(page)); } + if let Some(page) = try_on_demand_storyline_runs_page(&state, &query, &request_id).await? { + return Ok(Json(page)); + } let dataset_filter = query .dataset .as_deref() @@ -999,7 +1151,13 @@ async fn explorer_tree( .map(str::trim) .filter(|value| !value.is_empty()); let prefix = query.prefix.as_deref().unwrap_or(""); - let mut tree = explorer::catalog_tree(&summaries, dataset, prefix, explorer::MAX_TREE_CHILDREN); + let mut tree = explorer::catalog_tree_with_mounts( + &summaries, + runtime.snapshot.datasets(), + dataset, + prefix, + explorer::MAX_TREE_CHILDREN, + ); if let Some(name) = tree.dataset.clone() { if tree.prefix.is_empty() && let Some(dataset) = runtime.snapshot.dataset(&name) @@ -1007,6 +1165,30 @@ async fn explorer_tree( tree.ready_sources = Some(dataset.ready_source_count()); tree.error_sources = Some(dataset.error_source_count()); } + // Directory prefixes often have no run summaries yet; fill the next + // level from the live Dataset URI so import progress stays navigable. + if tree.children.is_empty() + && let Some(dataset) = runtime.snapshot.dataset(&name) + { + let under_directory = dataset.sources.iter().any(|source| { + source.kind == persisting_pchronicle::storage::CatalogSourceKind::Directory + && (tree.prefix == source.file + || tree.prefix.starts_with(&format!("{}/", source.file))) + }); + if under_directory + && let Ok(location) = + persisting_pchronicle::storage::DatasetLocation::parse(&dataset.mount.uri) + && let Ok(entries) = location.list_shallow_nav(&tree.prefix).await + { + let prefix = tree.prefix.clone(); + explorer::append_shallow_nav_children( + &mut tree, + &prefix, + &entries, + explorer::MAX_TREE_CHILDREN, + ); + } + } let (duration_ms, total_tokens) = tree_prefix_metrics(&runtime, &name, &tree.prefix).await; tree.duration_ms = duration_ms; tree.total_tokens = total_tokens; @@ -1021,15 +1203,17 @@ async fn tree_run_summaries( request_id: &RequestId, ) -> Result, ApiError> { let mut summaries = Vec::new(); - let mut compact_with_manifest = BTreeSet::new(); + let mut manifest_weighted = BTreeSet::new(); for dataset in runtime.snapshot.datasets() { for source in &dataset.sources { - if source.format.as_deref() != Some("compact-jsonl/v1") { - continue; - } let Some(record_count) = source.record_count else { continue; }; + let is_compact = source.format.as_deref() == Some("compact-jsonl/v1"); + let is_storyline = source.format.as_deref() == Some("storyline-lance"); + if !is_compact && !is_storyline { + continue; + } let weight = usize::try_from(record_count).unwrap_or(usize::MAX); let path = explorer::explorer_run_path( &dataset.mount.name, @@ -1044,7 +1228,11 @@ async fn tree_run_summaries( file: source.file.clone(), document_id: String::new(), run_id: None, - agent_id: "compact-jsonl".into(), + agent_id: if is_compact { + "compact-jsonl".into() + } else { + "storyline".into() + }, model_name: None, session_id: source.file.clone(), root_session_id: None, @@ -1055,12 +1243,20 @@ async fn tree_run_summaries( format: source.format.clone(), explorer_weight: Some(weight.max(1)), }); - compact_with_manifest.insert((dataset.mount.name.clone(), source.file.clone())); + manifest_weighted.insert((dataset.mount.name.clone(), source.file.clone())); } } if !runtime.snapshot.datasets().iter().any(|dataset| { dataset.sources.iter().any(|source| { - source.format.as_deref() != Some("compact-jsonl/v1") || source.record_count.is_none() + if source.kind + == persisting_pchronicle::storage::CatalogSourceKind::Directory + { + return false; + } + match source.format.as_deref() { + Some("compact-jsonl/v1") | Some("storyline-lance") => source.record_count.is_none(), + _ => true, + } }) }) { return Ok(summaries); @@ -1071,7 +1267,7 @@ async fn tree_run_summaries( .await .map_err(|error| fail(request_id, "explorer_tree", error))?; for summary in full.iter() { - if compact_with_manifest.contains(&(summary.dataset.clone(), summary.file.clone())) { + if manifest_weighted.contains(&(summary.dataset.clone(), summary.file.clone())) { continue; } summaries.push(summary.clone()); @@ -1240,6 +1436,11 @@ async fn resolve_run_summary( matches.retain(|run| run.root_session_id.as_ref() == Some(root)); } if matches.is_empty() { + if let Some(run) = + try_resolve_on_demand_storyline_run(state, query, request_id).await? + { + return Ok(run); + } return Err(ApiError::not_found("run was not found")); } if matches.len() > 1 { @@ -1250,6 +1451,197 @@ async fn resolve_run_summary( Ok(matches.into_iter().next().expect("one matching run")) } +/// Synthesize a RunSummary for a nested Storyline leaf that is reachable under +/// a Directory mount but absent from the catalog snapshot. +async fn try_resolve_on_demand_storyline_run( + state: &AppState, + query: &SessionQuery, + request_id: &RequestId, +) -> Result, ApiError> { + let Some(dataset_name) = query + .dataset + .as_deref() + .map(str::trim) + .filter(|value| !value.is_empty() && *value != "all") + else { + return Ok(None); + }; + let Some(file) = query + .file + .as_deref() + .map(str::trim) + .filter(|value| !value.is_empty()) + else { + return Ok(None); + }; + let session_id = query.session_id.trim(); + if session_id.is_empty() { + return Ok(None); + } + + let runtime = current_catalog(state, request_id).await?; + let Some(dataset) = runtime.snapshot.dataset(dataset_name) else { + return Ok(None); + }; + if dataset.sources.iter().any(|source| { + source.kind != persisting_pchronicle::storage::CatalogSourceKind::Directory + && (source.file == file || source.file.starts_with(&format!("{file}/"))) + }) { + return Ok(None); + } + let under_directory = dataset.sources.iter().any(|source| { + source.kind == persisting_pchronicle::storage::CatalogSourceKind::Directory + && (file == source.file || file.starts_with(&format!("{}/", source.file))) + }); + if !under_directory && !dataset.sources.is_empty() { + return Ok(None); + } + + let location = persisting_pchronicle::storage::DatasetLocation::parse(&dataset.mount.uri) + .map_err(|error| fail(request_id, "resolve_run", error))?; + if location + .probe_nav_dataset_kind(file) + .await + .map_err(|error| fail(request_id, "resolve_run", error))? + != Some("storyline") + { + return Ok(None); + } + let uri = format!( + "{}/{}", + dataset.mount.uri.trim_end_matches('/'), + file.trim_matches('/') + ); + let store = persisting_pchronicle::storage::StorylineLanceStore::open_uri(&uri) + .await + .map_err(|error| fail(request_id, "resolve_run", error))?; + let Some((_generation, ids)) = store + .document_ids_snapshot() + .await + .map_err(|error| fail(request_id, "resolve_run", error))? + else { + return Ok(None); + }; + if !ids.iter().any(|id| id == session_id) { + return Ok(None); + } + let path = explorer::explorer_run_path( + dataset_name, + file, + session_id, + session_id, + None, + None, + ); + Ok(Some(RunSummary { + dataset: dataset_name.to_string(), + file: file.to_string(), + document_id: session_id.to_string(), + run_id: query.run_id.clone().filter(|value| !value.is_empty()), + agent_id: if query.agent_id.trim().is_empty() { + "storyline".into() + } else { + query.agent_id.clone() + }, + model_name: None, + session_id: session_id.to_string(), + root_session_id: query.root_session_id.clone(), + path, + row_count: 1, + duplicate_event_ids: 0, + status: "completed".into(), + format: Some("storyline-lance".into()), + explorer_weight: None, + })) +} + +async fn load_on_demand_storyline_bundle( + state: &AppState, + run: &RunSummary, + request_id: &RequestId, + op: &'static str, +) -> Result, ApiError> { + let runtime = current_catalog(state, request_id).await?; + let Some(dataset) = runtime.snapshot.dataset(&run.dataset) else { + return Ok(None); + }; + // Never shadow a catalog-registered leaf source with a direct open. + if dataset.sources.iter().any(|source| { + source.kind != persisting_pchronicle::storage::CatalogSourceKind::Directory + && source.file == run.file + }) { + return Ok(None); + } + let under_directory = dataset.sources.iter().any(|source| { + source.kind == persisting_pchronicle::storage::CatalogSourceKind::Directory + && (run.file == source.file || run.file.starts_with(&format!("{}/", source.file))) + }); + if !under_directory { + return Ok(None); + } + let uri = format!( + "{}/{}", + dataset.mount.uri.trim_end_matches('/'), + run.file.trim_matches('/') + ); + let store = persisting_pchronicle::storage::StorylineLanceStore::open_uri(&uri) + .await + .map_err(|error| fail(request_id, op, error))?; + let document_id = if run.document_id.is_empty() { + run.session_id.clone() + } else { + run.document_id.clone() + }; + let stories = store + .get_storylines_by_document_ids(&[document_id.clone()]) + .await + .map_err(|error| fail(request_id, op, error))?; + let Some(Some(storyline)) = stories.into_iter().next() else { + return Ok(None); + }; + let document = persisting_pchronicle::document::storyline_to_events(&storyline) + .map_err(|error| fail(request_id, op, error))?; + Ok(Some( + persisting_pchronicle::storage::CatalogTrajectoryBundle { + storyline, + event_view: persisting_pchronicle::storage::CatalogEventView { + provenance: CatalogEventProvenance::SyntheticFromStoryline, + document, + }, + }, + )) +} + +async fn catalog_or_on_demand_trajectory_bundle( + state: &AppState, + run: &RunSummary, + request_id: &RequestId, + op: &'static str, +) -> Result { + let runtime = current_catalog(state, request_id).await?; + let key = catalog_storyline_key(run); + let catalog_result = if state.live_reads { + runtime.snapshot.load_live_trajectory_bundle(&key).await + } else { + runtime.snapshot.load_trajectory_bundle(&key).await + }; + match catalog_result { + Ok(Some(bundle)) => Ok(bundle), + Ok(None) => load_on_demand_storyline_bundle(state, run, request_id, op) + .await? + .ok_or_else(|| ApiError::not_found("run was not found")), + Err(error) => { + if let Some(bundle) = + load_on_demand_storyline_bundle(state, run, request_id, op).await? + { + Ok(bundle) + } else { + Err(fail(request_id, op, error)) + } + } + } +} + fn catalog_storyline_key(run: &RunSummary) -> CatalogStorylineKey { CatalogStorylineKey { dataset: run.dataset.clone(), @@ -1296,15 +1688,9 @@ async fn load_events( request_id: &RequestId, ) -> Result { let run = resolve_run_summary(state, query, request_id).await?; - let runtime = current_catalog(state, request_id).await?; - let key = catalog_storyline_key(&run); - let document = if state.live_reads { - runtime.snapshot.load_live_events(&key).await - } else { - runtime.snapshot.load_events(&key).await - } - .map_err(|error| fail(request_id, "load_events", error))? - .ok_or_else(|| ApiError::not_found("run was not found"))?; + let bundle = + catalog_or_on_demand_trajectory_bundle(state, &run, request_id, "load_events").await?; + let document = bundle.event_view; let offset = query .offset .unwrap_or(0) @@ -1365,17 +1751,10 @@ async fn storyline( ) -> Result, ApiError> { let query = api_query(query)?; let run = resolve_run_summary(&state, &query, &request_id).await?; - let runtime = current_catalog(&state, &request_id).await?; - let key = catalog_storyline_key(&run); - let document = if state.live_reads { - runtime.snapshot.load_live_storyline(&key).await - } else { - runtime.snapshot.load_storyline(&key).await - } - .map_err(|error| fail(&request_id, "storyline", error))? - .ok_or_else(|| ApiError::not_found("run was not found"))?; + let bundle = + catalog_or_on_demand_trajectory_bundle(&state, &run, &request_id, "storyline").await?; Ok(Json( - serde_json::to_value(document) + serde_json::to_value(bundle.storyline) .map_err(anyhow::Error::from) .map_err(|error| fail(&request_id, "storyline", error))?, )) @@ -1538,14 +1917,8 @@ async fn load_trajectory( turns: Vec::new(), }); } - let key = catalog_storyline_key(&run); - let bundle = if state.live_reads { - runtime.snapshot.load_live_trajectory_bundle(&key).await - } else { - runtime.snapshot.load_trajectory_bundle(&key).await - } - .map_err(|error| fail(request_id, "load_trajectory", error))? - .ok_or_else(|| ApiError::not_found("run was not found"))?; + let bundle = + catalog_or_on_demand_trajectory_bundle(state, &run, request_id, "load_trajectory").await?; let event_provenance = bundle.event_view.provenance; let records = bundle.event_view.document.events; let document = bundle.storyline; @@ -1728,10 +2101,13 @@ async fn explorer_turns( let session = query.session(); let loaded = load_trajectory(&state, &session, &request_id).await?; let runtime = current_catalog(&state, &request_id).await?; + // Nested Directory Storylines are opened on-demand and are absent from the + // prepared catalog; skip FTS path probing and keep in-memory turn pages. let paths = runtime .snapshot .storyline_table_paths(&loaded.run.dataset, &loaded.run.file) - .map_err(|error| fail(&request_id, "explorer_turns", error))?; + .ok() + .flatten(); let mut fts_available = if let Some(paths) = paths.as_ref() { match storyline_steps_fts_available(paths).await { Ok(available) => available, @@ -1758,6 +2134,12 @@ async fn explorer_turns( .map(str::trim) .filter(|value| !value.is_empty()) { + if paths.is_none() { + // On-demand nested Storylines are not registered in DuckDB; filter + // the already-loaded turns in memory instead of SQL FTS. + search_mode = "memory"; + (loaded.turns.clone(), Some(needle)) + } else { let expression = crate::combine_match_expressions(&[needle.to_owned()]) .map_err(|error| ApiError::invalid_request(error.to_string()))? .ok_or_else(|| ApiError::invalid_request("search query must not be empty"))?; @@ -1822,6 +2204,7 @@ async fn explorer_turns( Vec::new() }; (turns, None) + } } else { (loaded.turns.clone(), query.q.as_deref()) }; diff --git a/crates/persisting-pchronicle-cli/src/server/request_log.rs b/crates/persisting-pchronicle-cli/src/server/request_log.rs index 5f26efc2..82f93d56 100644 --- a/crates/persisting-pchronicle-cli/src/server/request_log.rs +++ b/crates/persisting-pchronicle-cli/src/server/request_log.rs @@ -194,13 +194,19 @@ fn inject_request_id_json(bytes: Vec, request_id: &str) -> (Vec, Option< } pub(crate) fn tracing_filter(level: crate::LogLevel) -> String { - let level = match level { - crate::LogLevel::Error => "error", - crate::LogLevel::Warn => "warn", - crate::LogLevel::Info => "info", - crate::LogLevel::Debug => "debug", - }; - format!("pchronicle.serve={level}") + match level { + crate::LogLevel::Error => "error".to_owned(), + crate::LogLevel::Warn => { + "warn,persisting_pchronicle=warn,persisting_pchronicle_cli=warn".to_owned() + } + crate::LogLevel::Info => { + // Keep CLI/import diagnostics readable: silence Lance/OpenDAL INFO + // spam (dataset load, FTS workers, If-Match noise) while still + // showing pChronicle warn for lease/CAS issues. + "info,persisting_pchronicle=warn,pchronicle.serve=info,lance=warn,lance_index=warn,opendal=warn,pchronicle.opendal=warn,object_store=warn,pchronicle.object_store_gate=warn".to_owned() + } + crate::LogLevel::Debug => "debug".to_owned(), + } } pub(crate) fn init_warehouse_tracing(level: crate::LogLevel) { @@ -214,6 +220,11 @@ pub(crate) fn init_warehouse_tracing(level: crate::LogLevel) { .try_init(); } +/// Initialize stderr tracing for non-serve commands (import lease diagnostics, etc.). +pub(crate) fn init_cli_tracing(level: crate::LogLevel) { + init_warehouse_tracing(level); +} + pub(crate) fn log_warehouse_startup(listen: &str, datasets: &[String], snapshot_id: Option<&str>) { let datasets = datasets.join(","); if let Some(snapshot_id) = snapshot_id { diff --git a/crates/persisting-pchronicle-cli/src/server/tests.rs b/crates/persisting-pchronicle-cli/src/server/tests.rs index f16ae642..e52632d5 100644 --- a/crates/persisting-pchronicle-cli/src/server/tests.rs +++ b/crates/persisting-pchronicle-cli/src/server/tests.rs @@ -540,11 +540,11 @@ async fn query_evidence_info_truncates_sql() { fn warehouse_tracing_filter_matches_log_level() { assert_eq!( super::request_log::tracing_filter(crate::LogLevel::Info), - "pchronicle.serve=info" + "info,persisting_pchronicle=warn,pchronicle.serve=info" ); assert_eq!( super::request_log::tracing_filter(crate::LogLevel::Error), - "pchronicle.serve=error" + "error" ); } diff --git a/crates/persisting-pchronicle-cli/src/sync.rs b/crates/persisting-pchronicle-cli/src/sync.rs index 26af8c3f..4f3fc4f9 100644 --- a/crates/persisting-pchronicle-cli/src/sync.rs +++ b/crates/persisting-pchronicle-cli/src/sync.rs @@ -276,6 +276,43 @@ fn changed_paths( mod tests { use super::*; + #[test] + fn prepare_destination_preserves_object_store_uri() { + assert_eq!( + prepare_destination("s3://bucket/prod/infra/agent/agentcompass", "Warehouse").unwrap(), + "s3://bucket/prod/infra/agent/agentcompass" + ); + } + + #[tokio::test] + async fn sync_pin_source_is_resolved_not_canonicalized() { + let mut stderr = Vec::new(); + let error = run( + SyncArgs { + from: "@origin/agentcompass".into(), + to: "/tmp/pchronicle-sync-warehouse".into(), + convert: "/tmp/pchronicle-sync-convert".into(), + input_format: ExchangeFormat::Auto, + columns: Vec::new(), + interval_seconds: 1, + once: true, + }, + None, + &mut stderr, + ) + .await + .expect_err("pin must expand through settings, not local canonicalize"); + let message = format!("{error:#}"); + assert!( + !message.contains("canonicalize sync source"), + "{message}" + ); + assert!( + message.contains("unknown Dataset pin") || message.contains("resolve sync source"), + "{message}" + ); + } + #[test] fn changed_paths_include_create_modify_and_delete() { let old = BTreeMap::from([( diff --git a/crates/persisting-pchronicle-cli/src/tests.rs b/crates/persisting-pchronicle-cli/src/tests.rs index 5d1760c0..e99bf80b 100644 --- a/crates/persisting-pchronicle-cli/src/tests.rs +++ b/crates/persisting-pchronicle-cli/src/tests.rs @@ -455,7 +455,9 @@ fn canonical_parser_surface_matches_the_cli_guide() -> Result<()> { assert_eq!(import.output.as_deref(), Some("./imported")); assert_eq!(import.format, ExchangeFormat::Atif); assert_eq!(import.output_format, Some(ImportOutputFormat::Preserve)); - assert_eq!(import.mode, ImportMode::Create); + assert_eq!(import.mode().unwrap(), ImportMode::Create); + assert!(!import.replace); + assert!(!import.append); assert_eq!(import.on_duplicate, None); assert!(!import.yes); @@ -466,17 +468,48 @@ fn canonical_parser_surface_matches_the_cli_guide() -> Result<()> { "input.json", "-t", "./imported", - "--mode", - "append", + "--append", "--on-duplicate", "skip", ])?; let Command::Import(import) = cli.command else { panic!("expected import command") }; - assert_eq!(import.mode, ImportMode::Append); + assert_eq!(import.mode().unwrap(), ImportMode::Append); + assert!(import.append); + assert!(!import.replace); assert_eq!(import.on_duplicate, Some(DuplicateIdPolicy::Skip)); + let cli = Cli::try_parse_from([ + "pchronicle", + "import", + "-f", + "input.json", + "-t", + "./imported", + "--replace", + "--yes", + ])?; + let Command::Import(import) = cli.command else { + panic!("expected import command") + }; + assert_eq!(import.mode().unwrap(), ImportMode::Replace); + assert!(import.replace); + assert!(!import.append); + assert!(import.yes); + + assert!(Cli::try_parse_from([ + "pchronicle", + "import", + "-f", + "input.json", + "-t", + "./imported", + "--replace", + "--append", + ]) + .is_err()); + let cli = Cli::try_parse_from(["pchronicle", "drop", "./imported", "--yes"])?; let Command::Drop(drop) = cli.command else { panic!("expected drop command") @@ -2789,15 +2822,17 @@ async fn object_store_replace_clears_existing_prefix_before_import() -> Result<( &output, "--output-format", "storyline", - "--mode", - "replace", + "--replace", "--yes", ])?; let mut stdout = Vec::new(); let mut stderr = Vec::new(); run(cli, false, &mut stdout, &mut stderr).await?; let stderr = String::from_utf8(stderr)?; - assert!(stderr.contains("status=replacing")); + assert!( + stderr.contains("deleted:total =") || stderr.contains("[deleting]"), + "replace should report delete progress, got: {stderr}" + ); let store = StorylineLanceStore::open_uri(&output).await?; let ids = store @@ -3262,8 +3297,7 @@ async fn append_storyline_import_suffixes_or_skips_existing_document_ids() -> Re duplicate.to_str().unwrap(), "--to", output.to_str().unwrap(), - "--mode", - "append", + "--append", ])?; let mut append_stdout = Vec::new(); let mut append_stderr = Vec::new(); @@ -3281,8 +3315,7 @@ async fn append_storyline_import_suffixes_or_skips_existing_document_ids() -> Re duplicate.to_str().unwrap(), "--to", output.to_str().unwrap(), - "--mode", - "append", + "--append", "--on-duplicate", "skip", ])?; @@ -3341,8 +3374,7 @@ async fn replace_and_drop_require_confirmation_and_accept_yes() -> Result<()> { output.to_str().unwrap(), "--output-format", "storyline", - "--mode", - "replace", + "--replace", ])?; let error = run(replace_without_yes, false, &mut Vec::new(), &mut Vec::new()) .await @@ -3360,8 +3392,7 @@ async fn replace_and_drop_require_confirmation_and_accept_yes() -> Result<()> { output.to_str().unwrap(), "--output-format", "storyline", - "--mode", - "replace", + "--replace", "--yes", ])?; assert!( @@ -3369,7 +3400,10 @@ async fn replace_and_drop_require_confirmation_and_accept_yes() -> Result<()> { .await .is_err() ); - assert!(output.join("old.marker").exists()); + // Storyline --replace clears the destination before import (not atomic). + assert!(!output.join("old.marker").exists()); + fs::create_dir_all(&output)?; + fs::write(output.join("old.marker"), "old")?; fs::write( &input, @@ -3387,8 +3421,7 @@ async fn replace_and_drop_require_confirmation_and_accept_yes() -> Result<()> { output.to_str().unwrap(), "--output-format", "storyline", - "--mode", - "replace", + "--replace", "--yes", ])?; run(replace, false, &mut Vec::new(), &mut Vec::new()).await?; @@ -3431,6 +3464,7 @@ async fn replace_and_drop_require_confirmation_and_accept_yes() -> Result<()> { interactive_drop, true, false, + false, &mut confirmation, &mut Vec::new(), &mut prompt, @@ -3512,7 +3546,9 @@ async fn directory_import_failure_does_not_publish_partial_output() -> Result<() .await .unwrap_err(); assert!(format!("{error:#}").contains("z-invalid.json"), "{error:#}"); - assert!(!output.exists()); + if output_format == ImportOutputFormat::Preserve { + assert!(!output.exists()); + } } assert!(!fs::read_dir(temp.path())?.any(|entry| { entry diff --git a/crates/persisting-pchronicle/src/search/storyline.rs b/crates/persisting-pchronicle/src/search/storyline.rs index dc8eca5a..8eaca396 100644 --- a/crates/persisting-pchronicle/src/search/storyline.rs +++ b/crates/persisting-pchronicle/src/search/storyline.rs @@ -68,9 +68,11 @@ pub(crate) async fn ensure_storyline_search_indexes(dataset: &mut Dataset) -> Re } ensure_default_jieba_model()?; + let table = crate::store::index_build_progress::table_label(dataset.uri()).to_owned(); + let mut jobs: Vec<(&str, &str)> = Vec::new(); for field in schema.fields() { if lance_arrow::json::is_json_field(field) { - ensure_storyline_search_index(dataset, field.name(), Some("json")).await?; + jobs.push((field.name(), "json")); } } for column in STORYLINE_FTS_COLUMNS { @@ -78,9 +80,18 @@ pub(crate) async fn ensure_storyline_search_indexes(dataset: &mut Dataset) -> Re .field_with_name(column) .is_ok_and(|field| !lance_arrow::json::is_json_field(field)) { - ensure_storyline_search_index(dataset, column, None).await?; + jobs.push((*column, "fts")); } } + let total = jobs.len(); + for (offset, (column, kind)) in jobs.into_iter().enumerate() { + crate::store::index_build_progress::note(format!( + "index {table}.{column} {kind} {}/{total}", + offset + 1 + )); + let tokenizer = if kind == "json" { Some("json") } else { None }; + ensure_storyline_search_index(dataset, column, tokenizer).await?; + } Ok(()) } diff --git a/crates/persisting-pchronicle/src/storage.rs b/crates/persisting-pchronicle/src/storage.rs index b524f303..0f63a95f 100644 --- a/crates/persisting-pchronicle/src/storage.rs +++ b/crates/persisting-pchronicle/src/storage.rs @@ -25,6 +25,11 @@ pub use crate::discovery::{ drop_lifecycle_run_partitions, expand_story_locations, expand_story_locations_blocking, }; +#[cfg(feature = "lance-store")] +pub use crate::store::index_build_progress::{ + Guard as IndexBuildProgressGuard, install as install_index_build_progress, +}; + #[cfg(feature = "lance-store")] pub use crate::store::{ AppendOutcome, AttemptRecord, AttemptRecordState, AttemptRegistry, CatalogDataset, @@ -36,6 +41,7 @@ pub use crate::store::{ DEFAULT_CONTENT_PREVIEW_BYTES, DEFAULT_DATASET_NAME, DEFAULT_MAX_EVENT_FALLBACK_BYTES, DEFAULT_MAX_EVENT_FALLBACK_ROWS, DEFAULT_PHYSICAL_PAGE_LIMIT, DatasetCatalogSnapshot, DatasetLocation, DatasetLocationKind, DatasetMount, DiscoveredSource, EventFactSnapshot, + ImportableObjectEvent, ShallowNavEntry, EventLogLayoutStats, EventWriterFence, ExportOutcome, LanceMaintenanceOptions, LanceMaintenanceReport, LeaseAcquireOutcome, ManifestKind, ManifestStats, NamespacePath, ObjectStoreManifestWriteMode, PhysicalColumn, PhysicalDataFile, PhysicalFileLayout, @@ -44,10 +50,11 @@ pub use crate::store::{ RawEventLanceStore, ReplayOutcome, RunControlStore, StorylineContentOptions, StorylineContentReadMode, StorylineDataSource, StorylineDataSourceOptions, StorylineLanceStore, StorylineMaintenanceReport, StorylineProjectionLineage, StorylineStreamImportReport, - StorylineTablePaths, TrajectoryStats, attempt_registry_now_ms, distinct_session_ids_in_run, + StorylineStreamOptions, StorylineTablePaths, TrajectoryStats, attempt_registry_now_ms, distinct_session_ids_in_run, export_source_dirs, export_story_bundle, inspect_physical_file, inspect_physical_layout, - inspect_physical_page, list_physical_sources, load_manifest, raw_event_lance_path, - write_compact_jsonl_manifest, + inspect_physical_page, list_physical_sources, load_manifest, load_manifest_at_uri, + raw_event_lance_path, write_compact_jsonl_manifest, write_storyline_manifest, + write_storyline_manifest_at_uri, }; // Compatibility exports; new callers should use `crate::search`. diff --git a/crates/persisting-pchronicle/src/store/catalog/discovery.rs b/crates/persisting-pchronicle/src/store/catalog/discovery.rs index 7d077796..73cfbf16 100644 --- a/crates/persisting-pchronicle/src/store/catalog/discovery.rs +++ b/crates/persisting-pchronicle/src/store/catalog/discovery.rs @@ -150,6 +150,14 @@ pub(super) async fn freeze_candidate( source_row.revision = Some(CatalogSourceRevision::Storyline { generation: paths.generation.clone(), }); + if let Ok(Some(manifest)) = + crate::store::chronicle_manifest::load_manifest_at_uri(&uri).await + && manifest.is_storyline_leaf() + && let Some(stats) = manifest.stats + { + source_row.record_count = Some(stats.record_count); + source_row.failed_count = Some(stats.failed_count); + } Ok(( source_row, Arc::new(LazySource::new( @@ -719,23 +727,38 @@ async fn collect_manifest_subtree( while let Some((current, current_manifest)) = stack.pop() { match current_manifest.kind { ManifestKind::Leaf => { - anyhow::ensure!( - current_manifest.is_compact_jsonl_leaf(), - "chronicle.manifest leaf format {:?} is not supported for discovery yet", - current_manifest.format - ); let metadata = fs::metadata(¤t)?; let file = if current == mount_root { ".".into() } else { relative_catalog_path(mount_root, ¤t, true)? }; - candidates.push(Candidate::Compact { - file, - uri: canonical_local_uri(¤t)?, - size_bytes: Some(metadata.len()), - last_modified: modified_string(&metadata), - }); + if current_manifest.is_compact_jsonl_leaf() { + candidates.push(Candidate::Compact { + file, + uri: canonical_local_uri(¤t)?, + size_bytes: Some(metadata.len()), + last_modified: modified_string(&metadata), + }); + } else if current_manifest.is_storyline_leaf() { + anyhow::ensure!( + current.join("CURRENT").is_file(), + "storyline chronicle.manifest requires CURRENT at {}", + current.display() + ); + let current_meta = fs::metadata(current.join("CURRENT"))?; + candidates.push(Candidate::Storyline { + file, + uri: canonical_local_uri(¤t)?, + size_bytes: Some(current_meta.len()), + last_modified: modified_string(¤t_meta), + }); + } else { + anyhow::bail!( + "chronicle.manifest leaf format {:?} is not supported for discovery yet", + current_manifest.format + ); + } } ManifestKind::Branch => { let mut entries = fs::read_dir(¤t) @@ -1014,18 +1037,36 @@ async fn probe_object_prefix( manifest.validate()?; match manifest.kind { ManifestKind::Leaf => { - anyhow::ensure!( - manifest.is_compact_jsonl_leaf(), + let meta = RemoteObjectMeta::from(entry); + if manifest.is_compact_jsonl_leaf() { + return Ok(Some(ObjectProbe::Source(Candidate::Compact { + file: source_file, + uri: child_uri(root_uri, relative), + size_bytes: Some(meta.size), + last_modified: Some(meta.last_modified), + }))); + } + if manifest.is_storyline_leaf() { + let current = store + .stat_file(&join("CURRENT")) + .await? + .ok_or_else(|| { + anyhow::anyhow!( + "storyline chronicle.manifest requires CURRENT under {relative}" + ) + })?; + let current_meta = RemoteObjectMeta::from(current); + return Ok(Some(ObjectProbe::Source(Candidate::Storyline { + file: source_file, + uri: child_uri(root_uri, relative), + size_bytes: Some(current_meta.size), + last_modified: Some(current_meta.last_modified), + }))); + } + anyhow::bail!( "chronicle.manifest leaf format {:?} is not supported for discovery yet", manifest.format ); - let meta = RemoteObjectMeta::from(entry); - return Ok(Some(ObjectProbe::Source(Candidate::Compact { - file: source_file, - uri: child_uri(root_uri, relative), - size_bytes: Some(meta.size), - last_modified: Some(meta.last_modified), - }))); } ManifestKind::Branch => return Ok(Some(ObjectProbe::Branch)), } diff --git a/crates/persisting-pchronicle/src/store/chronicle_manifest.rs b/crates/persisting-pchronicle/src/store/chronicle_manifest.rs index c96c6aa6..6aafd1ce 100644 --- a/crates/persisting-pchronicle/src/store/chronicle_manifest.rs +++ b/crates/persisting-pchronicle/src/store/chronicle_manifest.rs @@ -10,6 +10,8 @@ use serde::{Deserialize, Serialize}; pub const CHRONICLE_MANIFEST_FILE: &str = "chronicle.manifest"; pub const CHRONICLE_MANIFEST_SCHEMA_VERSION: u32 = 1; pub const COMPACT_JSONL_FORMAT: &str = "compact-jsonl/v1"; +/// Leaf format for a committed Storyline Lance store (RFC-0015 extension). +pub const STORYLINE_FORMAT: &str = "storyline/v1"; #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] #[serde(rename_all = "snake_case")] @@ -66,6 +68,28 @@ impl ChronicleManifest { } } + pub fn leaf_storyline( + fingerprint: impl Into, + record_count: u64, + failed_count: u64, + ) -> Self { + Self { + schema_version: CHRONICLE_MANIFEST_SCHEMA_VERSION, + kind: ManifestKind::Leaf, + format: Some(STORYLINE_FORMAT.into()), + identity: Some(ManifestIdentity { + fingerprint: fingerprint.into(), + }), + stats: Some(ManifestStats { + record_count, + failed_count, + min_timestamp: None, + max_timestamp: None, + total_tokens: None, + }), + } + } + pub fn branch() -> Self { Self { schema_version: CHRONICLE_MANIFEST_SCHEMA_VERSION, @@ -121,6 +145,12 @@ impl ChronicleManifest { && self.format.as_deref() == Some(COMPACT_JSONL_FORMAT) && self.validate().is_ok() } + + pub fn is_storyline_leaf(&self) -> bool { + self.kind == ManifestKind::Leaf + && self.format.as_deref() == Some(STORYLINE_FORMAT) + && self.validate().is_ok() + } } pub fn manifest_path(root: impl AsRef) -> PathBuf { @@ -131,6 +161,10 @@ pub fn lance_version_fingerprint(version: u64) -> String { format!("lance:version:{version}") } +pub fn storyline_generation_fingerprint(generation: impl AsRef) -> String { + format!("storyline:generation:{}", generation.as_ref()) +} + pub fn load_manifest(root: impl AsRef) -> Result> { let path = manifest_path(root); if !path.is_file() { @@ -201,6 +235,71 @@ pub fn write_compact_jsonl_manifest( atomic_write_manifest(root, &manifest) } +pub fn write_storyline_manifest( + root: impl AsRef, + generation: impl AsRef, + record_count: u64, + failed_count: u64, +) -> Result<()> { + let manifest = ChronicleManifest::leaf_storyline( + storyline_generation_fingerprint(generation), + record_count, + failed_count, + ); + atomic_write_manifest(root, &manifest) +} + +/// Publish a Storyline leaf manifesto at a local path or object-store Dataset URI. +pub async fn write_storyline_manifest_at_uri( + root_uri: &str, + generation: impl AsRef, + record_count: u64, + failed_count: u64, +) -> Result<()> { + let manifest = ChronicleManifest::leaf_storyline( + storyline_generation_fingerprint(generation), + record_count, + failed_count, + ); + manifest.validate()?; + let location = crate::store::location::DatasetLocation::parse(root_uri) + .with_context(|| format!("parse Dataset URI for chronicle.manifest ({root_uri})"))?; + if let Some(path) = location.local_path() { + return atomic_write_manifest(path, &manifest); + } + let encoded = toml::to_string_pretty(&manifest).context("encode chronicle.manifest")?; + location + .write_relative_bytes(CHRONICLE_MANIFEST_FILE, encoded.as_bytes()) + .await + .with_context(|| format!("write chronicle.manifest under {root_uri}")) +} + +/// Load a manifesto from a local path or object-store Dataset URI. +pub async fn load_manifest_at_uri(root_uri: &str) -> Result> { + let location = crate::store::location::DatasetLocation::parse(root_uri) + .with_context(|| format!("parse Dataset URI for chronicle.manifest ({root_uri})"))?; + if let Some(path) = location.local_path() { + return load_manifest(path); + } + match location.read_relative_bytes(CHRONICLE_MANIFEST_FILE).await { + Ok(bytes) => { + let text = std::str::from_utf8(&bytes).context("chronicle.manifest must be UTF-8")?; + let manifest: ChronicleManifest = + toml::from_str(text).context("parse chronicle.manifest")?; + manifest.validate()?; + Ok(Some(manifest)) + } + Err(error) => { + let message = error.to_string(); + if message.contains("not found") || message.contains("NotFound") { + Ok(None) + } else { + Err(error).with_context(|| format!("read chronicle.manifest under {root_uri}")) + } + } + } +} + /// True when a compact-jsonl leaf manifesto matches one Lance version. pub fn compact_jsonl_manifest_matches(manifest: &ChronicleManifest, lance_version: u64) -> bool { manifest.is_compact_jsonl_leaf() @@ -214,6 +313,17 @@ mod tests { use super::*; use tempfile::tempdir; + #[test] + fn storyline_leaf_round_trip() { + let manifest = ChronicleManifest::leaf_storyline("storyline:generation:abc", 7, 1); + manifest.validate().unwrap(); + assert!(manifest.is_storyline_leaf()); + assert!(!manifest.is_compact_jsonl_leaf()); + let encoded = toml::to_string_pretty(&manifest).unwrap(); + let decoded: ChronicleManifest = toml::from_str(&encoded).unwrap(); + assert_eq!(decoded, manifest); + } + #[test] fn leaf_round_trip_and_validation() { let manifest = ChronicleManifest::leaf_compact_jsonl("lance:version:3", 12); diff --git a/crates/persisting-pchronicle/src/store/index_build_progress.rs b/crates/persisting-pchronicle/src/store/index_build_progress.rs new file mode 100644 index 00000000..e549dad2 --- /dev/null +++ b/crates/persisting-pchronicle/src/store/index_build_progress.rs @@ -0,0 +1,54 @@ +//! Optional UI hook for long-running Lance index builds. +//! +//! Import / maintain callers can install a short-lived listener so progress stays +//! on the dense TTY surface instead of relying on Lance's INFO spam. + +use std::sync::{Arc, Mutex, OnceLock}; + +type Listener = Arc; + +fn slot() -> &'static Mutex> { + static SLOT: OnceLock>> = OnceLock::new(); + SLOT.get_or_init(|| Mutex::new(None)) +} + +/// Restores the previous listener when dropped. +pub struct Guard { + previous: Option, +} + +impl Drop for Guard { + fn drop(&mut self) { + if let Ok(mut slot) = slot().lock() { + *slot = self.previous.take(); + } + } +} + +/// Install a process-wide index-progress listener for the current scope. +pub fn install(listener: Arc) -> Guard { + let previous = match slot().lock() { + Ok(mut slot) => slot.replace(listener), + Err(_) => None, + }; + Guard { previous } +} + +/// Report a short, single-line index activity message (best-effort). +pub fn note(message: impl AsRef) { + let Ok(slot) = slot().lock() else { + return; + }; + if let Some(listener) = slot.as_ref() { + listener(message.as_ref()); + } +} + +pub(crate) fn table_label(uri: &str) -> &str { + let trimmed = uri.trim_end_matches('/'); + trimmed + .rsplit('/') + .next() + .unwrap_or(trimmed) + .trim_end_matches(".lance") +} diff --git a/crates/persisting-pchronicle/src/store/location.rs b/crates/persisting-pchronicle/src/store/location.rs index fcf46a67..113c3381 100644 --- a/crates/persisting-pchronicle/src/store/location.rs +++ b/crates/persisting-pchronicle/src/store/location.rs @@ -1,5 +1,6 @@ //! Dataset URI facade: one parse/exists/put path for local and object stores. +use std::collections::BTreeSet; use std::fs::{File, OpenOptions}; use std::io::{ErrorKind, Write}; use std::path::{Path, PathBuf}; @@ -9,6 +10,29 @@ use url::Url; use super::opendal_store::Store as OpendalStore; +/// One discovery event while walking importable JSON objects. +#[derive(Debug, Clone)] +pub enum ImportableObjectEvent { + /// Prefix currently being shallow-listed (`""` for the Dataset root). + Scanning { prefix: String }, + /// Importable `.json` / `.jsonl` / `.ndjson` object. + File { + key: String, + size: u64, + modified: Option, + }, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ShallowNavEntry { + pub name: String, + /// Navigational folder. Dataset leaves are never directories for explorer. + pub is_dir: bool, + /// Explorer data_type when this child is a Dataset leaf (`storyline`, + /// `compact-jsonl`, `other`, …). `None` for plain directories/files. + pub dataset_kind: Option, +} + #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum DatasetLocationKind { Local, @@ -211,6 +235,232 @@ impl DatasetLocation { Ok(bytes) } + /// Classify a Dataset-relative path as a navigable Dataset leaf, if markers + /// are present (`CURRENT`, leaf `chronicle.manifest`, events manifest). + pub async fn probe_nav_dataset_kind(&self, relative: &str) -> Result> { + let relative = relative.trim().trim_matches('/'); + anyhow::ensure!( + !relative.split('/').any(|part| part == ".."), + "relative object path must not contain '..'" + ); + if let Some(root) = &self.local_path { + let dir = if relative.is_empty() { + root.clone() + } else { + root.join(relative) + }; + if !dir.is_dir() { + return Ok(None); + } + if let Some(manifest) = crate::store::chronicle_manifest::try_load_manifest(&dir) { + if manifest.is_storyline_leaf() { + return Ok(Some("storyline")); + } + if manifest.is_compact_jsonl_leaf() { + return Ok(Some("compact-jsonl")); + } + if matches!(manifest.kind, crate::store::ManifestKind::Leaf) { + return Ok(Some("other")); + } + } + if dir.join("CURRENT").is_file() { + return Ok(Some("storyline")); + } + if dir.join("events.lance/_manifest.json").is_file() + || (dir.file_name().is_some_and(|name| name == "events.lance") + && dir.join("_manifest.json").is_file()) + { + return Ok(Some("other")); + } + return Ok(None); + } + + let store = OpendalStore::from_uri(&self.uri).await?; + let join = |name: &str| { + if relative.is_empty() { + name.to_string() + } else { + format!("{relative}/{name}") + } + }; + if let Some(entry) = store + .stat_file(&join(crate::store::CHRONICLE_MANIFEST_FILE)) + .await? + { + if let Some((bytes, _)) = store.read(&entry.path).await? { + if let Ok(text) = std::str::from_utf8(&bytes) + && let Ok(manifest) = + toml::from_str::(text) + && manifest.validate().is_ok() + { + if manifest.is_storyline_leaf() { + return Ok(Some("storyline")); + } + if manifest.is_compact_jsonl_leaf() { + return Ok(Some("compact-jsonl")); + } + if matches!(manifest.kind, crate::store::ManifestKind::Leaf) { + return Ok(Some("other")); + } + } + } + } + if store.stat_file(&join("CURRENT")).await?.is_some() { + return Ok(Some("storyline")); + } + if store + .stat_file(&join("events.lance/_manifest.json")) + .await? + .is_some() + || (relative.ends_with("events.lance") + && store.stat_file(&join("_manifest.json")).await?.is_some()) + { + return Ok(Some("other")); + } + Ok(None) + } + + /// Immediate children under a Dataset-relative prefix for explorer navigation. + /// + /// Returns directories and importable JSON files only. Hidden names, Lance + /// table interiors, and other leaf objects are skipped so the tree stays + /// useful while imports are still writing nested paths. + /// + /// If `relative` itself is already a Dataset leaf (Storyline / compact / + /// events), returns an empty list so callers treat it as a source file + /// instead of drilling into Lance internals like `generations/`. + pub async fn list_shallow_nav(&self, relative: &str) -> Result> { + let relative = relative.trim().trim_matches('/'); + anyhow::ensure!( + !relative.split('/').any(|part| part == ".."), + "relative object path must not contain '..'" + ); + if self.probe_nav_dataset_kind(relative).await?.is_some() { + return Ok(Vec::new()); + } + if let Some(root) = &self.local_path { + let dir = if relative.is_empty() { + root.clone() + } else { + root.join(relative) + }; + if !dir.is_dir() { + return Ok(Vec::new()); + } + let mut entries = std::fs::read_dir(&dir) + .with_context(|| format!("list {}", dir.display()))? + .collect::>>() + .with_context(|| format!("list {}", dir.display()))?; + entries.sort_by_key(|entry| entry.file_name()); + let mut out = Vec::new(); + for entry in entries { + let name = entry.file_name().to_string_lossy().into_owned(); + if !is_nav_child_name(&name) || is_storyline_interior_name(&name) { + continue; + } + let file_type = entry + .file_type() + .with_context(|| format!("stat {}", entry.path().display()))?; + if file_type.is_symlink() { + continue; + } + if file_type.is_dir() { + if name.ends_with(".lance") { + continue; + } + let child_rel = if relative.is_empty() { + name.clone() + } else { + format!("{relative}/{name}") + }; + if let Some(kind) = self.probe_nav_dataset_kind(&child_rel).await? { + out.push(ShallowNavEntry { + name, + is_dir: false, + dataset_kind: Some(kind.into()), + }); + } else { + out.push(ShallowNavEntry { + name, + is_dir: true, + dataset_kind: None, + }); + } + } else if file_type.is_file() && is_importable_json_name(&name) { + out.push(ShallowNavEntry { + name, + is_dir: false, + dataset_kind: None, + }); + } + } + return Ok(out); + } + + let store = OpendalStore::from_uri(&self.uri).await?; + let prefix = if relative.is_empty() { + String::new() + } else { + format!("{relative}/") + }; + let entries = store + .list_shallow(&prefix) + .await + .with_context(|| format!("list shallow children under {prefix}{}", self.uri))?; + let mut dirs = BTreeSet::new(); + let mut files = BTreeSet::new(); + for entry in entries { + let path = entry.path.trim_start_matches(&prefix).trim_matches('/'); + if path.is_empty() { + continue; + } + let child = path.split('/').next().unwrap_or(path); + if !is_nav_child_name(child) || is_storyline_interior_name(child) { + continue; + } + if entry.mode == opendal::EntryMode::FILE && !path.contains('/') { + if is_importable_json_name(child) { + files.insert(child.to_string()); + } + continue; + } + if child.ends_with(".lance") { + continue; + } + dirs.insert(child.to_string()); + } + let mut out = Vec::with_capacity(dirs.len() + files.len()); + for name in dirs { + let child_rel = if relative.is_empty() { + name.clone() + } else { + format!("{relative}/{name}") + }; + if let Some(kind) = self.probe_nav_dataset_kind(&child_rel).await? { + out.push(ShallowNavEntry { + name, + is_dir: false, + dataset_kind: Some(kind.into()), + }); + } else { + out.push(ShallowNavEntry { + name, + is_dir: true, + dataset_kind: None, + }); + } + } + for name in files { + out.push(ShallowNavEntry { + name, + is_dir: false, + dataset_kind: None, + }); + } + out.sort_by(|left, right| left.name.cmp(&right.name)); + Ok(out) + } + /// Recursively list importable `.json` / `.jsonl` / `.ndjson` object keys. /// Skips Lance table interiors (any path segment ending in `.lance`). pub async fn list_importable_json_objects(&self, max_files: usize) -> Result> { @@ -224,18 +474,45 @@ impl DatasetLocation { /// Like [`Self::list_importable_json_objects`], but also returns size and /// last-modified metadata for change detection (`sync`). + /// + /// Object-store discovery walks prefixes with shallow listings and skips + /// `.lance` / `_meta` directories so large Storyline/events trees are not + /// fully enumerated. Progress callbacks fire as prefixes are scanned and + /// as each importable object is found. pub async fn list_importable_json_object_stamps( &self, max_files: usize, ) -> Result)>> { + self.list_importable_json_object_stamps_with_progress(max_files, &mut |_, _| Ok(())) + .await + } + + /// Stream importable object-store (or local) JSON files without buffering the + /// full listing. Callers can overlap discovery with downstream work. + /// + /// `Scanning` events report the prefix currently being listed; `File` events + /// report each importable object as soon as it is found. Object-store order + /// follows BFS discovery (not lexicographic sort). + pub async fn for_each_importable_json_object_event( + &self, + max_files: usize, + mut on_event: F, + ) -> Result<()> + where + F: FnMut(ImportableObjectEvent) -> Fut, + Fut: std::future::Future>, + { anyhow::ensure!(max_files > 0, "import max_files must be positive"); if let Some(root) = &self.local_path { + on_event(ImportableObjectEvent::Scanning { + prefix: String::new(), + }) + .await?; let paths = list_local_importable_json_files(root)?; anyhow::ensure!( paths.len() <= max_files, "import input exceeds max_files limit of {max_files}" ); - let mut stamps = Vec::with_capacity(paths.len()); for path in paths { let relative = path .strip_prefix(root) @@ -244,51 +521,130 @@ impl DatasetLocation { .replace('\\', "/"); let metadata = std::fs::metadata(&path) .with_context(|| format!("stat importable file {}", path.display()))?; - stamps.push(( - relative, - metadata.len(), - metadata.modified().ok().and_then(|modified| { - modified - .duration_since(std::time::UNIX_EPOCH) - .ok() - .map(|duration| { - chrono::DateTime::::from_timestamp( - duration.as_secs() as i64, - duration.subsec_nanos(), - ) - .map(|value| value.to_rfc3339()) - }) - .flatten() - }), - )); + let size = metadata.len(); + let modified = metadata.modified().ok().and_then(|modified| { + modified + .duration_since(std::time::UNIX_EPOCH) + .ok() + .and_then(|duration| { + chrono::DateTime::::from_timestamp( + duration.as_secs() as i64, + duration.subsec_nanos(), + ) + .map(|value| value.to_rfc3339()) + }) + }); + on_event(ImportableObjectEvent::File { + key: relative, + size, + modified, + }) + .await?; } - return Ok(stamps); + return Ok(()); } let store = OpendalStore::from_uri(&self.uri).await?; - let entries = store - .list("") - .await - .with_context(|| format!("list importable objects under {}", self.uri))?; - let mut stamps = Vec::new(); - for entry in entries { - let key = entry.path.trim_matches('/').to_string(); - if key.is_empty() || !is_importable_json_object_key(&key) { - continue; + let mut pending = vec![String::new()]; + let mut found = 0usize; + while let Some(prefix) = pending.pop() { + on_event(ImportableObjectEvent::Scanning { + prefix: prefix.clone(), + }) + .await?; + let list_prefix = if prefix.is_empty() { + String::new() + } else { + format!("{prefix}/") + }; + let entries = store.list_shallow(&list_prefix).await.with_context(|| { + format!( + "list importable objects under {}{}", + self.uri, + if list_prefix.is_empty() { + String::new() + } else { + format!("/{prefix}") + } + ) + })?; + let mut child_dirs = BTreeSet::new(); + for entry in entries { + let path = entry.path.trim_start_matches(&list_prefix).trim_matches('/'); + if path.is_empty() { + continue; + } + let child = path.split('/').next().unwrap_or(path); + if !is_nav_child_name(child) { + continue; + } + let child_rel = if prefix.is_empty() { + child.to_string() + } else { + format!("{prefix}/{child}") + }; + if entry.mode == opendal::EntryMode::FILE && !path.contains('/') { + if !is_importable_json_name(child) { + continue; + } + anyhow::ensure!( + found < max_files, + "import input exceeds max_files limit of {max_files}" + ); + found = found.saturating_add(1); + on_event(ImportableObjectEvent::File { + key: child_rel, + size: entry.metadata.content_length(), + modified: entry.metadata.last_modified().map(|value| value.to_string()), + }) + .await?; + continue; + } + if child.ends_with(".lance") + || child == "_meta" + || is_storyline_interior_name(child) + { + continue; + } + child_dirs.insert(child_rel); } - anyhow::ensure!( - stamps.len() < max_files, - "import input exceeds max_files limit of {max_files}" - ); - stamps.push(( - key, - entry.metadata.content_length(), - entry - .metadata - .last_modified() - .map(|value| value.to_string()), - )); + pending.extend(child_dirs.into_iter().rev()); } + Ok(()) + } + + /// `on_progress(path, Some(size))` reports an importable file; `on_progress(prefix, None)` + /// reports the prefix currently being scanned. + pub async fn list_importable_json_object_stamps_with_progress( + &self, + max_files: usize, + on_progress: &mut F, + ) -> Result)>> + where + F: FnMut(&str, Option) -> Result<()> + Send, + { + let mut stamps = Vec::new(); + self.for_each_importable_json_object_event(max_files, |event| { + // Progress + collection run synchronously before the future is + // polled; for_each awaits each event immediately so this stays + // sequential and keeps `on_progress` / `stamps` as plain FnMut state. + let result = match event { + ImportableObjectEvent::Scanning { prefix } => on_progress(&prefix, None), + ImportableObjectEvent::File { + key, + size, + modified, + } => match on_progress(&key, Some(size)) { + Ok(()) => { + stamps.push((key, size, modified)); + Ok(()) + } + Err(error) => Err(error), + }, + }; + async move { result } + }) + .await?; stamps.sort_by(|left, right| left.0.cmp(&right.0)); Ok(stamps) } @@ -311,6 +667,18 @@ impl DatasetLocation { /// Remove the complete Dataset represented by this local directory or /// object-store prefix. pub async fn remove_all(&self) -> Result<()> { + self.remove_all_with_progress(|_, _, _| Ok(())).await + } + + /// Like [`Self::remove_all`], but reports progress for each deleted file. + /// + /// `on_progress` receives `(deleted, total, relative_path)` after each + /// successful file delete. `deleted` counts completed deletes; the final + /// call uses `deleted == total` with an empty path once the tree is gone. + pub async fn remove_all_with_progress(&self, mut on_progress: F) -> Result<()> + where + F: FnMut(u64, u64, &str) -> Result<()>, + { if let Some(path) = &self.local_path { anyhow::ensure!(path.exists(), "Dataset does not exist: {}", self.uri); anyhow::ensure!( @@ -318,8 +686,7 @@ impl DatasetLocation { "refusing to drop a filesystem root as a Dataset" ); anyhow::ensure!(path.is_dir(), "Dataset is not a directory: {}", self.uri); - std::fs::remove_dir_all(path) - .with_context(|| format!("drop local Dataset {}", path.display()))?; + remove_local_dir_with_progress(path, &mut on_progress)?; return Ok(()); } @@ -329,11 +696,99 @@ impl DatasetLocation { "refusing to drop an entire object-store bucket; name a Dataset prefix" ); let store = OpendalStore::from_uri(&self.uri).await?; + let entries = store + .list("") + .await + .with_context(|| format!("list objects under {}", self.uri))?; + let total = entries.len() as u64; + let mut deleted = 0_u64; + on_progress(deleted, total, "")?; + for entry in entries { + store + .remove(&entry.path) + .await + .with_context(|| format!("delete object {} under {}", entry.path, self.uri))?; + deleted = deleted.saturating_add(1); + on_progress(deleted, total, &entry.path)?; + } + // Clear any leftover prefix markers after individual object deletes. store.remove_all().await?; + on_progress(total, total, "")?; Ok(()) } } +fn remove_local_dir_with_progress(path: &Path, on_progress: &mut F) -> Result<()> +where + F: FnMut(u64, u64, &str) -> Result<()>, +{ + let files = list_local_files_recursive(path)?; + let total = files.len() as u64; + let mut deleted = 0_u64; + on_progress(deleted, total, "")?; + for file in files { + let relative = file + .strip_prefix(path) + .unwrap_or(file.as_path()) + .to_string_lossy() + .replace('\\', "/"); + std::fs::remove_file(&file) + .with_context(|| format!("delete file {}", file.display()))?; + deleted = deleted.saturating_add(1); + on_progress(deleted, total, &relative)?; + } + std::fs::remove_dir_all(path) + .with_context(|| format!("drop local Dataset {}", path.display()))?; + on_progress(total, total, "")?; + Ok(()) +} + +fn list_local_files_recursive(root: &Path) -> Result> { + let mut pending = vec![root.to_path_buf()]; + let mut files = Vec::new(); + while let Some(directory) = pending.pop() { + let mut entries = std::fs::read_dir(&directory) + .with_context(|| format!("read directory {}", directory.display()))? + .collect::>>()?; + entries.sort_by_key(std::fs::DirEntry::path); + for entry in entries { + let file_type = entry.file_type()?; + let path = entry.path(); + if file_type.is_dir() && !file_type.is_symlink() { + pending.push(path); + } else { + files.push(path); + } + } + } + files.sort(); + Ok(files) +} + +fn is_nav_child_name(name: &str) -> bool { + !name.is_empty() + && name != "." + && name != ".." + && !name.starts_with('.') + && name != "_meta" +} + +fn is_storyline_interior_name(name: &str) -> bool { + matches!(name, "generations" | "objects.lance" | "writer" | "leases") +} + +fn is_importable_json_name(name: &str) -> bool { + Path::new(name) + .extension() + .and_then(|extension| extension.to_str()) + .is_some_and(|extension| { + matches!( + extension.to_ascii_lowercase().as_str(), + "json" | "jsonl" | "ndjson" + ) + }) +} + fn is_importable_json_object_key(key: &str) -> bool { if key .split('/') diff --git a/crates/persisting-pchronicle/src/store/mod.rs b/crates/persisting-pchronicle/src/store/mod.rs index 76ee0980..098e3372 100644 --- a/crates/persisting-pchronicle/src/store/mod.rs +++ b/crates/persisting-pchronicle/src/store/mod.rs @@ -35,6 +35,8 @@ mod events; mod files; #[cfg(feature = "lance-store")] pub(crate) mod index_build_gate; +pub(crate) mod index_build_progress; +pub(crate) mod object_store_io_gate; #[cfg(feature = "lance-store")] mod inspect; #[cfg(feature = "lance-store")] @@ -75,8 +77,10 @@ pub use catalog::{ #[cfg(feature = "lance-store")] #[allow(unused_imports)] pub use chronicle_manifest::{ - CHRONICLE_MANIFEST_FILE, ChronicleManifest, ManifestKind, ManifestStats, atomic_write_manifest, - compact_jsonl_manifest_matches, load_manifest, try_load_manifest, write_compact_jsonl_manifest, + CHRONICLE_MANIFEST_FILE, ChronicleManifest, ManifestKind, ManifestStats, STORYLINE_FORMAT, + atomic_write_manifest, compact_jsonl_manifest_matches, load_manifest, load_manifest_at_uri, + try_load_manifest, write_compact_jsonl_manifest, write_storyline_manifest, + write_storyline_manifest_at_uri, }; #[cfg(feature = "lance-store")] pub use compact_jsonl::{ @@ -119,7 +123,9 @@ pub(crate) use local_query_manifest::{ LocalQueryInputFile, LocalQueryManifest, LocalQueryManifestOptions, }; #[cfg(feature = "lance-store")] -pub use location::{DatasetLocation, DatasetLocationKind}; +pub use location::{ + DatasetLocation, DatasetLocationKind, ImportableObjectEvent, ShallowNavEntry, +}; #[cfg(feature = "lance-store")] pub use query_engine::{ ChronicleQueryEngine, ChronicleQueryExecutionOptions, DEFAULT_QUERY_MEMORY_LIMIT_BYTES, @@ -137,9 +143,10 @@ pub use storyline::{ StorylineContentOptions, StorylineContentReadMode, StorylineDataFusionTableNames, StorylineDataSource, StorylineDataSourceOptions, StorylineLanceStore, StorylineMaintenanceReport, StorylineProjectionLineage, StorylineStreamImportReport, - StorylineTableKind, StorylineTablePaths, story_runs_arrow_schema, story_runs_from_batch, - story_runs_to_batch, story_steps_arrow_schema, story_steps_from_batch, story_steps_to_batch, - story_tool_calls_arrow_schema, story_tool_calls_from_batch, story_tool_calls_to_batch, + StorylineStreamOptions, StorylineTableKind, StorylineTablePaths, story_runs_arrow_schema, + story_runs_from_batch, story_runs_to_batch, story_steps_arrow_schema, story_steps_from_batch, + story_steps_to_batch, story_tool_calls_arrow_schema, story_tool_calls_from_batch, + story_tool_calls_to_batch, }; #[cfg(feature = "lance-store")] pub use storyline_model::{ diff --git a/crates/persisting-pchronicle/src/store/object_store_io_gate.rs b/crates/persisting-pchronicle/src/store/object_store_io_gate.rs new file mode 100644 index 00000000..5a515213 --- /dev/null +++ b/crates/persisting-pchronicle/src/store/object_store_io_gate.rs @@ -0,0 +1,232 @@ +//! Process-wide admission + AIMD backoff for remote object-store I/O. +//! +//! Lance opens and table writes against flaky S3-compatible gateways amplify +//! timeouts when several datasets race (list `_versions/`, retries, AIMD inside +//! object_store). This gate: +//! 1. caps concurrent remote Lance ops (default 1); +//! 2. after a transient failure, forces a shared cooldown + growing delay; +//! 3. decays the delay after a streak of successes. +//! +//! Local `file://` paths bypass the gate entirely. + +use std::sync::{Arc, Mutex, OnceLock}; +use std::time::{Duration, Instant}; + +use tokio::sync::{OwnedSemaphorePermit, Semaphore}; + +const DEFAULT_REMOTE_CONCURRENCY: usize = 1; +const MAX_REMOTE_CONCURRENCY: usize = 2; +const MAX_DELAY_MS: u64 = 30_000; +const SUCCESS_STREAK_TO_DECAY: u32 = 4; + +/// Whether the gated op is primarily reading metadata/objects or writing them. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum IoKind { + Read, + Write, +} + +impl IoKind { + fn as_str(self) -> &'static str { + match self { + Self::Read => "read", + Self::Write => "write", + } + } +} + +#[derive(Debug)] +struct AimdState { + /// Extra sleep applied before each remote acquire while degraded. + delay_ms: u64, + /// No new remote op starts until this instant. + cooldown_until: Option, + successes_since_backoff: u32, + failures: u64, + /// Last classified op that hit the gate (for progress UI). + last_kind: IoKind, +} + +impl Default for AimdState { + fn default() -> Self { + Self { + delay_ms: 0, + cooldown_until: None, + successes_since_backoff: 0, + failures: 0, + last_kind: IoKind::Read, + } + } +} + +struct Gate { + semaphore: Arc, + state: Mutex, +} + +fn gate() -> &'static Gate { + static GATE: OnceLock = OnceLock::new(); + GATE.get_or_init(|| { + let concurrency = std::env::var("PCHRONICLE_OBJECT_STORE_CONCURRENCY") + .ok() + .and_then(|value| value.parse::().ok()) + .unwrap_or(DEFAULT_REMOTE_CONCURRENCY) + .clamp(1, MAX_REMOTE_CONCURRENCY); + Gate { + semaphore: Arc::new(Semaphore::new(concurrency)), + state: Mutex::new(AimdState::default()), + } + }) +} + +/// True for s3/gs/az (and similar) URIs; false for local paths / file://. +pub(crate) fn is_remote_uri(uri: &str) -> bool { + let Some((scheme, _)) = uri.split_once("://") else { + return false; + }; + !matches!(scheme, "file" | "file+uring" | "memory" | "shared-memory") +} + +pub(crate) struct Permit { + _permit: Option, +} + +/// Acquire admission for a Lance/object-store operation on `uri`. +pub(crate) async fn acquire(uri: &str, kind: IoKind) -> Permit { + if !is_remote_uri(uri) { + return Permit { _permit: None }; + } + if let Ok(mut state) = gate().state.lock() { + state.last_kind = kind; + } + wait_out_degradation(kind).await; + let permit = gate() + .semaphore + .clone() + .acquire_owned() + .await + .expect("object-store I/O semaphore is never closed"); + wait_out_degradation(kind).await; + Permit { + _permit: Some(permit), + } +} + +async fn wait_out_degradation(kind: IoKind) { + let (sleep_for, delay_ms, failures) = { + let Ok(state) = gate().state.lock() else { + return; + }; + let cooldown = state + .cooldown_until + .and_then(|until| until.checked_duration_since(Instant::now())) + .unwrap_or_default(); + (cooldown, state.delay_ms, state.failures) + }; + if sleep_for.is_zero() { + return; + } + crate::store::index_build_progress::note(format!( + "s3 {} throttle wait {:.1}s (failures={failures}, delay={delay_ms}ms)", + kind.as_str(), + sleep_for.as_secs_f32() + )); + tracing::warn!( + target: "pchronicle.object_store_gate", + kind = kind.as_str(), + wait_ms = sleep_for.as_millis() as u64, + delay_ms, + failures, + "object-store I/O gate cooling down before next remote op" + ); + tokio::time::sleep(sleep_for).await; +} + +/// Publish the current I/O phase for progress UI without taking a permit. +/// Used around Lance writes that do not go through [`acquire`]. +pub(crate) fn mark_kind(kind: IoKind) { + if let Ok(mut state) = gate().state.lock() { + state.last_kind = kind; + } +} + +/// Record a successful remote op: decay shared delay after a streak. +pub(crate) fn note_success(uri: &str) { + if !is_remote_uri(uri) { + return; + } + let Ok(mut state) = gate().state.lock() else { + return; + }; + state.successes_since_backoff = state.successes_since_backoff.saturating_add(1); + if state.delay_ms == 0 { + return; + } + if state.successes_since_backoff >= SUCCESS_STREAK_TO_DECAY { + state.delay_ms /= 2; + state.successes_since_backoff = 0; + if state.delay_ms < 100 { + state.delay_ms = 0; + state.cooldown_until = None; + } + tracing::info!( + target: "pchronicle.object_store_gate", + delay_ms = state.delay_ms, + "object-store I/O gate recovered toward steady state" + ); + } +} + +/// Record a transient remote failure: grow shared delay and set a cooldown. +pub(crate) fn note_failure(uri: &str, kind: IoKind) { + if !is_remote_uri(uri) { + return; + } + let Ok(mut state) = gate().state.lock() else { + return; + }; + state.last_kind = kind; + state.failures = state.failures.saturating_add(1); + state.successes_since_backoff = 0; + state.delay_ms = if state.delay_ms == 0 { + 500 + } else { + state.delay_ms.saturating_mul(2).min(MAX_DELAY_MS) + }; + state.cooldown_until = Some(Instant::now() + Duration::from_millis(state.delay_ms)); + tracing::warn!( + target: "pchronicle.object_store_gate", + kind = kind.as_str(), + delay_ms = state.delay_ms, + failures = state.failures, + "object-store I/O gate backing off after transient failure" + ); + crate::store::index_build_progress::note(format!( + "s3 {} throttle backoff {}ms", + kind.as_str(), + state.delay_ms + )); +} + +#[cfg(test)] +pub(crate) fn debug_delay_ms() -> u64 { + gate() + .state + .lock() + .map(|state| state.delay_ms) + .unwrap_or(0) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn classifies_remote_uris() { + assert!(is_remote_uri("s3://bucket/prefix")); + assert!(is_remote_uri("gs://bucket/prefix")); + assert!(!is_remote_uri("/tmp/local")); + assert!(!is_remote_uri("file:///tmp/local")); + assert!(!is_remote_uri("shared-memory://x")); + } +} diff --git a/crates/persisting-pchronicle/src/store/opendal_store.rs b/crates/persisting-pchronicle/src/store/opendal_store.rs index bd15b216..f3e5a8c5 100644 --- a/crates/persisting-pchronicle/src/store/opendal_store.rs +++ b/crates/persisting-pchronicle/src/store/opendal_store.rs @@ -6,12 +6,38 @@ use anyhow::{Context, Result, anyhow}; use futures::TryStreamExt; +use opendal::layers::RetryLayer; use opendal::{EntryMode, ErrorKind, Metadata, Operator}; use std::collections::HashMap; use std::sync::Arc; use std::sync::{Mutex, OnceLock}; +use std::time::Duration; use url::Url; +/// Retries for transient object-store failures (DNS blips, connect resets, +/// 5xx, rate limits). Tuned for long imports over flaky endpoints: up to 8 +/// retries with exponential backoff + jitter, capped at 30s. +fn with_object_store_retries(operator: Operator) -> Operator { + operator.layer( + RetryLayer::new() + .with_notify(|event: opendal::layers::RetryEvent<'_>| { + tracing::warn!( + target: "pchronicle.opendal", + attempt = event.attempt, + retry_after_ms = event.retry_after.as_millis() as u64, + op = ?event.op, + error = %event.err, + "retrying temporary object-store error" + ); + }) + .with_jitter() + .with_factor(2.0) + .with_min_delay(Duration::from_millis(500)) + .with_max_delay(Duration::from_secs(30)) + .with_max_times(8), + ) +} + #[derive(Clone, Debug, PartialEq, Eq)] pub(crate) struct Version { pub(crate) etag: Option, @@ -63,14 +89,18 @@ impl Store { if let Some(operator) = map.get(uri) { operator.clone() } else { - let operator = Operator::from_uri(normalized.as_str()) - .with_context(|| format!("open OpenDAL store {uri}"))?; + let operator = with_object_store_retries( + Operator::from_uri(normalized.as_str()) + .with_context(|| format!("open OpenDAL store {uri}"))?, + ); map.insert(uri.to_string(), operator.clone()); operator } } else { - Operator::from_uri(normalized.as_str()) - .with_context(|| format!("open OpenDAL store {uri}"))? + with_object_store_retries( + Operator::from_uri(normalized.as_str()) + .with_context(|| format!("open OpenDAL store {uri}"))?, + ) }; let fallback_lock = if shared_memory { let locks = SHARED_LOCKS.get_or_init(|| Mutex::new(HashMap::new())); @@ -130,7 +160,19 @@ impl Store { .if_match(condition) .await .map(|_| ()) - .map_err(Into::into) + .map_err(|error| { + if is_conflict(&error) { + tracing::debug!( + target: "pchronicle.opendal", + path, + if_match = condition, + error = %error, + kind = ?error.kind(), + "conditional object write conflict (If-Match)" + ); + } + error.into() + }) } pub(crate) async fn write_overwrite(&self, path: &str, bytes: Vec) -> Result<()> { @@ -253,3 +295,19 @@ fn normalize_uri(uri: &str) -> Result { } Ok(parsed.to_string()) } + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn object_store_operator_accepts_retry_layer() -> Result<()> { + let store = Store::from_uri("shared-memory://pchronicle-retry-layer/root").await?; + store + .write_overwrite("probe.json", b"{\"ok\":true}".to_vec()) + .await?; + let loaded = store.read("probe.json").await?.context("probe missing")?; + assert_eq!(loaded.0, b"{\"ok\":true}"); + Ok(()) + } +} diff --git a/crates/persisting-pchronicle/src/store/storyline/content.rs b/crates/persisting-pchronicle/src/store/storyline/content.rs index 7c8d8d2a..00b33cf0 100644 --- a/crates/persisting-pchronicle/src/store/storyline/content.rs +++ b/crates/persisting-pchronicle/src/store/storyline/content.rs @@ -609,6 +609,7 @@ pub(crate) async fn commit_pending_content( snapshot_version: Option, pending: PendingContent, reopen_concurrent_create: bool, + build_indexes: bool, ) -> Result { let mut objects = pending.objects.into_values().collect::>(); objects.sort_by(|left, right| left.reference.content_id.cmp(&right.reference.content_id)); @@ -616,7 +617,7 @@ pub(crate) async fn commit_pending_content( let mut dataset = if let Some(snapshot_version) = snapshot_version { let mut dataset = open_objects(path, snapshot_version).await?; - let latest = Dataset::open(&uri).await?.version_id(); + let latest = super::open_dataset_uri(&uri).await?.version_id(); if latest != snapshot_version { dataset.restore().await.with_context(|| { format!( @@ -639,11 +640,15 @@ pub(crate) async fn commit_pending_content( .await { Ok(mut dataset) => { - ensure_content_index(&mut dataset).await?; + // Progressive imports defer indexes until a final maintain(); + // creating btree here would stall every first-batch commit. + if build_indexes { + ensure_content_index(&mut dataset).await?; + } return Ok(dataset.version_id()); } Err(lance::Error::DatasetAlreadyExists { .. }) if reopen_concurrent_create => { - Dataset::open(&uri).await.with_context(|| { + super::open_dataset_uri(&uri).await.with_context(|| { format!( "reopen concurrently created Storyline content store {}", path.display() @@ -679,6 +684,32 @@ pub(crate) async fn commit_pending_content( .execute_stream(reader) .await .with_context(|| format!("append Storyline content store {}", path.display()))?; + if build_indexes { + ensure_content_index(&mut dataset).await?; + dataset + .optimize_indices(&lance_index::optimize::OptimizeOptions::append()) + .await + .with_context(|| format!("extend Storyline content index {}", path.display()))?; + } + Ok(dataset.version_id()) +} + +/// Ensure + extend the objects.lance content_id btree (used by final maintain). +pub(crate) async fn ensure_optimize_objects_content_index( + path: &Path, + snapshot_version: u64, +) -> Result { + let uri = path.to_string_lossy().into_owned(); + let mut dataset = open_objects(path, snapshot_version).await?; + let latest = super::open_dataset_uri(&uri).await?.version_id(); + if latest != snapshot_version { + dataset.restore().await.with_context(|| { + format!( + "restore Storyline content store {} to version {snapshot_version}", + path.display() + ) + })?; + } ensure_content_index(&mut dataset).await?; dataset .optimize_indices(&lance_index::optimize::OptimizeOptions::append()) @@ -696,6 +727,11 @@ async fn ensure_content_index(dataset: &mut Dataset) -> Result<()> { { return Ok(()); } + crate::store::index_build_progress::note(format!( + "index {}.{} btree 1/1", + crate::store::index_build_progress::table_label(dataset.uri()), + CONTENT_ID_COLUMN + )); let _admission = super::super::index_build_gate::acquire().await; dataset .create_index( @@ -746,7 +782,7 @@ fn content_id_predicate<'a>(values: impl IntoIterator) -> String } pub(crate) async fn open_objects(path: &Path, version: u64) -> Result { - let dataset = Dataset::open(path.to_string_lossy().as_ref()) + let dataset = super::open_dataset_uri(path.to_string_lossy().as_ref()) .await .with_context(|| format!("open Storyline content store {}", path.display()))?; dataset.checkout_version(version).await.with_context(|| { diff --git a/crates/persisting-pchronicle/src/store/storyline/datafusion.rs b/crates/persisting-pchronicle/src/store/storyline/datafusion.rs index d4f455ab..42ef80d0 100644 --- a/crates/persisting-pchronicle/src/store/storyline/datafusion.rs +++ b/crates/persisting-pchronicle/src/store/storyline/datafusion.rs @@ -435,12 +435,24 @@ impl StorylineDataSource { paths: StorylineTablePaths, options: StorylineDataSourceOptions, ) -> Result { - let (runs, steps, tool_calls, objects) = tokio::try_join!( - open_dataset(&paths.runs, paths.runs_version), - open_dataset(&paths.steps, paths.steps_version), - open_dataset(&paths.tool_calls, paths.tool_calls_version), - open_objects(&paths.objects, paths.objects_version), - )?; + let remote = paths.runs.to_string_lossy().contains("://") + && !paths.runs.to_string_lossy().starts_with("file:"); + let (runs, steps, tool_calls, objects) = if remote { + // Avoid four concurrent Lance opens against flaky S3 gateways. + ( + open_dataset(&paths.runs, paths.runs_version).await?, + open_dataset(&paths.steps, paths.steps_version).await?, + open_dataset(&paths.tool_calls, paths.tool_calls_version).await?, + open_objects(&paths.objects, paths.objects_version).await?, + ) + } else { + tokio::try_join!( + open_dataset(&paths.runs, paths.runs_version), + open_dataset(&paths.steps, paths.steps_version), + open_dataset(&paths.tool_calls, paths.tool_calls_version), + open_objects(&paths.objects, paths.objects_version), + )? + }; let objects = Arc::new(objects); Ok(Self { paths, @@ -515,7 +527,7 @@ fn combine_filters(filters: &[Expr]) -> Option { } async fn open_dataset(path: &Path, version: u64) -> Result { - let dataset = Dataset::open(path.to_string_lossy().as_ref()) + let dataset = super::open_dataset_uri(path.to_string_lossy().as_ref()) .await .with_context(|| format!("open Storyline DataFusion table {}", path.display()))?; dataset.checkout_version(version).await.with_context(|| { diff --git a/crates/persisting-pchronicle/src/store/storyline/mod.rs b/crates/persisting-pchronicle/src/store/storyline/mod.rs index ffc02737..6244afcd 100644 --- a/crates/persisting-pchronicle/src/store/storyline/mod.rs +++ b/crates/persisting-pchronicle/src/store/storyline/mod.rs @@ -43,11 +43,12 @@ pub use rows::{ use std::collections::{HashMap, HashSet}; use std::fs::{File, OpenOptions}; +use std::future::Future; use std::io::Write; use std::path::{Path, PathBuf}; use std::sync::Arc; use std::sync::atomic::{AtomicU64, Ordering}; -use std::time::{SystemTime, UNIX_EPOCH}; +use std::time::{Duration, SystemTime, UNIX_EPOCH}; use anyhow::{Context, Result}; use fs2::FileExt; @@ -78,8 +79,8 @@ use crate::formats::unknown_fields::{compute_unknown_key_counts, validate_unknow use self::content::{ PendingContent, STORYLINE_OBJECTS_DATASET, collect_content_ids, commit_pending_content, - content_columns, externalize_batches, externalize_unknown_field_values, hydrate_batches, - open_objects, prune_unreferenced_objects, + content_columns, ensure_optimize_objects_content_index, externalize_batches, + externalize_unknown_field_values, hydrate_batches, open_objects, prune_unreferenced_objects, }; use super::AtifReader; use super::{LanceMaintenanceOptions, LanceMaintenanceReport, root_write_lock}; @@ -211,6 +212,9 @@ pub struct StorylineLanceStore { control_store: OpendalStore, write_lock: Arc>, control_lock: Arc>, + /// Some S3-compatible gateways always 412 on If-Match PUT. After the first + /// verified content-stable fallback, skip conditional writes for CURRENT. + current_if_match_unreliable: Arc, content_options: StorylineContentOptions, } @@ -443,6 +447,31 @@ fn release_waiting_content_create(root_uri: &str, first: bool) { } } +/// Options for streaming Storyline writes. +#[derive(Debug, Clone, Copy)] +pub struct StorylineStreamOptions { + /// When true, run Lance index ensure/optimize at the end of this stream. + /// Progressive imports set this false and call [`StorylineLanceStore::maintain`] + /// once after all batches land. + pub optimize_indices: bool, +} + +impl Default for StorylineStreamOptions { + fn default() -> Self { + Self { + optimize_indices: true, + } + } +} + +impl StorylineStreamOptions { + pub fn defer_index_optimize() -> Self { + Self { + optimize_indices: false, + } + } +} + impl StorylineLanceStore { pub async fn open(root: impl AsRef) -> Result { let root = root.as_ref().to_path_buf(); @@ -523,6 +552,7 @@ impl StorylineLanceStore { control_lock: Arc::new(tokio::sync::Mutex::new(())), root_uri, control_store, + current_if_match_unreliable: Arc::new(std::sync::atomic::AtomicBool::new(false)), content_options: StorylineContentOptions::default(), }) } @@ -585,19 +615,38 @@ impl StorylineLanceStore { let Some(paths) = self.resolve_current_table_paths().await? else { return Ok(None); }; - tokio::try_join!( - validate_table(&paths.generation, &paths.runs, paths.runs_version), - validate_table(&paths.generation, &paths.steps, paths.steps_version), + // Object-store gateways choke when Lance opens four datasets at once + // (each list/_versions + retries). Validate sequentially there; keep + // local try_join for speed. + if self.is_remote_object_store() { + validate_table(&paths.generation, &paths.runs, paths.runs_version).await?; + validate_table(&paths.generation, &paths.steps, paths.steps_version).await?; validate_table( &paths.generation, &paths.tool_calls, - paths.tool_calls_version - ), - validate_table(&paths.generation, &paths.objects, paths.objects_version), - )?; + paths.tool_calls_version, + ) + .await?; + validate_table(&paths.generation, &paths.objects, paths.objects_version).await?; + } else { + tokio::try_join!( + validate_table(&paths.generation, &paths.runs, paths.runs_version), + validate_table(&paths.generation, &paths.steps, paths.steps_version), + validate_table( + &paths.generation, + &paths.tool_calls, + paths.tool_calls_version + ), + validate_table(&paths.generation, &paths.objects, paths.objects_version), + )?; + } Ok(Some(paths)) } + fn is_remote_object_store(&self) -> bool { + self.root_uri.contains("://") && !matches!(self.storage_scheme(), "file" | "file+uring") + } + /// Return the generation and every stable per-document identity from one /// committed snapshot. pub async fn document_ids_snapshot(&self) -> Result)>> { @@ -657,6 +706,7 @@ impl StorylineLanceStore { Some(projection), StorylineStreamWriteMode::Replace, None, + StorylineStreamOptions::default(), ) .await?; published_storyline_report(outcome)?; @@ -698,6 +748,28 @@ impl StorylineLanceStore { None, StorylineStreamWriteMode::Replace, None, + StorylineStreamOptions::default(), + ) + .await?; + published_storyline_report(outcome) + } + + /// Like [`Self::replace_storyline_stream`], with explicit stream options. + pub async fn replace_storyline_stream_with_options( + &self, + stories: I, + options: StorylineStreamOptions, + ) -> Result + where + I: IntoIterator>, + { + let outcome = self + .replace_storyline_stream_with_projection( + stories, + None, + StorylineStreamWriteMode::Replace, + None, + options, ) .await?; published_storyline_report(outcome) @@ -710,6 +782,24 @@ impl StorylineLanceStore { stories: I, expected_generation: &str, ) -> Result + where + I: IntoIterator>, + { + self.append_storyline_stream_with_options( + stories, + expected_generation, + StorylineStreamOptions::default(), + ) + .await + } + + /// Like [`Self::append_storyline_stream`], with explicit stream options. + pub async fn append_storyline_stream_with_options( + &self, + stories: I, + expected_generation: &str, + options: StorylineStreamOptions, + ) -> Result where I: IntoIterator>, { @@ -719,6 +809,7 @@ impl StorylineLanceStore { None, StorylineStreamWriteMode::Replace, Some(expected_generation), + options, ) .await?; published_storyline_report(outcome) @@ -739,6 +830,7 @@ impl StorylineLanceStore { Some(projection), StorylineStreamWriteMode::Replace, None, + StorylineStreamOptions::default(), ) .await?; published_storyline_report(outcome) @@ -758,6 +850,7 @@ impl StorylineLanceStore { Some(projection), StorylineStreamWriteMode::CreateProjection, None, + StorylineStreamOptions::default(), ) .await } @@ -780,6 +873,7 @@ impl StorylineLanceStore { Some(projection), StorylineStreamWriteMode::Rebuild, None, + StorylineStreamOptions::default(), ) .await?; published_storyline_report(outcome) @@ -791,6 +885,7 @@ impl StorylineLanceStore { projection: Option, mode: StorylineStreamWriteMode, required_generation: Option<&str>, + stream_options: StorylineStreamOptions, ) -> Result where I: IntoIterator>, @@ -912,31 +1007,38 @@ impl StorylineLanceStore { original.as_ref().map(|paths| paths.objects_version), pending, mode == StorylineStreamWriteMode::CreateProjection, + stream_options.optimize_indices, ) .await; #[cfg(test)] release_waiting_content_create(&self.root_uri, first_content_create); let objects_version = objects_result?; - let (runs_version, steps_version, tool_calls_version) = tokio::try_join!( - write_batches( - &created.runs, - run_batches, - story_runs_arrow_schema(), - &RUN_INDEXES, - ), - write_batches( - &created.steps, - step_batches, - story_steps_arrow_schema(), - &STEP_INDEXES, - ), - write_batches( - &created.tool_calls, - tool_call_batches, - story_tool_calls_arrow_schema(), - &TOOL_CALL_INDEXES, - ), - )?; + let (runs_version, steps_version, tool_calls_version) = + join3_remote_aware( + self.is_remote_object_store(), + write_batches( + &created.runs, + run_batches, + story_runs_arrow_schema(), + &RUN_INDEXES, + stream_options.optimize_indices, + ), + write_batches( + &created.steps, + step_batches, + story_steps_arrow_schema(), + &STEP_INDEXES, + stream_options.optimize_indices, + ), + write_batches( + &created.tool_calls, + tool_call_batches, + story_tool_calls_arrow_schema(), + &TOOL_CALL_INDEXES, + stream_options.optimize_indices, + ), + ) + .await?; created.runs_version = runs_version; created.steps_version = steps_version; created.tool_calls_version = tool_calls_version; @@ -951,34 +1053,38 @@ impl StorylineLanceStore { Some(current.objects_version), pending, false, + stream_options.optimize_indices, ) .await?; - let (runs_version, steps_version, tool_calls_version) = tokio::try_join!( - replace_table_batches( - ¤t.runs, - current.runs_version, - &predicate, - &["document_id"], - run_batches, - story_runs_arrow_schema(), - ), - replace_table_batches( - ¤t.steps, - current.steps_version, - &predicate, - &["document_id", "step_id"], - step_batches, - story_steps_arrow_schema(), - ), - replace_table_batches( - ¤t.tool_calls, - current.tool_calls_version, - &predicate, - &["document_id", "step_id", "call_index"], - tool_call_batches, - story_tool_calls_arrow_schema(), - ), - )?; + let (runs_version, steps_version, tool_calls_version) = + join3_remote_aware( + self.is_remote_object_store(), + replace_table_batches( + ¤t.runs, + current.runs_version, + &predicate, + &["document_id"], + run_batches, + story_runs_arrow_schema(), + ), + replace_table_batches( + ¤t.steps, + current.steps_version, + &predicate, + &["document_id", "step_id"], + step_batches, + story_steps_arrow_schema(), + ), + replace_table_batches( + ¤t.tool_calls, + current.tool_calls_version, + &predicate, + &["document_id", "step_id", "call_index"], + tool_call_batches, + story_tool_calls_arrow_schema(), + ), + ) + .await?; current.runs_version = runs_version; current.steps_version = steps_version; current.tool_calls_version = tool_calls_version; @@ -993,12 +1099,12 @@ impl StorylineLanceStore { .as_ref() .context("missing streamed Storyline tables")?; let (runs_version, steps_version, tool_calls_version) = - // Build indexes for a new store (including a small import), - // and periodically after a large streamed import. Replacing - // one small region in an existing store must not rebuild and - // optimize every FTS/JSON index on every write; callers that - // need to catch up appended fragments can invoke `maintain`. - if original.is_none() || report.storylines > STREAM_IMPORT_STORIES { + // Build/optimize indexes for a brand-new store, or after a large + // one-shot streamed write. Progressive imports pass + // optimize_indices=false and call maintain() once at the end. + if stream_options.optimize_indices + && (original.is_none() || report.storylines > STREAM_IMPORT_STORIES) + { let maintenance = LanceMaintenanceOptions { // Extend scalar, FTS, and JSON indices once after // import, without putting compaction in the ingest @@ -1008,7 +1114,8 @@ impl StorylineLanceStore { vacuum_older_than: None, ..Default::default() }; - let (runs, steps, tool_calls) = tokio::try_join!( + let (runs, steps, tool_calls) = join3_remote_aware( + self.is_remote_object_store(), maintain_table_layout( ¤t.runs, current.runs_version, @@ -1027,7 +1134,8 @@ impl StorylineLanceStore { &TOOL_CALL_INDEXES, &maintenance, ), - )?; + ) + .await?; ( runs.final_version .context("missing imported runs version")?, @@ -1166,16 +1274,18 @@ impl StorylineLanceStore { } else { original.clone() }; - let (runs, steps, tool_calls) = tokio::try_join!( - maintain_table_layout(&paths.runs, paths.runs_version, &RUN_INDEXES, options,), - maintain_table_layout(&paths.steps, paths.steps_version, &STEP_INDEXES, options,), + let (runs, steps, tool_calls) = join3_remote_aware( + self.is_remote_object_store(), + maintain_table_layout(&paths.runs, paths.runs_version, &RUN_INDEXES, options), + maintain_table_layout(&paths.steps, paths.steps_version, &STEP_INDEXES, options), maintain_table_layout( &paths.tool_calls, paths.tool_calls_version, &TOOL_CALL_INDEXES, options, ), - )?; + ) + .await?; let runs_version = runs .final_version .context("missing maintained runs version")?; @@ -1208,9 +1318,15 @@ impl StorylineLanceStore { &tool_call_batches, StorylineTableKind::ToolCalls, )?); - let (objects_version, objects_removed) = + let (mut objects_version, objects_removed) = prune_unreferenced_objects(&paths.objects, paths.objects_version, &live_objects) .await?; + // Progressive imports defer objects.lance btree until here so + // mid-batch commits only write data. + if options.optimize_indices { + objects_version = + ensure_optimize_objects_content_index(&paths.objects, objects_version).await?; + } let generation = next_generation(); let snapshot = StorylineSnapshotPointer { schema_version: STORYLINE_LANCE_SCHEMA_VERSION, @@ -1322,6 +1438,7 @@ impl StorylineLanceStore { None, StorylineStreamWriteMode::Replace, None, + StorylineStreamOptions::default(), ) .await?; published_storyline_report(outcome)?; @@ -1476,18 +1593,21 @@ impl StorylineLanceStore { run_batches, story_runs_arrow_schema(), &RUN_INDEXES, + true, ), write_batches( &cloned.steps, step_batches, story_steps_arrow_schema(), &STEP_INDEXES, + true, ), write_batches( &cloned.tool_calls, tool_call_batches, story_tool_calls_arrow_schema(), &TOOL_CALL_INDEXES, + true, ), )?; cloned.generation.clone_from(&source.generation); @@ -1668,15 +1788,25 @@ async fn write_local_current(path: PathBuf, contents: Vec) -> Result<()> { } async fn validate_table(generation: &str, path: &Path, version: u64) -> Result<()> { - let dataset = Dataset::open(path.to_string_lossy().as_ref()) - .await - .with_context(|| { - format!( - "Storyline generation '{}' is incomplete: cannot open {}", - generation, + let uri = path.to_string_lossy(); + let dataset = open_dataset_uri(uri.as_ref()).await.map_err(|error| { + if is_not_found_storage_error(&error) { + error.context(format!( + "Storyline generation '{generation}' is incomplete: cannot open {}", path.display() - ) - })?; + )) + } else if is_transient_storage_error(&error) { + error.context(format!( + "Storyline generation '{generation}' could not be verified: object-store timeout opening {} (likely gateway overload, not a missing generation)", + path.display() + )) + } else { + error.context(format!( + "Storyline generation '{generation}' could not be verified: failed to open {}", + path.display() + )) + } + })?; dataset.checkout_version(version).await.with_context(|| { format!( "Storyline generation '{generation}' references missing version {version} of {}", @@ -1686,6 +1816,119 @@ async fn validate_table(generation: &str, path: &Path, version: u64) -> Result<( Ok(()) } +const DATASET_OPEN_MAX_ATTEMPTS: u32 = 8; + +fn error_chain_text(error: &anyhow::Error) -> String { + let mut parts = vec![error.to_string()]; + let mut source = error.source(); + while let Some(err) = source { + parts.push(err.to_string()); + source = err.source(); + } + parts.join(" | ").to_ascii_lowercase() +} + +fn is_transient_storage_error(error: &anyhow::Error) -> bool { + let text = error_chain_text(error); + [ + "timeout", + "timed out", + "error sending request", + "connection reset", + "connection refused", + "broken pipe", + "temporarily unavailable", + "slowdown", + "throttl", + "503", + "429", + "connect", + "tcp connect", + ] + .iter() + .any(|needle| text.contains(needle)) +} + +fn is_not_found_storage_error(error: &anyhow::Error) -> bool { + let text = error_chain_text(error); + [ + "not found", + "nosuchkey", + "no such key", + "404", + "does not exist", + ] + .iter() + .any(|needle| text.contains(needle)) + && !is_transient_storage_error(error) +} + +/// Open a Lance dataset with retries for flaky object-store gateways. +pub(super) async fn open_dataset_uri(uri: &str) -> Result { + let mut attempt = 0u32; + loop { + attempt += 1; + let _permit = crate::store::object_store_io_gate::acquire( + uri, + crate::store::object_store_io_gate::IoKind::Read, + ) + .await; + match Dataset::open(uri).await { + Ok(dataset) => { + crate::store::object_store_io_gate::note_success(uri); + return Ok(dataset); + } + Err(error) => { + let error = anyhow::Error::from(error); + if !is_transient_storage_error(&error) { + return Err(error).with_context(|| format!("open Lance dataset {uri}")); + } + crate::store::object_store_io_gate::note_failure( + uri, + crate::store::object_store_io_gate::IoKind::Read, + ); + if attempt >= DATASET_OPEN_MAX_ATTEMPTS { + return Err(error).with_context(|| format!("open Lance dataset {uri}")); + } + crate::store::index_build_progress::note(format!( + "retry open {} ({attempt}/{DATASET_OPEN_MAX_ATTEMPTS})", + crate::store::index_build_progress::table_label(uri) + )); + tracing::warn!( + uri = %uri, + attempt, + max_attempts = DATASET_OPEN_MAX_ATTEMPTS, + error = %error, + "transient object-store error opening Lance dataset; retrying under I/O gate" + ); + // Shared AIMD delay is applied on the next acquire(); keep a + // small per-attempt floor so we never spin. + tokio::time::sleep(Duration::from_millis(100)).await; + } + } + } +} + +/// Run three table futures in parallel locally, or sequentially on remote +/// object stores so we do not open/write three Lance datasets at once. +async fn join3_remote_aware( + remote: bool, + a: FA, + b: FB, + c: FC, +) -> Result<(A, B, C)> +where + FA: Future>, + FB: Future>, + FC: Future>, +{ + if remote { + Ok((a.await?, b.await?, c.await?)) + } else { + tokio::try_join!(a, b, c) + } +} + fn normalize_root_uri(value: &str) -> Result { let mut value = value.trim().to_string(); anyhow::ensure!(!value.is_empty(), "Storyline Lance root must not be empty"); @@ -1770,7 +2013,17 @@ async fn ensure_table_indexes(dataset: &mut Dataset, indexes: &[(&str, IndexType if dataset.count_rows(None).await? == 0 { return Ok(()); } - for (column, index_type) in indexes { + let table = crate::store::index_build_progress::table_label(dataset.uri()).to_string(); + let scalar_total = indexes.len(); + for (offset, (column, index_type)) in indexes.iter().enumerate() { + let kind = match index_type { + IndexType::Bitmap => "bitmap", + _ => "btree", + }; + crate::store::index_build_progress::note(format!( + "index {table}.{column} {kind} {}/{scalar_total}", + offset + 1 + )); let builtin = match index_type { IndexType::Bitmap => BuiltinIndexType::Bitmap, _ => BuiltinIndexType::BTree, @@ -1833,6 +2086,13 @@ async fn maintain_table_layout( })?; } if options.optimize_indices { + crate::store::object_store_io_gate::mark_kind( + crate::store::object_store_io_gate::IoKind::Write, + ); + crate::store::index_build_progress::note(format!( + "optimize indices {}", + crate::store::index_build_progress::table_label(path.to_string_lossy().as_ref()) + )); ensure_table_indexes(&mut dataset, indexes) .await .with_context(|| format!("ensure Storyline indices for {}", path.display()))?; @@ -1869,7 +2129,7 @@ async fn vacuum_table( let Some(retention) = retention else { return Ok(LanceMaintenanceReport::default()); }; - let dataset = Dataset::open(path.to_string_lossy().as_ref()) + let dataset = open_dataset_uri(path.to_string_lossy().as_ref()) .await .with_context(|| format!("open Storyline table {} for vacuum", path.display()))?; let retention = chrono::Duration::from_std(retention) @@ -1895,14 +2155,14 @@ fn merge_maintenance_reports( } async fn latest_table_version(path: &Path) -> Result { - Ok(Dataset::open(path.to_string_lossy().as_ref()) + Ok(open_dataset_uri(path.to_string_lossy().as_ref()) .await .with_context(|| format!("open Storyline Lance table {}", path.display()))? .version_id()) } async fn open_table_version(path: &Path, version: u64) -> Result { - let dataset = Dataset::open(path.to_string_lossy().as_ref()) + let dataset = open_dataset_uri(path.to_string_lossy().as_ref()) .await .with_context(|| format!("open Storyline Lance table {}", path.display()))?; dataset.checkout_version(version).await.with_context(|| { diff --git a/crates/persisting-pchronicle/src/store/storyline/mutation.rs b/crates/persisting-pchronicle/src/store/storyline/mutation.rs index 89d9b404..199d1fa2 100644 --- a/crates/persisting-pchronicle/src/store/storyline/mutation.rs +++ b/crates/persisting-pchronicle/src/store/storyline/mutation.rs @@ -251,8 +251,15 @@ pub(super) async fn write_batches( batches: Vec, schema: SchemaRef, indexes: &[(&str, IndexType)], + build_indexes: bool, ) -> Result { - write_record_batch_reader(path, Box::new(batch_reader(batches, schema)), indexes).await + write_record_batch_reader( + path, + Box::new(batch_reader(batches, schema)), + indexes, + build_indexes, + ) + .await } pub(super) async fn replace_table_batches( @@ -308,8 +315,10 @@ async fn write_record_batch_reader( path: &Path, reader: Box, indexes: &[(&str, IndexType)], + build_indexes: bool, ) -> Result { let uri = path.to_string_lossy().into_owned(); + crate::store::object_store_io_gate::mark_kind(crate::store::object_store_io_gate::IoKind::Write); let mut dataset = InsertBuilder::new(&uri) .with_params(&WriteParams { mode: WriteMode::Create, @@ -318,9 +327,12 @@ async fn write_record_batch_reader( .execute_stream(reader) .await .with_context(|| format!("stream ATIF into Storyline table {}", path.display()))?; - super::ensure_table_indexes(&mut dataset, indexes) - .await - .with_context(|| format!("ensure Storyline indexes for {}", path.display()))?; + if build_indexes { + super::ensure_table_indexes(&mut dataset, indexes) + .await + .with_context(|| format!("ensure Storyline indexes for {}", path.display()))?; + } + crate::store::object_store_io_gate::note_success(&uri); Ok(dataset.version_id()) } diff --git a/crates/persisting-pchronicle/src/store/storyline/writer_control.rs b/crates/persisting-pchronicle/src/store/storyline/writer_control.rs index 571ff08f..4a5ec50a 100644 --- a/crates/persisting-pchronicle/src/store/storyline/writer_control.rs +++ b/crates/persisting-pchronicle/src/store/storyline/writer_control.rs @@ -11,6 +11,16 @@ use super::{ }; const CONTROL_CAS_RETRIES: usize = 32; +/// Brief retries when CURRENT still shows a held lease after a prior release. +/// Object stores can lag on read-after-write for the control object. +#[cfg(not(test))] +const HELD_VISIBILITY_RETRIES: u32 = 30; +#[cfg(not(test))] +const HELD_RETRY_DELAY_MS: u64 = 200; +#[cfg(test)] +const HELD_VISIBILITY_RETRIES: u32 = 3; +#[cfg(test)] +const HELD_RETRY_DELAY_MS: u64 = 20; pub(super) const WRITER_LEASE_TTL_MS: u64 = 60_000; pub(super) const CURRENT_CONTROL_VERSION: u32 = 1; @@ -280,6 +290,34 @@ pub(super) fn unleased_publish_transition( Ok(Some(next)) } +fn format_lease_for_log(lease: &StorylineWriterLease, now_unix_ms: u64) -> String { + format!( + "owner={} epoch={} base_generation={} expires_in_ms={} issued_at_unix_ms={}", + lease.owner_id, + lease.epoch, + lease.base_generation.as_deref().unwrap_or(""), + lease.expires_at_unix_ms.saturating_sub(now_unix_ms), + lease.issued_at_unix_ms, + ) +} + +fn format_control_for_log(control: &StorylineCurrentControl, now_unix_ms: u64) -> String { + format!( + "revision={} committed={} lease={}", + control.revision, + control + .committed + .as_ref() + .map(|pointer| pointer.generation.as_str()) + .unwrap_or(""), + control + .lease + .as_ref() + .map(|lease| format_lease_for_log(lease, now_unix_ms)) + .unwrap_or_else(|| "".to_owned()), + ) +} + impl StorylineLanceStore { pub(super) async fn read_current_control(&self) -> Result { let result = if !self.root_uri.contains("://") { @@ -318,6 +356,7 @@ impl StorylineLanceStore { &self, control: &StorylineCurrentControl, expected: Option, + precondition: Option<&StorylineCurrentControl>, ) -> Result { validate_current_control(control)?; let contents = serde_json::to_vec(control).context("encode Storyline CURRENT control")?; @@ -325,30 +364,91 @@ impl StorylineLanceStore { write_local_current(self.root.join(CURRENT_FILE), contents).await?; return Ok(true); } - let result = match expected.as_ref() { - None => { - self.control_store - .write_create(CURRENT_FILE, contents) - .await - } - Some(version) => { - self.control_store - .write_match(CURRENT_FILE, contents, version) - .await - } - }; - match result { - Ok(_) => Ok(true), - Err(error) - if error - .downcast_ref::() - .is_some_and(opendal_store::is_conflict) => + + // Create still uses if_not_exists; updates may skip broken If-Match. + if expected.is_none() { + return match self + .control_store + .write_create(CURRENT_FILE, contents) + .await + { + Ok(()) => Ok(true), + Err(error) + if error + .downcast_ref::() + .is_some_and(opendal_store::is_conflict) => + { + Ok(false) + } + Err(error) => Err(error).with_context(|| { + format!("update Storyline CURRENT control for {}", self.root_uri) + }), + }; + } + + let skip_if_match = self + .current_if_match_unreliable + .load(std::sync::atomic::Ordering::Relaxed); + if !skip_if_match { + match self + .control_store + .write_match(CURRENT_FILE, contents.clone(), expected.as_ref().unwrap()) + .await { - Ok(false) + Ok(()) => return Ok(true), + Err(error) + if error + .downcast_ref::() + .is_some_and(opendal_store::is_conflict) => + { + let Some(precondition) = precondition else { + return Ok(false); + }; + let latest = self.read_current_control().await?; + if &latest.control != precondition { + tracing::debug!( + root_uri = %self.root_uri, + precondition = %format_control_for_log(precondition, unix_now_ms()), + latest = %format_control_for_log(&latest.control, unix_now_ms()), + "Storyline CURRENT conditional write conflict; control changed under us" + ); + return Ok(false); + } + // Remember for this store handle: avoid 412 spam on every commit. + let first = !self.current_if_match_unreliable.swap( + true, + std::sync::atomic::Ordering::Relaxed, + ); + if first { + tracing::warn!( + root_uri = %self.root_uri, + "Storyline CURRENT If-Match is unreliable on this object store; using content-checked overwrite for the rest of this writer (single-writer fallback)" + ); + } + } + Err(error) => { + return Err(error).with_context(|| { + format!("update Storyline CURRENT control for {}", self.root_uri) + }); + } + } + } else if let Some(precondition) = precondition { + let latest = self.read_current_control().await?; + if &latest.control != precondition { + return Ok(false); } - Err(error) => Err(error) - .with_context(|| format!("update Storyline CURRENT control for {}", self.root_uri)), } + + self.control_store + .write_overwrite(CURRENT_FILE, contents) + .await + .with_context(|| { + format!( + "overwrite Storyline CURRENT after If-Match fallback for {}", + self.root_uri + ) + })?; + Ok(true) } pub(super) async fn try_acquire_writer_lease( @@ -357,22 +457,57 @@ impl StorylineLanceStore { now_unix_ms: u64, ttl_ms: u64, ) -> Result { - let _control_guard = self.control_lock.lock().await; - for _ in 0..CONTROL_CAS_RETRIES { - let current = self.read_current_control().await?; - let (outcome, next) = - acquire_transition(¤t.control, owner_id, now_unix_ms, ttl_ms)?; - let Some(next) = next else { - return Ok(outcome); + let mut last_control = None; + for attempt in 1..=CONTROL_CAS_RETRIES { + let cas_result = { + let _control_guard = self.control_lock.lock().await; + let current = self.read_current_control().await?; + last_control = Some(current.control.clone()); + let (outcome, next) = + acquire_transition(¤t.control, owner_id, now_unix_ms, ttl_ms)?; + let Some(next) = next else { + return Ok(outcome); + }; + let expected_version = current.version.clone(); + let wrote = self + .try_write_current_control( + &next, + expected_version, + Some(¤t.control), + ) + .await?; + if wrote { + return Ok(outcome); + } + current }; - if self - .try_write_current_control(&next, current.version) - .await? - { - return Ok(outcome); - } + tracing::warn!( + root_uri = %self.root_uri, + owner_id, + attempt, + max_attempts = CONTROL_CAS_RETRIES, + expected_version = ?cas_result.version, + control = %format_control_for_log(&cas_result.control, now_unix_ms), + "Storyline CURRENT CAS conflict while acquiring writer lease; retrying" + ); + // Object-store backends may briefly reject conditional writes even + // when no other writer is active; back off before the next CAS. + // Sleep outside the control lock so renewals/other writers can proceed. + tokio::time::sleep(std::time::Duration::from_millis( + 20 + (attempt as u64).saturating_mul(15), + )) + .await; } - anyhow::bail!("Storyline commit conflict while acquiring writer lease") + anyhow::bail!( + "Storyline commit conflict while acquiring writer lease: CURRENT CAS exhausted after {} retries (root={}, owner={}, {})", + CONTROL_CAS_RETRIES, + self.root_uri, + owner_id, + last_control + .as_ref() + .map(|control| format_control_for_log(control, now_unix_ms)) + .unwrap_or_else(|| "control=".to_owned()), + ) } pub(super) async fn acquire_writer_lease_for_generation( @@ -380,29 +515,85 @@ impl StorylineLanceStore { owner_id: &str, expected_generation: Option<&str>, ) -> Result { - let acquired = match self - .try_acquire_writer_lease(owner_id, unix_now_ms(), WRITER_LEASE_TTL_MS) - .await? - { - LeaseAcquireOutcome::Held(_) => { - anyhow::bail!("Storyline commit conflict while acquiring writer lease") + let mut last_held: Option = None; + for attempt in 1..=HELD_VISIBILITY_RETRIES { + let now = unix_now_ms(); + match self + .try_acquire_writer_lease(owner_id, now, WRITER_LEASE_TTL_MS) + .await? + { + LeaseAcquireOutcome::Held(held) => { + tracing::warn!( + root_uri = %self.root_uri, + owner_id, + attempt, + max_attempts = HELD_VISIBILITY_RETRIES, + expected_generation = expected_generation.unwrap_or(""), + held = %format_lease_for_log(&held, now), + retry_delay_ms = HELD_RETRY_DELAY_MS, + "Storyline writer lease still held; retrying in case object-store CURRENT is stale" + ); + last_held = Some(held); + if attempt == HELD_VISIBILITY_RETRIES { + break; + } + tokio::time::sleep(std::time::Duration::from_millis(HELD_RETRY_DELAY_MS)).await; + } + LeaseAcquireOutcome::Acquired(acquired) => { + if acquired.lease.base_generation.as_deref() == expected_generation { + if attempt > 1 { + tracing::warn!( + root_uri = %self.root_uri, + owner_id, + attempt, + expected_generation = expected_generation.unwrap_or(""), + acquired = %format_lease_for_log(&acquired.lease, now), + "Storyline writer lease acquired after visibility/CAS retries" + ); + } + return Ok(acquired); + } + tracing::warn!( + root_uri = %self.root_uri, + owner_id, + attempt, + expected_generation = expected_generation.unwrap_or(""), + acquired = %format_lease_for_log(&acquired.lease, now), + "Storyline writer lease base_generation mismatch; releasing and failing" + ); + let conflict = anyhow::anyhow!( + "Storyline commit conflict while acquiring writer lease: base_generation mismatch (root={}, owner={}, expected={}, acquired={})", + self.root_uri, + owner_id, + expected_generation.unwrap_or(""), + format_lease_for_log(&acquired.lease, now), + ); + return match self + .release_writer_lease(owner_id, acquired.lease.epoch) + .await + { + Ok(true) => Err(conflict), + Ok(false) => Err(conflict + .context("mismatched writer lease was lost before release")), + Err(error) => Err(conflict.context(format!( + "failed to release mismatched writer lease: {error:#}" + ))), + }; + } } - LeaseAcquireOutcome::Acquired(acquired) => acquired, - }; - if acquired.lease.base_generation.as_deref() == expected_generation { - return Ok(acquired); - } - let conflict = anyhow::anyhow!("Storyline commit conflict while acquiring writer lease"); - match self - .release_writer_lease(owner_id, acquired.lease.epoch) - .await - { - Ok(true) => Err(conflict), - Ok(false) => Err(conflict.context("mismatched writer lease was lost before release")), - Err(error) => Err(conflict.context(format!( - "failed to release mismatched writer lease: {error:#}" - ))), } + let now = unix_now_ms(); + anyhow::bail!( + "Storyline commit conflict while acquiring writer lease: still held after {} visibility retries (root={}, owner={}, expected={}, {})", + HELD_VISIBILITY_RETRIES, + self.root_uri, + owner_id, + expected_generation.unwrap_or(""), + last_held + .as_ref() + .map(|lease| format_lease_for_log(lease, now)) + .unwrap_or_else(|| "held=".to_owned()), + ) } async fn transition_current_control( @@ -416,7 +607,7 @@ impl StorylineLanceStore { return Ok(false); }; if self - .try_write_current_control(&next, current.version) + .try_write_current_control(&next, current.version.clone(), Some(¤t.control)) .await? { return Ok(true); diff --git a/docs/src/en/pchronicle/guides/exchange.md b/docs/src/en/pchronicle/guides/exchange.md index 13b28abc..546a0a6d 100644 --- a/docs/src/en/pchronicle/guides/exchange.md +++ b/docs/src/en/pchronicle/guides/exchange.md @@ -20,13 +20,13 @@ Lance only to classify the tree ([RFC-0015](../../rfcs/0015-chronicle-manifest.m ```bash pchronicle import --from input.json \ - --to ./imported --input-format atif + --to ./imported --input-format atif ``` -The default `--mode create` refuses an existing target. Use `--mode append` +The default create behavior refuses an existing target. Use `--append` for an existing Storyline Dataset; duplicate `document_id` values receive a `#N` suffix by default, or can be skipped with `--on-duplicate skip`. Use -`--mode replace` to stage the complete import and atomically replace an existing +`--replace` to stage the complete import and atomically replace an existing local Dataset after confirmation; replacement requires interactive confirmation or `--yes`. Object-store Dataset replace clears the destination prefix before writing (not atomic; an interrupted replace may leave the target empty). @@ -48,7 +48,7 @@ output: ```bash pchronicle import --from ./corpus --to ./normalized \ - --output-format storyline + --output-format storyline ``` A validated, non-empty canonical Event Store is detected before JSON scanning @@ -68,7 +68,7 @@ In the squashed Dataset, `_file_` is `.` for all normalized rows: ```bash pchronicle query ./normalized \ - --sql 'SELECT _file_, COUNT(*) AS runs FROM dataset.runs GROUP BY _file_' + --sql 'SELECT _file_, COUNT(*) AS runs FROM dataset.runs GROUP BY _file_' ``` `document_id` is globally unique in Storyline output. Collisions receive a @@ -84,7 +84,7 @@ Stdin must be finite and explicit: ```bash cat input.json | pchronicle import --from - \ - --to ./imported --input-format openai-messages + --to ./imported --input-format openai-messages ``` After import, inspect the new boundary: @@ -98,14 +98,14 @@ pchronicle stats overview ./imported ```bash pchronicle export --from ./imported \ - --to restored.json --output-format atif + --to restored.json --output-format atif ``` Narrow the export with file path and external identity when needed: ```bash pchronicle export --from ./imported --to one.json --output-format actf \ - --source source.json --session-id session-42 --strict + --source source.json --session-id session-42 --strict ``` `--strict` fails when the target format cannot preserve the original exchange diff --git a/docs/src/en/pchronicle/reference/cases-self.md b/docs/src/en/pchronicle/reference/cases-self.md index c3d88006..113f2079 100644 --- a/docs/src/en/pchronicle/reference/cases-self.md +++ b/docs/src/en/pchronicle/reference/cases-self.md @@ -15,7 +15,7 @@ cd /tmp/pchronicle-cases ## S01: Browse a local Dataset ```bash -pchronicle import --from "$PCHRONICLE_CASE_FIXTURES/atif/support-ticket.json" --to ./trajectory-data --mode create +pchronicle import --from "$PCHRONICLE_CASE_FIXTURES/atif/support-ticket.json" --to ./trajectory-data pchronicle list ./trajectory-data pchronicle stats ./trajectory-data ``` @@ -25,9 +25,9 @@ Expected: the commands list runs, steps, and tool calls in the Dataset. ## S02: Run a SQL query ```bash -pchronicle import --from "$PCHRONICLE_CASE_FIXTURES/atif/support-ticket.json" --to ./trajectory-data --mode create +pchronicle import --from "$PCHRONICLE_CASE_FIXTURES/atif/support-ticket.json" --to ./trajectory-data pchronicle query ./trajectory-data \ - --sql 'SELECT COUNT(*) AS runs FROM dataset.runs' + --sql 'SELECT COUNT(*) AS runs FROM dataset.runs' ``` Expected: the query succeeds and returns a definite run count. @@ -35,7 +35,7 @@ Expected: the query succeeds and returns a definite run count. ## S03: Run a built-in analysis ```bash -pchronicle import --from "$PCHRONICLE_CASE_FIXTURES/atif/support-ticket.json" --to ./trajectory-data --mode create +pchronicle import --from "$PCHRONICLE_CASE_FIXTURES/atif/support-ticket.json" --to ./trajectory-data pchronicle stats overview ./trajectory-data ``` @@ -44,7 +44,7 @@ Expected: output includes run, step, and tool-call counts plus a time range. ## S04: Import and export ```bash -pchronicle import --from "$PCHRONICLE_CASE_FIXTURES/atif/support-ticket.json" --to ./trajectory-data --mode create +pchronicle import --from "$PCHRONICLE_CASE_FIXTURES/atif/support-ticket.json" --to ./trajectory-data pchronicle export --from ./trajectory-data --to ./output.atif.json --output-format atif test -s ./output.atif.json ``` diff --git a/docs/src/en/pchronicle/reference/cli.md b/docs/src/en/pchronicle/reference/cli.md index fe05c5ef..464c7ac9 100644 --- a/docs/src/en/pchronicle/reference/cli.md +++ b/docs/src/en/pchronicle/reference/cli.md @@ -8,14 +8,14 @@ line. New commands and scripts should use the syntax documented here. Start with the shortest path to a useful answer: - **Try the product:** `pchronicle onboard query` uses temporary example data - and needs no Dataset path. + and needs no Dataset path. - **Check a Dataset:** use `list`/`ls` and `stats overview` before writing SQL. - **Locate a run or phrase:** use `find --run-id`, `--session-id`, or - `--match`; inspect the returned identity before querying more data. + `--match`; inspect the returned identity before querying more data. - **Ask a repeatable question:** use `query --sql` or `query --file` and set - output and resource limits for automation. + output and resource limits for automation. - **Expose history:** use `serve` only after the read-only query works; the - [serve guide](../guides/serve.md) explains the lifecycle and shutdown path. + [serve guide](../guides/serve.md) explains the lifecycle and shutdown path. For a first interaction, copy this sequence: @@ -129,7 +129,7 @@ pchronicle list|ls [DATASET] [OPTIONS] pchronicle stats [DATASET] [OPTIONS] pchronicle stats [DATASET] [OPTIONS] pchronicle find [DATASET] - (--run-id ID|--document-id ID|--session-id ID|--match EXPRESSION) [OPTIONS] + (--run-id ID|--document-id ID|--session-id ID|--match EXPRESSION) [OPTIONS] ``` ```bash @@ -171,8 +171,8 @@ pchronicle query [DATASET|--mount NAME=DATASET ...] (--sql SQL|--file FILE_OR_ST ```bash pchronicle query ./dataset --sql 'SELECT COUNT(*) FROM dataset.runs' pchronicle query \ - --mount live=./live --mount archive=@archive \ - --file report.sql + --mount live=./live --mount archive=@archive \ + --file report.sql ``` Each invocation accepts one read-only statement with explicit resource limits. `--file -` reads SQL @@ -183,26 +183,26 @@ from stdin. Use `--format`, `--output`, `--max-output-rows`, ```text pchronicle import -f|--from SOURCE -t|--to NEW_DATASET - [-i|--input-format auto|atif|actf|openai-messages|storyline|codex|claude-code|compact-jsonl] - [-o|--output-format preserve|storyline|compact-jsonl] - [--mode create|append|replace] [--on-duplicate suffix|skip] [--yes] - [--column NAME=JSON_PATH]... [OPTIONS] + [-i|--input-format auto|atif|actf|openai-messages|storyline|codex|claude-code|compact-jsonl] + [-o|--output-format preserve|storyline|compact-jsonl] + [|--replace] [--append] [--on-duplicate suffix|skip] [--yes] + [--column NAME=JSON_PATH]... [OPTIONS] ``` ```bash pchronicle import -f input.json -t ./imported -i atif cat input.json | pchronicle import -f - -t ./imported -i openai-messages -pchronicle import -f more.json -t ./normalized --mode append --on-duplicate skip -pchronicle import -f rebuilt.json -t ./normalized --mode replace --yes +pchronicle import -f more.json -t ./normalized --append --on-duplicate skip +pchronicle import -f rebuilt.json -t ./normalized --replace --yes pchronicle import -f ./jsonl-root -t ./records.lance \ - -o compact-jsonl \ - --column id=$.event.id --column timestamp=$.event.time \ - --column model=$.payload.model + -o compact-jsonl \ + --column id=$.event.id --column timestamp=$.event.time \ + --column model=$.payload.model ``` -`-` means stdin. `create` is the default and requires a new destination. -`append` requires an existing Storyline Dataset and either suffixes colliding -`document_id` values with `#N` (the default) or skips them. `replace` moves the +`-` means stdin. Create is the default and requires a new destination. +`--append` requires an existing Storyline Dataset and either suffixes colliding +`document_id` values with `#N` (the default) or skips them. `--replace` moves the old local Dataset aside, publishes the fully imported Dataset with a rename transaction, and only then removes the old data. It requires interactive confirmation or `--yes`; an existing object-store Dataset cannot currently be @@ -224,8 +224,8 @@ local `create` and confirmed `replace`, but not stdin, object-store targets, or ```text pchronicle sync --from DIRECTORY --to DIRECTORY --convert DIRECTORY - [--input-format FORMAT] [--column NAME=JSON_PATH]... - [--interval DURATION] [--once] + [--input-format FORMAT] [--column NAME=JSON_PATH]... + [--interval DURATION] [--once] ``` `sync` is a resident polling worker for `.json`, `.jsonl`, and `.ndjson` files. @@ -259,7 +259,7 @@ filesystem roots or whole object-store buckets. ```text pchronicle export -f|--from DATASET -t|--to TARGET - -o|--output-format atif|actf|openai-messages|storyline|compact-jsonl [OPTIONS] + -o|--output-format atif|actf|openai-messages|storyline|compact-jsonl [OPTIONS] ``` ```bash @@ -274,7 +274,7 @@ unless `--overwrite` is explicit. ```text pchronicle agent [DATASET] - [--ask QUESTION|--ask-file FILE_OR_STDIN] [--no-overview] [--dry-run] + [--ask QUESTION|--ask-file FILE_OR_STDIN] [--no-overview] [--dry-run] ``` ```bash @@ -286,27 +286,27 @@ pchronicle agent claude @prod --ask 'Compare model latency' ```text pchronicle serve - [--listen LOOPBACK_ADDR] [--control LOOPBACK_ADDR] [--open] - [--gateway ADDRESS --gateway-dataset DATASET [--gateway-split TEMPLATE] - [--gateway-split-idle DURATION]] - [--gateway-config FILE --gateway-dataset DATASET [--gateway-state DIRECTORY]] - [--gateway-stream-markdown] [--gateway-debug] - [--catalog-config FILE] - [<[NAME=]DATASET> ...] -pchronicle serve catalog dataset add --catalog-config FILE NAME --uri URI [OPTIONS] + [--listen LOOPBACK_ADDR] [--control LOOPBACK_ADDR] [--open] + [--gateway ADDRESS --gateway-dataset DATASET [--gateway-split TEMPLATE] + [--gateway-split-idle DURATION]] + [--gateway-config FILE --gateway-dataset DATASET [--gateway-state DIRECTORY]] + [--gateway-stream-markdown] [--gateway-debug] + [--catalog-config FILE] + [<[NAME=]DATASET> ...] +pchronicle serve catalog dataset add --catalog-config FILE NAME --uri URI [OPTIONS] pchronicle serve catalog dataset remove --catalog-config FILE NAME... -pchronicle serve catalog dataset list --catalog-config FILE -pchronicle serve catalog issue --catalog-config FILE NAME -pchronicle serve catalog grant --catalog-config FILE NAME DATASET... +pchronicle serve catalog dataset list --catalog-config FILE +pchronicle serve catalog issue --catalog-config FILE NAME +pchronicle serve catalog grant --catalog-config FILE NAME DATASET... pchronicle serve catalog revoke --catalog-config FILE NAME DATASET... ``` ```bash pchronicle serve ./trajectory-data pchronicle serve \ - --gateway auto \ - --gateway-dataset ./trajectory-data \ - --gateway-split '{user}/{date}/{hour}' + --gateway auto \ + --gateway-dataset ./trajectory-data \ + --gateway-split '{user}/{date}/{hour}' ``` Every listener must use a loopback address. A bare single Dataset is mounted as @@ -351,13 +351,13 @@ The Directory ACL file contains users, datasets (libraries), and grants. Management commands create the file when it does not exist. ```text -pchronicle serve catalog issue --catalog-config FILE NAME -pchronicle serve catalog grant --catalog-config FILE NAME DATASET... +pchronicle serve catalog issue --catalog-config FILE NAME +pchronicle serve catalog grant --catalog-config FILE NAME DATASET... pchronicle serve catalog revoke --catalog-config FILE NAME DATASET... -pchronicle serve catalog dataset add --catalog-config FILE NAME --uri URI - [--endpoint URL] [--region REGION] [--access-key KEY] [--secret-key KEY] +pchronicle serve catalog dataset add --catalog-config FILE NAME --uri URI + [--endpoint URL] [--region REGION] [--access-key KEY] [--secret-key KEY] pchronicle serve catalog dataset remove --catalog-config FILE NAME... -pchronicle serve catalog dataset list --catalog-config FILE +pchronicle serve catalog dataset list --catalog-config FILE ``` `issue` generates a user AK/SK and prints the secret once. `dataset add` diff --git a/docs/src/en/rfcs/0015-chronicle-manifest.md b/docs/src/en/rfcs/0015-chronicle-manifest.md index 498a5a72..76a7f55e 100644 --- a/docs/src/en/rfcs/0015-chronicle-manifest.md +++ b/docs/src/en/rfcs/0015-chronicle-manifest.md @@ -155,7 +155,7 @@ source of truth that travels with the dataset. | Field | Type | Rules | |---|---|---| -| `format` | string | MUST be present for `kind = "leaf"`; v1 writers MUST use `compact-jsonl/v1` | +| `format` | string | MUST be present for `kind = "leaf"`; v1 writers MUST use `compact-jsonl/v1` or `storyline/v1` | Unknown `format` values MUST be preserved by generic readers; format-specific openers MAY reject unsupported values. diff --git a/docs/src/zh/pchronicle/guides/exchange.md b/docs/src/zh/pchronicle/guides/exchange.md index 4ae26f87..bb30b7a7 100644 --- a/docs/src/zh/pchronicle/guides/exchange.md +++ b/docs/src/zh/pchronicle/guides/exchange.md @@ -16,11 +16,11 @@ dataset 根写入 leaf `chronicle.manifest`，便于后续 discovery 不必仅 ```bash pchronicle import --from input.json \ - --to ./imported --input-format atif + --to ./imported --input-format atif ``` -默认 `--mode create` 会拒绝已有目标。`--mode append` 用于已有 Storyline Dataset；重复 -`document_id` 默认增加 `#N` 后缀，也可用 `--on-duplicate skip` 跳过。`--mode replace` 会先 +默认会拒绝已有目标。`--append` 用于已有 Storyline Dataset；重复 +`document_id` 默认增加 `#N` 后缀，也可用 `--on-duplicate skip` 跳过。`--replace` 会先 把完整导入写入临时路径，确认后以 rename 事务替换已有的本地 Dataset，最后才删除旧数据；要求 交互确认或传入 `--yes`。对象存储 Dataset 的 replace 会先清空目标前缀再写入（非原子；中断可能导致目标暂时为空）。普通文件可以自动识别。目录输入会递归扫描 `.json`、`.jsonl` 与 `.ndjson` 文件；默认输出会保留其相对 @@ -37,7 +37,7 @@ pchronicle import --from ./claude-sessions --to ./claude-ds --input-format claud ```bash pchronicle import --from ./corpus --to ./normalized \ - --output-format storyline + --output-format storyline ``` 经过验证且非空的 canonical Event Store 会在 JSON 扫描前被识别，并始终创建 @@ -56,7 +56,7 @@ squash 后，Dataset 所有规范化表中的 `_file_` 都是 `.`： ```bash pchronicle query ./normalized \ - --sql 'SELECT _file_, COUNT(*) AS runs FROM dataset.runs GROUP BY _file_' + --sql 'SELECT _file_, COUNT(*) AS runs FROM dataset.runs GROUP BY _file_' ``` Storyline 输出中的 `document_id` 全局唯一；冲突时会确定性地增加 `#N` 后缀，append 也可用 @@ -71,7 +71,7 @@ ATIF `.jsonl` 与 `.ndjson` 输入会逐条解码其中的非空记录。递归 ```bash cat input.json | pchronicle import --from - \ - --to ./imported --input-format openai-messages + --to ./imported --input-format openai-messages ``` 导入后检查新边界： @@ -85,14 +85,14 @@ pchronicle stats overview ./imported ```bash pchronicle export --from ./imported \ - --to restored.json --output-format atif + --to restored.json --output-format atif ``` 需要时使用文件路径与外部 ID 缩小导出范围： ```bash pchronicle export --from ./imported --to one.json --output-format actf \ - --source source.json --session-id session-42 --strict + --source source.json --session-id session-42 --strict ``` 目标格式无法保留原交换文档时，`--strict` 会失败。输出文件默认 create-only，覆盖必须显式 diff --git a/docs/src/zh/pchronicle/reference/cases-self.md b/docs/src/zh/pchronicle/reference/cases-self.md index 6d273728..b287f756 100644 --- a/docs/src/zh/pchronicle/reference/cases-self.md +++ b/docs/src/zh/pchronicle/reference/cases-self.md @@ -15,7 +15,7 @@ cd /tmp/pchronicle-cases ## S01：浏览本地 Dataset ```bash -pchronicle import --from "$PCHRONICLE_CASE_FIXTURES/atif/support-ticket.json" --to ./trajectory-data --mode create +pchronicle import --from "$PCHRONICLE_CASE_FIXTURES/atif/support-ticket.json" --to ./trajectory-data pchronicle list ./trajectory-data pchronicle stats ./trajectory-data ``` @@ -25,9 +25,9 @@ pchronicle stats ./trajectory-data ## S02：执行 SQL 查询 ```bash -pchronicle import --from "$PCHRONICLE_CASE_FIXTURES/atif/support-ticket.json" --to ./trajectory-data --mode create +pchronicle import --from "$PCHRONICLE_CASE_FIXTURES/atif/support-ticket.json" --to ./trajectory-data pchronicle query ./trajectory-data \ - --sql 'SELECT COUNT(*) AS runs FROM dataset.runs' + --sql 'SELECT COUNT(*) AS runs FROM dataset.runs' ``` 预期：查询成功并返回确定的 runs 数量。 @@ -35,7 +35,7 @@ pchronicle query ./trajectory-data \ ## S03：运行内建分析 ```bash -pchronicle import --from "$PCHRONICLE_CASE_FIXTURES/atif/support-ticket.json" --to ./trajectory-data --mode create +pchronicle import --from "$PCHRONICLE_CASE_FIXTURES/atif/support-ticket.json" --to ./trajectory-data pchronicle stats overview ./trajectory-data ``` @@ -44,7 +44,7 @@ pchronicle stats overview ./trajectory-data ## S04：导入和导出 ```bash -pchronicle import --from "$PCHRONICLE_CASE_FIXTURES/atif/support-ticket.json" --to ./trajectory-data --mode create +pchronicle import --from "$PCHRONICLE_CASE_FIXTURES/atif/support-ticket.json" --to ./trajectory-data pchronicle export --from ./trajectory-data --to ./output.atif.json --output-format atif test -s ./output.atif.json ``` diff --git a/docs/src/zh/pchronicle/reference/cli.md b/docs/src/zh/pchronicle/reference/cli.md index d9b1228c..8f01003e 100644 --- a/docs/src/zh/pchronicle/reference/cli.md +++ b/docs/src/zh/pchronicle/reference/cli.md @@ -53,8 +53,8 @@ Dataset 内部可以保存一种或多种受支持的运行数据格式。pChron `@NAME` 明确表示一个 dataset pin。裸字符串始终按路径或 URI 解释： ```text -prod 本地相对路径 ./prod -@prod 名为 prod 的 Dataset pin +prod 本地相对路径 ./prod +@prod 名为 prod 的 Dataset pin ``` 这种区分可以避免同名目录出现或消失时，命令突然解析到不同位置。 @@ -158,7 +158,7 @@ S3 凭证用 `--ak`/`--sk` 写在同一 pin 表中，不会被 `dataset list` / ```text pchronicle list [DATASET] [--physical] [--format auto|table|json] [--errors report|strict] - [--max-files N] [--max-entries N] + [--max-files N] [--max-entries N] ``` ```bash @@ -174,7 +174,7 @@ pchronicle list @prod --physical --format json --errors strict ```text pchronicle stats [DATASET] [--format auto|table|json] [--errors report|strict] [--timeout 30s] - [--max-files N] [--max-entries N] + [--max-files N] [--max-entries N] ``` ```bash @@ -190,8 +190,8 @@ canonical Event Store 的 Storyline projection 状态。还可以用 `--max-file ```text pchronicle stats [DATASET] - [--format auto|table|jsonl|csv|tsv] - [--limit 100] [--max-output-bytes 8MiB] [--timeout 30s] + [--format auto|table|jsonl|csv|tsv] + [--limit 100] [--max-output-bytes 8MiB] [--timeout 30s] ``` ```bash @@ -212,20 +212,20 @@ pchronicle stats tools @prod --format csv --limit 20 ```text pchronicle find [DATASET] - (--run-id ID|--document-id ID|--session-id ID|--match EXPRESSION) - [--source PATH] [--step-id N] [--match EXPRESSION ...] - [--format auto|table|json] [--max-results N] + (--run-id ID|--document-id ID|--session-id ID|--match EXPRESSION) + [--source PATH] [--step-id N] [--match EXPRESSION ...] + [--format auto|table|json] [--max-results N] ``` ```bash pchronicle find @prod --session-id session-42 pchronicle find ./dataset \ - --source nested/source.json \ - --session-id session-42 --step-id 7 + --source nested/source.json \ + --session-id session-42 --step-id 7 pchronicle find ./dataset \ - --match "timeout" --match "retry" --format json + --match "timeout" --match "retry" --format json pchronicle find ./dataset \ - --match '$.tags=important' --match '$.priority=2' --format json + --match '$.tags=important' --match '$.priority=2' --format json ``` 外部 ID 不保证在整个 Dataset 内唯一。没有 `--source` 时，同一个 ID 可以返回多个候选；结果中的 @@ -247,19 +247,19 @@ CLI 不一致时以 CLI 为准。 ```text pchronicle query [DATASET|--mount NAME=DATASET ...] (--sql SQL|--file FILE_OR_STDIN) - [--format auto|table|jsonl|csv] [--output PATH_OR_STDOUT] - [--max-output-rows N] [--max-output-bytes BYTES] [--timeout 30s] + [--format auto|table|jsonl|csv] [--output PATH_OR_STDOUT] + [--max-output-rows N] [--max-output-bytes BYTES] [--timeout 30s] ``` ```bash pchronicle query ./dataset \ - --sql 'SELECT COUNT(*) AS runs FROM dataset.runs' + --sql 'SELECT COUNT(*) AS runs FROM dataset.runs' pchronicle query \ - --mount live=./live \ - --mount archive=@archive \ - --sql 'SELECT * FROM live.runs - UNION ALL - SELECT * FROM archive.runs' + --mount live=./live \ + --mount archive=@archive \ + --sql 'SELECT * FROM live.runs + UNION ALL + SELECT * FROM archive.runs' ``` `--file` 从文件读取 SQL，`--file -` 从 stdin 读取；`--format`、`--output`、输出上限和 `--timeout` @@ -270,27 +270,27 @@ pchronicle query \ ```text pchronicle import -f|--from SOURCE -t|--to NEW_DATASET - [-i|--input-format FORMAT] [-o|--output-format preserve|storyline|compact-jsonl] - [--mode create|append|replace] [--on-duplicate suffix|skip] [--yes] - [--column NAME=JSON_PATH]... [--max-input-bytes BYTES] + [-i|--input-format FORMAT] [-o|--output-format preserve|storyline|compact-jsonl] + [--replace] [--append] [--on-duplicate suffix|skip] [--yes] + [--column NAME=JSON_PATH]... [--max-input-bytes BYTES] ``` ```bash pchronicle import \ - -f input.json -t ./imported -i atif + -f input.json -t ./imported -i atif pchronicle import \ - -f ./corpus \ - -t s3://bucket/normalized \ - -o storyline + -f ./corpus \ + -t s3://bucket/normalized \ + -o storyline pchronicle import \ - -f more.json -t ./normalized --mode append --on-duplicate skip + -f more.json -t ./normalized --append --on-duplicate skip pchronicle import \ - -f rebuilt.json -t ./normalized --mode replace --yes + -f rebuilt.json -t ./normalized --replace --yes pchronicle import \ - -f ./jsonl-root -t ./records.lance \ - -o compact-jsonl \ - --column id=$.event.id --column timestamp=$.event.time \ - --column model=$.payload.model + -f ./jsonl-root -t ./records.lance \ + -o compact-jsonl \ + --column id=$.event.id --column timestamp=$.event.time \ + --column model=$.payload.model ``` 长参数分别是 `--from`、`--to`、`--input-format` 和 `--output-format`。短 option 始终只有一个字符， @@ -309,8 +309,8 @@ pchronicle import \ | `compact-jsonl` | 是 | 是 | Codex 和 Claude Code session 是 decode-only 输入格式。Canonical Event Store 会自动识别并投影为 -Storyline Dataset。默认 `create` 模式要求目标不存在。`append` 要求目标是已有 Storyline Dataset； -重复 `document_id` 默认增加 `#N` 后缀，也可用 `--on-duplicate skip` 跳过。`replace` 会先将完整导入 +Storyline Dataset。默认创建要求目标不存在。`--append` 要求目标是已有 Storyline Dataset； +重复 `document_id` 默认增加 `#N` 后缀，也可用 `--on-duplicate skip` 跳过。`--replace` 会先将完整导入 写入临时路径，再将旧本地 Dataset rename 到备份路径、将新 Dataset rename 到正式路径，确认新路径 发布后才删除备份；因此必须交互确认或传入 `--yes`。对象存储 Dataset 的 replace 会先清空目标前缀再写入（非原子）。 @@ -327,8 +327,8 @@ Compact JSONL 是记录存储，不会转换或推断轨迹语义。指定 ```text pchronicle sync --from DIRECTORY --to DIRECTORY --convert DIRECTORY - [--input-format FORMAT] [--column NAME=JSON_PATH]... - [--interval DURATION] [--once] + [--input-format FORMAT] [--column NAME=JSON_PATH]... + [--interval DURATION] [--once] ``` `sync` 是常驻轮询器：监听源目录下的 `.json`、`.jsonl` 和 `.ndjson`。对于运行数据格式，它会将 @@ -355,16 +355,16 @@ pchronicle drop DATASET [--yes] ```text pchronicle export -f|--from DATASET -t|--to TARGET -o|--output-format FORMAT - [--source PATH] [--run-id ID|--document-id ID|--session-id ID] [--where EXPRESSION] - [--strict] [--overwrite] [--max-trajectories N] [--max-output-bytes BYTES] [--timeout 30s] + [--source PATH] [--run-id ID|--document-id ID|--session-id ID] [--where EXPRESSION] + [--strict] [--overwrite] [--max-trajectories N] [--max-output-bytes BYTES] [--timeout 30s] ``` ```bash pchronicle export \ - -f ./imported -t restored.json -o atif + -f ./imported -t restored.json -o atif pchronicle export \ - -f ./imported \ - -t - -o actf --session-id session-42 --strict + -f ./imported \ + -t - -o actf --session-id session-42 --strict ``` 长参数分别是 `--from`、`--to` 和 `--output-format`。过滤条件包括 `--source`、`--run-id`、 @@ -379,7 +379,7 @@ pchronicle export \ ```text pchronicle agent [DATASET] - [--ask QUESTION|--ask-file FILE_OR_STDIN] [--no-overview] [--dry-run] + [--ask QUESTION|--ask-file FILE_OR_STDIN] [--no-overview] [--dry-run] ``` ```bash @@ -395,27 +395,27 @@ Agent 注入是行为引导，不是 filesystem、network 或 tool permission ```text pchronicle serve - [--listen LOOPBACK_ADDR] [--control LOOPBACK_ADDR] [--open] - [--gateway ADDRESS --gateway-dataset DATASET [--gateway-split TEMPLATE] - [--gateway-split-idle DURATION]] - [--gateway-config FILE --gateway-dataset DATASET [--gateway-state DIRECTORY]] - [--gateway-stream-markdown] [--gateway-debug] - [--catalog-config FILE] - [<[NAME=]DATASET> ...] -pchronicle serve catalog dataset add --catalog-config FILE NAME --uri URI [OPTIONS] + [--listen LOOPBACK_ADDR] [--control LOOPBACK_ADDR] [--open] + [--gateway ADDRESS --gateway-dataset DATASET [--gateway-split TEMPLATE] + [--gateway-split-idle DURATION]] + [--gateway-config FILE --gateway-dataset DATASET [--gateway-state DIRECTORY]] + [--gateway-stream-markdown] [--gateway-debug] + [--catalog-config FILE] + [<[NAME=]DATASET> ...] +pchronicle serve catalog dataset add --catalog-config FILE NAME --uri URI [OPTIONS] pchronicle serve catalog dataset remove --catalog-config FILE NAME... -pchronicle serve catalog dataset list --catalog-config FILE -pchronicle serve catalog issue --catalog-config FILE NAME -pchronicle serve catalog grant --catalog-config FILE NAME DATASET... +pchronicle serve catalog dataset list --catalog-config FILE +pchronicle serve catalog issue --catalog-config FILE NAME +pchronicle serve catalog grant --catalog-config FILE NAME DATASET... pchronicle serve catalog revoke --catalog-config FILE NAME DATASET... ``` ```bash pchronicle serve ./trajectory-data pchronicle serve \ - --gateway auto \ - --gateway-dataset ./trajectory-data \ - --gateway-split '{user}/{date}/{hour}' + --gateway auto \ + --gateway-dataset ./trajectory-data \ + --gateway-split '{user}/{date}/{hour}' ``` 未指定服务 flag 时，只读 Web/API 默认监听 `127.0.0.1:0`。多个 Dataset 使用 @@ -442,13 +442,13 @@ loopback；服务准备完成后，stdout 输出一行版本化 readiness JSON Directory ACL 文件包含用户、datasets（libraries）和 grants。配置文件不存在时，管理命令会自动创建。 ```text -pchronicle serve catalog issue --catalog-config FILE NAME -pchronicle serve catalog grant --catalog-config FILE NAME DATASET... +pchronicle serve catalog issue --catalog-config FILE NAME +pchronicle serve catalog grant --catalog-config FILE NAME DATASET... pchronicle serve catalog revoke --catalog-config FILE NAME DATASET... -pchronicle serve catalog dataset add --catalog-config FILE NAME --uri URI - [--endpoint URL] [--region REGION] [--access-key KEY] [--secret-key KEY] +pchronicle serve catalog dataset add --catalog-config FILE NAME --uri URI + [--endpoint URL] [--region REGION] [--access-key KEY] [--secret-key KEY] pchronicle serve catalog dataset remove --catalog-config FILE NAME... -pchronicle serve catalog dataset list --catalog-config FILE +pchronicle serve catalog dataset list --catalog-config FILE ``` `issue` 生成用户 AK/SK 并只显示一次 secret；`dataset add` 只登记 URI 与可选后端存储凭据， @@ -485,9 +485,9 @@ pchronicle dataset pin local ./trajectory-data pchronicle dataset pin default @local pchronicle import \ - -f ./training.json \ - -t ./trajectory-data/training \ - -i openai-messages + -f ./training.json \ + -t ./trajectory-data/training \ + -i openai-messages pchronicle list pchronicle stats @@ -501,16 +501,16 @@ pchronicle dataset pin live s3://bucket/live pchronicle dataset pin archive s3://bucket/archive pchronicle query \ - --mount live=@live \ - --mount archive=@archive \ - --sql 'SELECT model_name, COUNT(*) AS steps - FROM ( - SELECT model_name FROM live.steps - UNION ALL - SELECT model_name FROM archive.steps - ) - GROUP BY model_name - ORDER BY steps DESC' + --mount live=@live \ + --mount archive=@archive \ + --sql 'SELECT model_name, COUNT(*) AS steps + FROM ( + SELECT model_name FROM live.steps + UNION ALL + SELECT model_name FROM archive.steps + ) + GROUP BY model_name + ORDER BY steps DESC' ``` ### 找到并严格导出一条 Run @@ -519,29 +519,29 @@ pchronicle query \ pchronicle find @prod --session-id session-42 --format json pchronicle export \ - -f @prod \ - -t session-42.actf.json \ - -o actf \ - --source nested/source.json \ - --session-id session-42 \ - --strict + -f @prod \ + -t session-42.actf.json \ + -o actf \ + --source nested/source.json \ + --session-id session-42 \ + --strict ``` ### 在 CI 中使用 ```bash pchronicle \ - -c ./ci-config.toml \ - --log-level error \ - status ./fixtures \ - --format json > status.json + -c ./ci-config.toml \ + --log-level error \ + status ./fixtures \ + --format json > status.json pchronicle \ - -c ./ci-config.toml \ - --log-level error \ - query ./fixtures \ - --file checks.sql \ - --format jsonl > checks.jsonl + -c ./ci-config.toml \ + --log-level error \ + query ./fixtures \ + --file checks.sql \ + --format jsonl > checks.jsonl ``` 定位后再写 SQL 见 [发现并查询](../guides/discover-and-query.md)，交换见 diff --git a/docs/src/zh/rfcs/0015-chronicle-manifest.md b/docs/src/zh/rfcs/0015-chronicle-manifest.md index 313aad2b..4d7639d0 100644 --- a/docs/src/zh/rfcs/0015-chronicle-manifest.md +++ b/docs/src/zh/rfcs/0015-chronicle-manifest.md @@ -129,7 +129,7 @@ Warehouse / Catalog MAY 在进程内缓存已发现的 leaf stats 与前缀聚 | 字段 | 类型 | 规则 | |---|---|---| -| `format` | string | `kind = "leaf"` 时 MUST 存在；v1 写入方 MUST 使用 `compact-jsonl/v1` | +| `format` | string | `kind = "leaf"` 时 MUST 存在；v1 写入方 MUST 使用 `compact-jsonl/v1` 或 `storyline/v1` | 未知 `format` 值 MUST 被通用读者保留；特定格式 opener MAY 拒绝不支持的值。 @@ -185,6 +185,9 @@ kind = "branch" MUST 只通过该 store API，不得在上层另写并行 sidecar。 - Compact JSONL `import` / 成功 republish / `sync` snapshot MUST 在输出 dataset 根写入 `chronicle.manifest`。 +- Storyline `import`（`--output-format storyline`）MUST 在每次分批 commit 后更新输出根上的 + leaf `chronicle.manifest`（`format = "storyline/v1"`，`record_count` 为已提交累计条数）， + 以便 Warehouse catalog / explorer 在导入过程中观察到进展。 - 本机文件系统上的写入 MUST 原子（写临时文件再 rename）。 - 物理写入成功后，`fingerprint` MUST 匹配已发布修订，且 `[stats].record_count` MUST 等于已发布行数。 - 若 dataset 写成功但 manifest 写失败，`import_path` MUST 失败（不发布半成品契约）；对 From 2f9c5b265c1d17a863ab7d5e7f30cee3cc3e315c Mon Sep 17 00:00:00 2001 From: Reiase Date: Thu, 10 Sep 2026 20:04:32 +0800 Subject: [PATCH 7/8] feat(timestamp): enhance timestamp handling and introduce JSON sanitization Improved the handling of timestamps in the Storyline format by adding lenient parsing methods for various timestamp string formats. Introduced a new module for sanitizing JSON input to handle non-standard tokens like NaN and Infinity, ensuring compatibility with scientific and Python-generated data. Updated the StorylineTimestamp struct to support optional timestamps and refined the deserialization process. This update aims to enhance robustness and flexibility in timestamp processing across different formats. --- crates/persisting-pchronicle-cli/README.md | 13 +- .../persisting-pchronicle-cli/src/exchange.rs | 3593 ----------------- .../src/exchange/decode.rs | 917 +++++ .../src/exchange/drop.rs | 99 + .../src/exchange/export.rs | 402 ++ .../src/exchange/import.rs | 2298 +++++++++++ .../src/exchange/mod.rs | 23 + .../src/exchange/pipeline.rs | 770 ++++ .../src/exchange/progress.rs | 959 +++++ .../src/exchange/staging.rs | 153 + .../src/exchange/sync.rs | 102 + .../src/exchange/wal.rs | 368 ++ crates/persisting-pchronicle-cli/src/lib.rs | 72 +- .../persisting-pchronicle-cli/src/onboard.rs | 16 +- .../src/server/explorer.rs | 57 +- .../src/server/mod.rs | 143 +- .../persisting-pchronicle-cli/src/settings.rs | 42 + crates/persisting-pchronicle-cli/src/sync.rs | 266 +- crates/persisting-pchronicle-cli/src/tests.rs | 76 +- .../src/formats/actf/convert.rs | 117 +- .../src/formats/actf/mod.rs | 395 +- .../persisting-pchronicle/src/formats/atif.rs | 3 +- .../src/formats/common/json_sanitize.rs | 102 + .../src/formats/common/mod.rs | 3 + .../src/formats/detect.rs | 99 + .../src/formats/openai_corpus.rs | 4 +- .../src/formats/storyline.rs | 68 +- .../src/formats/timestamp.rs | 175 +- .../src/formats/unknown_fields.rs | 3 + crates/persisting-pchronicle/src/storage.rs | 44 +- .../src/store/catalog/discovery.rs | 19 +- .../src/store/compact_jsonl.rs | 205 +- .../src/store/location.rs | 54 +- crates/persisting-pchronicle/src/store/mod.rs | 21 +- .../src/store/object_store_io_gate.rs | 396 +- .../src/store/storyline/content.rs | 57 +- .../src/store/storyline/mod.rs | 130 +- .../src/store/storyline/mutation.rs | 78 +- .../src/store/storyline/rows.rs | 116 +- .../src/store/storyline/tests.rs | 2 + .../src/store/storyline/writer_control.rs | 23 +- docs/src/en/pchronicle/reference/cli.md | 45 +- docs/src/en/rfcs/0014-compact-jsonl.md | 10 +- docs/src/zh/pchronicle/reference/cli.md | 37 +- docs/src/zh/rfcs/0014-compact-jsonl.md | 5 +- 45 files changed, 8354 insertions(+), 4226 deletions(-) delete mode 100644 crates/persisting-pchronicle-cli/src/exchange.rs create mode 100644 crates/persisting-pchronicle-cli/src/exchange/decode.rs create mode 100644 crates/persisting-pchronicle-cli/src/exchange/drop.rs create mode 100644 crates/persisting-pchronicle-cli/src/exchange/export.rs create mode 100644 crates/persisting-pchronicle-cli/src/exchange/import.rs create mode 100644 crates/persisting-pchronicle-cli/src/exchange/mod.rs create mode 100644 crates/persisting-pchronicle-cli/src/exchange/pipeline.rs create mode 100644 crates/persisting-pchronicle-cli/src/exchange/progress.rs create mode 100644 crates/persisting-pchronicle-cli/src/exchange/staging.rs create mode 100644 crates/persisting-pchronicle-cli/src/exchange/sync.rs create mode 100644 crates/persisting-pchronicle-cli/src/exchange/wal.rs create mode 100644 crates/persisting-pchronicle/src/formats/common/json_sanitize.rs diff --git a/crates/persisting-pchronicle-cli/README.md b/crates/persisting-pchronicle-cli/README.md index f57d64f0..91158215 100644 --- a/crates/persisting-pchronicle-cli/README.md +++ b/crates/persisting-pchronicle-cli/README.md @@ -18,13 +18,12 @@ Current commands include `onboard`, `dataset` (pin/unpin/list/show/set/rename), `agent` sessions, Source-local `find`, create/append/replace `import`, destructive `drop`, complete-trajectory `export`, directory `sync`, `echo`, and `serve`. Import and export support ATIF, OpenAI Messages, ACTF, -Storyline JSON, and record-level Compact JSONL. `sync --from SOURCE --to -WAREHOUSE --convert OUTPUT` polls a local source directory, atomically mirrors -supported JSON files into a local Warehouse Dataset byte-for-byte, and rebuilds -a Storyline Lance Dataset at the conversion output on each coalesced batch. -With `--input-format compact-jsonl`, each batch instead replaces a compact Lance -snapshot at `OUTPUT`; `--to` remains required but is not written. Use `--once` -for a finite run. +Storyline JSON, and record-level Compact JSONL. `sync --from SOURCE [--mirror +MIRROR] [--to OUTPUT]` polls a source directory and, on each coalesced batch, +optionally rebuilds a Compact JSONL Lance Dataset at `--mirror` and/or a +Storyline Lance Dataset at `--to`. Provide at least one destination. With +`--input-format compact-jsonl`, only `--mirror` is valid. Use `--once` for a +finite run. `pchronicle serve --control 127.0.0.1:0 URI` is normally launched by pPilot or pVisor. `serve --listen` is the read-only Warehouse. Public bind addresses are diff --git a/crates/persisting-pchronicle-cli/src/exchange.rs b/crates/persisting-pchronicle-cli/src/exchange.rs deleted file mode 100644 index 826a3679..00000000 --- a/crates/persisting-pchronicle-cli/src/exchange.rs +++ /dev/null @@ -1,3593 +0,0 @@ -use super::*; - -#[derive(Serialize)] -struct DropResponse { - dataset_uri: String, - dropped: bool, -} - -pub(super) async fn run_drop( - args: DropArgs, - settings_override: Option<&Path>, - stdin_is_terminal: bool, - stdin: &mut dyn Read, - stdout: &mut dyn Write, - stderr: &mut dyn Write, -) -> Result<()> { - let dataset_uri = expand_dataset_reference(&args.dataset_uri, settings_override, false)?; - let mut location = DatasetLocation::parse(&dataset_uri)?; - if !location.exists().await? { - return Err(cli_boundary_error( - BoundaryCode::NotFound, - format!("Dataset does not exist: {}", location.as_str()), - )); - } - if location.local_path().is_some() { - location = location.into_existing()?; - } - confirm_destructive_dataset( - "drop", - location.as_str(), - args.yes, - stdin_is_terminal, - stdin, - stderr, - )?; - location.remove_all().await?; - let response = DropResponse { - dataset_uri: location.as_str().to_string(), - dropped: true, - }; - serde_json::to_writer_pretty(&mut *stdout, &response).context("encode pChronicle drop JSON")?; - writeln!(stdout).context("write pChronicle drop JSON")?; - writeln!( - stderr, - "dataset_uri={} status=dropped", - response.dataset_uri - ) - .context("write pChronicle drop metadata")?; - Ok(()) -} - -async fn prepare_import_destination( - args: &ImportArgs, - output_arg: &str, - stdin_is_terminal: bool, - stdin: &mut dyn Read, - stderr: &mut dyn Write, -) -> Result { - let parsed = DatasetLocation::parse(output_arg)?; - let exists = parsed.exists().await?; - match args.mode()? { - ImportMode::Create => { - if parsed.is_object_store() { - anyhow::ensure!(!exists, "import output already exists"); - Ok(PreparedImportDestination { - location: parsed, - replace_existing: false, - }) - } else { - Ok(PreparedImportDestination { - location: parsed.into_create_target()?, - replace_existing: false, - }) - } - } - ImportMode::Append => { - if !exists { - return Err(cli_boundary_error( - BoundaryCode::NotFound, - format!("append target Dataset does not exist: {}", parsed.as_str()), - )); - } - let location = if parsed.local_path().is_some() { - parsed.into_existing()? - } else { - parsed - }; - Ok(PreparedImportDestination { - location, - replace_existing: false, - }) - } - ImportMode::Replace => { - if !exists { - return if parsed.is_object_store() { - Ok(PreparedImportDestination { - location: parsed, - replace_existing: false, - }) - } else { - Ok(PreparedImportDestination { - location: parsed.into_create_target()?, - replace_existing: false, - }) - }; - } - let existing = parsed.into_existing()?; - ensure_import_source_outside_destination(args, &existing)?; - confirm_destructive_dataset( - "replace", - existing.as_str(), - args.yes, - stdin_is_terminal, - stdin, - stderr, - )?; - Ok(PreparedImportDestination { - location: existing, - replace_existing: true, - }) - } - } -} - -struct PreparedImportDestination { - location: DatasetLocation, - replace_existing: bool, -} - -fn ensure_import_source_outside_destination( - args: &ImportArgs, - destination: &DatasetLocation, -) -> Result<()> { - let (Some(source), Some(target)) = ( - (args.from != "-").then(|| Path::new(&args.from)), - destination.local_path(), - ) else { - return Ok(()); - }; - let source = std::fs::canonicalize(source).context("canonicalize replace import source")?; - anyhow::ensure!( - !source.starts_with(target), - "replace import source is inside the Dataset that would be replaced" - ); - Ok(()) -} - -fn confirm_destructive_dataset( - action: &str, - dataset_uri: &str, - yes: bool, - stdin_is_terminal: bool, - stdin: &mut dyn Read, - stderr: &mut dyn Write, -) -> Result<()> { - if yes { - return Ok(()); - } - if !stdin_is_terminal { - return Err(cli_boundary_error( - BoundaryCode::InvalidRequest, - format!("{action} requires confirmation; rerun with --yes"), - )); - } - write!( - stderr, - "Permanently {action} Dataset '{dataset_uri}'? [y/N] " - ) - .context("write Dataset confirmation prompt")?; - stderr - .flush() - .context("flush Dataset confirmation prompt")?; - let mut answer = Vec::new(); - let mut byte = [0u8; 1]; - while answer.len() <= 16 && stdin.read(&mut byte).context("read Dataset confirmation")? == 1 { - if byte[0] == b'\n' { - break; - } - answer.push(byte[0]); - } - let answer = std::str::from_utf8(&answer) - .context("Dataset confirmation is not UTF-8")? - .trim(); - if matches!(answer.to_ascii_lowercase().as_str(), "y" | "yes") { - return Ok(()); - } - Err(cli_boundary_error( - BoundaryCode::InvalidRequest, - format!("{action} cancelled"), - )) -} - -pub(super) async fn run_import( - mut args: ImportArgs, - settings_override: Option<&Path>, - stdin_is_terminal: bool, - stderr_is_terminal: bool, - stdin: &mut dyn Read, - stdout: &mut dyn Write, - stderr: &mut dyn Write, -) -> Result<()> { - args.stream = args.from == "-" || args.stream; - let max_input_bytes = match args.max_input_bytes { - Some(0) => { - return Err(anyhow!("--max-input-bytes must be greater than zero")); - } - Some(limit) => limit, - None => usize::MAX, - }; - anyhow::ensure!( - args.from == "-" || !args.stream, - "--stream requires --from -" - ); - if args.stream { - anyhow::ensure!( - args.format != ExchangeFormat::Auto, - "stdin import requires an explicit --input-format" - ); - } - let mode = args.mode()?; - anyhow::ensure!( - mode == ImportMode::Append || args.on_duplicate.is_none(), - "--on-duplicate is only valid with --append" - ); - anyhow::ensure!( - mode == ImportMode::Replace || !args.yes, - "--yes is only valid with --replace" - ); - anyhow::ensure!( - !(args.stream && mode == ImportMode::Replace && !args.yes), - "stdin replace import requires --yes because stdin carries the import data" - ); - if args.from != "-" { - args.from = expand_dataset_reference(&args.from, settings_override, true)?; - } - let from_location = (!args.stream) - .then(|| DatasetLocation::parse(&args.from)) - .transpose()?; - let canonical = if let Some(location) = &from_location { - let looks_like_store = location.is_object_store() - || location.local_path().is_some_and(std::path::Path::is_dir); - if looks_like_store { - probe_canonical_event_store(location.as_str()).await? - } else { - None - } - } else { - None - }; - let output_arg = match args.output.as_deref() { - Some(output) => expand_dataset_reference(output, settings_override, false)?, - None => default_import_output(&args, settings_override)?, - }; - if args.format == ExchangeFormat::CompactJsonl - || args.output_format == Some(ImportOutputFormat::CompactJsonl) - { - args.format = ExchangeFormat::CompactJsonl; - return run_compact_jsonl_import(args, &output_arg, stdout, stderr).await; - } - let requested_destination = DatasetLocation::parse(&output_arg)?; - if canonical.is_none() - && requested_destination.is_object_store() - && args.output_format != Some(ImportOutputFormat::Storyline) - { - anyhow::ensure!( - mode == ImportMode::Append && args.output_format.is_none(), - "object-store import requires --output-format storyline" - ); - } - let prepared = - prepare_import_destination(&args, &output_arg, stdin_is_terminal, stdin, stderr).await?; - let destination = prepared.location; - let replace_existing = prepared.replace_existing; - if let Some(snapshot) = canonical { - anyhow::ensure!( - mode != ImportMode::Append, - "canonical event import does not support --append" - ); - return run_canonical_event_import( - args, - snapshot, - destination, - replace_existing, - stdout, - stderr, - ) - .await; - } - let mut progress = ImportProgress::new(stderr_is_terminal); - let object_store_from = from_location - .as_ref() - .filter(|location| location.is_object_store() && !args.stream) - .cloned(); - let (directory_input, candidates) = if args.stream { - progress.set_discovered(1, 0)?; - (false, Vec::new()) - } else if object_store_from.is_some() { - // Object-store Sources are discovered inside the Storyline pipeline so - // listing overlaps read/parse/write instead of buffering the full tree. - (true, Vec::new()) - } else if from_location.is_some() { - let (directory_input, candidates) = collect_import_candidates(Path::new(&args.from))?; - let discovered_bytes = candidates.iter().try_fold(0u64, |total, candidate| { - total - .checked_add(candidate.size_hint) - .context("import discovered byte count overflow") - })?; - progress.set_discovered(candidates.len() as u64, discovered_bytes)?; - (directory_input, candidates) - } else { - (false, Vec::new()) - }; - anyhow::ensure!( - mode != ImportMode::Append || args.output_format != Some(ImportOutputFormat::Preserve), - "append import requires --output-format storyline (or omit it)" - ); - let output_format = args - .output_format - .unwrap_or(if mode == ImportMode::Append { - ImportOutputFormat::Storyline - } else { - ImportOutputFormat::Preserve - }); - let duplicate_policy = args.on_duplicate.unwrap_or(DuplicateIdPolicy::Suffix); - let (dataset_uri, imported_sources, unknown_field_warnings, skipped_warnings) = if mode - == ImportMode::Append - { - let store = StorylineLanceStore::open_uri(destination.as_str()) - .await - .context("open append target as a Storyline Lance Dataset")?; - anyhow::ensure!( - store.current_table_paths().await?.is_some(), - "append target is not a committed Storyline Dataset" - ); - let (append_generation, existing_document_ids) = store - .document_ids_snapshot() - .await? - .context("append target has no committed Storyline snapshot")?; - let existing_document_ids = existing_document_ids.into_iter().collect(); - let (imported_sources, unknown_field_warnings, skipped_warnings) = - squash_storyline_into_store( - &store, - &args, - stdin, - &mut progress, - &candidates, - object_store_from.clone(), - StorylineImportOptions { - max_input_bytes, - directory_input, - seen_document_ids: existing_document_ids, - duplicate_policy, - allow_empty: true, - append_generation: Some(append_generation), - }, - ) - .await?; - ( - destination.as_str().to_string(), - imported_sources, - unknown_field_warnings, - skipped_warnings, - ) - } else if destination.is_object_store() || output_format == ImportOutputFormat::Storyline { - // Storyline imports commit in place so progressive CURRENT + - // chronicle.manifest updates are visible to a live catalog mount. - // Remote object-store targets stage locally first: Lance index builds - // on S3 are extremely slow, so we write+index on disk then upload. - if destination.exists().await? { - if replace_existing { - destination - .remove_all_with_progress(|deleted, total, path| { - progress.note_deleted(deleted, total, path) - }) - .await - .with_context(|| { - format!("remove existing Dataset {}", destination.as_str()) - })?; - progress.finish()?; - // Delete progress reuses the paint lines but must not wipe discovery - // totals collected before replace (local candidates only). - progress.reset_import_counters(); - if object_store_from.is_none() { - let discovered_bytes = candidates.iter().try_fold(0u64, |total, candidate| { - total - .checked_add(candidate.size_hint) - .context("import discovered byte count overflow") - })?; - progress.set_discovered(candidates.len() as u64, discovered_bytes)?; - } - } else { - return Err(cli_boundary_error( - BoundaryCode::Conflict, - "import output already exists", - )); - } - } - let (imported_sources, unknown_field_warnings, skipped_warnings) = - if destination.is_object_store() { - progress.set_phase(ImportPhase::Writing, "local staging (indexes on disk)")?; - let staging = tempfile::Builder::new() - .prefix("pchronicle-storyline-stage-") - .tempdir() - .context("create local Storyline staging directory")?; - let store = StorylineLanceStore::open(staging.path()) - .await - .context("open local Storyline staging Dataset")?; - let result = squash_storyline_into_store( - &store, - &args, - stdin, - &mut progress, - &candidates, - object_store_from.clone(), - StorylineImportOptions::create(max_input_bytes, directory_input), - ) - .await?; - upload_local_storyline_dataset(staging.path(), &destination, &mut progress) - .await - .with_context(|| { - format!( - "upload staged Storyline Dataset to {}", - destination.as_str() - ) - })?; - result - } else { - let store = StorylineLanceStore::open_uri(destination.as_str()) - .await - .context("create squashed Storyline Lance Dataset")?; - squash_storyline_into_store( - &store, - &args, - stdin, - &mut progress, - &candidates, - object_store_from.clone(), - StorylineImportOptions::create(max_input_bytes, directory_input), - ) - .await? - }; - ( - destination.as_str().to_string(), - imported_sources, - unknown_field_warnings, - skipped_warnings, - ) - } else { - let output = destination - .local_path() - .context("local import output must be a filesystem path")? - .to_path_buf(); - let parent = output - .parent() - .context("import output must have a parent directory")?; - let staging = tempfile::Builder::new() - .prefix(".pchronicle-import-") - .tempdir_in(parent) - .with_context(|| format!("create import staging directory in {}", parent.display()))?; - let (imported_sources, unknown_field_warnings, skipped_warnings) = match output_format { - ImportOutputFormat::Preserve => { - let mut unknown_field_warnings = - persisting_pchronicle::model::UnknownFieldImportWarnings::default(); - let mut imported_sources = Vec::new(); - let mut skipped_warnings = Vec::new(); - if args.stream { - progress.set_phase(ImportPhase::Reading, "stdin")?; - let input = read_bounded(stdin, max_input_bytes, "stdin")?; - progress.set_phase(ImportPhase::Parsing, "stdin")?; - if let Some(source) = stage_preserved_import_source( - args.format, - None, - None, - None, - &input, - staging.path(), - &mut unknown_field_warnings, - &mut skipped_warnings, - )? { - progress.set_phase(ImportPhase::Writing, &source.source_path)?; - progress.note_imported(source.input_bytes as u64)?; - imported_sources.push(source); - } else { - progress.note_imported(input.len() as u64)?; - } - } else { - for candidate in &candidates { - let name = candidate.relative_path.to_string_lossy(); - let label = format!("import source {name}"); - progress.set_phase(ImportPhase::Reading, &name)?; - let input = - load_import_candidate_bytes(candidate, max_input_bytes, &label).await?; - progress.set_phase(ImportPhase::Parsing, &name)?; - if let Some(source) = stage_preserved_import_source( - args.format, - Some(&candidate.path), - Some(&candidate.relative_path), - candidate.output_relative_path.as_deref(), - &input, - staging.path(), - &mut unknown_field_warnings, - &mut skipped_warnings, - )? { - progress.set_phase(ImportPhase::Writing, &source.source_path)?; - progress.note_imported(source.input_bytes as u64)?; - imported_sources.push(source); - } else { - progress.note_imported(input.len() as u64)?; - } - } - } - (imported_sources, unknown_field_warnings, skipped_warnings) - } - ImportOutputFormat::Storyline => { - unreachable!("storyline import commits in place above") - } - ImportOutputFormat::CompactJsonl => unreachable!("compact import handled above"), - }; - if imported_sources.is_empty() { - return Err(empty_auto_directory_import_error(directory_input)); - } - - std::fs::File::open(staging.path()) - .and_then(|directory| directory.sync_all()) - .context("sync import staging directory")?; - - let staging_path = staging.keep(); - let mut cleanup = StagingPathGuard::new(staging_path.clone()); - publish_staged_dataset(&staging_path, &output, replace_existing, Some(&mut progress)).await?; - cleanup.disarm(); - ( - output.to_string_lossy().into_owned(), - imported_sources, - unknown_field_warnings, - skipped_warnings, - ) - }; - if imported_sources.is_empty() { - return Err(empty_auto_directory_import_error(directory_input)); - } - let trajectories = imported_sources.iter().try_fold(0usize, |total, source| { - total - .checked_add(source.trajectories) - .context("import trajectory count overflow") - })?; - let input_bytes = imported_sources.iter().try_fold(0usize, |total, source| { - total - .checked_add(source.input_bytes) - .context("import input byte count overflow") - })?; - - let single_source = (!directory_input).then(|| { - imported_sources - .first() - .expect("stdin and regular-file imports have one Source") - }); - let response = ImportResponse { - dataset_uri, - source_path: single_source.map(|source| source.source_path.clone()), - format: single_source.map(|source| source.format.as_str().to_owned()), - output_format: output_format.response_name().into(), - sources: imported_sources.len(), - trajectories, - fact_rows: None, - input_bytes: Some(input_bytes), - }; - serde_json::to_writer_pretty(&mut *stdout, &response) - .context("encode pChronicle import JSON")?; - writeln!(stdout).context("write pChronicle import JSON")?; - progress.finish()?; - if let (Some(source_path), Some(format)) = (&response.source_path, &response.format) { - progress.notice(&format!( - "dataset_uri={} source={} format={} output_format={} trajectories={} input_bytes={}", - response.dataset_uri, - source_path, - format, - response.output_format, - response.trajectories, - response - .input_bytes - .expect("JSON imports always report input bytes"), - ))?; - } else { - progress.notice(&format!( - "dataset_uri={} sources={} output_format={} trajectories={} input_bytes={}", - response.dataset_uri, - response.sources, - response.output_format, - response.trajectories, - response - .input_bytes - .expect("JSON imports always report input bytes"), - ))?; - } - for line in skipped_warnings { - progress.notice(&line)?; - } - for line in unknown_field_warnings.warning_lines() { - progress.notice(&line)?; - } - progress.flush_log(stderr)?; - Ok(()) -} - -async fn run_compact_jsonl_import( - args: ImportArgs, - output_arg: &str, - stdout: &mut dyn Write, - stderr: &mut dyn Write, -) -> Result<()> { - anyhow::ensure!( - args.mode()? != ImportMode::Append, - "compact JSONL append is not supported; use sync or replace" - ); - anyhow::ensure!( - args.from != "-", - "compact JSONL import does not support stdin" - ); - let input = Path::new(&args.from); - let output = Path::new(output_arg); - anyhow::ensure!( - !output_arg.starts_with("s3://") && !output_arg.starts_with("oss://"), - "compact JSONL currently requires local paths" - ); - if args.mode()? == ImportMode::Create { - anyhow::ensure!(!output.exists(), "import output already exists"); - } - let columns = args - .columns - .iter() - .map(|item| { - let (name, path) = item - .split_once('=') - .context("--column must be NAME=JSON_PATH")?; - persisting_pchronicle::storage::CompactJsonlColumn::new(name.trim(), path.trim()) - }) - .collect::>>()?; - let options = persisting_pchronicle::storage::CompactJsonlOptions { - columns, - offload_threshold: 4 * 1024 * 1024, - }; - let parent = output - .parent() - .filter(|path| !path.as_os_str().is_empty()) - .unwrap_or_else(|| Path::new(".")); - let staging = tempfile::Builder::new() - .prefix(".pchronicle-compact-jsonl-") - .tempdir_in(parent)?; - let rows = persisting_pchronicle::storage::CompactJsonlStore::import_path( - input, - staging.path(), - &options, - ) - .await?; - std::fs::File::open(staging.path())?.sync_all()?; - let staging_path = staging.keep(); - let mut cleanup = StagingPathGuard::new(staging_path.clone()); - publish_staged_dataset(&staging_path, output, output.exists(), None).await?; - cleanup.disarm(); - serde_json::to_writer_pretty( - &mut *stdout, - &serde_json::json!({"dataset_uri": output_arg, "output_format": "compact-jsonl", "rows": rows}), - )?; - writeln!(stdout)?; - writeln!( - stderr, - "dataset_uri={} output_format=compact-jsonl rows={rows}", - output_arg - )?; - Ok(()) -} - -/// Run one full snapshot import for the resident sync worker. -/// -/// The existing import path already stages local outputs atomically, mirrors -/// deletions, and rebuilds a Storyline Lance destination from the same source -/// directory. Keeping the orchestration here avoids a second decoder or -/// Dataset publication protocol in the sync command. -pub(crate) async fn sync_snapshot( - source: &str, - warehouse: &str, - storyline: &str, - input_format: ExchangeFormat, - columns: &[String], -) -> Result<()> { - if input_format == ExchangeFormat::CompactJsonl { - let mut stdout = std::io::sink(); - let mut stderr = std::io::sink(); - return run_compact_jsonl_import( - ImportArgs { - from: source.to_owned(), - output: Some(storyline.to_owned()), - format: ExchangeFormat::CompactJsonl, - output_format: Some(ImportOutputFormat::CompactJsonl), - replace: true, - append: false, - mode: None, - on_duplicate: None, - yes: true, - stream: false, - max_input_bytes: Some(256 * 1024 * 1024), - commit_every: None, - columns: columns.to_vec(), - }, - storyline, - &mut stdout, - &mut stderr, - ) - .await; - } - // ponytail: rebuild one atomic snapshot per coalesced batch; add affected-document mutation - // when profiling shows full-directory rebuilds are the bottleneck. - let mut stdout = std::io::sink(); - let mut stderr = std::io::sink(); - let mut stdin = std::io::empty(); - run_import( - ImportArgs { - from: source.to_owned(), - output: Some(warehouse.to_owned()), - format: input_format, - output_format: Some(ImportOutputFormat::Preserve), - replace: true, - append: false, - mode: None, - on_duplicate: None, - yes: true, - stream: false, - max_input_bytes: Some(256 * 1024 * 1024), - commit_every: None, - columns: Vec::new(), - }, - None, - false, - false, - &mut stdin, - &mut stdout, - &mut stderr, - ) - .await - .context("sync source into Warehouse")?; - run_import( - ImportArgs { - from: source.to_owned(), - output: Some(storyline.to_owned()), - format: input_format, - output_format: Some(ImportOutputFormat::Storyline), - replace: true, - append: false, - mode: None, - on_duplicate: None, - yes: true, - stream: false, - max_input_bytes: Some(256 * 1024 * 1024), - commit_every: None, - columns: Vec::new(), - }, - None, - false, - false, - &mut stdin, - &mut stdout, - &mut stderr, - ) - .await - .context("sync source into Storyline Lance")?; - Ok(()) -} - -struct StorylineImportOptions { - max_input_bytes: usize, - directory_input: bool, - seen_document_ids: HashSet, - duplicate_policy: DuplicateIdPolicy, - allow_empty: bool, - append_generation: Option, -} - -impl StorylineImportOptions { - fn create(max_input_bytes: usize, directory_input: bool) -> Self { - Self { - max_input_bytes, - directory_input, - seen_document_ids: HashSet::new(), - duplicate_policy: DuplicateIdPolicy::Suffix, - allow_empty: false, - append_generation: None, - } - } -} - -/// How many Sources the reader may prefetch ahead of parse/write. -/// Bounded so large object-store imports do not buffer unbounded memory. -const IMPORT_READ_AHEAD: usize = 3; -/// Pipeline channel capacity for object-store discover/read events. Listing -/// emits Discovered first; this buffer only absorbs Loaded messages while a -/// commit is in flight. -const IMPORT_PIPELINE_CHANNEL: usize = 16; - -struct PipelineLoadedSource { - candidate: ImportFileCandidate, - bytes: Vec, -} - -enum PipelineMsg { - Scanning(String), - Discovered { path: String, bytes: u64 }, - Loaded(PipelineLoadedSource), -} - -fn spawn_candidates_load_producer( - candidates: Vec, - max_input_bytes: usize, - reading_ahead: Arc>, -) -> ( - tokio::sync::mpsc::Receiver>, - tokio::task::JoinHandle<()>, -) { - let (tx, rx) = tokio::sync::mpsc::channel::>(IMPORT_READ_AHEAD); - let producer = tokio::spawn(async move { - for candidate in candidates { - let name = candidate.relative_path.to_string_lossy().into_owned(); - if let Ok(mut guard) = reading_ahead.lock() { - *guard = name.clone(); - } - let label = format!("import source {name}"); - let loaded = match load_import_candidate_bytes(&candidate, max_input_bytes, &label).await - { - Ok(bytes) => Ok(PipelineMsg::Loaded(PipelineLoadedSource { candidate, bytes })), - Err(error) => Err(error), - }; - if tx.send(loaded).await.is_err() { - return; - } - } - if let Ok(mut guard) = reading_ahead.lock() { - guard.clear(); - } - }); - (rx, producer) -} - -fn spawn_object_store_discover_load_producer( - location: DatasetLocation, - max_input_bytes: usize, - reading_ahead: Arc>, -) -> ( - tokio::sync::mpsc::Receiver>, - tokio::task::JoinHandle<()>, -) { - let (tx, rx) = tokio::sync::mpsc::channel::>(IMPORT_PIPELINE_CHANNEL); - let producer = tokio::spawn(async move { - let remote_root = location.as_str().to_owned(); - // List completely before any Load so discovery totals keep moving even - // when a later commit/index stalls the consumer. - let pending_files = Arc::new(std::sync::Mutex::new(Vec::<(String, u64)>::new())); - let list_result = location - .for_each_importable_json_object_event( - persisting_pchronicle::storage::DEFAULT_MAX_LOCAL_QUERY_FILES, - |event| { - let tx = tx.clone(); - let pending_files = Arc::clone(&pending_files); - async move { - match event { - persisting_pchronicle::storage::ImportableObjectEvent::Scanning { - prefix, - } => { - let _ = tx.send(Ok(PipelineMsg::Scanning(prefix))).await; - Ok(()) - } - persisting_pchronicle::storage::ImportableObjectEvent::File { - key, - size, - .. - } => { - if tx - .send(Ok(PipelineMsg::Discovered { - path: key.clone(), - bytes: size, - })) - .await - .is_err() - { - return Ok(()); - } - if let Ok(mut guard) = pending_files.lock() { - guard.push((key, size)); - } - Ok(()) - } - } - } - }, - ) - .await; - if let Err(error) = list_result { - let _ = tx.send(Err(error)).await; - if let Ok(mut guard) = reading_ahead.lock() { - guard.clear(); - } - return; - } - let files = match pending_files.lock() { - Ok(mut guard) => std::mem::take(&mut *guard), - Err(_) => Vec::new(), - }; - for (key, size) in files { - if let Ok(mut guard) = reading_ahead.lock() { - *guard = key.clone(); - } - let relative_path = PathBuf::from(&key); - let candidate = ImportFileCandidate { - path: relative_path.clone(), - output_relative_path: Some(relative_path.clone()), - relative_path, - content: None, - remote_root: Some(remote_root.clone()), - size_hint: size, - }; - let label = format!("import source {key}"); - match load_import_candidate_bytes(&candidate, max_input_bytes, &label).await { - Ok(bytes) => { - if tx - .send(Ok(PipelineMsg::Loaded(PipelineLoadedSource { - candidate, - bytes, - }))) - .await - .is_err() - { - break; - } - } - Err(error) => { - let _ = tx.send(Err(error)).await; - break; - } - } - } - if let Ok(mut guard) = reading_ahead.lock() { - guard.clear(); - } - }); - (rx, producer) -} - -async fn squash_storyline_into_store( - store: &StorylineLanceStore, - args: &ImportArgs, - stdin: &mut dyn Read, - progress: &mut ImportProgress, - candidates: &[ImportFileCandidate], - object_store_from: Option, - options: StorylineImportOptions, -) -> Result<( - Vec, - persisting_pchronicle::model::UnknownFieldImportWarnings, - Vec, -)> { - let StorylineImportOptions { - max_input_bytes, - directory_input, - seen_document_ids, - duplicate_policy, - allow_empty, - append_generation, - } = options; - if args.stream { - return squash_storyline_stdin_into_store( - store, - args.format, - max_input_bytes, - stdin, - progress, - seen_document_ids, - duplicate_policy, - allow_empty, - directory_input, - append_generation, - commit_batch_schedule(args), - ) - .await; - } - let source = match object_store_from { - Some(location) => ObjectStoreImportSource::Location(location), - None => ObjectStoreImportSource::Candidates(candidates.to_vec()), - }; - squash_storyline_files_pipeline( - store, - args.format, - max_input_bytes, - progress, - source, - seen_document_ids, - duplicate_policy, - allow_empty, - directory_input, - append_generation, - commit_batch_schedule(args), - ) - .await -} - -const DEFAULT_COMMIT_BATCH_START: usize = 64; -const DEFAULT_COMMIT_BATCH_MAX: usize = 4096; - -#[derive(Debug, Clone)] -struct CommitBatchSchedule { - next: usize, - max: usize, - fixed: bool, -} - -impl CommitBatchSchedule { - fn adaptive() -> Self { - Self { - next: DEFAULT_COMMIT_BATCH_START, - max: DEFAULT_COMMIT_BATCH_MAX, - fixed: false, - } - } - - fn fixed(n: usize) -> Self { - let n = n.max(1); - Self { - next: n, - max: n, - fixed: true, - } - } - - fn current(&self) -> usize { - self.next - } - - fn after_commit(&mut self) { - if self.fixed { - return; - } - self.next = self.next.saturating_mul(2).min(self.max); - } -} - -fn commit_batch_schedule(args: &ImportArgs) -> CommitBatchSchedule { - match args.commit_every { - Some(n) => CommitBatchSchedule::fixed(n), - None => CommitBatchSchedule::adaptive(), - } -} - -enum ObjectStoreImportSource { - Candidates(Vec), - Location(DatasetLocation), -} - -#[allow(clippy::too_many_arguments)] -async fn squash_storyline_files_pipeline( - store: &StorylineLanceStore, - requested_format: ExchangeFormat, - max_input_bytes: usize, - progress: &mut ImportProgress, - source: ObjectStoreImportSource, - mut seen_document_ids: HashSet, - duplicate_policy: DuplicateIdPolicy, - allow_empty: bool, - directory_input: bool, - mut append_generation: Option, - mut commit_schedule: CommitBatchSchedule, -) -> Result<( - Vec, - persisting_pchronicle::model::UnknownFieldImportWarnings, - Vec, -)> { - let reading_ahead = Arc::new(std::sync::Mutex::new(String::new())); - let (mut rx, producer) = match source { - ObjectStoreImportSource::Candidates(candidates) => { - spawn_candidates_load_producer(candidates, max_input_bytes, Arc::clone(&reading_ahead)) - } - ObjectStoreImportSource::Location(location) => { - progress.set_phase(ImportPhase::Discovering, location.as_str())?; - spawn_object_store_discover_load_producer( - location, - max_input_bytes, - Arc::clone(&reading_ahead), - ) - } - }; - - let mut unknown_field_warnings = - persisting_pchronicle::model::UnknownFieldImportWarnings::default(); - let mut skipped_warnings = Vec::new(); - let mut imported_sources: Vec = Vec::new(); - let mut batch = Vec::with_capacity(commit_schedule.current()); - let mut committed_storylines = 0u64; - let mut skipped_commit_storylines = 0usize; - let mut saw_any = false; - let mut current_storylines = Vec::new().into_iter(); - let mut producer_done = false; - let mut discovered_any = false; - - loop { - if let Some(mut storyline) = current_storylines.next() { - saw_any = true; - if let Some(warning) = - apply_duplicate_document_policy(&mut storyline, &mut seen_document_ids, duplicate_policy) - { - if warning.contains("skipped") { - skipped_warnings.push(warning); - continue; - } - skipped_warnings.push(warning); - } - let metadata = imported_sources - .last_mut() - .expect("decoded Storyline has source metadata"); - metadata.trajectories = metadata - .trajectories - .checked_add(1) - .context("import trajectory count overflow")?; - batch.push(storyline); - if batch.len() >= commit_schedule.current() { - match commit_or_skip_storyline_import_batch( - store, - progress, - std::mem::take(&mut batch), - &mut append_generation, - committed_storylines, - &mut commit_schedule, - ) - .await? - { - StorylineBatchCommit::Committed(total) => { - committed_storylines = total; - } - StorylineBatchCommit::Skipped { batch_len, warning } => { - skipped_commit_storylines = skipped_commit_storylines - .saturating_add(batch_len as usize); - skipped_warnings.push(warning); - retract_imported_trajectories( - &mut imported_sources, - batch_len as usize, - ); - } - } - batch.reserve(commit_schedule.current()); - } - continue; - } - - if producer_done { - break; - } - - // Surface producer read activity while waiting for the next Source. - let msg = loop { - if let Ok(guard) = reading_ahead.lock() { - progress.set_reading_ahead(guard.as_str())?; - } - tokio::select! { - item = rx.recv() => break item, - _ = tokio::time::sleep(std::time::Duration::from_millis(100)) => {} - } - }; - match msg { - Some(Ok(PipelineMsg::Scanning(prefix))) => { - progress.note_scanning(&prefix)?; - } - Some(Ok(PipelineMsg::Discovered { path, bytes })) => { - discovered_any = true; - progress.note_discovered(&path, bytes)?; - } - Some(Ok(PipelineMsg::Loaded(loaded))) => { - let name = loaded.candidate.relative_path.to_string_lossy().into_owned(); - if let Ok(guard) = reading_ahead.lock() { - progress.set_reading_ahead(guard.as_str())?; - } - progress.set_phase(ImportPhase::Parsing, &name)?; - match decode_import_source( - requested_format, - ImportOutputFormat::Storyline, - Some(&loaded.candidate.path), - Some(&loaded.candidate.relative_path), - loaded.candidate.output_relative_path.as_deref(), - &loaded.bytes, - &mut unknown_field_warnings, - )? { - DecodeImportOutcome::Imported(decoded) => { - progress.set_phase( - ImportPhase::Writing, - &decoded.diagnostic_path.to_string_lossy(), - )?; - progress.note_imported(decoded.metadata.input_bytes as u64)?; - let mut metadata = decoded.metadata; - metadata.trajectories = 0; - imported_sources.push(metadata); - current_storylines = decoded.storylines.into_iter(); - } - DecodeImportOutcome::Skipped { path, reason } => { - progress.note_imported(0)?; - skipped_warnings.push(skipped_import_warning(&path, &reason)); - } - } - } - Some(Err(error)) => { - producer.abort(); - return Err(error); - } - None => { - producer_done = true; - progress.clear_reading_ahead()?; - } - } - } - - if !batch.is_empty() { - match commit_or_skip_storyline_import_batch( - store, - progress, - std::mem::take(&mut batch), - &mut append_generation, - committed_storylines, - &mut commit_schedule, - ) - .await? - { - StorylineBatchCommit::Committed(total) => { - committed_storylines = total; - } - StorylineBatchCommit::Skipped { batch_len, warning } => { - skipped_commit_storylines = - skipped_commit_storylines.saturating_add(batch_len as usize); - skipped_warnings.push(warning); - retract_imported_trajectories(&mut imported_sources, batch_len as usize); - } - } - } - progress.clear_reading_ahead()?; - - match producer.await { - Ok(()) => {} - Err(error) if error.is_cancelled() => {} - Err(error) => return Err(anyhow!("import reader task failed: {error}")), - } - - if imported_sources.is_empty() { - if allow_empty && !saw_any { - return Ok((imported_sources, unknown_field_warnings, skipped_warnings)); - } - if !discovered_any { - return Err(cli_boundary_error( - BoundaryCode::InvalidRequest, - "import object prefix contains no .json, .jsonl, or .ndjson files", - )); - } - return Err(empty_auto_directory_import_error(directory_input)); - } - // Drop Sources that lost every trajectory to skipped commits so empty - // placeholders do not inflate the import summary. - if skipped_commit_storylines > 0 { - imported_sources.retain(|source| source.trajectories > 0); - } - if imported_sources.is_empty() { - return Err(anyhow!( - "storyline import committed no trajectories after skipping failed batches" - )); - } - anyhow::ensure!( - store.current_table_paths().await?.is_some(), - "squashed Storyline Lance Dataset has no committed snapshot" - ); - let imported_trajectories = imported_sources.iter().try_fold(0usize, |total, source| { - total - .checked_add(source.trajectories) - .context("import trajectory count overflow") - })?; - anyhow::ensure!( - committed_storylines as usize == imported_trajectories, - "squashed Storyline import report does not match decoded trajectory count" - ); - finalize_storyline_import_indexes(store, progress).await?; - Ok((imported_sources, unknown_field_warnings, skipped_warnings)) -} - -#[allow(clippy::too_many_arguments)] -async fn squash_storyline_stdin_into_store( - store: &StorylineLanceStore, - requested_format: ExchangeFormat, - max_input_bytes: usize, - stdin: &mut dyn Read, - progress: &mut ImportProgress, - seen_document_ids: HashSet, - duplicate_policy: DuplicateIdPolicy, - allow_empty: bool, - directory_input: bool, - append_generation: Option, - commit_schedule: CommitBatchSchedule, -) -> Result<( - Vec, - persisting_pchronicle::model::UnknownFieldImportWarnings, - Vec, -)> { - let import = StorylineImportIterator::stdin( - requested_format, - max_input_bytes, - stdin, - progress, - seen_document_ids, - duplicate_policy, - ); - drain_storyline_import_batches( - store, - import, - append_generation, - commit_schedule, - allow_empty, - directory_input, - ) - .await -} - -fn apply_duplicate_document_policy( - storyline: &mut StorylineDocument, - seen_document_ids: &mut HashSet, - duplicate_policy: DuplicateIdPolicy, -) -> Option { - let original = storyline.document_id().to_string(); - match duplicate_policy { - DuplicateIdPolicy::Suffix => { - uniquify_storyline_document_id(storyline, seen_document_ids).map( - |(original, renamed)| { - format!("warning: duplicate document_id '{original}' renamed to '{renamed}'") - }, - ) - } - DuplicateIdPolicy::Skip => { - if !seen_document_ids.insert(original.clone()) { - Some(format!( - "warning: duplicate document_id '{original}' skipped" - )) - } else { - None - } - } - } -} - -async fn drain_storyline_import_batches( - store: &StorylineLanceStore, - mut import: StorylineImportIterator<'_>, - mut append_generation: Option, - mut commit_schedule: CommitBatchSchedule, - allow_empty: bool, - directory_input: bool, -) -> Result<( - Vec, - persisting_pchronicle::model::UnknownFieldImportWarnings, - Vec, -)> { - let mut batch = Vec::with_capacity(commit_schedule.current()); - let mut committed_storylines = 0u64; - let mut skipped_commit_storylines = 0usize; - let mut commit_skip_warnings = Vec::new(); - let mut saw_any = false; - - loop { - match import.next_document().await { - Some(item) => { - saw_any = true; - batch.push(item?); - if batch.len() < commit_schedule.current() { - continue; - } - match commit_or_skip_storyline_import_batch( - store, - import.progress, - std::mem::take(&mut batch), - &mut append_generation, - committed_storylines, - &mut commit_schedule, - ) - .await? - { - StorylineBatchCommit::Committed(total) => { - committed_storylines = total; - } - StorylineBatchCommit::Skipped { batch_len, warning } => { - skipped_commit_storylines = skipped_commit_storylines - .saturating_add(batch_len as usize); - commit_skip_warnings.push(warning); - } - } - batch.reserve(commit_schedule.current()); - } - None if batch.is_empty() => break, - None => { - match commit_or_skip_storyline_import_batch( - store, - import.progress, - std::mem::take(&mut batch), - &mut append_generation, - committed_storylines, - &mut commit_schedule, - ) - .await? - { - StorylineBatchCommit::Committed(total) => { - committed_storylines = total; - } - StorylineBatchCommit::Skipped { batch_len, warning } => { - skipped_commit_storylines = skipped_commit_storylines - .saturating_add(batch_len as usize); - commit_skip_warnings.push(warning); - } - } - break; - } - } - } - - let (mut imported_sources, unknown_field_warnings, mut skipped_warnings) = - import.into_result_parts(); - skipped_warnings.extend(commit_skip_warnings); - retract_imported_trajectories(&mut imported_sources, skipped_commit_storylines); - if skipped_commit_storylines > 0 { - imported_sources.retain(|source| source.trajectories > 0); - } - if imported_sources.is_empty() { - if allow_empty && !saw_any { - return Ok((imported_sources, unknown_field_warnings, skipped_warnings)); - } - if skipped_commit_storylines > 0 { - return Err(anyhow!( - "storyline import committed no trajectories after skipping failed batches" - )); - } - return Err(empty_auto_directory_import_error(directory_input)); - } - anyhow::ensure!( - store.current_table_paths().await?.is_some(), - "squashed Storyline Lance Dataset has no committed snapshot" - ); - let imported_trajectories = imported_sources.iter().try_fold(0usize, |total, source| { - total - .checked_add(source.trajectories) - .context("import trajectory count overflow") - })?; - anyhow::ensure!( - committed_storylines as usize == imported_trajectories, - "squashed Storyline import report does not match decoded trajectory count" - ); - finalize_storyline_import_indexes(store, import.progress).await?; - Ok((imported_sources, unknown_field_warnings, skipped_warnings)) -} - -enum StorylineBatchCommit { - Committed(u64), - Skipped { batch_len: u64, warning: String }, -} - -fn is_skippable_storyline_commit_error(error: &anyhow::Error) -> bool { - let text = format!("{error:#}").to_ascii_lowercase(); - text.contains("timeout") - || text.contains("timed out") - || text.contains("error sending request") - || text.contains("conditionnotmatch") - || text.contains("preconditionfailed") - || text.contains("precondition failed") - || text.contains("throttle") - || text.contains("slow down") - || text.contains("503") - || text.contains("429") - || text.contains("connection reset") - || text.contains("broken pipe") - || text.contains("lanceerror(io)") - || text.contains("generic s3 error") - || text.contains("client error (connect)") -} - -fn retract_imported_trajectories(sources: &mut [ImportedSource], mut count: usize) { - for source in sources.iter_mut().rev() { - if count == 0 { - break; - } - let take = source.trajectories.min(count); - source.trajectories -= take; - count -= take; - } -} - -async fn refresh_append_generation_after_skip( - store: &StorylineLanceStore, - append_generation: &mut Option, -) { - match store.current_table_paths().await { - Ok(Some(paths)) => { - *append_generation = Some(paths.generation); - } - Ok(None) => {} - Err(error) => { - tracing::warn!( - root = %store.root_uri(), - error = %error, - "failed to refresh Storyline generation after skipped commit batch" - ); - } - } -} - -async fn commit_or_skip_storyline_import_batch( - store: &StorylineLanceStore, - progress: &mut ImportProgress, - batch: Vec, - append_generation: &mut Option, - committed_storylines: u64, - commit_schedule: &mut CommitBatchSchedule, -) -> Result { - let batch_len = batch.len() as u64; - let sample_ids = batch - .iter() - .take(8) - .map(|storyline| storyline.document_id().to_string()) - .collect::>(); - match commit_storyline_import_batch( - store, - progress, - batch, - append_generation, - committed_storylines, - ) - .await - { - Ok(total) => { - commit_schedule.after_commit(); - Ok(StorylineBatchCommit::Committed(total)) - } - Err(error) if is_skippable_storyline_commit_error(&error) => { - tracing::warn!( - committed_before = committed_storylines, - batch_len, - root = %store.root_uri(), - sample_document_ids = ?sample_ids, - error = %format!("{error:#}"), - "skipping storyline commit batch after transient storage failure; continuing import" - ); - refresh_append_generation_after_skip(store, append_generation).await; - if !commit_schedule.fixed { - commit_schedule.next = DEFAULT_COMMIT_BATCH_START; - } - let warning = format!( - "warning: skipped storyline commit batch of {batch_len} trajectories (committed_before={committed_storylines}, sample_document_ids={sample_ids:?}): {error:#}" - ); - let _ = progress.notice(&warning); - Ok(StorylineBatchCommit::Skipped { batch_len, warning }) - } - Err(error) => Err(error), - } -} - -async fn finalize_storyline_import_indexes( - store: &StorylineLanceStore, - progress: &mut ImportProgress, -) -> Result<()> { - progress.set_phase(ImportPhase::Writing, "optimize indices (final)")?; - let _index_progress = progress.attach_index_progress(); - store - .maintain(&persisting_pchronicle::storage::LanceMaintenanceOptions { - compact: false, - optimize_indices: true, - vacuum_older_than: None, - ..Default::default() - }) - .await - .context("finalize Storyline indexes after progressive import")?; - progress.set_phase(ImportPhase::Writing, "optimize indices done")?; - Ok(()) -} - -fn collect_local_relative_files(root: &Path) -> Result> { - fn walk(root: &Path, dir: &Path, out: &mut Vec) -> Result<()> { - for entry in std::fs::read_dir(dir) - .with_context(|| format!("read staging directory {}", dir.display()))? - { - let entry = entry?; - let path = entry.path(); - if path.is_dir() { - walk(root, &path, out)?; - continue; - } - let relative = path - .strip_prefix(root) - .with_context(|| format!("strip staging root from {}", path.display()))? - .to_string_lossy() - .replace('\\', "/"); - if !relative.is_empty() { - out.push(relative); - } - } - Ok(()) - } - let mut files = Vec::new(); - walk(root, root, &mut files)?; - files.sort(); - Ok(files) -} - -fn is_deferred_storyline_publish_key(relative: &str) -> bool { - matches!( - relative, - "CURRENT" | "chronicle.manifest" | ".storyline-write.lock" - ) || relative.ends_with("/CURRENT") - || relative.ends_with("/chronicle.manifest") -} - -async fn upload_local_storyline_dataset( - local_root: &Path, - destination: &DatasetLocation, - progress: &mut ImportProgress, -) -> Result<()> { - let files = collect_local_relative_files(local_root)?; - anyhow::ensure!( - files.iter().any(|path| path == "CURRENT"), - "staged Storyline Dataset is missing CURRENT" - ); - let (deferred, eager): (Vec<_>, Vec<_>) = files - .into_iter() - .partition(|path| is_deferred_storyline_publish_key(path)); - let total = eager.len().saturating_add(deferred.len()) as u64; - let mut uploaded = 0u64; - for relative in eager.into_iter().chain(deferred) { - if relative == ".storyline-write.lock" { - continue; - } - uploaded = uploaded.saturating_add(1); - progress.set_phase( - ImportPhase::Writing, - &format!( - "upload {uploaded}/{total} {}", - truncate_middle(&relative, 56) - ), - )?; - let bytes = tokio::fs::read(local_root.join(&relative)) - .await - .with_context(|| format!("read staged file {relative}"))?; - destination - .write_relative_bytes(&relative, &bytes) - .await - .with_context(|| format!("upload staged file {relative}"))?; - } - progress.set_phase(ImportPhase::Writing, "upload complete")?; - Ok(()) -} - -async fn commit_storyline_import_batch( - store: &StorylineLanceStore, - progress: &mut ImportProgress, - batch: Vec, - append_generation: &mut Option, - committed_storylines: u64, -) -> Result { - anyhow::ensure!(!batch.is_empty(), "storyline import commit batch is empty"); - let batch_len = batch.len() as u64; - progress.set_phase( - ImportPhase::Writing, - &format!("commit {batch_len} trajectories"), - )?; - let report = match append_generation.as_deref() { - Some(generation) => { - tracing::info!( - committed_before = committed_storylines, - batch_len, - expected_generation = generation, - root = %store.root_uri(), - "storyline progressive append commit starting" - ); - store - .append_storyline_stream_with_options( - batch.into_iter().map(Ok), - generation, - persisting_pchronicle::storage::StorylineStreamOptions::defer_index_optimize(), - ) - .await - .with_context(|| { - format!( - "storyline progressive append commit failed (committed_before={committed_storylines}, batch={batch_len}, expected_generation={generation}, root={})", - store.root_uri() - ) - })? - } - None => { - tracing::info!( - batch_len, - root = %store.root_uri(), - "storyline progressive replace commit starting" - ); - store - .replace_storyline_stream_with_options( - batch.into_iter().map(Ok), - persisting_pchronicle::storage::StorylineStreamOptions::defer_index_optimize(), - ) - .await - .with_context(|| { - format!( - "storyline progressive replace commit failed (batch={batch_len}, root={})", - store.root_uri() - ) - })? - } - }; - anyhow::ensure!( - report.storylines as u64 == batch_len, - "storyline import batch report does not match batch size" - ); - let paths = store - .current_table_paths() - .await? - .context("storyline import batch produced no committed snapshot")?; - let total = committed_storylines - .checked_add(batch_len) - .context("import trajectory count overflow")?; - persisting_pchronicle::storage::write_storyline_manifest_at_uri( - store.root_uri(), - &paths.generation, - total, - 0, - ) - .await - .context("write progressive chronicle.manifest after storyline commit")?; - *append_generation = Some(paths.generation.clone()); - progress.note_committed(total)?; - Ok(total) -} - -async fn run_canonical_event_import( - args: ImportArgs, - _snapshot: EventFactSnapshot, - destination: DatasetLocation, - replace_existing: bool, - stdout: &mut dyn Write, - stderr: &mut dyn Write, -) -> Result<()> { - anyhow::ensure!( - args.format == ExchangeFormat::Auto, - "canonical event import does not accept a JSON exchange --format" - ); - anyhow::ensure!( - args.output_format != Some(ImportOutputFormat::Preserve), - "canonical event import cannot preserve an existing canonical event Store" - ); - if destination.exists().await? && !replace_existing { - return Err(cli_boundary_error( - BoundaryCode::Conflict, - "import output already exists", - )); - } - let output_uri = destination.as_str().to_string(); - - let (report, staged_path) = if replace_existing { - let output = destination - .local_path() - .context("replace import output must be a local Dataset path")?; - let parent = output - .parent() - .context("replace import output must have a parent directory")?; - let staging = tempfile::Builder::new() - .prefix(".pchronicle-import-") - .tempdir_in(parent) - .with_context(|| format!("create import staging directory in {}", parent.display()))?; - let staging_uri = staging.path().to_string_lossy().into_owned(); - let report = - match build_storyline_projection(&args.from, &staging_uri, "events.lance").await? { - StorylineProjectionBuildOutcome::Built(report) => report, - StorylineProjectionBuildOutcome::OutputNotEmpty => { - return Err(cli_boundary_error( - BoundaryCode::Conflict, - "import staging Dataset already exists", - )); - } - }; - std::fs::File::open(staging.path()) - .and_then(|directory| directory.sync_all()) - .context("sync import staging directory")?; - (report, Some((staging.keep(), output.to_path_buf()))) - } else { - let report = - match build_storyline_projection(&args.from, &output_uri, "events.lance").await? { - StorylineProjectionBuildOutcome::Built(report) => report, - StorylineProjectionBuildOutcome::OutputNotEmpty => { - return Err(cli_boundary_error( - BoundaryCode::Conflict, - "import output already exists", - )); - } - }; - (report, None) - }; - if let Some((staging_path, output)) = staged_path { - let mut cleanup = StagingPathGuard::new(staging_path.clone()); - publish_staged_dataset(&staging_path, &output, true, None).await?; - cleanup.disarm(); - } - let response = ImportResponse { - dataset_uri: output_uri, - source_path: Some("events.lance".into()), - format: Some("events".into()), - output_format: ImportOutputFormat::Storyline.response_name().into(), - sources: 1, - trajectories: report.storylines, - fact_rows: Some(report.fact_rows), - input_bytes: None, - }; - serde_json::to_writer_pretty(&mut *stdout, &response) - .context("encode canonical event import JSON")?; - writeln!(stdout).context("write canonical event import JSON")?; - writeln!( - stderr, - "dataset_uri={} source=events.lance format=events output_format={} trajectories={} fact_rows={}", - response.dataset_uri, - response.output_format, - response.trajectories, - report.fact_rows, - ) - .context("write canonical event import metadata")?; - Ok(()) -} - -pub(super) async fn run_export( - mut args: ExportArgs, - settings_override: Option<&Path>, - stdout: &mut dyn Write, - stderr: &mut dyn Write, -) -> Result<()> { - anyhow::ensure!( - args.max_trajectories > 0, - "--max-trajectories must be greater than zero" - ); - anyhow::ensure!( - args.max_output_bytes > 0, - "--max-output-bytes must be greater than zero" - ); - anyhow::ensure!( - args.timeout_seconds > 0, - "--timeout must be greater than zero" - ); - args.stream = args.output == "-" || args.stream; - anyhow::ensure!( - args.output == "-" || !args.stream, - "--stream requires --to -" - ); - anyhow::ensure!( - !(args.output == "-" && args.overwrite), - "--overwrite cannot be used with stdout" - ); - if let Some(source) = &args.source { - validate_source_path(source)?; - } - if let Some(run_id) = &args.run_id { - validate_find_id("--run-id", run_id)?; - } - if let Some(document_id) = &args.document_id { - validate_find_id("--document-id", document_id)?; - } - if let Some(session_id) = &args.session_id { - validate_find_id("--session-id", session_id)?; - } - if let Some(expression) = &args.r#where { - anyhow::ensure!(!expression.trim().is_empty(), "--where must not be empty"); - anyhow::ensure!( - expression.len() <= 16 * 1024, - "--where exceeds the 16384-byte limit" - ); - } - - let format = ExchangeFormat::from(args.format); - let dataset = resolve_dataset_uri(args.from.as_deref(), settings_override)?; - if args.output != "-" { - args.output = expand_dataset_reference(&args.output, settings_override, false)?; - } - if format == ExchangeFormat::CompactJsonl { - anyhow::ensure!( - args.source.is_none() - && args.run_id.is_none() - && args.document_id.is_none() - && args.session_id.is_none() - && args.r#where.is_none(), - "compact JSONL export does not support filters" - ); - anyhow::ensure!( - args.output != "-", - "compact JSONL export requires a directory output" - ); - anyhow::ensure!( - args.overwrite || !Path::new(&args.output).exists(), - "export output already exists; pass --overwrite" - ); - let rows = - persisting_pchronicle::storage::CompactJsonlStore::export_path(&dataset, &args.output) - .await?; - writeln!( - stderr, - "format=compact-jsonl rows={} output={}", - rows, args.output - )?; - return Ok(()); - } - let (_, dataset_uris, snapshot) = - discover_query_snapshot(Some(&dataset), &[], args.max_files, args.max_entries).await?; - let dataset_uri = dataset_uris - .first() - .cloned() - .context("export Dataset URI missing after discovery")?; - let snapshot = Arc::new(snapshot); - let snapshot_id = snapshot.snapshot_id().to_string(); - let deadline = Duration::from_secs(args.timeout_seconds); - let export = tokio::time::timeout( - deadline, - export_from_snapshot(&args, format, &dataset_uri, snapshot.clone()), - ) - .await - .with_context(|| { - format!( - "Dataset export timed out after {} seconds", - args.timeout_seconds - ) - })??; - ensure_export_trajectory_budget(export.trajectories, args.max_trajectories)?; - ensure_output_byte_budget(export.bytes.len(), args.max_output_bytes, "encoded export")?; - write_export_output(&args.output, &export.bytes, args.overwrite, stdout).await?; - writeln!( - stderr, - "snapshot_id={} format={} trajectories={} output_bytes={} exact={}", - snapshot_id, - format.as_str(), - export.trajectories, - export.bytes.len(), - export.exact, - ) - .context("write pChronicle export metadata")?; - Ok(()) -} - -struct EncodedExport { - bytes: Vec, - trajectories: usize, - exact: bool, -} - -async fn export_from_snapshot( - args: &ExportArgs, - format: ExchangeFormat, - dataset_uri: &str, - snapshot: Arc, -) -> Result { - if let Some(export) = exact_local_file_export(args, format, dataset_uri, &snapshot).await? { - return Ok(export); - } - anyhow::ensure!( - !args.strict, - "strict export requires an unfiltered source file already stored in the requested format" - ); - - let sql = export_address_sql(args)?; - let engine = snapshot.clone().query_engine(Default::default()).await?; - let row_limit = args - .max_trajectories - .checked_add(1) - .context("--max-trajectories is too large")?; - let mut addresses = LimitedBuffer::new(args.max_output_bytes); - let write_result = engine - .write_query_jsonl_bounded(&sql, &mut addresses, Some(row_limit)) - .await; - let address_bytes = match addresses.finish(write_result)? { - QueryOutputBudgetOutcome::Complete(bytes) => bytes, - QueryOutputBudgetOutcome::RowLimitExceeded => { - return Err(cli_boundary_error( - BoundaryCode::ResourceExhausted, - format!( - "export exceeds max_trajectories limit of {}", - args.max_trajectories - ), - )); - } - QueryOutputBudgetOutcome::ByteLimitExceeded => { - return Err(cli_boundary_error( - BoundaryCode::ResourceExhausted, - format!( - "export address selection exceeds max_output_bytes limit of {}", - args.max_output_bytes - ), - )); - } - }; - let mut addresses = address_bytes - .split(|byte| *byte == b'\n') - .filter(|line| !line.is_empty()) - .map(|line| serde_json::from_slice(line).context("decode export run address")) - .collect::>>()?; - ensure_export_trajectory_budget(addresses.len(), args.max_trajectories)?; - anyhow::ensure!(!addresses.is_empty(), "export selection matched no runs"); - addresses.sort_by(|left, right| { - (&left.source_path, &left.document_id, &left.session_id).cmp(&( - &right.source_path, - &right.document_id, - &right.session_id, - )) - }); - let mut stories = Vec::with_capacity(addresses.len()); - let mut normalized_bytes = 0usize; - for address in &addresses { - let key = CatalogStorylineKey { - dataset: DEFAULT_DATASET_NAME.into(), - file: address.source_path.clone(), - document_id: address.document_id.clone(), - session_id: address.session_id.clone(), - }; - let story = snapshot - .load_storyline(&key) - .await - .with_context(|| { - format!( - "load export run {}/{}", - address.source_path, address.session_id - ) - })? - .with_context(|| { - format!( - "export run disappeared from snapshot: {}/{}", - address.source_path, address.session_id - ) - })?; - anyhow::ensure!( - story.trajectory_id.as_deref().unwrap_or(&story.session_id) == address.document_id, - "export run document ID changed within the snapshot" - ); - anyhow::ensure!( - story.run_id == address.run_id, - "export run runtime ID changed within the snapshot" - ); - normalized_bytes = normalized_bytes.saturating_add(serde_json::to_vec(&story)?.len()); - ensure_output_byte_budget(normalized_bytes, args.max_output_bytes, "normalized export")?; - stories.push(story); - } - let bytes = encode_export(format, &stories)?; - Ok(EncodedExport { - bytes, - trajectories: stories.len(), - exact: false, - }) -} - -async fn exact_local_file_export( - args: &ExportArgs, - format: ExchangeFormat, - dataset_uri: &str, - snapshot: &DatasetCatalogSnapshot, -) -> Result> { - if args.document_id.is_some() - || args.run_id.is_some() - || args.session_id.is_some() - || args.r#where.is_some() - { - return Ok(None); - } - let Some(dataset) = snapshot.dataset(DEFAULT_DATASET_NAME) else { - return Ok(None); - }; - let sources = dataset - .sources - .iter() - .filter(|source| source.status == CatalogSourceStatus::Ready) - .filter(|source| { - args.source - .as_deref() - .is_none_or(|selected| selected == source.file) - }) - .collect::>(); - if sources.len() != 1 || sources[0].kind != CatalogSourceKind::File { - return Ok(None); - } - let root = Path::new(dataset_uri); - if !root.is_dir() { - return Ok(None); - } - let source_path = root.join(&sources[0].file); - let source_path = std::fs::canonicalize(&source_path).context("canonicalize export Source")?; - anyhow::ensure!( - source_path.starts_with(root), - "export Source resolves outside the local Dataset" - ); - let input = std::fs::read(&source_path).context("read exact export Source")?; - ensure_output_byte_budget(input.len(), args.max_output_bytes, "exact export")?; - let text = std::str::from_utf8(&input).context("exact export Source must be UTF-8")?; - let detected = detect_format(Some(&source_path), Some(text))?; - if detected != exchange_document_format(format) { - return Ok(None); - } - let trajectories = validate_import_source(format, &source_path).await?; - anyhow::ensure!( - sources[0].size_bytes == Some(input.len() as u64) - && sources[0].snapshot_ref().as_deref() == Some(&local_file_snapshot_ref(&source_path)), - "export Source changed after the Snapshot was created" - ); - Ok(Some(EncodedExport { - bytes: input, - trajectories, - exact: true, - })) -} - -fn ensure_export_trajectory_budget(trajectories: usize, max_trajectories: u64) -> Result<()> { - if usize::try_from(max_trajectories).is_ok_and(|limit| trajectories > limit) { - return Err(cli_boundary_error( - BoundaryCode::ResourceExhausted, - format!("export exceeds max_trajectories limit of {max_trajectories}"), - )); - } - Ok(()) -} - -fn export_address_sql(args: &ExportArgs) -> Result { - let mut predicates = Vec::new(); - if let Some(source) = &args.source { - predicates.push(format!("_file_ = {}", sql_string(source))); - } - if let Some(run_id) = &args.run_id { - predicates.push(format!("run_id = {}", sql_string(run_id))); - } - if let Some(document_id) = &args.document_id { - predicates.push(format!("document_id = {}", sql_string(document_id))); - } - if let Some(session_id) = &args.session_id { - predicates.push(format!("session_id = {}", sql_string(session_id))); - } - if let Some(expression) = &args.r#where { - predicates.push(format!("({expression})")); - } - let predicate = if predicates.is_empty() { - String::new() - } else { - format!(" WHERE {}", predicates.join(" AND ")) - }; - let limit = args - .max_trajectories - .checked_add(1) - .context("--max-trajectories is too large")?; - Ok(format!( - "SELECT _file_ AS source_path, document_id, run_id, session_id \ - FROM dataset.trajectories{predicate} \ - ORDER BY _file_, document_id, session_id LIMIT {limit}" - )) -} - -fn encode_export(format: ExchangeFormat, stories: &[StorylineDocument]) -> Result> { - let value = match format { - ExchangeFormat::Atif => encode_json_storylines(DocumentFormat::Atif, stories)?, - ExchangeFormat::Actf => encode_json_storylines(DocumentFormat::Actf, stories)?, - ExchangeFormat::OpenaiMessages => { - encode_json_storylines(DocumentFormat::OpenaiMsg, stories)? - } - ExchangeFormat::Storyline => encode_json_storylines(DocumentFormat::Storyline, stories)?, - ExchangeFormat::Codex | ExchangeFormat::ClaudeCode => { - bail!("{format} is decode-only and cannot be exported") - } - ExchangeFormat::CompactJsonl | ExchangeFormat::Auto => { - unreachable!("exchange export format was validated") - } - }; - let mut output = serde_json::to_vec_pretty(&value).context("encode export JSON")?; - output.push(b'\n'); - Ok(output) -} - -fn exchange_document_format(format: ExchangeFormat) -> Option { - match format { - ExchangeFormat::Atif => Some(DocumentFormat::Atif), - ExchangeFormat::Actf => Some(DocumentFormat::Actf), - ExchangeFormat::OpenaiMessages => Some(DocumentFormat::OpenaiMsg), - ExchangeFormat::Storyline => Some(DocumentFormat::Storyline), - ExchangeFormat::Codex => Some(DocumentFormat::Codex), - ExchangeFormat::ClaudeCode => Some(DocumentFormat::ClaudeCode), - ExchangeFormat::CompactJsonl | ExchangeFormat::Auto => None, - } -} - -async fn write_export_output( - output: &str, - bytes: &[u8], - overwrite: bool, - stdout: &mut dyn Write, -) -> Result<()> { - if output == "-" { - stdout.write_all(bytes).context("write export stream")?; - return Ok(()); - } - DatasetLocation::parse(output)? - .put_bytes(bytes, overwrite) - .await -} - -fn local_file_snapshot_ref(path: &Path) -> String { - let mut hash = blake3::Hasher::new(); - hash.update(path.to_string_lossy().as_bytes()); - if let Ok(metadata) = std::fs::metadata(path) { - hash.update(&metadata.len().to_le_bytes()); - if let Ok(modified) = metadata.modified() - && let Ok(duration) = modified.duration_since(std::time::UNIX_EPOCH) - { - hash.update(&duration.as_nanos().to_le_bytes()); - } - #[cfg(unix)] - { - use std::os::unix::fs::MetadataExt; - hash.update(&metadata.dev().to_le_bytes()); - hash.update(&metadata.ino().to_le_bytes()); - } - } - format!("local:{}", hash.finalize().to_hex()) -} - -#[derive(Debug, Clone)] -struct ImportFileCandidate { - path: PathBuf, - relative_path: PathBuf, - output_relative_path: Option, - /// Prefetched bytes (tests / rare callers). Normal imports leave this empty - /// and read local paths or object-store keys on demand. - content: Option>, - /// Object-store Dataset root URI; when set, bytes are fetched lazily. - remote_root: Option, - /// Size from discovery (`stat` / object metadata) for progress totals. - size_hint: u64, -} - -#[derive(Debug)] -struct ImportedSource { - source_path: String, - format: DocumentFormat, - trajectories: usize, - input_bytes: usize, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -enum ImportPhase { - Discovering, - Deleting, - Reading, - Parsing, - Writing, -} - -impl ImportPhase { - fn as_str(self) -> &'static str { - match self { - Self::Discovering => "discovering", - Self::Deleting => "deleting", - Self::Reading => "reading", - Self::Parsing => "parsing", - Self::Writing => "writing", - } - } -} - -/// Dense import progress: TTY paints three in-place lines; redirected stderr gets -/// one summary line per completed source (buffered, flushed at the end). -struct ImportProgress { - tty: bool, - discovered_files: u64, - discovered_bytes: u64, - imported_files: u64, - imported_bytes: u64, - /// Storyline trajectories successfully committed so far. - committed: u64, - /// Replace/drop delete progress (separate from discovery totals). - deleted_files: u64, - delete_total: u64, - phase: ImportPhase, - file: String, - /// Producer side of the read→parse pipeline (empty when idle). - reading_ahead: String, - painted: bool, - log_lines: Vec, - last_paint: Option, - /// Shared with index-build callbacks so Lance work updates line 3 in place. - surface: Arc>, -} - -#[derive(Debug, Clone)] -struct ImportProgressSurface { - tty: bool, - painted: bool, - deleting: bool, - line1: String, - line2: String, - reading_ahead: String, - phase: String, - file: String, -} - -impl ImportProgressSurface { - fn paint_activity(&mut self, activity: &str) -> Result<()> { - if !self.tty { - return Ok(()); - } - let file = if activity.is_empty() { - if self.file.is_empty() { - "-".to_owned() - } else { - truncate_middle(&self.file, 72) - } - } else { - truncate_middle(activity, 96) - }; - let line3 = if !self.reading_ahead.is_empty() && !activity.is_empty() { - format!( - "[reading] {} | [writing] {file}", - truncate_middle(&self.reading_ahead, 40), - ) - } else if !self.reading_ahead.is_empty() && self.phase == "reading" { - format!( - "[reading] {}", - truncate_middle(&self.reading_ahead, 96) - ) - } else if !self.reading_ahead.is_empty() { - format!( - "[reading] {} | [{}] {file}", - truncate_middle(&self.reading_ahead, 40), - self.phase, - ) - } else { - format!("[{}] {file}", if activity.is_empty() { self.phase.as_str() } else { "writing" }) - }; - - let mut err = std::io::stderr(); - if self.painted { - write!(err, "\x1b[2A").context("move import progress cursor")?; - } - if self.deleting { - write!(err, "\r\x1b[2K{}\n\r\x1b[2K\n\r\x1b[2K{line3}", self.line1) - .context("paint delete progress")?; - } else { - write!( - err, - "\r\x1b[2K{}\n\r\x1b[2K{}\n\r\x1b[2K{line3}", - self.line1, self.line2 - ) - .context("paint import progress")?; - } - err.flush().context("flush import progress")?; - self.painted = true; - Ok(()) - } -} - -impl ImportProgress { - fn new(tty: bool) -> Self { - Self { - tty, - discovered_files: 0, - discovered_bytes: 0, - imported_files: 0, - imported_bytes: 0, - committed: 0, - deleted_files: 0, - delete_total: 0, - phase: ImportPhase::Discovering, - file: String::new(), - reading_ahead: String::new(), - painted: false, - log_lines: Vec::new(), - last_paint: None, - surface: Arc::new(std::sync::Mutex::new(ImportProgressSurface { - tty, - painted: false, - deleting: false, - line1: String::new(), - line2: String::new(), - reading_ahead: String::new(), - phase: ImportPhase::Discovering.as_str().to_owned(), - file: String::new(), - })), - } - } - - fn attach_index_progress(&self) -> persisting_pchronicle::storage::IndexBuildProgressGuard { - let surface = Arc::clone(&self.surface); - persisting_pchronicle::storage::install_index_build_progress(Arc::new(move |message| { - if let Ok(mut surface) = surface.lock() { - let _ = surface.paint_activity(message); - } - })) - } - - fn reset_import_counters(&mut self) { - self.imported_files = 0; - self.imported_bytes = 0; - self.committed = 0; - self.deleted_files = 0; - self.delete_total = 0; - self.reading_ahead.clear(); - self.file.clear(); - } - - fn set_discovered(&mut self, files: u64, bytes: u64) -> Result<()> { - self.discovered_files = files; - self.discovered_bytes = bytes; - self.phase = ImportPhase::Discovering; - self.file.clear(); - self.paint(false) - } - - fn note_discovered(&mut self, file: &str, bytes: u64) -> Result<()> { - self.discovered_files = self.discovered_files.saturating_add(1); - self.discovered_bytes = self.discovered_bytes.saturating_add(bytes); - self.phase = ImportPhase::Discovering; - self.file = file.to_owned(); - // Throttle TTY paints during large listings so discovery stays responsive. - let should_paint = !self.tty - || self - .last_paint - .map(|at| at.elapsed() >= std::time::Duration::from_millis(100)) - .unwrap_or(true) - || self.discovered_files == 1 - || self.discovered_files % 64 == 0; - if should_paint { - self.paint(true)?; - } - Ok(()) - } - - fn note_scanning(&mut self, prefix: &str) -> Result<()> { - self.phase = ImportPhase::Discovering; - self.file = if prefix.is_empty() { - "/".to_owned() - } else { - format!("{prefix}/") - }; - let should_paint = !self.tty - || self - .last_paint - .map(|at| at.elapsed() >= std::time::Duration::from_millis(100)) - .unwrap_or(true); - if should_paint { - self.paint(true)?; - } - Ok(()) - } - - fn note_deleted(&mut self, deleted: u64, total: u64, path: &str) -> Result<()> { - self.deleted_files = deleted; - self.delete_total = total; - self.phase = ImportPhase::Deleting; - self.file = path.to_owned(); - if deleted == total { - // Always emit a final summary line for non-TTY logs. - return self.paint(false); - } - let should_paint = !self.tty - || path.is_empty() - || self - .last_paint - .map(|at| at.elapsed() >= std::time::Duration::from_millis(100)) - .unwrap_or(true) - || deleted == 1 - || deleted % 64 == 0; - if should_paint { - self.paint(true)?; - } - Ok(()) - } - - fn set_phase(&mut self, phase: ImportPhase, file: &str) -> Result<()> { - self.phase = phase; - self.file = file.to_owned(); - self.paint(true) - } - - fn set_reading_ahead(&mut self, file: &str) -> Result<()> { - self.reading_ahead = file.to_owned(); - let should_paint = !self.tty - || self - .last_paint - .map(|at| at.elapsed() >= std::time::Duration::from_millis(100)) - .unwrap_or(true); - if should_paint { - self.paint(true)?; - } - Ok(()) - } - - fn clear_reading_ahead(&mut self) -> Result<()> { - if self.reading_ahead.is_empty() { - return Ok(()); - } - self.reading_ahead.clear(); - self.paint(true) - } - - fn note_imported(&mut self, bytes: u64) -> Result<()> { - self.imported_files = self.imported_files.saturating_add(1); - self.imported_bytes = self.imported_bytes.saturating_add(bytes); - self.paint(false) - } - - fn note_committed(&mut self, committed: u64) -> Result<()> { - self.committed = committed; - self.phase = ImportPhase::Writing; - self.file = format!("commit trajectories={committed}"); - self.paint(false) - } - - fn finish(&mut self) -> Result<()> { - if let Ok(surface) = self.surface.lock() { - self.painted = surface.painted; - } - if self.tty && self.painted { - let mut err = std::io::stderr(); - writeln!(err).context("finish import progress")?; - err.flush().context("flush import progress")?; - self.painted = false; - if let Ok(mut surface) = self.surface.lock() { - surface.painted = false; - } - } - Ok(()) - } - - fn notice(&mut self, message: &str) -> Result<()> { - self.finish()?; - if self.tty { - let mut err = std::io::stderr(); - writeln!(err, "{message}").context("write import notice")?; - err.flush().context("flush import notice")?; - } else { - self.log_lines.push(message.to_owned()); - } - Ok(()) - } - - fn flush_log(self, out: &mut dyn Write) -> Result<()> { - for line in self.log_lines { - writeln!(out, "{line}").context("flush import progress log")?; - } - Ok(()) - } - - fn paint(&mut self, phase_only: bool) -> Result<()> { - if let Ok(surface) = self.surface.lock() { - self.painted = surface.painted; - } - let deleting = self.phase == ImportPhase::Deleting; - let line1 = if deleting { - format!( - "deleted:total = {}/{}", - self.deleted_files, self.delete_total - ) - } else { - format!( - "imported:discovered = {}/{}", - self.imported_files, self.discovered_files - ) - }; - let line2 = if deleting { - String::new() - } else { - format!( - "committed = {} ; size = {}:{}", - self.committed, - format_byte_count(self.imported_bytes), - format_byte_count(self.discovered_bytes) - ) - }; - let file = if self.file.is_empty() { - "-".to_owned() - } else { - truncate_middle(&self.file, 72) - }; - let line3 = if !self.reading_ahead.is_empty() - && matches!( - self.phase, - ImportPhase::Parsing | ImportPhase::Writing - ) - { - format!( - "[reading] {} | [{}] {file}", - truncate_middle(&self.reading_ahead, 48), - self.phase.as_str(), - ) - } else if !self.reading_ahead.is_empty() && self.phase == ImportPhase::Reading { - format!( - "[reading] {}", - truncate_middle(&self.reading_ahead, 96) - ) - } else { - format!("[{}] {file}", self.phase.as_str()) - }; - - if let Ok(mut surface) = self.surface.lock() { - surface.tty = self.tty; - surface.deleting = deleting; - surface.line1 = line1.clone(); - surface.line2 = line2.clone(); - surface.reading_ahead = self.reading_ahead.clone(); - surface.phase = self.phase.as_str().to_owned(); - surface.file = self.file.clone(); - surface.painted = self.painted; - } - - if self.tty { - let mut err = std::io::stderr(); - if self.painted { - write!(err, "\x1b[2A").context("move import progress cursor")?; - } - if deleting { - write!(err, "\r\x1b[2K{line1}\n\r\x1b[2K\n\r\x1b[2K{line3}") - .context("paint delete progress")?; - } else { - write!(err, "\r\x1b[2K{line1}\n\r\x1b[2K{line2}\n\r\x1b[2K{line3}") - .context("paint import progress")?; - } - err.flush().context("flush import progress")?; - self.painted = true; - if let Ok(mut surface) = self.surface.lock() { - surface.painted = true; - } - self.last_paint = Some(std::time::Instant::now()); - return Ok(()); - } - - if phase_only { - return Ok(()); - } - if deleting { - self.log_lines - .push(format!("{line1}; {line3}")); - } else { - self.log_lines - .push(format!("{line1}; {line2}; {line3}")); - } - Ok(()) - } -} - -fn format_byte_count(bytes: u64) -> String { - const KIB: f64 = 1024.0; - const MIB: f64 = 1024.0 * 1024.0; - const GIB: f64 = 1024.0 * 1024.0 * 1024.0; - let value = bytes as f64; - if value >= GIB { - format!("{:.1}GiB", value / GIB) - } else if value >= MIB { - format!("{:.1}MiB", value / MIB) - } else if value >= KIB { - format!("{:.1}KiB", value / KIB) - } else { - format!("{bytes}B") - } -} - -fn truncate_middle(value: &str, max_chars: usize) -> String { - let chars: Vec = value.chars().collect(); - if chars.len() <= max_chars { - return value.to_owned(); - } - if max_chars <= 3 { - return chars.into_iter().take(max_chars).collect(); - } - let head = (max_chars - 1) / 2; - let tail = max_chars - 1 - head; - let mut out: String = chars.iter().take(head).collect(); - out.push('…'); - out.extend(chars.iter().skip(chars.len() - tail)); - out -} - -#[cfg(test)] -mod import_progress_tests { - use super::*; - - #[test] - fn commit_batch_schedule_grows_to_cap() { - let mut schedule = CommitBatchSchedule::adaptive(); - assert_eq!(schedule.current(), 64); - schedule.after_commit(); - assert_eq!(schedule.current(), 128); - schedule.after_commit(); - assert_eq!(schedule.current(), 256); - schedule.after_commit(); - assert_eq!(schedule.current(), 512); - schedule.after_commit(); - assert_eq!(schedule.current(), 1024); - schedule.after_commit(); - assert_eq!(schedule.current(), 2048); - schedule.after_commit(); - assert_eq!(schedule.current(), 4096); - schedule.after_commit(); - assert_eq!(schedule.current(), 4096); - } - - #[test] - fn skippable_commit_errors_cover_s3_timeouts_and_preconditions() { - assert!(is_skippable_storyline_commit_error(&anyhow!( - "LanceError(IO): Generic S3 error: operation timed out" - ))); - assert!(is_skippable_storyline_commit_error(&anyhow!( - "ConditionNotMatch (persistent) PreconditionFailed" - ))); - assert!(!is_skippable_storyline_commit_error(&anyhow!( - "duplicate document_id policy rejected payload" - ))); - } - - #[test] - fn retract_imported_trajectories_from_tail_sources() { - let mut sources = vec![ - ImportedSource { - source_path: "a.json".into(), - format: DocumentFormat::Atif, - trajectories: 3, - input_bytes: 10, - }, - ImportedSource { - source_path: "b.json".into(), - format: DocumentFormat::Atif, - trajectories: 2, - input_bytes: 10, - }, - ]; - retract_imported_trajectories(&mut sources, 3); - assert_eq!(sources[0].trajectories, 2); - assert_eq!(sources[1].trajectories, 0); - } - - #[test] - fn commit_batch_schedule_fixed_stays_put() { - let mut schedule = CommitBatchSchedule::fixed(50); - assert_eq!(schedule.current(), 50); - schedule.after_commit(); - assert_eq!(schedule.current(), 50); - } - - #[test] - fn format_byte_count_uses_binary_units() { - assert_eq!(format_byte_count(512), "512B"); - assert_eq!(format_byte_count(1536), "1.5KiB"); - assert_eq!(format_byte_count(2 * 1024 * 1024), "2.0MiB"); - } - - #[test] - fn non_tty_progress_emits_dense_completed_lines() { - let mut progress = ImportProgress::new(false); - progress.set_discovered(2, 300).unwrap(); - progress.set_phase(ImportPhase::Reading, "a/long.json").unwrap(); - progress.set_phase(ImportPhase::Parsing, "a/long.json").unwrap(); - progress.note_imported(100).unwrap(); - progress.set_phase(ImportPhase::Writing, "b.json").unwrap(); - progress.note_imported(200).unwrap(); - progress.note_committed(3).unwrap(); - let mut out = Vec::new(); - progress.flush_log(&mut out).unwrap(); - let text = String::from_utf8(out).unwrap(); - assert!(text.contains("imported:discovered = 1/2"), "{text}"); - assert!(text.contains("imported:discovered = 2/2"), "{text}"); - assert!(text.contains("committed = 3"), "{text}"); - assert!(text.contains("size ="), "{text}"); - assert!(text.contains("[writing] commit trajectories=3") || text.contains("[writing] b.json") || text.contains("[parsing] a/long.json"), "{text}"); - assert!(!text.contains("status=fetching"), "{text}"); - } -} - -fn collect_import_candidates(input: &Path) -> Result<(bool, Vec)> { - let metadata = std::fs::symlink_metadata(input) - .with_context(|| format!("inspect import input {}", input.display()))?; - let explicit_file = if metadata.file_type().is_symlink() { - std::fs::metadata(input) - .with_context(|| format!("inspect import input target {}", input.display()))? - .is_file() - } else { - metadata.is_file() - }; - if explicit_file { - let relative_path = input - .file_name() - .map(PathBuf::from) - .context("import input file has no filename")?; - let size_hint = metadata.len(); - return Ok(( - false, - vec![ImportFileCandidate { - path: input.to_path_buf(), - relative_path, - output_relative_path: None, - content: None, - remote_root: None, - size_hint, - }], - )); - } - anyhow::ensure!( - metadata.is_dir(), - "import input must be a regular file or directory" - ); - - let paths = collect_visible_json_files(input)?; - let mut candidates = Vec::with_capacity(paths.len()); - for path in paths { - let relative_path = path - .strip_prefix(input) - .context("derive Dataset-relative import source path")? - .to_path_buf(); - let size_hint = std::fs::metadata(&path) - .map(|meta| meta.len()) - .unwrap_or(0); - candidates.push(ImportFileCandidate { - path, - output_relative_path: Some(relative_path.clone()), - relative_path, - content: None, - remote_root: None, - size_hint, - }); - } - candidates.sort_by(|left, right| left.relative_path.cmp(&right.relative_path)); - if candidates.is_empty() { - return Err(cli_boundary_error( - BoundaryCode::InvalidRequest, - "import directory contains no .json, .jsonl, or .ndjson files", - )); - } - Ok((true, candidates)) -} - -/// Recursively collect absolute paths of visible `.json` / `.jsonl` / `.ndjson` -/// files under `root`. Shared by `import` and `sync`; not Catalog Directory -/// discovery (which is one-level and skips loose files). -pub(crate) fn collect_visible_json_files(root: &Path) -> Result> { - let mut pending = vec![root.to_path_buf()]; - let mut files = Vec::new(); - while let Some(directory) = pending.pop() { - let mut entries = std::fs::read_dir(&directory) - .with_context(|| format!("read directory {}", directory.display()))? - .collect::>>()?; - entries.sort_by_key(std::fs::DirEntry::path); - for entry in entries { - let file_type = entry.file_type()?; - if file_type.is_symlink() { - continue; - } - let path = entry.path(); - if file_type.is_dir() { - pending.push(path); - } else if file_type.is_file() && is_visible_json_file(&path) { - let relative = path - .strip_prefix(root) - .unwrap_or(path.as_path()) - .to_string_lossy() - .replace('\\', "/"); - if relative.split('/').any(|part| part == "_meta") { - continue; - } - files.push(path); - } - } - } - files.sort(); - Ok(files) -} - -fn is_visible_json_file(path: &Path) -> bool { - path.extension() - .and_then(|extension| extension.to_str()) - .is_some_and(|extension| { - matches!( - extension.to_ascii_lowercase().as_str(), - "json" | "jsonl" | "ndjson" - ) - }) -} - -async fn load_import_candidate_bytes( - candidate: &ImportFileCandidate, - max_input_bytes: usize, - label: &str, -) -> Result> { - if let Some(content) = &candidate.content { - anyhow::ensure!( - content.len() <= max_input_bytes, - "{label} exceeds max_input_bytes limit of {max_input_bytes}" - ); - return Ok(content.clone()); - } - if let Some(remote_root) = &candidate.remote_root { - let key = candidate.relative_path.to_string_lossy().replace('\\', "/"); - let location = DatasetLocation::parse(remote_root)?; - let bytes = location - .read_relative_bytes(&key) - .await - .with_context(|| format!("read import object {key} under {remote_root}"))?; - anyhow::ensure!( - bytes.len() <= max_input_bytes, - "{label} exceeds max_input_bytes limit of {max_input_bytes}" - ); - return Ok(bytes); - } - let file = std::fs::File::open(&candidate.path).with_context(|| format!("open {label}"))?; - read_bounded(file, max_input_bytes, label) -} - -fn scope_import_source_error(error: anyhow::Error, source_path: &Path) -> anyhow::Error { - if let Some(boundary) = error.downcast_ref::() { - return cli_boundary_error( - boundary.code, - format!("{}: {}", source_path.display(), boundary.message), - ); - } - error.context(format!("import source {}", source_path.display())) -} - -struct DecodedImportSource { - diagnostic_path: PathBuf, - metadata: ImportedSource, - storylines: Vec, -} - -enum DecodeImportOutcome { - Imported(DecodedImportSource), - Skipped { path: PathBuf, reason: String }, -} - -enum ImportFormatResolution { - Format(ExchangeFormat), - Skip(String), -} - -enum StorylineImportInputs<'a> { - Stdin(Option<&'a mut dyn Read>), -} - -struct StorylineImportIterator<'a> { - requested_format: ExchangeFormat, - max_input_bytes: usize, - progress: &'a mut ImportProgress, - inputs: StorylineImportInputs<'a>, - current: std::vec::IntoIter, - imported_sources: Vec, - unknown_field_warnings: persisting_pchronicle::model::UnknownFieldImportWarnings, - skipped_warnings: Vec, - seen_document_ids: HashSet, - duplicate_policy: DuplicateIdPolicy, - failed: bool, -} - -impl<'a> StorylineImportIterator<'a> { - fn stdin( - requested_format: ExchangeFormat, - max_input_bytes: usize, - stdin: &'a mut dyn Read, - progress: &'a mut ImportProgress, - seen_document_ids: HashSet, - duplicate_policy: DuplicateIdPolicy, - ) -> Self { - Self { - requested_format, - max_input_bytes, - progress, - inputs: StorylineImportInputs::Stdin(Some(stdin)), - current: Vec::new().into_iter(), - imported_sources: Vec::new(), - unknown_field_warnings: - persisting_pchronicle::model::UnknownFieldImportWarnings::default(), - skipped_warnings: Vec::new(), - seen_document_ids, - duplicate_policy, - failed: false, - } - } - - async fn decode_next_source(&mut self) -> Result> { - loop { - let outcome = match &mut self.inputs { - StorylineImportInputs::Stdin(stdin) => { - let Some(stdin) = stdin.take() else { - return Ok(None); - }; - self.progress.set_phase(ImportPhase::Reading, "stdin")?; - let input = read_bounded(stdin, self.max_input_bytes, "stdin")?; - self.progress.set_phase(ImportPhase::Parsing, "stdin")?; - decode_import_source( - self.requested_format, - ImportOutputFormat::Storyline, - None, - None, - None, - &input, - &mut self.unknown_field_warnings, - )? - } - }; - match outcome { - DecodeImportOutcome::Imported(decoded) => { - self.progress.set_phase( - ImportPhase::Writing, - &decoded.diagnostic_path.to_string_lossy(), - )?; - self.progress - .note_imported(decoded.metadata.input_bytes as u64)?; - return Ok(Some(decoded)); - } - DecodeImportOutcome::Skipped { path, reason } => { - self.progress.note_imported(0)?; - self.skipped_warnings - .push(skipped_import_warning(&path, &reason)); - } - } - } - } - - fn into_result_parts( - self, - ) -> ( - Vec, - persisting_pchronicle::model::UnknownFieldImportWarnings, - Vec, - ) { - ( - self.imported_sources, - self.unknown_field_warnings, - self.skipped_warnings, - ) - } - - async fn next_document(&mut self) -> Option> { - loop { - if let Some(mut storyline) = self.current.next() { - let original = storyline.document_id().to_string(); - match self.duplicate_policy { - DuplicateIdPolicy::Suffix => { - if let Some((original, renamed)) = uniquify_storyline_document_id( - &mut storyline, - &mut self.seen_document_ids, - ) { - self.skipped_warnings.push(format!( - "warning: duplicate document_id '{original}' renamed to '{renamed}'" - )); - } - } - DuplicateIdPolicy::Skip => { - if !self.seen_document_ids.insert(original.clone()) { - self.skipped_warnings.push(format!( - "warning: duplicate document_id '{original}' skipped" - )); - continue; - } - } - } - let metadata = self - .imported_sources - .last_mut() - .expect("decoded Storyline has source metadata"); - metadata.trajectories = metadata - .trajectories - .checked_add(1) - .expect("import trajectory count overflow"); - return Some(Ok(storyline)); - } - if self.failed { - return None; - } - match self.decode_next_source().await { - Ok(Some(decoded)) => { - let mut metadata = decoded.metadata; - metadata.trajectories = 0; - self.imported_sources.push(metadata); - self.current = decoded.storylines.into_iter(); - } - Ok(None) => return None, - Err(error) => { - self.failed = true; - return Some(Err(error)); - } - } - } - } -} - -fn uniquify_storyline_document_id( - story: &mut StorylineDocument, - seen: &mut HashSet, -) -> Option<(String, String)> { - let preferred = story.document_id().to_string(); - if seen.insert(preferred.clone()) { - return None; - } - let mut suffix = 1u64; - let renamed = loop { - let candidate = format!("{preferred}#{suffix}"); - if seen.insert(candidate.clone()) { - break candidate; - } - suffix = suffix - .checked_add(1) - .expect("document_id disambiguation suffix overflow"); - }; - if story - .trajectory_id - .as_deref() - .is_some_and(|id| !id.is_empty()) - { - story.trajectory_id = Some(renamed.clone()); - } else { - story.session_id = renamed.clone(); - } - Some((preferred, renamed)) -} - -#[allow(clippy::too_many_arguments)] -fn decode_import_source( - requested_format: ExchangeFormat, - output_format: ImportOutputFormat, - input_path: Option<&Path>, - decode_relative_path: Option<&Path>, - logical_source_path: Option<&Path>, - input: &[u8], - unknown_field_warnings: &mut persisting_pchronicle::model::UnknownFieldImportWarnings, -) -> Result { - let diagnostic_path = decode_relative_path - .unwrap_or_else(|| Path::new("stdin")) - .to_path_buf(); - let text = std::str::from_utf8(input).map_err(|error| { - cli_boundary_error( - BoundaryCode::InvalidRequest, - format!("{} is not UTF-8: {error}", diagnostic_path.display()), - ) - })?; - let allow_skip = requested_format == ExchangeFormat::Auto && logical_source_path.is_some(); - let format = match resolve_import_format(requested_format, input_path, text, allow_skip) - .map_err(|error| { - if logical_source_path.is_some() { - scope_import_source_error(error, &diagnostic_path) - } else { - error - } - })? { - ImportFormatResolution::Format(format) => format, - ImportFormatResolution::Skip(reason) => { - return Ok(DecodeImportOutcome::Skipped { - path: diagnostic_path, - reason, - }); - } - }; - let document_format = exchange_document_format(format) - .context("supported import format must map to a physical document format")?; - let source_path = logical_source_path - .map(PathBuf::from) - .unwrap_or_else(|| single_import_source_path(format, output_format, input_path)); - let decode_relative_path = decode_relative_path.unwrap_or(&source_path); - let storylines = - decode_json_storylines(document_format, text, decode_relative_path).map_err(|issue| { - let code = match issue.kind() { - InputIssueKind::Invalid => BoundaryCode::InvalidRequest, - InputIssueKind::Unsupported => BoundaryCode::Unsupported, - }; - cli_boundary_error( - code, - import_input_issue_message(&issue, decode_relative_path), - ) - }); - let storylines = match storylines { - Ok(storylines) => storylines, - Err(error) if allow_skip => { - return Ok(DecodeImportOutcome::Skipped { - path: diagnostic_path, - reason: error.to_string(), - }); - } - Err(error) => return Err(error), - }; - unknown_field_warnings - .observe_storylines(&storylines) - .map_err(|issue| { - cli_boundary_error( - BoundaryCode::InvalidRequest, - import_input_issue_message(&issue, decode_relative_path), - ) - })?; - - let metadata = ImportedSource { - source_path: source_path - .to_str() - .context("Dataset-relative import Source path is not UTF-8")? - .to_owned(), - format: document_format, - trajectories: storylines.len(), - input_bytes: input.len(), - }; - Ok(DecodeImportOutcome::Imported(DecodedImportSource { - diagnostic_path, - metadata, - storylines, - })) -} - -#[allow(clippy::too_many_arguments)] -fn stage_preserved_import_source( - requested_format: ExchangeFormat, - input_path: Option<&Path>, - decode_relative_path: Option<&Path>, - logical_source_path: Option<&Path>, - input: &[u8], - staging_root: &Path, - unknown_field_warnings: &mut persisting_pchronicle::model::UnknownFieldImportWarnings, - skipped_warnings: &mut Vec, -) -> Result> { - let decoded = match decode_import_source( - requested_format, - ImportOutputFormat::Preserve, - input_path, - decode_relative_path, - logical_source_path, - input, - unknown_field_warnings, - )? { - DecodeImportOutcome::Imported(decoded) => decoded, - DecodeImportOutcome::Skipped { path, reason } => { - skipped_warnings.push(skipped_import_warning(&path, &reason)); - return Ok(None); - } - }; - validate_import_storylines(&decoded.storylines).map_err(|error| { - if logical_source_path.is_some() { - scope_import_source_error(error, &decoded.diagnostic_path) - } else { - error - } - })?; - - let staged_source = staging_root.join(&decoded.metadata.source_path); - let staged_parent = staged_source - .parent() - .context("staged import Source has no parent")?; - std::fs::create_dir_all(staged_parent) - .with_context(|| format!("create staged Source parent {}", staged_parent.display()))?; - let mut file = std::fs::OpenOptions::new() - .write(true) - .create_new(true) - .open(&staged_source) - .with_context(|| format!("create staged Source {}", decoded.metadata.source_path))?; - file.write_all(input) - .with_context(|| format!("write staged Source {}", decoded.metadata.source_path))?; - file.sync_all() - .with_context(|| format!("sync staged Source {}", decoded.metadata.source_path))?; - Ok(Some(decoded.metadata)) -} - -fn read_bounded(mut reader: impl Read, max_bytes: usize, label: &str) -> Result> { - let mut input = Vec::new(); - if max_bytes == usize::MAX { - reader - .read_to_end(&mut input) - .with_context(|| format!("read {label}"))?; - } else { - let limit = u64::try_from(max_bytes) - .ok() - .and_then(|limit| limit.checked_add(1)) - .ok_or_else(|| { - cli_boundary_error( - BoundaryCode::InvalidRequest, - "--max-input-bytes is too large", - ) - })?; - reader - .by_ref() - .take(limit) - .read_to_end(&mut input) - .with_context(|| format!("read {label}"))?; - if input.len() > max_bytes { - return Err(cli_boundary_error( - BoundaryCode::ResourceExhausted, - format!("{label} exceeds max_input_bytes limit of {max_bytes}"), - )); - } - } - if input.is_empty() { - return Err(cli_boundary_error( - BoundaryCode::InvalidRequest, - format!("{label} is empty"), - )); - } - Ok(input) -} - -fn resolve_import_format( - requested: ExchangeFormat, - input_path: Option<&Path>, - input: &str, - allow_skip: bool, -) -> Result { - let format = match requested { - ExchangeFormat::Auto => match detect_format(input_path, Some(input))? { - Some(DocumentFormat::Atif) => ExchangeFormat::Atif, - Some(DocumentFormat::Actf) => ExchangeFormat::Actf, - Some(DocumentFormat::OpenaiMsg) => ExchangeFormat::OpenaiMessages, - Some(DocumentFormat::Storyline) => ExchangeFormat::Storyline, - Some(DocumentFormat::Codex) => ExchangeFormat::Codex, - Some(DocumentFormat::ClaudeCode) => ExchangeFormat::ClaudeCode, - Some(format) if allow_skip => { - return Ok(ImportFormatResolution::Skip(format!( - "detected import format '{format}' is not a queryable JSON format" - ))); - } - Some(format) => { - return Err(cli_boundary_error( - BoundaryCode::Unsupported, - format!("detected import format '{format}' is not a queryable JSON format"), - )); - } - None if allow_skip && looks_like_json_document(input) => { - return Ok(ImportFormatResolution::Skip( - "cannot detect import format".into(), - )); - } - None => { - return Err(cli_boundary_error( - BoundaryCode::InvalidRequest, - "cannot detect import format; pass --format explicitly", - )); - } - }, - ExchangeFormat::Atif => ExchangeFormat::Atif, - ExchangeFormat::Actf => ExchangeFormat::Actf, - ExchangeFormat::OpenaiMessages => ExchangeFormat::OpenaiMessages, - ExchangeFormat::Storyline => ExchangeFormat::Storyline, - ExchangeFormat::Codex => ExchangeFormat::Codex, - ExchangeFormat::ClaudeCode => ExchangeFormat::ClaudeCode, - ExchangeFormat::CompactJsonl => ExchangeFormat::CompactJsonl, - }; - if !matches!( - format, - ExchangeFormat::Atif - | ExchangeFormat::Actf - | ExchangeFormat::OpenaiMessages - | ExchangeFormat::Storyline - | ExchangeFormat::Codex - | ExchangeFormat::ClaudeCode - | ExchangeFormat::CompactJsonl - ) { - return Err(cli_boundary_error( - BoundaryCode::Unsupported, - format!( - "import format '{format}' is not supported by the first queryable import increment" - ), - )); - } - Ok(ImportFormatResolution::Format(format)) -} - -fn looks_like_json_document(input: &str) -> bool { - let trimmed = input.trim_start(); - if !(trimmed.starts_with('{') || trimmed.starts_with('[')) { - return false; - } - if serde_json::from_str::(trimmed).is_ok() { - return true; - } - trimmed - .lines() - .find(|line| !line.trim().is_empty()) - .is_some_and(|line| serde_json::from_str::(line).is_ok()) -} - -fn skipped_import_warning(path: &Path, reason: &str) -> String { - format!( - "warning: skipped import source {}: {reason}", - path.display() - ) -} - -fn empty_auto_directory_import_error(directory_input: bool) -> anyhow::Error { - cli_boundary_error( - BoundaryCode::InvalidRequest, - if directory_input { - "import directory contains no detectable trajectory files" - } else { - "cannot detect import format; pass --format explicitly" - }, - ) -} - -fn import_source_name(format: ExchangeFormat) -> &'static str { - match format { - ExchangeFormat::Atif => "trajectories.atif.json", - ExchangeFormat::Actf => "trajectories.actf.json", - ExchangeFormat::OpenaiMessages => "session_steps.json", - ExchangeFormat::Storyline => "trajectories.storyline.json", - ExchangeFormat::Codex => "session.codex.jsonl", - ExchangeFormat::ClaudeCode => "session.claude-code.jsonl", - ExchangeFormat::CompactJsonl => "compact.jsonl", - _ => unreachable!("unsupported import format was rejected"), - } -} - -fn single_import_source_path( - format: ExchangeFormat, - output_format: ImportOutputFormat, - input_path: Option<&Path>, -) -> PathBuf { - if format == ExchangeFormat::Atif && output_format == ImportOutputFormat::Preserve { - let line_extension = input_path - .and_then(Path::extension) - .and_then(|extension| extension.to_str()) - .map(str::to_ascii_lowercase) - .filter(|extension| matches!(extension.as_str(), "jsonl" | "ndjson")); - if let Some(extension) = line_extension { - return PathBuf::from(format!("trajectories.atif.{extension}")); - } - } - PathBuf::from(import_source_name(format)) -} - -fn import_input_issue_message(issue: &InputIssue, source_path: &Path) -> String { - match issue.location() { - Some(location) => format!("{} {location}: {}", source_path.display(), issue.message()), - None => format!("{}: {}", source_path.display(), issue.message()), - } -} - -fn validate_import_storylines(storylines: &[StorylineDocument]) -> Result { - Ok(storylines.len()) -} - -pub(super) async fn validate_import_source(format: ExchangeFormat, path: &Path) -> Result { - let format = exchange_document_format(format) - .context("supported import format must map to a physical document format")?; - let source = open_document(format, path).await?; - let mut seen = HashSet::new(); - let mut document_count = 0usize; - source - .for_each_storyline(|story| { - let document_id = story.document_id(); - if !seen.insert(document_id.to_string()) { - return Err(cli_boundary_error( - BoundaryCode::InvalidRequest, - "import contains duplicate document_id", - )); - } - document_count = document_count - .checked_add(1) - .ok_or_else(|| anyhow::anyhow!("import document count overflow"))?; - Ok(()) - }) - .await?; - Ok(document_count) -} - -struct StagingPathGuard { - path: Option, -} - -impl StagingPathGuard { - fn new(path: PathBuf) -> Self { - Self { path: Some(path) } - } - - fn disarm(&mut self) { - self.path = None; - } -} - -impl Drop for StagingPathGuard { - fn drop(&mut self) { - if let Some(path) = &self.path { - let _ = std::fs::remove_dir_all(path); - } - } -} - -async fn publish_staged_dataset( - staging: &Path, - output: &Path, - replace_existing: bool, - progress: Option<&mut ImportProgress>, -) -> Result<()> { - let parent = output - .parent() - .context("Dataset output must have a parent directory")?; - if !replace_existing { - rename_noreplace(staging, output) - .with_context(|| format!("publish new Dataset {}", output.display()))?; - sync_dataset_parent(parent)?; - return Ok(()); - } - - let backup = parent.join(format!( - ".pchronicle-replace-{}-{}", - output - .file_name() - .map(|name| name.to_string_lossy()) - .unwrap_or_else(|| std::borrow::Cow::Borrowed("dataset")), - uuid::Uuid::new_v4().simple() - )); - rename_noreplace(output, &backup) - .with_context(|| format!("move existing Dataset to {}", backup.display()))?; - if let Err(error) = sync_dataset_parent(parent) { - return Err(rollback_replacement(output, &backup, error)); - } - if let Err(error) = rename_noreplace(staging, output) - .with_context(|| format!("publish replacement Dataset {}", output.display())) - { - return Err(rollback_replacement(output, &backup, error)); - } - sync_dataset_parent(parent).with_context(|| { - format!( - "sync replacement Dataset parent {}; old Dataset remains at {}", - parent.display(), - backup.display() - ) - })?; - let backup_location = DatasetLocation::parse( - backup - .to_str() - .context("replaced Dataset backup path is not valid UTF-8")?, - )?; - if let Some(progress) = progress { - backup_location - .remove_all_with_progress(|deleted, total, path| { - progress.note_deleted(deleted, total, path) - }) - .await - .with_context(|| format!("delete replaced Dataset backup {}", backup.display()))?; - progress.finish()?; - } else { - backup_location - .remove_all() - .await - .with_context(|| format!("delete replaced Dataset backup {}", backup.display()))?; - } - sync_dataset_parent(parent)?; - Ok(()) -} - -fn rollback_replacement(output: &Path, backup: &Path, error: anyhow::Error) -> anyhow::Error { - match rename_noreplace(backup, output) { - Ok(()) => error, - Err(rollback_error) => anyhow!( - "{error}; failed to restore old Dataset from {} to {}: {rollback_error}", - backup.display(), - output.display() - ), - } -} - -fn sync_dataset_parent(parent: &Path) -> Result<()> { - std::fs::File::open(parent) - .and_then(|directory| directory.sync_all()) - .with_context(|| format!("sync Dataset parent {}", parent.display()))?; - Ok(()) -} - -#[cfg(any(target_os = "linux", target_os = "macos"))] -pub(super) fn rename_noreplace(from: &Path, to: &Path) -> std::io::Result<()> { - use std::os::unix::ffi::OsStrExt; - - let from = CString::new(from.as_os_str().as_bytes())?; - let to = CString::new(to.as_os_str().as_bytes())?; - #[cfg(target_os = "linux")] - // SAFETY: both pointers come from live CString values and are NUL-terminated. - // Call SYS_renameat2 directly so the binary still links on manylinux2014 - // (glibc 2.17). The renameat2() wrapper only exists in glibc 2.28+. - let result = unsafe { - libc::syscall( - libc::SYS_renameat2, - libc::AT_FDCWD, - from.as_ptr(), - libc::AT_FDCWD, - to.as_ptr(), - libc::RENAME_NOREPLACE, - ) - }; - #[cfg(target_os = "macos")] - // SAFETY: both pointers come from live CString values and are NUL-terminated. - let result = unsafe { libc::renamex_np(from.as_ptr(), to.as_ptr(), libc::RENAME_EXCL) }; - if result == 0 { - Ok(()) - } else { - Err(std::io::Error::last_os_error()) - } -} - -#[cfg(not(any(target_os = "linux", target_os = "macos")))] -pub(super) fn rename_noreplace(_from: &Path, _to: &Path) -> std::io::Result<()> { - Err(std::io::Error::new( - std::io::ErrorKind::Unsupported, - "atomic create-only Dataset publish is unsupported on this platform", - )) -} diff --git a/crates/persisting-pchronicle-cli/src/exchange/decode.rs b/crates/persisting-pchronicle-cli/src/exchange/decode.rs new file mode 100644 index 00000000..2c4c75ff --- /dev/null +++ b/crates/persisting-pchronicle-cli/src/exchange/decode.rs @@ -0,0 +1,917 @@ +//! Import candidates, decode, format resolution, and validation. + +use super::super::*; +use super::progress::{CliProgress, StageId}; +use anyhow::{Context, Result}; +use persisting_pchronicle::document::{ + DocumentFormat, InputIssue, InputIssueKind, decode_json_storylines, detect_format, + open_document, +}; +use persisting_pchronicle::model::StorylineDocument; +use std::collections::HashSet; +use std::io::Read; +use std::path::{Path, PathBuf}; + +#[derive(Debug, Clone)] +pub(crate) struct ImportFileCandidate { + pub(crate) path: PathBuf, + pub(crate) relative_path: PathBuf, + pub(crate) output_relative_path: Option, + /// Prefetched bytes (tests / rare callers). Normal imports leave this empty + /// and read local paths or object-store keys on demand. + pub(crate) content: Option>, + /// Object-store Dataset root URI; when set, bytes are fetched lazily. + pub(crate) remote_root: Option, + /// Size from discovery (`stat` / object metadata) for progress totals. + pub(crate) size_hint: u64, +} + +#[derive(Debug)] +pub(crate) struct ImportedSource { + pub(crate) source_path: String, + pub(crate) format: DocumentFormat, + pub(crate) trajectories: usize, + pub(crate) input_bytes: usize, +} + +pub(crate) fn exchange_document_format(format: ExchangeFormat) -> Option { + match format { + ExchangeFormat::Atif => Some(DocumentFormat::Atif), + ExchangeFormat::Actf => Some(DocumentFormat::Actf), + ExchangeFormat::OpenaiMessages => Some(DocumentFormat::OpenaiMsg), + ExchangeFormat::Storyline => Some(DocumentFormat::Storyline), + ExchangeFormat::Codex => Some(DocumentFormat::Codex), + ExchangeFormat::ClaudeCode => Some(DocumentFormat::ClaudeCode), + ExchangeFormat::CompactJsonl | ExchangeFormat::Auto => None, + } +} + +pub(crate) fn apply_duplicate_document_policy( + storyline: &mut StorylineDocument, + seen_document_ids: &mut HashSet, + duplicate_policy: DuplicateIdPolicy, +) -> Option { + let original = storyline.document_id().to_string(); + match duplicate_policy { + DuplicateIdPolicy::Suffix => uniquify_storyline_document_id(storyline, seen_document_ids) + .map(|(original, renamed)| { + format!("warning: duplicate document_id '{original}' renamed to '{renamed}'") + }), + DuplicateIdPolicy::Skip => { + if !seen_document_ids.insert(original.clone()) { + Some(format!( + "warning: duplicate document_id '{original}' skipped" + )) + } else { + None + } + } + } +} + +pub(crate) fn collect_import_candidates(input: &Path) -> Result<(bool, Vec)> { + let metadata = std::fs::symlink_metadata(input) + .with_context(|| format!("inspect import input {}", input.display()))?; + let explicit_file = if metadata.file_type().is_symlink() { + std::fs::metadata(input) + .with_context(|| format!("inspect import input target {}", input.display()))? + .is_file() + } else { + metadata.is_file() + }; + if explicit_file { + let relative_path = input + .file_name() + .map(PathBuf::from) + .context("import input file has no filename")?; + let size_hint = metadata.len(); + return Ok(( + false, + vec![ImportFileCandidate { + path: input.to_path_buf(), + relative_path, + output_relative_path: None, + content: None, + remote_root: None, + size_hint, + }], + )); + } + anyhow::ensure!( + metadata.is_dir(), + "import input must be a regular file or directory" + ); + + let paths = collect_visible_json_files(input)?; + let mut candidates = Vec::with_capacity(paths.len()); + for path in paths { + let relative_path = path + .strip_prefix(input) + .context("derive Dataset-relative import source path")? + .to_path_buf(); + let size_hint = std::fs::metadata(&path).map(|meta| meta.len()).unwrap_or(0); + candidates.push(ImportFileCandidate { + path, + output_relative_path: Some(relative_path.clone()), + relative_path, + content: None, + remote_root: None, + size_hint, + }); + } + candidates.sort_by(|left, right| left.relative_path.cmp(&right.relative_path)); + if candidates.is_empty() { + return Err(cli_boundary_error( + BoundaryCode::InvalidRequest, + "import directory contains no .json, .jsonl, or .ndjson files", + )); + } + Ok((true, candidates)) +} + +/// Recursively collect absolute paths of visible `.json` / `.jsonl` / `.ndjson` +/// files under `root`. Shared by `import` and `sync`; not Catalog Directory +/// discovery (which is one-level and skips loose files). +pub(crate) fn collect_visible_json_files(root: &Path) -> Result> { + let mut pending = vec![root.to_path_buf()]; + let mut files = Vec::new(); + while let Some(directory) = pending.pop() { + let mut entries = std::fs::read_dir(&directory) + .with_context(|| format!("read directory {}", directory.display()))? + .collect::>>()?; + entries.sort_by_key(std::fs::DirEntry::path); + for entry in entries { + let file_type = entry.file_type()?; + if file_type.is_symlink() { + continue; + } + let path = entry.path(); + if file_type.is_dir() { + pending.push(path); + } else if file_type.is_file() && is_visible_json_file(&path) { + let relative = path + .strip_prefix(root) + .unwrap_or(path.as_path()) + .to_string_lossy() + .replace('\\', "/"); + if relative.split('/').any(|part| part == "_meta") { + continue; + } + files.push(path); + } + } + } + files.sort(); + Ok(files) +} + +pub(crate) fn is_visible_json_file(path: &Path) -> bool { + path.extension() + .and_then(|extension| extension.to_str()) + .is_some_and(|extension| { + matches!( + extension.to_ascii_lowercase().as_str(), + "json" | "jsonl" | "ndjson" + ) + }) +} + +pub(crate) async fn load_import_candidate_bytes( + candidate: &ImportFileCandidate, + max_input_bytes: usize, + label: &str, +) -> Result> { + if let Some(content) = &candidate.content { + anyhow::ensure!( + content.len() <= max_input_bytes, + "{label} exceeds max_input_bytes limit of {max_input_bytes}" + ); + return Ok(content.clone()); + } + if let Some(remote_root) = &candidate.remote_root { + let key = candidate.relative_path.to_string_lossy().replace('\\', "/"); + let location = DatasetLocation::parse(remote_root)?; + let bytes = location + .read_relative_bytes(&key) + .await + .with_context(|| format!("read import object {key} under {remote_root}"))?; + anyhow::ensure!( + bytes.len() <= max_input_bytes, + "{label} exceeds max_input_bytes limit of {max_input_bytes}" + ); + return Ok(bytes); + } + let file = std::fs::File::open(&candidate.path).with_context(|| format!("open {label}"))?; + read_bounded(file, max_input_bytes, label) +} + +pub(crate) fn scope_import_source_error(error: anyhow::Error, source_path: &Path) -> anyhow::Error { + if let Some(boundary) = error.downcast_ref::() { + return cli_boundary_error( + boundary.code, + format!("{}: {}", source_path.display(), boundary.message), + ); + } + error.context(format!("import source {}", source_path.display())) +} + +pub(crate) struct DecodedImportSource { + pub(crate) diagnostic_path: PathBuf, + pub(crate) metadata: ImportedSource, + pub(crate) storylines: Vec, +} + +pub(crate) enum DecodeImportOutcome { + Imported(DecodedImportSource), + Skipped { path: PathBuf, reason: String }, +} + +#[derive(Debug, PartialEq, Eq)] +pub(crate) enum ImportFormatResolution { + Format(ExchangeFormat), + Skip(String), +} + +pub(crate) enum StorylineImportInputs<'a> { + Stdin(Option<&'a mut dyn Read>), +} + +pub(crate) struct StorylineImportIterator<'a> { + pub(crate) requested_format: ExchangeFormat, + pub(crate) suggested_format: Option, + pub(crate) max_input_bytes: usize, + pub(crate) progress: &'a mut CliProgress, + pub(crate) inputs: StorylineImportInputs<'a>, + pub(crate) current: std::vec::IntoIter, + pub(crate) imported_sources: Vec, + pub(crate) unknown_field_warnings: persisting_pchronicle::model::UnknownFieldImportWarnings, + pub(crate) skipped_warnings: Vec, + pub(crate) seen_document_ids: HashSet, + pub(crate) duplicate_policy: DuplicateIdPolicy, + pub(crate) failed: bool, +} + +impl<'a> StorylineImportIterator<'a> { + pub(crate) fn stdin( + requested_format: ExchangeFormat, + suggested_format: Option, + max_input_bytes: usize, + stdin: &'a mut dyn Read, + progress: &'a mut CliProgress, + seen_document_ids: HashSet, + duplicate_policy: DuplicateIdPolicy, + ) -> Self { + Self { + requested_format, + suggested_format, + max_input_bytes, + progress, + inputs: StorylineImportInputs::Stdin(Some(stdin)), + current: Vec::new().into_iter(), + imported_sources: Vec::new(), + unknown_field_warnings: + persisting_pchronicle::model::UnknownFieldImportWarnings::default(), + skipped_warnings: Vec::new(), + seen_document_ids, + duplicate_policy, + failed: false, + } + } + + pub(crate) async fn decode_next_source(&mut self) -> Result> { + loop { + let outcome = match &mut self.inputs { + StorylineImportInputs::Stdin(stdin) => { + let Some(stdin) = stdin.take() else { + return Ok(None); + }; + self.progress.stage(StageId::Fetch).set_current("stdin"); + let input = read_bounded(stdin, self.max_input_bytes, "stdin")?; + self.progress.note_fetched("stdin", input.len() as u64)?; + self.progress.stage(StageId::Parse).set_current("stdin"); + decode_import_source( + self.requested_format, + self.suggested_format, + ImportOutputFormat::Storyline, + None, + None, + None, + &input, + &mut self.unknown_field_warnings, + )? + } + }; + match outcome { + DecodeImportOutcome::Imported(decoded) => { + self.progress.note_parsed( + &decoded.diagnostic_path.to_string_lossy(), + decoded.metadata.input_bytes as u64, + )?; + return Ok(Some(decoded)); + } + DecodeImportOutcome::Skipped { path, reason } => { + self.progress.note_parsed(&path.to_string_lossy(), 0)?; + self.skipped_warnings + .push(skipped_import_warning(&path, &reason)); + } + } + } + } + + pub(crate) fn into_result_parts( + self, + ) -> ( + Vec, + persisting_pchronicle::model::UnknownFieldImportWarnings, + Vec, + &'a mut CliProgress, + ) { + ( + self.imported_sources, + self.unknown_field_warnings, + self.skipped_warnings, + self.progress, + ) + } + + pub(crate) async fn next_document(&mut self) -> Option> { + loop { + if let Some(mut storyline) = self.current.next() { + let original = storyline.document_id().to_string(); + match self.duplicate_policy { + DuplicateIdPolicy::Suffix => { + if let Some((original, renamed)) = uniquify_storyline_document_id( + &mut storyline, + &mut self.seen_document_ids, + ) { + self.skipped_warnings.push(format!( + "warning: duplicate document_id '{original}' renamed to '{renamed}'" + )); + } + } + DuplicateIdPolicy::Skip => { + if !self.seen_document_ids.insert(original.clone()) { + self.skipped_warnings.push(format!( + "warning: duplicate document_id '{original}' skipped" + )); + continue; + } + } + } + let metadata = self + .imported_sources + .last_mut() + .expect("decoded Storyline has source metadata"); + metadata.trajectories = metadata + .trajectories + .checked_add(1) + .expect("import trajectory count overflow"); + return Some(Ok(storyline)); + } + if self.failed { + return None; + } + match self.decode_next_source().await { + Ok(Some(decoded)) => { + let mut metadata = decoded.metadata; + metadata.trajectories = 0; + self.imported_sources.push(metadata); + self.current = decoded.storylines.into_iter(); + } + Ok(None) => return None, + Err(error) => { + self.failed = true; + return Some(Err(error)); + } + } + } + } +} + +pub(crate) fn uniquify_storyline_document_id( + story: &mut StorylineDocument, + seen: &mut HashSet, +) -> Option<(String, String)> { + let preferred = story.document_id().to_string(); + if seen.insert(preferred.clone()) { + return None; + } + let mut suffix = 1u64; + let renamed = loop { + let candidate = format!("{preferred}#{suffix}"); + if seen.insert(candidate.clone()) { + break candidate; + } + suffix = suffix + .checked_add(1) + .expect("document_id disambiguation suffix overflow"); + }; + if story + .trajectory_id + .as_deref() + .is_some_and(|id| !id.is_empty()) + { + story.trajectory_id = Some(renamed.clone()); + } else { + story.session_id = renamed.clone(); + } + Some((preferred, renamed)) +} + +#[allow(clippy::too_many_arguments)] +pub(crate) fn decode_import_source( + requested_format: ExchangeFormat, + suggested_format: Option, + output_format: ImportOutputFormat, + input_path: Option<&Path>, + decode_relative_path: Option<&Path>, + logical_source_path: Option<&Path>, + input: &[u8], + unknown_field_warnings: &mut persisting_pchronicle::model::UnknownFieldImportWarnings, +) -> Result { + let diagnostic_path = decode_relative_path + .unwrap_or_else(|| Path::new("stdin")) + .to_path_buf(); + let text = std::str::from_utf8(input).map_err(|error| { + cli_boundary_error( + BoundaryCode::InvalidRequest, + format!("{} is not UTF-8: {error}", diagnostic_path.display()), + ) + })?; + let allow_skip = requested_format == ExchangeFormat::Auto && logical_source_path.is_some(); + let format = match resolve_import_format( + requested_format, + suggested_format, + input_path, + text, + allow_skip, + ) + .map_err(|error| { + if logical_source_path.is_some() { + scope_import_source_error(error, &diagnostic_path) + } else { + error + } + })? { + ImportFormatResolution::Format(format) => format, + ImportFormatResolution::Skip(reason) => { + return Ok(DecodeImportOutcome::Skipped { + path: diagnostic_path, + reason, + }); + } + }; + let document_format = exchange_document_format(format) + .context("supported import format must map to a physical document format")?; + let source_path = logical_source_path + .map(PathBuf::from) + .unwrap_or_else(|| single_import_source_path(format, output_format, input_path)); + let decode_relative_path = decode_relative_path.unwrap_or(&source_path); + let storylines = + decode_json_storylines(document_format, text, decode_relative_path).map_err(|issue| { + let code = match issue.kind() { + InputIssueKind::Invalid => BoundaryCode::InvalidRequest, + InputIssueKind::Unsupported => BoundaryCode::Unsupported, + }; + cli_boundary_error( + code, + import_input_issue_message(&issue, decode_relative_path), + ) + }); + let storylines = match storylines { + Ok(storylines) => storylines, + Err(error) if allow_skip => { + return Ok(DecodeImportOutcome::Skipped { + path: diagnostic_path, + reason: error.to_string(), + }); + } + Err(error) => return Err(error), + }; + unknown_field_warnings + .observe_storylines(&storylines) + .map_err(|issue| { + cli_boundary_error( + BoundaryCode::InvalidRequest, + import_input_issue_message(&issue, decode_relative_path), + ) + })?; + + let metadata = ImportedSource { + source_path: source_path + .to_str() + .context("Dataset-relative import Source path is not UTF-8")? + .to_owned(), + format: document_format, + trajectories: storylines.len(), + input_bytes: input.len(), + }; + Ok(DecodeImportOutcome::Imported(DecodedImportSource { + diagnostic_path, + metadata, + storylines, + })) +} + +#[allow(clippy::too_many_arguments)] +pub(crate) fn stage_preserved_import_source( + requested_format: ExchangeFormat, + suggested_format: Option, + input_path: Option<&Path>, + decode_relative_path: Option<&Path>, + logical_source_path: Option<&Path>, + input: &[u8], + staging_root: &Path, + unknown_field_warnings: &mut persisting_pchronicle::model::UnknownFieldImportWarnings, + skipped_warnings: &mut Vec, +) -> Result> { + let decoded = match decode_import_source( + requested_format, + suggested_format, + ImportOutputFormat::Preserve, + input_path, + decode_relative_path, + logical_source_path, + input, + unknown_field_warnings, + )? { + DecodeImportOutcome::Imported(decoded) => decoded, + DecodeImportOutcome::Skipped { path, reason } => { + skipped_warnings.push(skipped_import_warning(&path, &reason)); + return Ok(None); + } + }; + validate_import_storylines(&decoded.storylines).map_err(|error| { + if logical_source_path.is_some() { + scope_import_source_error(error, &decoded.diagnostic_path) + } else { + error + } + })?; + + let staged_source = staging_root.join(&decoded.metadata.source_path); + let staged_parent = staged_source + .parent() + .context("staged import Source has no parent")?; + std::fs::create_dir_all(staged_parent) + .with_context(|| format!("create staged Source parent {}", staged_parent.display()))?; + let mut file = std::fs::OpenOptions::new() + .write(true) + .create_new(true) + .open(&staged_source) + .with_context(|| format!("create staged Source {}", decoded.metadata.source_path))?; + file.write_all(input) + .with_context(|| format!("write staged Source {}", decoded.metadata.source_path))?; + file.sync_all() + .with_context(|| format!("sync staged Source {}", decoded.metadata.source_path))?; + Ok(Some(decoded.metadata)) +} + +pub(crate) fn read_bounded( + mut reader: impl Read, + max_bytes: usize, + label: &str, +) -> Result> { + let mut input = Vec::new(); + if max_bytes == usize::MAX { + reader + .read_to_end(&mut input) + .with_context(|| format!("read {label}"))?; + } else { + let limit = u64::try_from(max_bytes) + .ok() + .and_then(|limit| limit.checked_add(1)) + .ok_or_else(|| { + cli_boundary_error( + BoundaryCode::InvalidRequest, + "--max-input-bytes is too large", + ) + })?; + reader + .by_ref() + .take(limit) + .read_to_end(&mut input) + .with_context(|| format!("read {label}"))?; + if input.len() > max_bytes { + return Err(cli_boundary_error( + BoundaryCode::ResourceExhausted, + format!("{label} exceeds max_input_bytes limit of {max_bytes}"), + )); + } + } + if input.is_empty() { + return Err(cli_boundary_error( + BoundaryCode::InvalidRequest, + format!("{label} is empty"), + )); + } + Ok(input) +} + +pub(crate) fn resolve_import_format( + requested: ExchangeFormat, + suggested: Option, + input_path: Option<&Path>, + input: &str, + allow_skip: bool, +) -> Result { + let format = match requested { + ExchangeFormat::Auto => match detect_format(input_path, Some(input))? { + Some(DocumentFormat::Atif) => ExchangeFormat::Atif, + Some(DocumentFormat::Actf) => ExchangeFormat::Actf, + Some(DocumentFormat::OpenaiMsg) => ExchangeFormat::OpenaiMessages, + Some(DocumentFormat::Storyline) => ExchangeFormat::Storyline, + Some(DocumentFormat::Codex) => ExchangeFormat::Codex, + Some(DocumentFormat::ClaudeCode) => ExchangeFormat::ClaudeCode, + Some(format) if allow_skip => { + return Ok(ImportFormatResolution::Skip(format!( + "detected import format '{format}' is not a queryable JSON format" + ))); + } + Some(format) => { + return Err(cli_boundary_error( + BoundaryCode::Unsupported, + format!("detected import format '{format}' is not a queryable JSON format"), + )); + } + None => { + if let Some(hint) = suggested.filter(|format| *format != ExchangeFormat::Auto) + && suggested_format_compatible(hint, input_path, input) + { + hint + } else if allow_skip && looks_like_json_document(input) { + return Ok(ImportFormatResolution::Skip( + "cannot detect import format".into(), + )); + } else { + return Err(cli_boundary_error( + BoundaryCode::InvalidRequest, + if suggested.is_some() { + "cannot detect import format; --suggested-format did not match this file (pass --format to force)" + } else { + "cannot detect import format; pass --format explicitly or --suggested-format to assist" + }, + )); + } + } + }, + ExchangeFormat::Atif => ExchangeFormat::Atif, + ExchangeFormat::Actf => ExchangeFormat::Actf, + ExchangeFormat::OpenaiMessages => ExchangeFormat::OpenaiMessages, + ExchangeFormat::Storyline => ExchangeFormat::Storyline, + ExchangeFormat::Codex => ExchangeFormat::Codex, + ExchangeFormat::ClaudeCode => ExchangeFormat::ClaudeCode, + ExchangeFormat::CompactJsonl => ExchangeFormat::CompactJsonl, + }; + if !matches!( + format, + ExchangeFormat::Atif + | ExchangeFormat::Actf + | ExchangeFormat::OpenaiMessages + | ExchangeFormat::Storyline + | ExchangeFormat::Codex + | ExchangeFormat::ClaudeCode + | ExchangeFormat::CompactJsonl + ) { + return Err(cli_boundary_error( + BoundaryCode::Unsupported, + format!( + "import format '{format}' is not supported by the first queryable import increment" + ), + )); + } + Ok(ImportFormatResolution::Format(format)) +} + +/// Weak compatibility check used only with `--suggested-format`. +/// +/// Stronger than blind force, weaker than auto fingerprint: the file must still +/// look like the suggested family before we accept the hint. +pub(crate) fn suggested_format_compatible( + suggested: ExchangeFormat, + input_path: Option<&Path>, + input: &str, +) -> bool { + match suggested { + ExchangeFormat::Actf => weakly_compatible_actf(input), + ExchangeFormat::Atif => weakly_compatible_json_keys(input, &["agent", "steps"]), + ExchangeFormat::Storyline => { + weakly_compatible_json_keys(input, &["schema_version", "session", "turns"]) + || weakly_compatible_json_keys(input, &["schema_version", "session", "agent"]) + } + ExchangeFormat::OpenaiMessages => { + weakly_compatible_json_keys(input, &["session_id", "messages"]) + || weakly_compatible_json_keys(input, &["messages", "step_id"]) + } + ExchangeFormat::Codex | ExchangeFormat::ClaudeCode => { + looks_like_json_document(input) + && input_path.is_some_and(|path| { + path.extension() + .and_then(|ext| ext.to_str()) + .is_some_and(|ext| { + matches!(ext.to_ascii_lowercase().as_str(), "jsonl" | "ndjson") + }) + }) + } + ExchangeFormat::CompactJsonl | ExchangeFormat::Auto => false, + } +} + +fn weakly_compatible_actf(input: &str) -> bool { + // Assist only: root shape, not trajectory schema fingerprint. + // Avoid full JSON parse so Python NaN dumps still qualify. + let trimmed = input.trim_start(); + (trimmed.starts_with('{') || trimmed.starts_with('[')) + && trimmed.contains("\"task_id\"") + && trimmed.contains("\"attempts\"") +} + +fn weakly_compatible_json_keys(input: &str, required: &[&str]) -> bool { + let trimmed = input.trim_start(); + let Ok(value) = serde_json::from_str::(trimmed) else { + return false; + }; + let Some(object) = value.as_object() else { + return false; + }; + required.iter().all(|key| object.contains_key(*key)) +} + +pub(crate) fn looks_like_json_document(input: &str) -> bool { + let trimmed = input.trim_start(); + if !(trimmed.starts_with('{') || trimmed.starts_with('[')) { + return false; + } + if serde_json::from_str::(trimmed).is_ok() { + return true; + } + trimmed + .lines() + .find(|line| !line.trim().is_empty()) + .is_some_and(|line| serde_json::from_str::(line).is_ok()) +} + +pub(crate) fn skipped_import_warning(path: &Path, reason: &str) -> String { + format!( + "warning: skipped import source {}: {reason}", + path.display() + ) +} + +pub(crate) fn empty_auto_directory_import_error(directory_input: bool) -> anyhow::Error { + cli_boundary_error( + BoundaryCode::InvalidRequest, + if directory_input { + "import directory contains no detectable trajectory files" + } else { + "cannot detect import format; pass --format explicitly" + }, + ) +} + +pub(crate) fn import_source_name(format: ExchangeFormat) -> &'static str { + match format { + ExchangeFormat::Atif => "trajectories.atif.json", + ExchangeFormat::Actf => "trajectories.actf.json", + ExchangeFormat::OpenaiMessages => "session_steps.json", + ExchangeFormat::Storyline => "trajectories.storyline.json", + ExchangeFormat::Codex => "session.codex.jsonl", + ExchangeFormat::ClaudeCode => "session.claude-code.jsonl", + ExchangeFormat::CompactJsonl => "compact.jsonl", + _ => unreachable!("unsupported import format was rejected"), + } +} + +pub(crate) fn single_import_source_path( + format: ExchangeFormat, + output_format: ImportOutputFormat, + input_path: Option<&Path>, +) -> PathBuf { + if format == ExchangeFormat::Atif && output_format == ImportOutputFormat::Preserve { + let line_extension = input_path + .and_then(Path::extension) + .and_then(|extension| extension.to_str()) + .map(str::to_ascii_lowercase) + .filter(|extension| matches!(extension.as_str(), "jsonl" | "ndjson")); + if let Some(extension) = line_extension { + return PathBuf::from(format!("trajectories.atif.{extension}")); + } + } + PathBuf::from(import_source_name(format)) +} + +pub(crate) fn import_input_issue_message(issue: &InputIssue, source_path: &Path) -> String { + match issue.location() { + Some(location) => format!("{} {location}: {}", source_path.display(), issue.message()), + None => format!("{}: {}", source_path.display(), issue.message()), + } +} + +pub(crate) fn validate_import_storylines(storylines: &[StorylineDocument]) -> Result { + Ok(storylines.len()) +} + +pub(crate) async fn validate_import_source(format: ExchangeFormat, path: &Path) -> Result { + let format = exchange_document_format(format) + .context("supported import format must map to a physical document format")?; + let source = open_document(format, path).await?; + let mut seen = HashSet::new(); + let mut document_count = 0usize; + source + .for_each_storyline(|story| { + let document_id = story.document_id(); + if !seen.insert(document_id.to_string()) { + return Err(cli_boundary_error( + BoundaryCode::InvalidRequest, + "import contains duplicate document_id", + )); + } + document_count = document_count + .checked_add(1) + .ok_or_else(|| anyhow::anyhow!("import document count overflow"))?; + Ok(()) + }) + .await?; + Ok(document_count) +} + +#[cfg(test)] +mod tests { + use super::*; + use std::path::Path; + + #[test] + fn visible_json_extensions() { + assert!(is_visible_json_file(Path::new("a.json"))); + assert!(is_visible_json_file(Path::new("a.JSONL"))); + assert!(is_visible_json_file(Path::new("a.ndjson"))); + assert!(!is_visible_json_file(Path::new("a.txt"))); + } + + #[test] + fn looks_like_json_document_smoke() { + assert!(looks_like_json_document(r#"{"a":1}"#)); + assert!(looks_like_json_document("\n[1,2]\n")); + assert!(!looks_like_json_document("not json")); + } + + #[test] + fn exchange_document_format_maps_known() { + assert_eq!( + exchange_document_format(ExchangeFormat::Atif), + Some(DocumentFormat::Atif) + ); + assert!(exchange_document_format(ExchangeFormat::Auto).is_none()); + } + + #[test] + fn suggested_actf_assists_when_auto_fingerprint_misses() { + // Object trajectory with steps but no ACTF_ schema_version: auto stays None. + let input = r#"{ + "task_id":"travel-planning", + "attempts":{"1":{ + "correct":false, + "trajectory":{ + "steps":[], + "started_at":"2026-06-17T07:26:27Z", + "finished_at":"2026-06-17T07:26:28Z" + } + }} + }"#; + let err = resolve_import_format(ExchangeFormat::Auto, None, None, input, false) + .unwrap_err() + .to_string(); + assert!(err.contains("cannot detect import format")); + assert_eq!( + resolve_import_format( + ExchangeFormat::Auto, + Some(ExchangeFormat::Actf), + None, + input, + false + ) + .unwrap(), + ImportFormatResolution::Format(ExchangeFormat::Actf) + ); + } + + #[test] + fn suggested_actf_rejects_incompatible_shape() { + let input = r#"{"error":"boom","message":"no pe"}"#; + assert!(!suggested_format_compatible( + ExchangeFormat::Actf, + None, + input + )); + let err = resolve_import_format( + ExchangeFormat::Auto, + Some(ExchangeFormat::Actf), + None, + input, + false, + ) + .unwrap_err() + .to_string(); + assert!(err.contains("--suggested-format did not match")); + } +} diff --git a/crates/persisting-pchronicle-cli/src/exchange/drop.rs b/crates/persisting-pchronicle-cli/src/exchange/drop.rs new file mode 100644 index 00000000..a9feb146 --- /dev/null +++ b/crates/persisting-pchronicle-cli/src/exchange/drop.rs @@ -0,0 +1,99 @@ +use super::super::*; +use anyhow::{Context, Result}; +use serde::Serialize; +use std::io::{Read, Write}; +use std::path::Path; + +#[derive(Serialize)] +struct DropResponse { + dataset_uri: String, + dropped: bool, +} + +pub(crate) async fn run_drop( + args: DropArgs, + settings_override: Option<&Path>, + stdin_is_terminal: bool, + stdin: &mut dyn Read, + stdout: &mut dyn Write, + stderr: &mut dyn Write, +) -> Result<()> { + let dataset_uri = expand_dataset_reference(&args.dataset_uri, settings_override, false)?; + let mut location = DatasetLocation::parse(&dataset_uri)?; + if !location.exists().await? { + return Err(cli_boundary_error( + BoundaryCode::NotFound, + format!("Dataset does not exist: {}", location.as_str()), + )); + } + if location.local_path().is_some() { + location = location.into_existing()?; + } + confirm_destructive_dataset( + "drop", + location.as_str(), + args.yes, + stdin_is_terminal, + stdin, + stderr, + )?; + location.remove_all().await?; + let response = DropResponse { + dataset_uri: location.as_str().to_string(), + dropped: true, + }; + serde_json::to_writer_pretty(&mut *stdout, &response).context("encode pChronicle drop JSON")?; + writeln!(stdout).context("write pChronicle drop JSON")?; + writeln!( + stderr, + "dataset_uri={} status=dropped", + response.dataset_uri + ) + .context("write pChronicle drop metadata")?; + Ok(()) +} + +pub(crate) fn confirm_destructive_dataset( + action: &str, + dataset_uri: &str, + yes: bool, + stdin_is_terminal: bool, + stdin: &mut dyn Read, + stderr: &mut dyn Write, +) -> Result<()> { + if yes { + return Ok(()); + } + if !stdin_is_terminal { + return Err(cli_boundary_error( + BoundaryCode::InvalidRequest, + format!("{action} requires confirmation; rerun with --yes"), + )); + } + write!( + stderr, + "Permanently {action} Dataset '{dataset_uri}'? [y/N] " + ) + .context("write Dataset confirmation prompt")?; + stderr + .flush() + .context("flush Dataset confirmation prompt")?; + let mut answer = Vec::new(); + let mut byte = [0u8; 1]; + while answer.len() <= 16 && stdin.read(&mut byte).context("read Dataset confirmation")? == 1 { + if byte[0] == b'\n' { + break; + } + answer.push(byte[0]); + } + let answer = std::str::from_utf8(&answer) + .context("Dataset confirmation is not UTF-8")? + .trim(); + if matches!(answer.to_ascii_lowercase().as_str(), "y" | "yes") { + return Ok(()); + } + Err(cli_boundary_error( + BoundaryCode::InvalidRequest, + format!("{action} cancelled"), + )) +} diff --git a/crates/persisting-pchronicle-cli/src/exchange/export.rs b/crates/persisting-pchronicle-cli/src/exchange/export.rs new file mode 100644 index 00000000..a5dc839a --- /dev/null +++ b/crates/persisting-pchronicle-cli/src/exchange/export.rs @@ -0,0 +1,402 @@ +//! Dataset export command. + +use super::super::*; +use super::decode::{exchange_document_format, validate_import_source}; +use anyhow::{Context, Result, bail}; +use persisting_pchronicle::document::{DocumentFormat, detect_format, encode_json_storylines}; +use persisting_pchronicle::model::StorylineDocument; +use persisting_pchronicle::storage::{ + CatalogSourceKind, CatalogSourceStatus, CatalogStorylineKey, DEFAULT_DATASET_NAME, + DatasetCatalogSnapshot, +}; +use std::io::Write; +use std::path::Path; +use std::sync::Arc; +use std::time::Duration; + +pub(crate) async fn run_export( + mut args: ExportArgs, + settings_override: Option<&Path>, + stdout: &mut dyn Write, + stderr: &mut dyn Write, +) -> Result<()> { + anyhow::ensure!( + args.max_trajectories > 0, + "--max-trajectories must be greater than zero" + ); + anyhow::ensure!( + args.max_output_bytes > 0, + "--max-output-bytes must be greater than zero" + ); + anyhow::ensure!( + args.timeout_seconds > 0, + "--timeout must be greater than zero" + ); + args.stream = args.output == "-" || args.stream; + anyhow::ensure!( + args.output == "-" || !args.stream, + "--stream requires --to -" + ); + anyhow::ensure!( + !(args.output == "-" && args.overwrite), + "--overwrite cannot be used with stdout" + ); + if let Some(source) = &args.source { + validate_source_path(source)?; + } + if let Some(run_id) = &args.run_id { + validate_find_id("--run-id", run_id)?; + } + if let Some(document_id) = &args.document_id { + validate_find_id("--document-id", document_id)?; + } + if let Some(session_id) = &args.session_id { + validate_find_id("--session-id", session_id)?; + } + if let Some(expression) = &args.r#where { + anyhow::ensure!(!expression.trim().is_empty(), "--where must not be empty"); + anyhow::ensure!( + expression.len() <= 16 * 1024, + "--where exceeds the 16384-byte limit" + ); + } + + let format = ExchangeFormat::from(args.format); + let dataset = resolve_dataset_uri(args.from.as_deref(), settings_override)?; + if args.output != "-" { + args.output = expand_dataset_reference(&args.output, settings_override, false)?; + } + if format == ExchangeFormat::CompactJsonl { + anyhow::ensure!( + args.source.is_none() + && args.run_id.is_none() + && args.document_id.is_none() + && args.session_id.is_none() + && args.r#where.is_none(), + "compact JSONL export does not support filters" + ); + anyhow::ensure!( + args.output != "-", + "compact JSONL export requires a directory output" + ); + anyhow::ensure!( + args.overwrite || !Path::new(&args.output).exists(), + "export output already exists; pass --overwrite" + ); + let rows = + persisting_pchronicle::storage::CompactJsonlStore::export_path(&dataset, &args.output) + .await?; + writeln!( + stderr, + "format=compact-jsonl rows={} output={}", + rows, args.output + )?; + return Ok(()); + } + let (_, dataset_uris, snapshot) = + discover_query_snapshot(Some(&dataset), &[], args.max_files, args.max_entries).await?; + let dataset_uri = dataset_uris + .first() + .cloned() + .context("export Dataset URI missing after discovery")?; + let snapshot = Arc::new(snapshot); + let snapshot_id = snapshot.snapshot_id().to_string(); + let deadline = Duration::from_secs(args.timeout_seconds); + let export = tokio::time::timeout( + deadline, + export_from_snapshot(&args, format, &dataset_uri, snapshot.clone()), + ) + .await + .with_context(|| { + format!( + "Dataset export timed out after {} seconds", + args.timeout_seconds + ) + })??; + ensure_export_trajectory_budget(export.trajectories, args.max_trajectories)?; + ensure_output_byte_budget(export.bytes.len(), args.max_output_bytes, "encoded export")?; + write_export_output(&args.output, &export.bytes, args.overwrite, stdout).await?; + writeln!( + stderr, + "snapshot_id={} format={} trajectories={} output_bytes={} exact={}", + snapshot_id, + format.as_str(), + export.trajectories, + export.bytes.len(), + export.exact, + ) + .context("write pChronicle export metadata")?; + Ok(()) +} + +pub(crate) struct EncodedExport { + bytes: Vec, + trajectories: usize, + exact: bool, +} + +pub(crate) async fn export_from_snapshot( + args: &ExportArgs, + format: ExchangeFormat, + dataset_uri: &str, + snapshot: Arc, +) -> Result { + if let Some(export) = exact_local_file_export(args, format, dataset_uri, &snapshot).await? { + return Ok(export); + } + anyhow::ensure!( + !args.strict, + "strict export requires an unfiltered source file already stored in the requested format" + ); + + let sql = export_address_sql(args)?; + let engine = snapshot.clone().query_engine(Default::default()).await?; + let row_limit = args + .max_trajectories + .checked_add(1) + .context("--max-trajectories is too large")?; + let mut addresses = LimitedBuffer::new(args.max_output_bytes); + let write_result = engine + .write_query_jsonl_bounded(&sql, &mut addresses, Some(row_limit)) + .await; + let address_bytes = match addresses.finish(write_result)? { + QueryOutputBudgetOutcome::Complete(bytes) => bytes, + QueryOutputBudgetOutcome::RowLimitExceeded => { + return Err(cli_boundary_error( + BoundaryCode::ResourceExhausted, + format!( + "export exceeds max_trajectories limit of {}", + args.max_trajectories + ), + )); + } + QueryOutputBudgetOutcome::ByteLimitExceeded => { + return Err(cli_boundary_error( + BoundaryCode::ResourceExhausted, + format!( + "export address selection exceeds max_output_bytes limit of {}", + args.max_output_bytes + ), + )); + } + }; + let mut addresses = address_bytes + .split(|byte| *byte == b'\n') + .filter(|line| !line.is_empty()) + .map(|line| serde_json::from_slice(line).context("decode export run address")) + .collect::>>()?; + ensure_export_trajectory_budget(addresses.len(), args.max_trajectories)?; + anyhow::ensure!(!addresses.is_empty(), "export selection matched no runs"); + addresses.sort_by(|left, right| { + (&left.source_path, &left.document_id, &left.session_id).cmp(&( + &right.source_path, + &right.document_id, + &right.session_id, + )) + }); + let mut stories = Vec::with_capacity(addresses.len()); + let mut normalized_bytes = 0usize; + for address in &addresses { + let key = CatalogStorylineKey { + dataset: DEFAULT_DATASET_NAME.into(), + file: address.source_path.clone(), + document_id: address.document_id.clone(), + session_id: address.session_id.clone(), + }; + let story = snapshot + .load_storyline(&key) + .await + .with_context(|| { + format!( + "load export run {}/{}", + address.source_path, address.session_id + ) + })? + .with_context(|| { + format!( + "export run disappeared from snapshot: {}/{}", + address.source_path, address.session_id + ) + })?; + anyhow::ensure!( + story.trajectory_id.as_deref().unwrap_or(&story.session_id) == address.document_id, + "export run document ID changed within the snapshot" + ); + anyhow::ensure!( + story.run_id == address.run_id, + "export run runtime ID changed within the snapshot" + ); + normalized_bytes = normalized_bytes.saturating_add(serde_json::to_vec(&story)?.len()); + ensure_output_byte_budget(normalized_bytes, args.max_output_bytes, "normalized export")?; + stories.push(story); + } + let bytes = encode_export(format, &stories)?; + Ok(EncodedExport { + bytes, + trajectories: stories.len(), + exact: false, + }) +} + +pub(crate) async fn exact_local_file_export( + args: &ExportArgs, + format: ExchangeFormat, + dataset_uri: &str, + snapshot: &DatasetCatalogSnapshot, +) -> Result> { + if args.document_id.is_some() + || args.run_id.is_some() + || args.session_id.is_some() + || args.r#where.is_some() + { + return Ok(None); + } + let Some(dataset) = snapshot.dataset(DEFAULT_DATASET_NAME) else { + return Ok(None); + }; + let sources = dataset + .sources + .iter() + .filter(|source| source.status == CatalogSourceStatus::Ready) + .filter(|source| { + args.source + .as_deref() + .is_none_or(|selected| selected == source.file) + }) + .collect::>(); + if sources.len() != 1 || sources[0].kind != CatalogSourceKind::File { + return Ok(None); + } + let root = Path::new(dataset_uri); + if !root.is_dir() { + return Ok(None); + } + let source_path = root.join(&sources[0].file); + let source_path = std::fs::canonicalize(&source_path).context("canonicalize export Source")?; + anyhow::ensure!( + source_path.starts_with(root), + "export Source resolves outside the local Dataset" + ); + let input = std::fs::read(&source_path).context("read exact export Source")?; + ensure_output_byte_budget(input.len(), args.max_output_bytes, "exact export")?; + let text = std::str::from_utf8(&input).context("exact export Source must be UTF-8")?; + let detected = detect_format(Some(&source_path), Some(text))?; + if detected != exchange_document_format(format) { + return Ok(None); + } + let trajectories = validate_import_source(format, &source_path).await?; + anyhow::ensure!( + sources[0].size_bytes == Some(input.len() as u64) + && sources[0].snapshot_ref().as_deref() == Some(&local_file_snapshot_ref(&source_path)), + "export Source changed after the Snapshot was created" + ); + Ok(Some(EncodedExport { + bytes: input, + trajectories, + exact: true, + })) +} + +pub(crate) fn ensure_export_trajectory_budget( + trajectories: usize, + max_trajectories: u64, +) -> Result<()> { + if usize::try_from(max_trajectories).is_ok_and(|limit| trajectories > limit) { + return Err(cli_boundary_error( + BoundaryCode::ResourceExhausted, + format!("export exceeds max_trajectories limit of {max_trajectories}"), + )); + } + Ok(()) +} + +pub(crate) fn export_address_sql(args: &ExportArgs) -> Result { + let mut predicates = Vec::new(); + if let Some(source) = &args.source { + predicates.push(format!("_file_ = {}", sql_string(source))); + } + if let Some(run_id) = &args.run_id { + predicates.push(format!("run_id = {}", sql_string(run_id))); + } + if let Some(document_id) = &args.document_id { + predicates.push(format!("document_id = {}", sql_string(document_id))); + } + if let Some(session_id) = &args.session_id { + predicates.push(format!("session_id = {}", sql_string(session_id))); + } + if let Some(expression) = &args.r#where { + predicates.push(format!("({expression})")); + } + let predicate = if predicates.is_empty() { + String::new() + } else { + format!(" WHERE {}", predicates.join(" AND ")) + }; + let limit = args + .max_trajectories + .checked_add(1) + .context("--max-trajectories is too large")?; + Ok(format!( + "SELECT _file_ AS source_path, document_id, run_id, session_id \ + FROM dataset.trajectories{predicate} \ + ORDER BY _file_, document_id, session_id LIMIT {limit}" + )) +} + +pub(crate) fn encode_export( + format: ExchangeFormat, + stories: &[StorylineDocument], +) -> Result> { + let value = match format { + ExchangeFormat::Atif => encode_json_storylines(DocumentFormat::Atif, stories)?, + ExchangeFormat::Actf => encode_json_storylines(DocumentFormat::Actf, stories)?, + ExchangeFormat::OpenaiMessages => { + encode_json_storylines(DocumentFormat::OpenaiMsg, stories)? + } + ExchangeFormat::Storyline => encode_json_storylines(DocumentFormat::Storyline, stories)?, + ExchangeFormat::Codex | ExchangeFormat::ClaudeCode => { + bail!("{format} is decode-only and cannot be exported") + } + ExchangeFormat::CompactJsonl | ExchangeFormat::Auto => { + unreachable!("exchange export format was validated") + } + }; + let mut output = serde_json::to_vec_pretty(&value).context("encode export JSON")?; + output.push(b'\n'); + Ok(output) +} + +pub(crate) async fn write_export_output( + output: &str, + bytes: &[u8], + overwrite: bool, + stdout: &mut dyn Write, +) -> Result<()> { + if output == "-" { + stdout.write_all(bytes).context("write export stream")?; + return Ok(()); + } + DatasetLocation::parse(output)? + .put_bytes(bytes, overwrite) + .await +} + +pub(crate) fn local_file_snapshot_ref(path: &Path) -> String { + let mut hash = blake3::Hasher::new(); + hash.update(path.to_string_lossy().as_bytes()); + if let Ok(metadata) = std::fs::metadata(path) { + hash.update(&metadata.len().to_le_bytes()); + if let Ok(modified) = metadata.modified() + && let Ok(duration) = modified.duration_since(std::time::UNIX_EPOCH) + { + hash.update(&duration.as_nanos().to_le_bytes()); + } + #[cfg(unix)] + { + use std::os::unix::fs::MetadataExt; + hash.update(&metadata.dev().to_le_bytes()); + hash.update(&metadata.ino().to_le_bytes()); + } + } + format!("local:{}", hash.finalize().to_hex()) +} diff --git a/crates/persisting-pchronicle-cli/src/exchange/import.rs b/crates/persisting-pchronicle-cli/src/exchange/import.rs new file mode 100644 index 00000000..6d69ae37 --- /dev/null +++ b/crates/persisting-pchronicle-cli/src/exchange/import.rs @@ -0,0 +1,2298 @@ +//! Import command: storyline squash/commit/finalize, compact, and event paths. + +use super::super::*; +use super::decode::*; +use super::drop::confirm_destructive_dataset; +use super::pipeline::*; +use super::progress::{CliProgress, StageHandle, StageId, format_byte_count}; +use super::staging::*; +use anyhow::{Context, Result, anyhow}; +use persisting_pchronicle::model::StorylineDocument; +use persisting_pchronicle::storage::StorylineLanceStore; +use std::collections::{HashSet, VecDeque}; +use std::fs::OpenOptions; +use std::io::{Read, Write}; +use std::path::Path; +use std::sync::Arc; + +pub(crate) async fn prepare_import_destination( + args: &ImportArgs, + output_arg: &str, + stdin_is_terminal: bool, + stdin: &mut dyn Read, + stderr: &mut dyn Write, +) -> Result { + let parsed = DatasetLocation::parse(output_arg)?; + let exists = parsed.exists().await?; + match args.mode()? { + ImportMode::Create => { + if parsed.is_object_store() { + anyhow::ensure!(!exists, "import output already exists"); + Ok(PreparedImportDestination { + location: parsed, + replace_existing: false, + }) + } else { + Ok(PreparedImportDestination { + location: parsed.into_create_target()?, + replace_existing: false, + }) + } + } + ImportMode::Append => { + if !exists { + return Err(cli_boundary_error( + BoundaryCode::NotFound, + format!("append target Dataset does not exist: {}", parsed.as_str()), + )); + } + let location = if parsed.local_path().is_some() { + parsed.into_existing()? + } else { + parsed + }; + Ok(PreparedImportDestination { + location, + replace_existing: false, + }) + } + ImportMode::Replace => { + if !exists { + return if parsed.is_object_store() { + Ok(PreparedImportDestination { + location: parsed, + replace_existing: false, + }) + } else { + Ok(PreparedImportDestination { + location: parsed.into_create_target()?, + replace_existing: false, + }) + }; + } + let existing = parsed.into_existing()?; + ensure_import_source_outside_destination(args, &existing)?; + confirm_destructive_dataset( + "replace", + existing.as_str(), + args.yes, + stdin_is_terminal, + stdin, + stderr, + )?; + Ok(PreparedImportDestination { + location: existing, + replace_existing: true, + }) + } + } +} + +pub(crate) struct PreparedImportDestination { + location: DatasetLocation, + replace_existing: bool, +} + +pub(crate) fn ensure_import_source_outside_destination( + args: &ImportArgs, + destination: &DatasetLocation, +) -> Result<()> { + let (Some(source), Some(target)) = ( + (args.from != "-").then(|| Path::new(&args.from)), + destination.local_path(), + ) else { + return Ok(()); + }; + let source = std::fs::canonicalize(source).context("canonicalize replace import source")?; + anyhow::ensure!( + !source.starts_with(target), + "replace import source is inside the Dataset that would be replaced" + ); + Ok(()) +} + +pub(crate) async fn run_import( + mut args: ImportArgs, + settings_override: Option<&Path>, + stdin_is_terminal: bool, + stderr_is_terminal: bool, + stdin: &mut dyn Read, + stdout: &mut dyn Write, + stderr: &mut dyn Write, +) -> Result<()> { + args.stream = args.from == "-" || args.stream; + reset_import_log()?; + let max_input_bytes = match args.max_input_bytes { + Some(0) => { + return Err(anyhow!("--max-input-bytes must be greater than zero")); + } + Some(limit) => limit, + None => usize::MAX, + }; + anyhow::ensure!( + args.from == "-" || !args.stream, + "--stream requires --from -" + ); + if args.stream { + anyhow::ensure!( + args.format != ExchangeFormat::Auto, + "stdin import requires an explicit --input-format" + ); + } + if let Some(suggested) = args.suggested_format { + anyhow::ensure!( + args.format == ExchangeFormat::Auto, + "--suggested-format is only valid with --format auto" + ); + anyhow::ensure!( + suggested != ExchangeFormat::Auto, + "--suggested-format cannot be auto" + ); + anyhow::ensure!( + suggested != ExchangeFormat::CompactJsonl, + "--suggested-format cannot be compact-jsonl; pass --format compact-jsonl instead" + ); + } + let mode = args.mode()?; + anyhow::ensure!( + mode == ImportMode::Append || args.on_duplicate.is_none(), + "--on-duplicate is only valid with --append" + ); + anyhow::ensure!( + mode == ImportMode::Replace || !args.yes, + "--yes is only valid with --replace" + ); + anyhow::ensure!( + !(args.stream && mode == ImportMode::Replace && !args.yes), + "stdin replace import requires --yes because stdin carries the import data" + ); + if args.from != "-" { + args.from = expand_dataset_reference(&args.from, settings_override, true)?; + } + let from_location = (!args.stream) + .then(|| DatasetLocation::parse(&args.from)) + .transpose()?; + let canonical = if let Some(location) = &from_location { + let looks_like_store = location.is_object_store() + || location.local_path().is_some_and(std::path::Path::is_dir); + if looks_like_store { + probe_canonical_event_store(location.as_str()).await? + } else { + None + } + } else { + None + }; + let output_arg = match args.output.as_deref() { + Some(output) => expand_dataset_reference(output, settings_override, false)?, + None => default_import_output(&args, settings_override)?, + }; + if args.format == ExchangeFormat::CompactJsonl + || args.output_format == Some(ImportOutputFormat::CompactJsonl) + { + args.format = ExchangeFormat::CompactJsonl; + return run_compact_jsonl_import(args, &output_arg, stdout, stderr, stderr_is_terminal) + .await; + } + let requested_destination = DatasetLocation::parse(&output_arg)?; + if canonical.is_none() + && requested_destination.is_object_store() + && args.output_format != Some(ImportOutputFormat::Storyline) + { + anyhow::ensure!( + mode == ImportMode::Append && args.output_format.is_none(), + "object-store import requires --output-format storyline" + ); + } + let prepared = + prepare_import_destination(&args, &output_arg, stdin_is_terminal, stdin, stderr).await?; + let destination = prepared.location; + let replace_existing = prepared.replace_existing; + if let Some(snapshot) = canonical { + anyhow::ensure!( + mode != ImportMode::Append, + "canonical event import does not support --append" + ); + return run_canonical_event_import( + args, + snapshot, + destination, + replace_existing, + stdout, + stderr, + ) + .await; + } + let mut progress = CliProgress::new(stderr_is_terminal); + let _s3_throttle_ui = progress.attach_object_store_throttle(); + let object_store_from = from_location + .as_ref() + .filter(|location| location.is_object_store() && !args.stream) + .cloned(); + let (directory_input, candidates) = if args.stream { + progress.set_discovered(1, 0)?; + (false, Vec::new()) + } else if object_store_from.is_some() { + // Object-store Sources are discovered inside the Storyline pipeline so + // listing overlaps read/parse/write instead of buffering the full tree. + (true, Vec::new()) + } else if from_location.is_some() { + let (directory_input, candidates) = collect_import_candidates(Path::new(&args.from))?; + let discovered_bytes = candidates.iter().try_fold(0u64, |total, candidate| { + total + .checked_add(candidate.size_hint) + .context("import discovered byte count overflow") + })?; + progress.set_discovered(candidates.len() as u64, discovered_bytes)?; + (directory_input, candidates) + } else { + (false, Vec::new()) + }; + anyhow::ensure!( + mode != ImportMode::Append || args.output_format != Some(ImportOutputFormat::Preserve), + "append import requires --output-format storyline (or omit it)" + ); + let output_format = args.output_format.unwrap_or(if mode == ImportMode::Append { + ImportOutputFormat::Storyline + } else { + ImportOutputFormat::Preserve + }); + let duplicate_policy = args.on_duplicate.unwrap_or(DuplicateIdPolicy::Suffix); + let (wal, skip_paths) = open_import_wal( + &args, + &args.from, + destination.as_str(), + output_format, + )?; + if let Some(wal) = &wal + && let Ok(guard) = wal.lock() + { + progress.notice(&format!( + "import_wal={} job_id={} done={} failed={} resume={}", + guard.dir().display(), + guard.job().job_id, + guard.done_count(), + guard.failed_count(), + args.resume, + ))?; + } + let (dataset_uri, imported_sources, unknown_field_warnings, skipped_warnings) = if mode + == ImportMode::Append + { + let store = StorylineLanceStore::open_uri(destination.as_str()) + .await + .context("open append target as a Storyline Lance Dataset")?; + anyhow::ensure!( + store.current_table_paths().await?.is_some(), + "append target is not a committed Storyline Dataset" + ); + let (append_generation, existing_document_ids) = store + .document_ids_snapshot() + .await? + .context("append target has no committed Storyline snapshot")?; + let existing_storyline_count = existing_document_ids.len() as u64; + let existing_document_ids = existing_document_ids.into_iter().collect(); + let (imported_sources, unknown_field_warnings, skipped_warnings) = + squash_storyline_into_store( + &store, + &args, + stdin, + &mut progress, + &candidates, + object_store_from.clone(), + StorylineImportOptions { + max_input_bytes, + directory_input, + seen_document_ids: existing_document_ids, + duplicate_policy, + allow_empty: true, + append_generation: Some(append_generation), + initial_storyline_count: existing_storyline_count, + wal: wal.clone(), + skip_paths: Arc::clone(&skip_paths), + }, + ) + .await?; + ( + destination.as_str().to_string(), + imported_sources, + unknown_field_warnings, + skipped_warnings, + ) + } else if destination.is_object_store() || output_format == ImportOutputFormat::Storyline { + // Storyline imports commit in place so progressive CURRENT + + // chronicle.manifest updates are visible to a live catalog mount. + if destination.exists().await? && !replace_existing { + return Err(cli_boundary_error( + BoundaryCode::Conflict, + "import output already exists", + )); + } + let (imported_sources, unknown_field_warnings, skipped_warnings) = if destination + .is_object_store() + { + // Write directly to the remote Dataset. Progressive commits must be + // visible on the destination during long imports; local staging + + // final upload hides all progress until the job finishes. + if replace_existing { + destination + .remove_all_with_progress(|deleted, total, path| { + progress.note_deleted(deleted, total, path) + }) + .await + .with_context(|| { + format!("delete replaced Dataset prefix {}", destination.as_str()) + })?; + } + let store = StorylineLanceStore::open_uri(destination.as_str()) + .await + .with_context(|| { + format!( + "open remote Storyline Dataset for import at {}", + destination.as_str() + ) + })?; + squash_storyline_into_store( + &store, + &args, + stdin, + &mut progress, + &candidates, + object_store_from.clone(), + StorylineImportOptions::create(max_input_bytes, directory_input) + .with_wal(wal.clone(), Arc::clone(&skip_paths)), + ) + .await? + } else { + let output = destination + .local_path() + .context("local Storyline output must be a filesystem path")?; + let staging = tempfile::Builder::new() + .prefix(".pchronicle-storyline-stage-") + .tempdir_in(output.parent().context("Storyline output has no parent")?) + .context("create local Storyline staging directory")?; + let store = StorylineLanceStore::open(staging.path()) + .await + .context("create staged Storyline Lance Dataset")?; + let result = squash_storyline_into_store( + &store, + &args, + stdin, + &mut progress, + &candidates, + object_store_from.clone(), + StorylineImportOptions::create(max_input_bytes, directory_input) + .with_wal(wal.clone(), Arc::clone(&skip_paths)), + ) + .await?; + let staging_path = staging.keep(); + let mut cleanup = StagingPathGuard::new(staging_path.clone()); + publish_staged_dataset(&staging_path, output, replace_existing, Some(&mut progress)) + .await?; + cleanup.disarm(); + result + }; + ( + destination.as_str().to_string(), + imported_sources, + unknown_field_warnings, + skipped_warnings, + ) + } else { + let output = destination + .local_path() + .context("local import output must be a filesystem path")? + .to_path_buf(); + let parent = output + .parent() + .context("import output must have a parent directory")?; + let staging = tempfile::Builder::new() + .prefix(".pchronicle-import-") + .tempdir_in(parent) + .with_context(|| format!("create import staging directory in {}", parent.display()))?; + let (imported_sources, unknown_field_warnings, skipped_warnings) = match output_format { + ImportOutputFormat::Preserve => { + let mut unknown_field_warnings = + persisting_pchronicle::model::UnknownFieldImportWarnings::default(); + let mut imported_sources = Vec::new(); + let mut skipped_warnings = Vec::new(); + if args.stream { + progress.stage(StageId::Fetch).set_current("stdin"); + let input = read_bounded(stdin, max_input_bytes, "stdin")?; + progress.note_fetched("stdin", input.len() as u64)?; + progress.stage(StageId::Parse).set_current("stdin"); + if let Some(source) = stage_preserved_import_source( + args.format, + args.suggested_format, + None, + None, + None, + &input, + staging.path(), + &mut unknown_field_warnings, + &mut skipped_warnings, + )? { + progress.note_parsed(&source.source_path, source.input_bytes as u64)?; + imported_sources.push(source); + } else { + progress.note_parsed("stdin", input.len() as u64)?; + } + } else { + progress + .stage(StageId::Discover) + .set_total_items(candidates.len() as u64); + for candidate in &candidates { + let name = candidate.relative_path.to_string_lossy().into_owned(); + let label = format!("import source {name}"); + progress.note_discovered(&name, candidate.size_hint)?; + progress.stage(StageId::Fetch).set_current(&name); + let input = + load_import_candidate_bytes(candidate, max_input_bytes, &label).await?; + progress.note_fetched(&name, input.len() as u64)?; + progress.stage(StageId::Parse).set_current(&name); + match stage_preserved_import_source( + args.format, + args.suggested_format, + Some(&candidate.path), + Some(&candidate.relative_path), + candidate.output_relative_path.as_deref(), + &input, + staging.path(), + &mut unknown_field_warnings, + &mut skipped_warnings, + ) { + Ok(Some(source)) => { + progress + .note_parsed(&source.source_path, source.input_bytes as u64)?; + imported_sources.push(source); + } + Ok(None) => { + progress.note_parsed(&name, input.len() as u64)?; + } + Err(error) => { + let warning = skipped_import_warning( + Path::new(&name), + &format!("{error:#}"), + ); + let _ = append_import_log(&name, &error); + skipped_warnings.push(warning); + progress.note_parsed(&name, input.len() as u64)?; + } + } + } + } + (imported_sources, unknown_field_warnings, skipped_warnings) + } + ImportOutputFormat::Storyline => { + unreachable!("storyline import commits in place above") + } + ImportOutputFormat::CompactJsonl => unreachable!("compact import handled above"), + }; + if imported_sources.is_empty() { + return Err(empty_auto_directory_import_error(directory_input)); + } + + std::fs::File::open(staging.path()) + .and_then(|directory| directory.sync_all()) + .context("sync import staging directory")?; + + let staging_path = staging.keep(); + let mut cleanup = StagingPathGuard::new(staging_path.clone()); + publish_staged_dataset( + &staging_path, + &output, + replace_existing, + Some(&mut progress), + ) + .await?; + cleanup.disarm(); + ( + output.to_string_lossy().into_owned(), + imported_sources, + unknown_field_warnings, + skipped_warnings, + ) + }; + if imported_sources.is_empty() { + return Err(empty_auto_directory_import_error(directory_input)); + } + let trajectories = imported_sources.iter().try_fold(0usize, |total, source| { + total + .checked_add(source.trajectories) + .context("import trajectory count overflow") + })?; + let input_bytes = imported_sources.iter().try_fold(0usize, |total, source| { + total + .checked_add(source.input_bytes) + .context("import input byte count overflow") + })?; + let on_disk_bytes = measure_storyline_on_disk_bytes(&dataset_uri, output_format).await; + if let Some(bytes) = on_disk_bytes { + progress.stage(StageId::Commit).set_bytes(bytes); + progress + .stage(StageId::Commit) + .set_current(format!("on_disk={}", format_byte_count(bytes))); + } + + let single_source = (!directory_input).then(|| { + imported_sources + .first() + .expect("stdin and regular-file imports have one Source") + }); + let response = ImportResponse { + dataset_uri, + source_path: single_source.map(|source| source.source_path.clone()), + format: single_source.map(|source| source.format.as_str().to_owned()), + output_format: output_format.response_name().into(), + sources: imported_sources.len(), + trajectories, + fact_rows: None, + input_bytes: Some(input_bytes), + on_disk_bytes, + }; + serde_json::to_writer_pretty(&mut *stdout, &response) + .context("encode pChronicle import JSON")?; + writeln!(stdout).context("write pChronicle import JSON")?; + progress.finish()?; + if let (Some(source_path), Some(format)) = (&response.source_path, &response.format) { + progress.notice(&format!( + "dataset_uri={} source={} format={} output_format={} trajectories={} input_bytes={}{}", + response.dataset_uri, + source_path, + format, + response.output_format, + response.trajectories, + response + .input_bytes + .expect("JSON imports always report input bytes"), + on_disk_bytes_suffix(response.on_disk_bytes), + ))?; + } else { + progress.notice(&format!( + "dataset_uri={} sources={} output_format={} trajectories={} input_bytes={}{}", + response.dataset_uri, + response.sources, + response.output_format, + response.trajectories, + response + .input_bytes + .expect("JSON imports always report input bytes"), + on_disk_bytes_suffix(response.on_disk_bytes), + ))?; + } + for line in skipped_warnings { + progress.notice(&line)?; + } + for line in unknown_field_warnings.warning_lines() { + progress.notice(&line)?; + } + progress.flush_log(stderr)?; + Ok(()) +} + +/// Keep per-source failures durable while allowing a large import to continue. +pub(crate) fn reset_import_log() -> Result<()> { + OpenOptions::new() + .create(true) + .write(true) + .truncate(true) + .open("import.log") + .context("reset import.log")?; + Ok(()) +} + +pub(crate) fn append_import_log(path: &str, error: &anyhow::Error) -> Result<()> { + let mut log = OpenOptions::new() + .create(true) + .append(true) + .open("import.log") + .context("open import.log")?; + writeln!(log, "source={path}\terror={error:#}").context("append import.log") +} + +async fn measure_storyline_on_disk_bytes( + dataset_uri: &str, + output_format: ImportOutputFormat, +) -> Option { + if output_format != ImportOutputFormat::Storyline { + // Preserve / other modes may leave non-Storyline trees; skip. + // Object-store imports always write Storyline even when the CLI + // defaulted output_format from the destination kind. + let Ok(location) = DatasetLocation::parse(dataset_uri) else { + return None; + }; + if !location.is_object_store() { + return None; + } + } + match StorylineLanceStore::open_uri(dataset_uri).await { + Ok(store) => match store.on_disk_bytes().await { + Ok(bytes) => Some(bytes), + Err(error) => { + tracing::warn!( + dataset_uri, + error = %error, + "failed to measure Storyline on-disk bytes after import" + ); + None + } + }, + Err(error) => { + tracing::warn!( + dataset_uri, + error = %error, + "failed to reopen Storyline Dataset to measure on-disk bytes" + ); + None + } + } +} + +fn on_disk_bytes_suffix(on_disk_bytes: Option) -> String { + match on_disk_bytes { + Some(bytes) => format!(" on_disk_bytes={bytes} ({})", format_byte_count(bytes)), + None => String::new(), + } +} + +pub(crate) async fn run_compact_jsonl_import( + args: ImportArgs, + output_arg: &str, + stdout: &mut dyn Write, + stderr: &mut dyn Write, + stderr_is_terminal: bool, +) -> Result<()> { + anyhow::ensure!( + args.mode()? != ImportMode::Append, + "compact JSONL append is not supported; use sync or replace" + ); + anyhow::ensure!( + args.from != "-", + "compact JSONL import does not support stdin" + ); + let input = Path::new(&args.from); + let output = Path::new(output_arg); + anyhow::ensure!( + !output_arg.starts_with("s3://") && !output_arg.starts_with("oss://"), + "compact JSONL currently requires local paths" + ); + if args.mode()? == ImportMode::Create { + anyhow::ensure!(!output.exists(), "import output already exists"); + } + let columns = args + .columns + .iter() + .map(|item| { + let (name, path) = item + .split_once('=') + .context("--column must be NAME=JSON_PATH")?; + persisting_pchronicle::storage::CompactJsonlColumn::new(name.trim(), path.trim()) + }) + .collect::>>()?; + let options = persisting_pchronicle::storage::CompactJsonlOptions { + columns, + offload_threshold: 4 * 1024 * 1024, + }; + let parent = output + .parent() + .filter(|path| !path.as_os_str().is_empty()) + .unwrap_or_else(|| Path::new(".")); + let staging = tempfile::Builder::new() + .prefix(".pchronicle-compact-jsonl-") + .tempdir_in(parent)?; + let mut progress = CliProgress::new(stderr_is_terminal); + let _index_progress = progress.attach_index_progress(); + let rows = { + let progress = &mut progress; + persisting_pchronicle::storage::CompactJsonlStore::import_path_with_progress( + input, + staging.path(), + &options, + |event| match event { + persisting_pchronicle::storage::CompactJsonlImportEvent::Listed { + files, + bytes, + } => progress.set_discovered(files, bytes), + persisting_pchronicle::storage::CompactJsonlImportEvent::Reading { + relative, + file_bytes, + file_rows, + total_rows, + done, + } => { + let label = format!("{relative} rows={file_rows} total={total_rows}"); + progress.stage(StageId::Fetch).set_current(label.clone()); + progress.stage(StageId::Parse).set_current(&label); + if done { + progress.note_fetched(&relative, file_bytes)?; + progress.note_parsed(&relative, file_bytes)?; + } + Ok(()) + } + persisting_pchronicle::storage::CompactJsonlImportEvent::Building { + phase, + rows, + processed, + } => { + let commit = progress.stage(StageId::Commit); + commit.set_queue_cap(rows); + if let Some(processed) = processed { + commit.set_queue(processed); + commit.set_current(format!("{} {processed}/{rows}", phase.as_str())); + } else { + commit.set_queue(rows); + commit.set_current(format!("{} rows={rows}", phase.as_str())); + } + Ok(()) + } + persisting_pchronicle::storage::CompactJsonlImportEvent::Written { rows } => { + progress.note_committed(rows, 0) + } + }, + ) + .await? + }; + std::fs::File::open(staging.path())?.sync_all()?; + let staging_path = staging.keep(); + let mut cleanup = StagingPathGuard::new(staging_path.clone()); + publish_staged_dataset(&staging_path, output, output.exists(), Some(&mut progress)).await?; + cleanup.disarm(); + progress.finish()?; + serde_json::to_writer_pretty( + &mut *stdout, + &serde_json::json!({"dataset_uri": output_arg, "output_format": "compact-jsonl", "rows": rows}), + )?; + writeln!(stdout)?; + writeln!( + stderr, + "dataset_uri={} output_format=compact-jsonl rows={rows}", + output_arg + )?; + progress.flush_log(stderr)?; + Ok(()) +} + +pub(crate) struct StorylineImportOptions { + max_input_bytes: usize, + directory_input: bool, + seen_document_ids: HashSet, + duplicate_policy: DuplicateIdPolicy, + allow_empty: bool, + append_generation: Option, + initial_storyline_count: u64, + wal: Option>>, + skip_paths: std::sync::Arc>, +} + +impl StorylineImportOptions { + pub(crate) fn create(max_input_bytes: usize, directory_input: bool) -> Self { + Self { + max_input_bytes, + directory_input, + seen_document_ids: HashSet::new(), + duplicate_policy: DuplicateIdPolicy::Suffix, + allow_empty: false, + append_generation: None, + initial_storyline_count: 0, + wal: None, + skip_paths: std::sync::Arc::new(HashSet::new()), + } + } + + pub(crate) fn with_wal( + mut self, + wal: Option>>, + skip_paths: std::sync::Arc>, + ) -> Self { + self.wal = wal; + self.skip_paths = skip_paths; + self + } +} + +pub(crate) async fn squash_storyline_into_store( + store: &StorylineLanceStore, + args: &ImportArgs, + stdin: &mut dyn Read, + progress: &mut CliProgress, + candidates: &[ImportFileCandidate], + object_store_from: Option, + options: StorylineImportOptions, +) -> Result<( + Vec, + persisting_pchronicle::model::UnknownFieldImportWarnings, + Vec, +)> { + let StorylineImportOptions { + max_input_bytes, + directory_input, + seen_document_ids, + duplicate_policy, + allow_empty, + append_generation, + initial_storyline_count, + wal, + skip_paths, + } = options; + if args.stream { + return squash_storyline_stdin_into_store( + store, + args.format, + args.suggested_format, + max_input_bytes, + stdin, + progress, + seen_document_ids, + duplicate_policy, + allow_empty, + directory_input, + append_generation, + initial_storyline_count, + commit_batch_schedule(args), + ) + .await; + } + let source = match object_store_from { + Some(location) => ObjectStoreImportSource::Location(location), + None => ObjectStoreImportSource::Candidates(candidates.to_vec()), + }; + squash_storyline_files_pipeline( + store, + args.format, + args.suggested_format, + max_input_bytes, + progress, + source, + seen_document_ids, + duplicate_policy, + allow_empty, + directory_input, + append_generation, + initial_storyline_count, + commit_batch_schedule(args), + wal, + skip_paths, + ) + .await +} + +type SharedImportWal = std::sync::Arc>; +type ImportWalSkipSet = std::sync::Arc>; + +pub(crate) fn open_import_wal( + args: &ImportArgs, + from: &str, + to: &str, + output_format: ImportOutputFormat, +) -> Result<(Option, ImportWalSkipSet)> { + let output_name = output_format.response_name(); + let suggested = args + .suggested_format + .map(|format| format.as_str().to_string()); + let root = args + .wal_dir + .clone() + .unwrap_or_else(super::wal::ImportWal::default_root); + let wal = super::wal::ImportWal::open_or_create( + &root, + from, + to, + output_name, + suggested.as_deref(), + args.resume, + args.reset, + )?; + let skip = if args.resume { + std::sync::Arc::new(wal.skip_paths()) + } else { + std::sync::Arc::new(HashSet::new()) + }; + Ok(( + Some(std::sync::Arc::new(std::sync::Mutex::new(wal))), + skip, + )) +} + +pub(crate) const DEFAULT_COMMIT_BATCH_START: usize = 64; +pub(crate) const DEFAULT_COMMIT_BATCH_MAX: usize = 4096; + +#[derive(Debug, Clone)] +pub(crate) struct CommitBatchSchedule { + pub(crate) next: usize, + pub(crate) max: usize, + pub(crate) fixed: bool, +} + +impl CommitBatchSchedule { + pub(crate) fn adaptive() -> Self { + Self { + next: DEFAULT_COMMIT_BATCH_START, + max: DEFAULT_COMMIT_BATCH_MAX, + fixed: false, + } + } + + pub(crate) fn fixed(n: usize) -> Self { + let n = n.max(1); + Self { + next: n, + max: n, + fixed: true, + } + } + + pub(crate) fn current(&self) -> usize { + self.next + } + + pub(crate) fn after_commit(&mut self) { + if self.fixed { + return; + } + self.next = self.next.saturating_mul(2).min(self.max); + } +} + +pub(crate) fn commit_batch_schedule(args: &ImportArgs) -> CommitBatchSchedule { + match args.commit_every { + Some(n) => CommitBatchSchedule::fixed(n), + None => CommitBatchSchedule::adaptive(), + } +} + +pub(crate) enum ObjectStoreImportSource { + Candidates(Vec), + Location(DatasetLocation), +} + +#[derive(Debug)] +pub(crate) struct CommitStageOutcome { + pub(crate) imported_sources: Vec, + pub(crate) skipped_warnings: Vec, + pub(crate) committed_storylines: u64, + pub(crate) skipped_commit_storylines: usize, + pub(crate) saw_any: bool, + pub(crate) discovered_any: bool, +} + +pub(crate) struct CommitStageConfig { + pub(crate) store: StorylineLanceStore, + pub(crate) commit: StageHandle, + pub(crate) fetch: StageHandle, + pub(crate) parse: StageHandle, + pub(crate) seen_document_ids: HashSet, + pub(crate) duplicate_policy: DuplicateIdPolicy, + pub(crate) append_generation: Option, + pub(crate) initial_storyline_count: u64, + pub(crate) commit_schedule: CommitBatchSchedule, + pub(crate) unknown_field_warnings: std::sync::Arc< + tokio::sync::Mutex, + >, + pub(crate) wal: Option>>, +} + +struct BatchEntry { + storyline: StorylineDocument, + source_path: String, +} + +struct SourceCommitTracker { + /// Remaining storylines not yet successfully committed for each source. + remaining: std::collections::HashMap, + totals: std::collections::HashMap, +} + +impl SourceCommitTracker { + fn new() -> Self { + Self { + remaining: std::collections::HashMap::new(), + totals: std::collections::HashMap::new(), + } + } + + fn register(&mut self, path: &str, count: u64) { + if count == 0 { + return; + } + *self.remaining.entry(path.to_owned()).or_insert(0) += count; + *self.totals.entry(path.to_owned()).or_insert(0) += count; + } + + fn note_committed( + &mut self, + paths: &[String], + wal: &Option>>, + ) { + let mut completed = Vec::new(); + for path in paths { + if let Some(left) = self.remaining.get_mut(path) { + *left = left.saturating_sub(1); + if *left == 0 { + completed.push(path.clone()); + } + } + } + if let Some(wal) = wal + && let Ok(mut guard) = wal.lock() + { + for path in &completed { + let total = self.totals.remove(path).unwrap_or(1); + self.remaining.remove(path); + let _ = guard.mark_done(path, total); + } + } else { + for path in &completed { + self.remaining.remove(path); + self.totals.remove(path); + } + } + } + + fn note_failed_paths( + &mut self, + paths: &[String], + error: &str, + wal: &Option>>, + ) { + let unique = paths.iter().cloned().collect::>(); + for path in &unique { + self.remaining.remove(path); + self.totals.remove(path); + } + if let Some(wal) = wal + && let Ok(mut guard) = wal.lock() + { + for path in unique { + let _ = guard.mark_failed(&path, error); + } + } + } +} + +/// Single-worker commit stage running on its own tokio task. +/// +/// Double-buffers batches: while one batch is writing to storage, keep draining +/// `parsed_rx` into the next batch so parse→commit backpressure does not stall +/// the whole pipeline for the full remote commit latency. +pub(crate) fn spawn_commit_stage( + mut parsed_rx: tokio::sync::mpsc::Receiver>, + config: CommitStageConfig, +) -> tokio::task::JoinHandle> { + let CommitStageConfig { + store, + commit, + fetch, + parse, + mut seen_document_ids, + duplicate_policy, + mut append_generation, + initial_storyline_count, + mut commit_schedule, + unknown_field_warnings, + wal, + } = config; + tokio::spawn(async move { + let mut skipped_warnings = Vec::new(); + let mut imported_sources: Vec = Vec::new(); + let mut batch: Vec = Vec::with_capacity(commit_schedule.current()); + let mut source_bytes_left = 0u64; + let mut source_storylines_left = 0u64; + let mut committed_storylines = 0u64; + let mut skipped_commit_storylines = 0usize; + let mut saw_any = false; + let mut current_source_path = String::new(); + let mut current_storylines = Vec::new().into_iter(); + let mut producer_done = false; + let mut discovered_any = false; + let mut inflight: Option = None; + let mut lookahead: VecDeque> = VecDeque::new(); + let mut sources = SourceCommitTracker::new(); + refresh_commit_queue(&commit, &commit_schedule, batch.len(), &inflight); + + loop { + if let Some(mut storyline) = current_storylines.next() { + saw_any = true; + if let Some(warning) = apply_duplicate_document_policy( + &mut storyline, + &mut seen_document_ids, + duplicate_policy, + ) { + if warning.contains("skipped") { + skipped_warnings.push(warning); + let share = take_source_byte_share( + &mut source_bytes_left, + &mut source_storylines_left, + ); + commit.record_skipped(1, share); + sources.note_committed( + std::slice::from_ref(¤t_source_path), + &wal, + ); + continue; + } + skipped_warnings.push(warning); + } + let metadata = imported_sources + .last_mut() + .expect("decoded Storyline has source metadata"); + metadata.trajectories = metadata + .trajectories + .checked_add(1) + .context("import trajectory count overflow")?; + let share = + take_source_byte_share(&mut source_bytes_left, &mut source_storylines_left); + commit.record_bytes(share); + batch.push(BatchEntry { + storyline, + source_path: current_source_path.clone(), + }); + refresh_commit_queue(&commit, &commit_schedule, batch.len(), &inflight); + if batch.len() >= commit_schedule.current() { + join_inflight_commit_batch_draining( + &store, + &mut inflight, + &mut parsed_rx, + &mut lookahead, + &mut producer_done, + &mut append_generation, + &mut committed_storylines, + &mut commit_schedule, + &mut skipped_commit_storylines, + &mut skipped_warnings, + &mut sources, + &wal, + ) + .await?; + inflight = Some(spawn_inflight_commit_batch( + store.clone(), + commit.clone(), + std::mem::take(&mut batch), + append_generation.clone(), + committed_storylines, + initial_storyline_count, + )); + batch.reserve(commit_schedule.current()); + refresh_commit_queue(&commit, &commit_schedule, batch.len(), &inflight); + } + continue; + } + + if producer_done && lookahead.is_empty() { + break; + } + + let received = if let Some(item) = lookahead.pop_front() { + Some(item) + } else { + commit.enter_upstream_wait(); + let received = parsed_rx.recv().await; + commit.leave_upstream_wait(); + received + }; + match received { + Some(Ok(ParsedItem::Imported { + diagnostic_path: _, + mut metadata, + storylines, + warnings, + })) => { + unknown_field_warnings.lock().await.merge(&warnings); + discovered_any = true; + let storyline_count = storylines.len() as u64; + source_bytes_left = metadata.input_bytes as u64; + source_storylines_left = storyline_count; + current_source_path = metadata.source_path.clone(); + sources.register(¤t_source_path, storyline_count); + if storyline_count > 0 { + commit.record_inbound(storyline_count); + } else if let Some(wal) = &wal + && let Ok(mut guard) = wal.lock() + { + let _ = guard.mark_done(¤t_source_path, 0); + } + metadata.trajectories = 0; + imported_sources.push(metadata); + current_storylines = storylines.into_iter(); + } + Some(Ok(ParsedItem::Skipped { + path, + reason, + bytes: _, + })) => { + discovered_any = true; + let path_key = path.to_string_lossy().into_owned(); + let warning = skipped_import_warning(&path, &reason); + let _ = append_import_log(&path_key, &anyhow!("{reason}")); + if let Some(wal) = &wal + && let Ok(mut guard) = wal.lock() + { + let _ = guard.mark_failed(&path_key, &reason); + } + skipped_warnings.push(warning); + } + Some(Err(error)) => { + let _ = join_inflight_commit_batch( + &store, + &mut inflight, + &mut append_generation, + &mut committed_storylines, + &mut commit_schedule, + &mut skipped_commit_storylines, + &mut skipped_warnings, + &mut sources, + &wal, + ) + .await; + return Err(error); + } + None => { + producer_done = true; + fetch.clear_current(); + parse.clear_current(); + } + } + } + + join_inflight_commit_batch( + &store, + &mut inflight, + &mut append_generation, + &mut committed_storylines, + &mut commit_schedule, + &mut skipped_commit_storylines, + &mut skipped_warnings, + &mut sources, + &wal, + ) + .await?; + + if !batch.is_empty() { + let paths = batch + .iter() + .map(|entry| entry.source_path.clone()) + .collect::>(); + let storylines = batch + .into_iter() + .map(|entry| entry.storyline) + .collect::>(); + let mut state = StorylineCommitState { + append_generation: &mut append_generation, + committed_storylines, + initial_storyline_count, + commit_schedule: &mut commit_schedule, + }; + match commit_or_skip_storyline_import_batch(&store, &commit, storylines, &mut state) + .await + { + Ok(total) => { + committed_storylines = total; + sources.note_committed(&paths, &wal); + } + Err(error) if is_skippable_storyline_commit_error(&error) => { + skipped_commit_storylines = + skipped_commit_storylines.saturating_add(paths.len()); + let message = format!("{error:#}"); + skipped_warnings.push(message.clone()); + sources.note_failed_paths(&paths, &message, &wal); + refresh_append_generation_after_skip(&store, &mut append_generation).await; + } + Err(error) => return Err(error), + } + refresh_commit_queue(&commit, &commit_schedule, 0, &None); + } + + commit.clear_current(); + Ok(CommitStageOutcome { + imported_sources, + skipped_warnings, + committed_storylines, + skipped_commit_storylines, + saw_any, + discovered_any, + }) + }) +} + +struct InflightCommitBatch { + handle: tokio::task::JoinHandle>, + batch_len: usize, + source_paths: Vec, +} + +enum InflightCommitOutcome { + Committed { + total: u64, + generation: Option, + }, + Skipped { + error: String, + batch_len: usize, + }, +} + +/// Parsed-item lookahead while both storyline buffers are occupied. +const COMMIT_LOOKAHEAD_CAP: usize = PARSE_TO_COMMIT_BUFFER.saturating_mul(2); + +fn refresh_commit_queue( + commit: &StageHandle, + schedule: &CommitBatchSchedule, + filling: usize, + inflight: &Option, +) { + let inflight_len = inflight.as_ref().map(|job| job.batch_len).unwrap_or(0); + // Double-buffer capacity: one batch writing + one batch filling. + let cap = schedule.current().saturating_mul(2) as u64; + commit.set_queue_cap(cap); + commit.set_queue((filling + inflight_len) as u64); +} + +fn spawn_inflight_commit_batch( + store: StorylineLanceStore, + commit: StageHandle, + batch: Vec, + mut append_generation: Option, + committed_storylines: u64, + initial_storyline_count: u64, +) -> InflightCommitBatch { + let batch_len = batch.len(); + let source_paths = batch + .iter() + .map(|entry| entry.source_path.clone()) + .collect::>(); + let sample_ids = batch + .iter() + .take(8) + .map(|entry| entry.storyline.document_id().to_string()) + .collect::>(); + let storylines = batch + .into_iter() + .map(|entry| entry.storyline) + .collect::>(); + let handle = tokio::spawn(async move { + match commit_storyline_import_batch( + &store, + &commit, + storylines, + &mut append_generation, + committed_storylines, + initial_storyline_count, + ) + .await + { + Ok(total) => Ok(InflightCommitOutcome::Committed { + total, + generation: append_generation, + }), + Err(error) if is_skippable_storyline_commit_error(&error) => { + Ok(InflightCommitOutcome::Skipped { + error: format!( + "storyline commit batch skipped (batch={batch_len}, committed_before={committed_storylines}, sample_document_ids={sample_ids:?}): {error:#}" + ), + batch_len, + }) + } + Err(error) => Err(error), + } + }); + InflightCommitBatch { + handle, + batch_len, + source_paths, + } +} + +#[allow(clippy::too_many_arguments)] +fn apply_inflight_outcome( + outcome: InflightCommitOutcome, + source_paths: &[String], + append_generation: &mut Option, + committed_storylines: &mut u64, + commit_schedule: &mut CommitBatchSchedule, + skipped_commit_storylines: &mut usize, + skipped_warnings: &mut Vec, + sources: &mut SourceCommitTracker, + wal: &Option, +) -> bool { + match outcome { + InflightCommitOutcome::Committed { total, generation } => { + *append_generation = generation; + *committed_storylines = total; + commit_schedule.after_commit(); + sources.note_committed(source_paths, wal); + false + } + InflightCommitOutcome::Skipped { error, batch_len } => { + *skipped_commit_storylines = skipped_commit_storylines.saturating_add(batch_len); + skipped_warnings.push(error.clone()); + sources.note_failed_paths(source_paths, &error, wal); + true + } + } +} + +#[allow(clippy::too_many_arguments)] +async fn join_inflight_commit_batch( + store: &StorylineLanceStore, + inflight: &mut Option, + append_generation: &mut Option, + committed_storylines: &mut u64, + commit_schedule: &mut CommitBatchSchedule, + skipped_commit_storylines: &mut usize, + skipped_warnings: &mut Vec, + sources: &mut SourceCommitTracker, + wal: &Option, +) -> Result<()> { + let Some(job) = inflight.take() else { + return Ok(()); + }; + let outcome = job + .handle + .await + .context("storyline commit batch task join failed")??; + let needs_refresh = apply_inflight_outcome( + outcome, + &job.source_paths, + append_generation, + committed_storylines, + commit_schedule, + skipped_commit_storylines, + skipped_warnings, + sources, + wal, + ); + if needs_refresh { + refresh_append_generation_after_skip(store, append_generation).await; + } + Ok(()) +} + +/// Join the in-flight write, draining parse→commit into `lookahead` meanwhile. +#[allow(clippy::too_many_arguments)] +async fn join_inflight_commit_batch_draining( + store: &StorylineLanceStore, + inflight: &mut Option, + parsed_rx: &mut tokio::sync::mpsc::Receiver>, + lookahead: &mut VecDeque>, + producer_done: &mut bool, + append_generation: &mut Option, + committed_storylines: &mut u64, + commit_schedule: &mut CommitBatchSchedule, + skipped_commit_storylines: &mut usize, + skipped_warnings: &mut Vec, + sources: &mut SourceCommitTracker, + wal: &Option, +) -> Result<()> { + let Some(mut job) = inflight.take() else { + return Ok(()); + }; + loop { + tokio::select! { + biased; + joined = &mut job.handle => { + let outcome = joined + .context("storyline commit batch task join failed")??; + let needs_refresh = apply_inflight_outcome( + outcome, + &job.source_paths, + append_generation, + committed_storylines, + commit_schedule, + skipped_commit_storylines, + skipped_warnings, + sources, + wal, + ); + if needs_refresh { + refresh_append_generation_after_skip(store, append_generation).await; + } + return Ok(()); + } + item = parsed_rx.recv(), if !*producer_done && lookahead.len() < COMMIT_LOOKAHEAD_CAP => { + match item { + Some(parsed) => { + lookahead.push_back(parsed); + } + None => { + *producer_done = true; + } + } + } + } + } +} + +#[allow(clippy::too_many_arguments)] +pub(crate) async fn squash_storyline_files_pipeline( + store: &StorylineLanceStore, + requested_format: ExchangeFormat, + suggested_format: Option, + max_input_bytes: usize, + progress: &mut CliProgress, + source: ObjectStoreImportSource, + seen_document_ids: HashSet, + duplicate_policy: DuplicateIdPolicy, + allow_empty: bool, + directory_input: bool, + append_generation: Option, + initial_storyline_count: u64, + commit_schedule: CommitBatchSchedule, + wal: Option>>, + skip_paths: std::sync::Arc>, +) -> Result<( + Vec, + persisting_pchronicle::model::UnknownFieldImportWarnings, + Vec, +)> { + let ImportPipelineHandles { + parsed_rx, + joins, + unknown_field_warnings, + } = match source { + ObjectStoreImportSource::Candidates(candidates) => spawn_candidates_fetch_pipeline( + candidates, + ImportPipelineConfig { + max_input_bytes, + requested_format, + suggested_format, + discover: progress.stage(StageId::Discover), + fetch: progress.stage(StageId::Fetch), + parse: progress.stage(StageId::Parse), + commit: progress.stage(StageId::Commit), + skip_paths: Arc::clone(&skip_paths), + }, + ), + ObjectStoreImportSource::Location(location) => { + progress + .stage(StageId::Discover) + .set_current(location.as_str()); + spawn_location_fetch_pipeline( + location, + ImportPipelineConfig { + max_input_bytes, + requested_format, + suggested_format, + discover: progress.stage(StageId::Discover), + fetch: progress.stage(StageId::Fetch), + parse: progress.stage(StageId::Parse), + commit: progress.stage(StageId::Commit), + skip_paths, + }, + ) + } + }; + + let commit_join = spawn_commit_stage( + parsed_rx, + CommitStageConfig { + store: store.clone(), + commit: progress.stage(StageId::Commit), + fetch: progress.stage(StageId::Fetch), + parse: progress.stage(StageId::Parse), + seen_document_ids, + duplicate_policy, + append_generation, + initial_storyline_count, + commit_schedule, + unknown_field_warnings: Arc::clone(&unknown_field_warnings), + wal, + }, + ); + + let commit_result = match commit_join.await { + Ok(result) => result, + Err(error) if error.is_cancelled() => Err(anyhow!("commit stage cancelled")), + Err(error) => Err(anyhow!("commit stage task failed: {error}")), + }; + + let CommitStageOutcome { + mut imported_sources, + skipped_warnings, + committed_storylines, + skipped_commit_storylines, + saw_any, + discovered_any, + } = match commit_result { + Ok(outcome) => { + join_pipeline_stages(joins).await?; + outcome + } + Err(error) => { + for join in &joins { + join.abort(); + } + let _ = join_pipeline_stages(joins).await; + return Err(error); + } + }; + + let unknown_field_warnings = unknown_field_warnings.lock().await.clone(); + + retract_imported_trajectories(&mut imported_sources, skipped_commit_storylines); + if skipped_commit_storylines > 0 { + imported_sources.retain(|source| source.trajectories > 0); + } + if imported_sources.is_empty() { + if allow_empty && !saw_any { + return Ok((imported_sources, unknown_field_warnings, skipped_warnings)); + } + if skipped_commit_storylines > 0 { + return Err(anyhow!( + "storyline import committed no trajectories after skipping failed batches" + )); + } + if !discovered_any { + return Err(cli_boundary_error( + BoundaryCode::InvalidRequest, + "import object prefix contains no .json, .jsonl, or .ndjson files", + )); + } + return Err(empty_auto_directory_import_error(directory_input)); + } + anyhow::ensure!( + store.current_table_paths().await?.is_some(), + "squashed Storyline Lance Dataset has no committed snapshot" + ); + let imported_trajectories = imported_sources.iter().try_fold(0usize, |total, source| { + total + .checked_add(source.trajectories) + .context("import trajectory count overflow") + })?; + anyhow::ensure!( + committed_storylines as usize == imported_trajectories, + "squashed Storyline import report does not match decoded trajectory count" + ); + finalize_storyline_import_indexes(store, progress).await?; + Ok((imported_sources, unknown_field_warnings, skipped_warnings)) +} + +#[allow(clippy::too_many_arguments)] +pub(crate) async fn squash_storyline_stdin_into_store( + store: &StorylineLanceStore, + requested_format: ExchangeFormat, + suggested_format: Option, + max_input_bytes: usize, + stdin: &mut dyn Read, + progress: &mut CliProgress, + seen_document_ids: HashSet, + duplicate_policy: DuplicateIdPolicy, + allow_empty: bool, + directory_input: bool, + append_generation: Option, + initial_storyline_count: u64, + commit_schedule: CommitBatchSchedule, +) -> Result<( + Vec, + persisting_pchronicle::model::UnknownFieldImportWarnings, + Vec, +)> { + let import = StorylineImportIterator::stdin( + requested_format, + suggested_format, + max_input_bytes, + stdin, + progress, + seen_document_ids, + duplicate_policy, + ); + drain_storyline_import_batches( + store, + import, + append_generation, + commit_schedule, + allow_empty, + directory_input, + initial_storyline_count, + ) + .await +} + +pub(crate) async fn drain_storyline_import_batches( + store: &StorylineLanceStore, + mut import: StorylineImportIterator<'_>, + mut append_generation: Option, + mut commit_schedule: CommitBatchSchedule, + allow_empty: bool, + directory_input: bool, + initial_storyline_count: u64, +) -> Result<( + Vec, + persisting_pchronicle::model::UnknownFieldImportWarnings, + Vec, +)> { + let mut batch = Vec::with_capacity(commit_schedule.current()); + let mut committed_storylines = 0u64; + let mut skipped_commit_storylines = 0usize; + let mut commit_skip_warnings = Vec::new(); + let mut saw_any = false; + + loop { + match import.next_document().await { + Some(item) => { + saw_any = true; + batch.push(item?); + if batch.len() < commit_schedule.current() { + continue; + } + let batch_len = batch.len() as u64; + let commit = import.progress.stage(StageId::Commit); + let mut state = StorylineCommitState { + append_generation: &mut append_generation, + committed_storylines, + initial_storyline_count, + commit_schedule: &mut commit_schedule, + }; + match commit_or_skip_storyline_import_batch( + store, + &commit, + std::mem::take(&mut batch), + &mut state, + ) + .await + { + Ok(total) => { + committed_storylines = total; + } + Err(error) if is_skippable_storyline_commit_error(&error) => { + skipped_commit_storylines = + skipped_commit_storylines.saturating_add(batch_len as usize); + commit_skip_warnings.push(format!("{error:#}")); + refresh_append_generation_after_skip(store, &mut append_generation).await; + } + Err(error) => return Err(error), + } + batch.reserve(commit_schedule.current()); + } + None if batch.is_empty() => break, + None => { + let batch_len = batch.len() as u64; + let commit = import.progress.stage(StageId::Commit); + let mut state = StorylineCommitState { + append_generation: &mut append_generation, + committed_storylines, + initial_storyline_count, + commit_schedule: &mut commit_schedule, + }; + match commit_or_skip_storyline_import_batch( + store, + &commit, + std::mem::take(&mut batch), + &mut state, + ) + .await + { + Ok(total) => { + committed_storylines = total; + } + Err(error) if is_skippable_storyline_commit_error(&error) => { + skipped_commit_storylines = + skipped_commit_storylines.saturating_add(batch_len as usize); + commit_skip_warnings.push(format!("{error:#}")); + refresh_append_generation_after_skip(store, &mut append_generation).await; + } + Err(error) => return Err(error), + } + break; + } + } + } + + let (mut imported_sources, unknown_field_warnings, mut skipped_warnings, progress) = + import.into_result_parts(); + skipped_warnings.extend(commit_skip_warnings); + retract_imported_trajectories(&mut imported_sources, skipped_commit_storylines); + if skipped_commit_storylines > 0 { + imported_sources.retain(|source| source.trajectories > 0); + } + if imported_sources.is_empty() { + if allow_empty && !saw_any { + return Ok((imported_sources, unknown_field_warnings, skipped_warnings)); + } + if skipped_commit_storylines > 0 { + return Err(anyhow!( + "storyline import committed no trajectories after skipping failed batches" + )); + } + return Err(empty_auto_directory_import_error(directory_input)); + } + anyhow::ensure!( + store.current_table_paths().await?.is_some(), + "squashed Storyline Lance Dataset has no committed snapshot" + ); + let imported_trajectories = imported_sources.iter().try_fold(0usize, |total, source| { + total + .checked_add(source.trajectories) + .context("import trajectory count overflow") + })?; + anyhow::ensure!( + committed_storylines as usize == imported_trajectories, + "squashed Storyline import report does not match decoded trajectory count" + ); + finalize_storyline_import_indexes(store, progress).await?; + Ok((imported_sources, unknown_field_warnings, skipped_warnings)) +} + +pub(crate) fn is_skippable_storyline_commit_error(error: &anyhow::Error) -> bool { + let text = format!("{error:#}").to_ascii_lowercase(); + text.contains("timeout") + || text.contains("timed out") + || text.contains("error sending request") + || text.contains("conditionnotmatch") + || text.contains("preconditionfailed") + || text.contains("precondition failed") + || text.contains("throttle") + || text.contains("slow down") + || text.contains("503") + || text.contains("429") + || text.contains("connection reset") + || text.contains("broken pipe") + || text.contains("lanceerror(io)") + || text.contains("generic s3 error") + || text.contains("client error (connect)") + || text.contains("byte array offset overflow") + || text.contains("arrow encode panicked") + || text.contains("max_chunk_bytes") + || text.contains("max_document_bytes") + || text.contains("max_chunk_rows") + || text.contains("max_document_rows") +} + +pub(crate) fn retract_imported_trajectories(sources: &mut [ImportedSource], mut count: usize) { + for source in sources.iter_mut().rev() { + if count == 0 { + break; + } + let take = source.trajectories.min(count); + source.trajectories -= take; + count -= take; + } +} + +pub(crate) async fn refresh_append_generation_after_skip( + store: &StorylineLanceStore, + append_generation: &mut Option, +) { + match store.current_table_paths().await { + Ok(Some(paths)) => { + *append_generation = Some(paths.generation); + } + Ok(None) => {} + Err(error) => { + tracing::warn!( + root = %store.root_uri(), + error = %error, + "failed to refresh Storyline generation after skipped commit batch" + ); + } + } +} + +pub(crate) fn take_source_byte_share(bytes_left: &mut u64, storylines_left: &mut u64) -> u64 { + if *storylines_left == 0 { + return 0; + } + let share = if *storylines_left == 1 { + *bytes_left + } else { + *bytes_left / *storylines_left + }; + *bytes_left = bytes_left.saturating_sub(share); + *storylines_left = storylines_left.saturating_sub(1); + share +} + +pub(crate) struct StorylineCommitState<'a> { + pub(crate) append_generation: &'a mut Option, + pub(crate) committed_storylines: u64, + pub(crate) initial_storyline_count: u64, + pub(crate) commit_schedule: &'a mut CommitBatchSchedule, +} + +pub(crate) async fn commit_or_skip_storyline_import_batch( + store: &StorylineLanceStore, + commit: &StageHandle, + batch: Vec, + state: &mut StorylineCommitState<'_>, +) -> Result { + let batch_len = batch.len() as u64; + let sample_ids = batch + .iter() + .take(8) + .map(|storyline| storyline.document_id().to_string()) + .collect::>(); + match commit_storyline_import_batch( + store, + commit, + batch, + state.append_generation, + state.committed_storylines, + state.initial_storyline_count, + ) + .await + { + Ok(total) => { + state.commit_schedule.after_commit(); + Ok(total) + } + Err(error) if is_skippable_storyline_commit_error(&error) => Err(error).context( + format!( + "storyline commit batch failed after transient storage error (batch={batch_len}, committed_before={}, sample_document_ids={sample_ids:?})", + state.committed_storylines + ), + ), + Err(error) => Err(error), + } +} + +pub(crate) async fn finalize_storyline_import_indexes( + store: &StorylineLanceStore, + progress: &mut CliProgress, +) -> Result<()> { + progress + .stage(StageId::Commit) + .set_current("optimize indices (final)"); + let _index_progress = progress.attach_index_progress(); + store + .maintain(&persisting_pchronicle::storage::LanceMaintenanceOptions { + compact: false, + optimize_indices: true, + vacuum_older_than: None, + ..Default::default() + }) + .await + .context("finalize Storyline indexes after progressive import")?; + progress + .stage(StageId::Commit) + .set_current("optimize indices done"); + Ok(()) +} + +pub(crate) async fn commit_storyline_import_batch( + store: &StorylineLanceStore, + commit: &StageHandle, + batch: Vec, + append_generation: &mut Option, + committed_storylines: u64, + initial_storyline_count: u64, +) -> Result { + anyhow::ensure!(!batch.is_empty(), "storyline import commit batch is empty"); + let batch_len = batch.len() as u64; + commit.set_current(format!("batch={batch_len}")); + let report = match append_generation.as_deref() { + Some(generation) => { + tracing::info!( + committed_before = committed_storylines, + batch_len, + expected_generation = generation, + root = %store.root_uri(), + "storyline progressive append commit starting" + ); + store + .append_storyline_stream_with_options( + batch.into_iter().map(Ok), + generation, + persisting_pchronicle::storage::StorylineStreamOptions::defer_index_optimize(), + ) + .await + .with_context(|| { + format!( + "storyline progressive append commit failed (committed_before={committed_storylines}, batch={batch_len}, expected_generation={generation}, root={})", + store.root_uri() + ) + })? + } + None => { + tracing::info!( + batch_len, + root = %store.root_uri(), + "storyline progressive replace commit starting" + ); + store + .replace_storyline_stream_with_options( + batch.into_iter().map(Ok), + persisting_pchronicle::storage::StorylineStreamOptions::defer_index_optimize(), + ) + .await + .with_context(|| { + format!( + "storyline progressive replace commit failed (batch={batch_len}, root={})", + store.root_uri() + ) + })? + } + }; + anyhow::ensure!( + report.storylines as u64 == batch_len, + "storyline import batch report does not match batch size" + ); + let paths = store + .current_table_paths() + .await? + .context("storyline import batch produced no committed snapshot")?; + let imported_total = committed_storylines + .checked_add(batch_len) + .context("import trajectory count overflow")?; + let manifest_total = initial_storyline_count + .checked_add(imported_total) + .context("import manifest record count overflow")?; + persisting_pchronicle::storage::write_storyline_manifest_at_uri( + store.root_uri(), + &paths.generation, + manifest_total, + 0, + ) + .await + .context("write progressive chronicle.manifest after storyline commit")?; + *append_generation = Some(paths.generation.clone()); + // Bytes were already attributed when trajectories entered the batch. + commit.note_committed(imported_total, 0); + Ok(imported_total) +} + +pub(crate) async fn run_canonical_event_import( + args: ImportArgs, + _snapshot: EventFactSnapshot, + destination: DatasetLocation, + replace_existing: bool, + stdout: &mut dyn Write, + stderr: &mut dyn Write, +) -> Result<()> { + anyhow::ensure!( + args.format == ExchangeFormat::Auto, + "canonical event import does not accept a JSON exchange --format" + ); + anyhow::ensure!( + args.output_format != Some(ImportOutputFormat::Preserve), + "canonical event import cannot preserve an existing canonical event Store" + ); + if destination.exists().await? && !replace_existing { + return Err(cli_boundary_error( + BoundaryCode::Conflict, + "import output already exists", + )); + } + let output_uri = destination.as_str().to_string(); + + let (report, staged_path) = if replace_existing { + let output = destination + .local_path() + .context("replace import output must be a local Dataset path")?; + let parent = output + .parent() + .context("replace import output must have a parent directory")?; + let staging = tempfile::Builder::new() + .prefix(".pchronicle-import-") + .tempdir_in(parent) + .with_context(|| format!("create import staging directory in {}", parent.display()))?; + let staging_uri = staging.path().to_string_lossy().into_owned(); + let report = + match build_storyline_projection(&args.from, &staging_uri, "events.lance").await? { + StorylineProjectionBuildOutcome::Built(report) => report, + StorylineProjectionBuildOutcome::OutputNotEmpty => { + return Err(cli_boundary_error( + BoundaryCode::Conflict, + "import staging Dataset already exists", + )); + } + }; + std::fs::File::open(staging.path()) + .and_then(|directory| directory.sync_all()) + .context("sync import staging directory")?; + (report, Some((staging.keep(), output.to_path_buf()))) + } else { + let report = + match build_storyline_projection(&args.from, &output_uri, "events.lance").await? { + StorylineProjectionBuildOutcome::Built(report) => report, + StorylineProjectionBuildOutcome::OutputNotEmpty => { + return Err(cli_boundary_error( + BoundaryCode::Conflict, + "import output already exists", + )); + } + }; + (report, None) + }; + if let Some((staging_path, output)) = staged_path { + let mut cleanup = StagingPathGuard::new(staging_path.clone()); + publish_staged_dataset(&staging_path, &output, true, None).await?; + cleanup.disarm(); + } + let response = ImportResponse { + dataset_uri: output_uri, + source_path: Some("events.lance".into()), + format: Some("events".into()), + output_format: ImportOutputFormat::Storyline.response_name().into(), + sources: 1, + trajectories: report.storylines, + fact_rows: Some(report.fact_rows), + input_bytes: None, + on_disk_bytes: None, + }; + serde_json::to_writer_pretty(&mut *stdout, &response) + .context("encode canonical event import JSON")?; + writeln!(stdout).context("write canonical event import JSON")?; + writeln!( + stderr, + "dataset_uri={} source=events.lance format=events output_format={} trajectories={} fact_rows={}", + response.dataset_uri, + response.output_format, + response.trajectories, + report.fact_rows, + ) + .context("write canonical event import metadata")?; + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + use persisting_pchronicle::document::DocumentFormat; + + #[test] + fn commit_batch_schedule_grows_to_cap() { + let mut schedule = CommitBatchSchedule::adaptive(); + assert_eq!(schedule.current(), 64); + schedule.after_commit(); + assert_eq!(schedule.current(), 128); + schedule.after_commit(); + assert_eq!(schedule.current(), 256); + schedule.after_commit(); + assert_eq!(schedule.current(), 512); + schedule.after_commit(); + assert_eq!(schedule.current(), 1024); + schedule.after_commit(); + assert_eq!(schedule.current(), 2048); + schedule.after_commit(); + assert_eq!(schedule.current(), 4096); + schedule.after_commit(); + assert_eq!(schedule.current(), 4096); + } + + #[test] + fn skippable_commit_errors_cover_s3_timeouts_and_preconditions() { + assert!(is_skippable_storyline_commit_error(&anyhow!( + "LanceError(IO): Generic S3 error: operation timed out" + ))); + assert!(is_skippable_storyline_commit_error(&anyhow!( + "ConditionNotMatch (persistent) PreconditionFailed" + ))); + assert!(is_skippable_storyline_commit_error(&anyhow!( + "arrow encode panicked: byte array offset overflow" + ))); + assert!(is_skippable_storyline_commit_error(&anyhow!( + "document exceeds max_chunk_bytes" + ))); + assert!(!is_skippable_storyline_commit_error(&anyhow!( + "duplicate document_id policy rejected payload" + ))); + } + + #[test] + fn open_import_wal_skips_only_on_resume() { + let root = tempfile::tempdir().unwrap(); + let mut base = ImportArgs { + from: "s3://bucket/from".into(), + output: Some("s3://bucket/to".into()), + format: ExchangeFormat::Auto, + suggested_format: None, + output_format: Some(ImportOutputFormat::Storyline), + replace: false, + append: false, + on_duplicate: None, + yes: true, + stream: false, + max_input_bytes: None, + commit_every: None, + resume: false, + wal_dir: Some(root.path().to_path_buf()), + reset: false, + columns: Vec::new(), + }; + let (wal, skip) = open_import_wal( + &base, + &base.from, + "s3://bucket/to", + ImportOutputFormat::Storyline, + ) + .unwrap(); + assert!(skip.is_empty()); + { + let mut guard = wal.as_ref().unwrap().lock().unwrap(); + guard.mark_done("done.json", 1).unwrap(); + guard.mark_failed("fail.json", "parse").unwrap(); + } + + let (_, skip_again) = open_import_wal( + &base, + &base.from, + "s3://bucket/to", + ImportOutputFormat::Storyline, + ) + .unwrap(); + assert!( + skip_again.is_empty(), + "without --resume, prior WAL entries must not be skipped" + ); + + base.resume = true; + let (_, skip_resume) = open_import_wal( + &base, + &base.from, + "s3://bucket/to", + ImportOutputFormat::Storyline, + ) + .unwrap(); + assert!(skip_resume.contains("done.json")); + assert!(skip_resume.contains("fail.json")); + } + + #[test] + fn source_commit_tracker_marks_done_when_all_storylines_commit() { + let root = tempfile::tempdir().unwrap(); + let wal = super::super::wal::ImportWal::open_or_create( + root.path(), + "from", + "to", + "storyline-lance", + None, + false, + false, + ) + .unwrap(); + let wal = std::sync::Arc::new(std::sync::Mutex::new(wal)); + let mut tracker = SourceCommitTracker::new(); + tracker.register("a.json", 2); + tracker.note_committed(&[String::from("a.json")], &Some(wal.clone())); + assert!(!wal.lock().unwrap().should_skip("a.json")); + tracker.note_committed(&[String::from("a.json")], &Some(wal.clone())); + assert!(wal.lock().unwrap().should_skip("a.json")); + } + + #[test] + fn retract_imported_trajectories_from_tail_sources() { + let mut sources = vec![ + ImportedSource { + source_path: "a.json".into(), + format: DocumentFormat::Atif, + trajectories: 3, + input_bytes: 10, + }, + ImportedSource { + source_path: "b.json".into(), + format: DocumentFormat::Atif, + trajectories: 2, + input_bytes: 10, + }, + ]; + retract_imported_trajectories(&mut sources, 3); + assert_eq!(sources[0].trajectories, 2); + assert_eq!(sources[1].trajectories, 0); + } + + #[test] + fn commit_batch_schedule_fixed_stays_put() { + let mut schedule = CommitBatchSchedule::fixed(50); + assert_eq!(schedule.current(), 50); + schedule.after_commit(); + assert_eq!(schedule.current(), 50); + } +} diff --git a/crates/persisting-pchronicle-cli/src/exchange/mod.rs b/crates/persisting-pchronicle-cli/src/exchange/mod.rs new file mode 100644 index 00000000..65881f91 --- /dev/null +++ b/crates/persisting-pchronicle-cli/src/exchange/mod.rs @@ -0,0 +1,23 @@ +//! Dataset exchange: import, export, drop, and sync snapshot helpers. + +mod decode; +mod drop; +mod export; +mod import; +mod pipeline; +mod progress; +mod staging; +mod sync; +mod wal; + +pub(crate) use decode::collect_visible_json_files; +pub(crate) use drop::run_drop; +pub(crate) use export::run_export; +pub(crate) use import::run_import; +pub(crate) use sync::sync_snapshot; + +// Re-exported for lib/tests; production call sites often go through sibling modules. +#[allow(unused_imports)] +pub(crate) use decode::validate_import_source; +#[allow(unused_imports)] +pub(crate) use staging::rename_noreplace; diff --git a/crates/persisting-pchronicle-cli/src/exchange/pipeline.rs b/crates/persisting-pchronicle-cli/src/exchange/pipeline.rs new file mode 100644 index 00000000..40b10e60 --- /dev/null +++ b/crates/persisting-pchronicle-cli/src/exchange/pipeline.rs @@ -0,0 +1,770 @@ +//! Generic multi-stage producer/consumer pipeline with bounded buffers. +//! +//! Stages communicate through `tokio::sync::mpsc` channels: a full buffer +//! applies backpressure to the upstream producer. Each stage reports through a +//! [`StageHandle`](super::progress::StageHandle). +//! +//! Import shape: +//! `discover (1) → fetch (N) →[8]→ parse (N) →[8]→ commit (1 task)` + +use super::decode::{DecodeImportOutcome, DecodedImportSource, ImportedSource}; +use super::progress::StageHandle; +use anyhow::{Result, anyhow}; +use persisting_pchronicle::model::StorylineDocument; +use std::collections::{BTreeMap, HashSet}; +use std::future::Future; +use std::path::PathBuf; +use std::sync::Arc; +use tokio::sync::mpsc; +use tokio::task::JoinHandle; + +/// Discover → fetch buffer (listing can run ahead of I/O). +pub(crate) const DISCOVER_TO_FETCH_BUFFER: usize = 64; +/// Fetch → parse buffer. Keep modest: each slot holds a full source payload. +pub(crate) const FETCH_TO_PARSE_BUFFER: usize = 8; +/// Parse → commit buffer. Small on purpose — decoded Storylines are heavy, and +/// remote commit (often S3 with concurrency 1) is the usual bottleneck; a large +/// backlog only burns RAM. Keep enough headroom that parse workers do not thrash +/// on every commit AIMD pause / full-batch flush. +pub(crate) const PARSE_TO_COMMIT_BUFFER: usize = 8; +/// Parallel fetch workers. +pub(crate) const FETCH_STAGE_CONCURRENCY: usize = 4; +/// Parallel parse workers. +pub(crate) const PARSE_STAGE_CONCURRENCY: usize = 4; + +/// One item flowing out of the discover stage. +#[derive(Debug, Clone)] +pub(crate) struct DiscoveredItem { + pub(crate) path: String, + pub(crate) bytes: u64, + /// Object-store Dataset root when the path is a remote key (kept for diagnostics). + #[allow(dead_code)] + pub(crate) remote_root: Option, +} + +/// Bytes loaded for one discovered source. +#[derive(Debug)] +pub(crate) struct FetchedItem { + pub(crate) path: String, + pub(crate) relative_path: PathBuf, + pub(crate) output_relative_path: Option, + pub(crate) bytes: Vec, + #[allow(dead_code)] + pub(crate) size_hint: u64, +} + +/// Decode result ready for the single-worker commit stage. +#[derive(Debug)] +pub(crate) enum ParsedItem { + Imported { + #[allow(dead_code)] + diagnostic_path: PathBuf, + metadata: ImportedSource, + storylines: Vec, + warnings: persisting_pchronicle::model::UnknownFieldImportWarnings, + }, + Skipped { + path: PathBuf, + reason: String, + #[allow(dead_code)] + bytes: u64, + }, +} + +/// A bounded link between two stages (backpressure when full). +pub(crate) struct StageChannel { + pub(crate) tx: mpsc::Sender>, + pub(crate) rx: mpsc::Receiver>, +} + +impl StageChannel { + pub(crate) fn bounded(capacity: usize) -> Self { + let (tx, rx) = mpsc::channel(capacity.max(1)); + Self { tx, rx } + } +} + +pub(crate) struct ParallelMapOptions { + pub(crate) capacity: usize, + pub(crate) workers: usize, + pub(crate) outbound: StageHandle, + pub(crate) downstream: &'static str, + pub(crate) track_outbound_queue: bool, +} + +/// Send into a bounded stage channel, surfacing backpressure on the progress line. +/// +/// When `track_inbound_queue` is true, `inbound`'s `queue_depth` is incremented on a +/// successful enqueue so the UI shows the real channel length. Commit uses batch +/// fill instead, so parse→commit passes `false`. +pub(crate) async fn send_with_flow_control( + tx: &mpsc::Sender>, + item: Result, + sender: &StageHandle, + inbound: &StageHandle, + downstream: &'static str, + track_inbound_queue: bool, +) -> bool { + match tx.try_reserve() { + Ok(permit) => { + permit.send(item); + if track_inbound_queue { + inbound.queue_push(); + } + true + } + Err(mpsc::error::TrySendError::Full(_)) => { + sender.enter_flow_wait(format!("pending→{downstream}")); + let ok = tx.send(item).await.is_ok(); + sender.leave_flow_wait(); + if ok && track_inbound_queue { + inbound.queue_push(); + } + ok + } + Err(mpsc::error::TrySendError::Closed(_)) => false, + } +} + +/// Spawn a source stage that only produces items (no upstream). +/// +/// `downstream` labels the next stage for backpressure UI (e.g. `"fetch"`). +pub(crate) fn spawn_source_stage( + capacity: usize, + progress: StageHandle, + body: F, +) -> (mpsc::Receiver>, JoinHandle<()>) +where + T: Send + 'static, + F: FnOnce(mpsc::Sender>, StageHandle) -> Fut + Send + 'static, + Fut: Future + Send + 'static, +{ + let StageChannel { tx, rx } = StageChannel::bounded(capacity); + let handle = progress.clone(); + let join = tokio::spawn(async move { + body(tx, handle).await; + }); + (rx, join) +} + +/// Spawn a 1:1 map stage: recv `In` → process → send `Out`. +#[cfg(test)] +pub(crate) fn spawn_map_stage( + mut rx: mpsc::Receiver>, + capacity: usize, + progress: StageHandle, + outbound: StageHandle, + downstream: &'static str, + mut map: F, +) -> (mpsc::Receiver>, JoinHandle<()>) +where + In: Send + 'static, + Out: Send + 'static, + F: FnMut(In, StageHandle) -> Fut + Send + 'static, + Fut: Future> + Send + 'static, +{ + let StageChannel { tx, rx: out_rx } = StageChannel::bounded(capacity); + let join = tokio::spawn(async move { + loop { + progress.enter_upstream_wait(); + let item = rx.recv().await; + progress.leave_upstream_wait(); + let Some(item) = item else { + break; + }; + progress.queue_pop(); + match item { + Ok(input) => match map(input, progress.clone()).await { + Ok(output) => { + if !send_with_flow_control( + &tx, + Ok(output), + &progress, + &outbound, + downstream, + true, + ) + .await + { + return; + } + } + Err(error) => { + progress.record_error(format!("{error:#}")); + let _ = send_with_flow_control( + &tx, + Err(error), + &progress, + &outbound, + downstream, + true, + ) + .await; + return; + } + }, + Err(error) => { + progress.record_error(format!("{error:#}")); + let _ = send_with_flow_control( + &tx, + Err(error), + &progress, + &outbound, + downstream, + true, + ) + .await; + return; + } + } + } + }); + (out_rx, join) +} + +/// Spawn a bounded parallel map stage (multi-worker). +pub(crate) fn spawn_parallel_map_stage( + mut rx: mpsc::Receiver>, + progress: StageHandle, + options: ParallelMapOptions, + map: F, +) -> (mpsc::Receiver>, JoinHandle<()>) +where + In: Send + 'static, + Out: Send + 'static, + F: Fn(In, StageHandle) -> Fut + Clone + Send + Sync + 'static, + Fut: Future> + Send + 'static, +{ + let StageChannel { tx, rx: out_rx } = StageChannel::bounded(options.capacity); + let workers = options.workers.max(1); + let outbound = options.outbound; + let downstream = options.downstream; + let track_outbound_queue = options.track_outbound_queue; + let join = tokio::spawn(async move { + let mut tasks = tokio::task::JoinSet::new(); + let mut pending = BTreeMap::new(); + let mut next_input = 0usize; + let mut next_output = 0usize; + let mut input_closed = false; + loop { + while !input_closed && tasks.len() < workers { + progress.enter_upstream_wait(); + let received = rx.recv().await; + progress.leave_upstream_wait(); + match received { + Some(item) => { + progress.queue_pop(); + let sequence = next_input; + next_input += 1; + let progress = progress.clone(); + let map = map.clone(); + if item.is_err() { + input_closed = true; + } + tasks.spawn(async move { + let result = match item { + Ok(input) => map(input, progress.clone()).await, + Err(error) => Err(error), + }; + if let Err(error) = &result { + progress.record_error(format!("{error:#}")); + } + (sequence, result) + }); + } + None => input_closed = true, + } + } + if tasks.is_empty() { + break; + } + let Some(joined) = tasks.join_next().await else { + break; + }; + let (sequence, result) = match joined { + Ok(result) => result, + Err(error) => { + progress.record_error(format!("parallel map worker failed: {error}")); + return; + } + }; + pending.insert(sequence, result); + while let Some(result) = pending.remove(&next_output) { + if !send_with_flow_control( + &tx, + result, + &progress, + &outbound, + downstream, + track_outbound_queue, + ) + .await + { + return; + } + next_output += 1; + } + } + }); + (out_rx, join) +} + +fn spawn_parse_stage( + fetched_rx: mpsc::Receiver>, + requested_format: crate::ExchangeFormat, + suggested_format: Option, + parse: StageHandle, + commit: StageHandle, + _unknown_field_warnings: Arc< + tokio::sync::Mutex, + >, +) -> (mpsc::Receiver>, JoinHandle<()>) { + spawn_parallel_map_stage( + fetched_rx, + parse, + ParallelMapOptions { + capacity: PARSE_TO_COMMIT_BUFFER, + workers: PARSE_STAGE_CONCURRENCY, + outbound: commit, + downstream: "commit", + track_outbound_queue: false, + }, + move |fetched, parse| { + async move { + let name = fetched.path.clone(); + parse.set_current(name.clone()); + let mut warnings = + persisting_pchronicle::model::UnknownFieldImportWarnings::default(); + let parse_result = super::decode::decode_import_source( + requested_format, + suggested_format, + crate::ImportOutputFormat::Storyline, + Some(std::path::Path::new(&fetched.path)), + Some(&fetched.relative_path), + fetched.output_relative_path.as_deref(), + &fetched.bytes, + &mut warnings, + ); + match parse_result { + Ok(DecodeImportOutcome::Imported(DecodedImportSource { + diagnostic_path, + metadata, + storylines, + })) => { + let bytes = metadata.input_bytes as u64; + if storylines.is_empty() { + parse.record_empty(1, bytes); + } else { + parse.record(1, bytes); + } + Ok(ParsedItem::Imported { + diagnostic_path, + metadata, + storylines, + warnings, + }) + } + Ok(DecodeImportOutcome::Skipped { path, reason }) => { + parse.record_skipped(1, fetched.bytes.len() as u64); + Ok(ParsedItem::Skipped { + path, + reason, + bytes: fetched.bytes.len() as u64, + }) + } + Err(error) => { + // Soft-skip: keep large imports moving; commit worker logs. + parse.record_error(format!("{error:#}")); + Ok(ParsedItem::Skipped { + path: PathBuf::from(&name), + reason: format!("{error:#}"), + bytes: fetched.bytes.len() as u64, + }) + } + } + } + }, + ) +} + +/// Wire helpers for import: discover → fetch → parse (commit is a separate task). +pub(crate) struct ImportPipelineHandles { + pub(crate) parsed_rx: mpsc::Receiver>, + pub(crate) joins: Vec>, + pub(crate) unknown_field_warnings: + Arc>, +} + +pub(crate) struct ImportPipelineConfig { + pub(crate) max_input_bytes: usize, + pub(crate) requested_format: crate::ExchangeFormat, + pub(crate) suggested_format: Option, + pub(crate) discover: StageHandle, + pub(crate) fetch: StageHandle, + pub(crate) parse: StageHandle, + pub(crate) commit: StageHandle, + /// Relative source paths already completed or failed in a prior run. + pub(crate) skip_paths: Arc>, +} + +/// Build discover→fetch→parse for a prelisted candidate set. +pub(crate) fn spawn_candidates_fetch_pipeline( + candidates: Vec, + config: ImportPipelineConfig, +) -> ImportPipelineHandles { + let ImportPipelineConfig { + max_input_bytes, + requested_format, + suggested_format, + discover, + fetch, + parse, + commit, + skip_paths, + } = config; + let unknown_field_warnings = Arc::new(tokio::sync::Mutex::new( + persisting_pchronicle::model::UnknownFieldImportWarnings::default(), + )); + let total = candidates.len() as u64; + discover.set_total_items(total); + fetch.set_queue_cap(DISCOVER_TO_FETCH_BUFFER as u64); + parse.set_queue_cap(FETCH_TO_PARSE_BUFFER as u64); + // Commit queue shows batch fill, configured in the commit task. + let fetch_for_discover = fetch.clone(); + let (discovered_rx, discover_join) = spawn_source_stage( + DISCOVER_TO_FETCH_BUFFER, + discover.clone(), + move |tx, discover| async move { + for candidate in candidates { + let path = candidate.relative_path.to_string_lossy().into_owned(); + if skip_paths.contains(&path) { + discover.record_skipped(1, candidate.size_hint); + continue; + } + let bytes = candidate.size_hint; + // Discover totals were already set via set_discovered; only + // refresh the activity label while feeding the fetch stage. + discover.set_current(path.clone()); + let item = DiscoveredItem { + path, + bytes, + remote_root: candidate.remote_root.clone(), + }; + if !send_with_flow_control( + &tx, + Ok((item, candidate)), + &discover, + &fetch_for_discover, + "reading", + true, + ) + .await + { + return; + } + } + discover.clear_current(); + }, + ); + + let parse_for_fetch = parse.clone(); + let (fetched_rx, fetch_join) = spawn_parallel_map_stage( + discovered_rx, + fetch, + ParallelMapOptions { + capacity: FETCH_TO_PARSE_BUFFER, + workers: FETCH_STAGE_CONCURRENCY, + outbound: parse_for_fetch, + downstream: "parsing", + track_outbound_queue: true, + }, + move |(item, candidate), fetch| async move { + fetch.set_current(item.path.clone()); + let label = format!("import source {}", item.path); + let bytes = + super::decode::load_import_candidate_bytes(&candidate, max_input_bytes, &label) + .await?; + let fetched = FetchedItem { + path: item.path, + relative_path: candidate.relative_path, + output_relative_path: candidate.output_relative_path, + size_hint: item.bytes, + bytes, + }; + fetch.record(1, fetched.bytes.len() as u64); + Ok(fetched) + }, + ); + + let (parsed_rx, parse_join) = spawn_parse_stage( + fetched_rx, + requested_format, + suggested_format, + parse, + commit, + Arc::clone(&unknown_field_warnings), + ); + + ImportPipelineHandles { + parsed_rx, + joins: vec![discover_join, fetch_join, parse_join], + unknown_field_warnings, + } +} + +/// Build discover→fetch→parse for an object-store (or local tree) location. +pub(crate) fn spawn_location_fetch_pipeline( + location: persisting_pchronicle::storage::DatasetLocation, + config: ImportPipelineConfig, +) -> ImportPipelineHandles { + let ImportPipelineConfig { + max_input_bytes, + requested_format, + suggested_format, + discover, + fetch, + parse, + commit, + skip_paths, + } = config; + let unknown_field_warnings = Arc::new(tokio::sync::Mutex::new( + persisting_pchronicle::model::UnknownFieldImportWarnings::default(), + )); + fetch.set_queue_cap(DISCOVER_TO_FETCH_BUFFER as u64); + parse.set_queue_cap(FETCH_TO_PARSE_BUFFER as u64); + // Commit queue shows batch fill, configured in the commit task. + let remote_root = location.as_str().to_owned(); + let fetch_for_discover = fetch.clone(); + let (discovered_rx, discover_join) = spawn_source_stage( + DISCOVER_TO_FETCH_BUFFER, + discover.clone(), + move |tx, discover| async move { + let list_result = location + .for_each_importable_json_object_event( + persisting_pchronicle::storage::DEFAULT_MAX_LOCAL_QUERY_FILES, + |event| { + let tx = tx.clone(); + let discover = discover.clone(); + let fetch = fetch_for_discover.clone(); + let skip_paths = Arc::clone(&skip_paths); + async move { + match event { + persisting_pchronicle::storage::ImportableObjectEvent::Scanning { + prefix, + } => { + let label = if prefix.is_empty() { + "/".to_owned() + } else { + format!("{prefix}/") + }; + discover.set_current(label); + Ok(()) + } + persisting_pchronicle::storage::ImportableObjectEvent::File { + key, + size, + .. + } => { + if skip_paths.contains(&key) { + discover.record_skipped(1, size); + return Ok(()); + } + discover.set_current(key.clone()); + discover.record(1, size); + if !send_with_flow_control( + &tx, + Ok(DiscoveredItem { + path: key, + bytes: size, + remote_root: None, + }), + &discover, + &fetch, + "reading", + true, + ) + .await + { + return Ok(()); + } + Ok(()) + } + } + } + }, + ) + .await; + if let Err(error) = list_result { + discover.record_error(format!("{error:#}")); + if tx.send(Err(error)).await.is_ok() { + fetch_for_discover.queue_push(); + } + return; + } + discover.clear_current(); + }, + ); + + let remote_root_for_fetch = remote_root; + let parse_for_fetch = parse.clone(); + let (fetched_rx, fetch_join) = spawn_parallel_map_stage( + discovered_rx, + fetch, + ParallelMapOptions { + capacity: FETCH_TO_PARSE_BUFFER, + workers: FETCH_STAGE_CONCURRENCY, + outbound: parse_for_fetch, + downstream: "parsing", + track_outbound_queue: true, + }, + move |item, fetch| { + let remote_root = remote_root_for_fetch.clone(); + async move { + fetch.set_current(item.path.clone()); + let relative_path = PathBuf::from(&item.path); + let candidate = super::decode::ImportFileCandidate { + path: relative_path.clone(), + output_relative_path: Some(relative_path.clone()), + relative_path: relative_path.clone(), + content: None, + remote_root: Some(remote_root), + size_hint: item.bytes, + }; + let label = format!("import source {}", item.path); + let bytes = + super::decode::load_import_candidate_bytes(&candidate, max_input_bytes, &label) + .await?; + fetch.record(1, bytes.len() as u64); + Ok(FetchedItem { + path: item.path, + relative_path, + output_relative_path: candidate.output_relative_path, + size_hint: item.bytes, + bytes, + }) + } + }, + ); + + let (parsed_rx, parse_join) = spawn_parse_stage( + fetched_rx, + requested_format, + suggested_format, + parse, + commit, + Arc::clone(&unknown_field_warnings), + ); + + ImportPipelineHandles { + parsed_rx, + joins: vec![discover_join, fetch_join, parse_join], + unknown_field_warnings, + } +} + +pub(crate) async fn join_pipeline_stages(joins: Vec>) -> Result<()> { + for join in joins { + match join.await { + Ok(()) => {} + Err(error) if error.is_cancelled() => {} + Err(error) => return Err(anyhow!("pipeline stage task failed: {error}")), + } + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::exchange::progress::{CliProgress, StageId}; + + #[tokio::test] + async fn map_stage_applies_backpressure_and_transforms() { + let progress = CliProgress::new(false); + let (rx, join) = + spawn_source_stage(1, progress.stage(StageId::Discover), |tx, _| async move { + for i in 0..5u64 { + tx.send(Ok(i)).await.unwrap(); + } + }); + let (mut out_rx, map_join) = spawn_map_stage( + rx, + 1, + progress.stage(StageId::Fetch), + progress.stage(StageId::Parse), + "parsing", + |n, _| async move { Ok(n * 10) }, + ); + let mut got = Vec::new(); + while let Some(item) = out_rx.recv().await { + got.push(item.unwrap()); + } + join_pipeline_stages(vec![join, map_join]).await.unwrap(); + assert_eq!(got, vec![0, 10, 20, 30, 40]); + } + + #[tokio::test] + async fn parallel_map_stage_uses_multiple_workers() { + let progress = CliProgress::new(false); + let (rx, join) = + spawn_source_stage(8, progress.stage(StageId::Discover), |tx, _| async move { + for i in 0..8u64 { + tx.send(Ok(i)).await.unwrap(); + } + }); + let (mut out_rx, map_join) = spawn_parallel_map_stage( + rx, + progress.stage(StageId::Fetch), + ParallelMapOptions { + capacity: FETCH_TO_PARSE_BUFFER, + workers: 4, + outbound: progress.stage(StageId::Parse), + downstream: "parsing", + track_outbound_queue: true, + }, + |n, _| async move { + tokio::time::sleep(std::time::Duration::from_millis(40 - n * 5)).await; + Ok(n) + }, + ); + let mut got = Vec::new(); + while let Some(item) = out_rx.recv().await { + got.push(item.unwrap()); + } + join_pipeline_stages(vec![join, map_join]).await.unwrap(); + assert_eq!(got, (0..8).collect::>()); + } + + #[tokio::test] + async fn map_stage_records_error_on_failure() { + let progress = CliProgress::new(false); + let fetch = progress.stage(StageId::Fetch); + let (rx, join) = + spawn_source_stage(2, progress.stage(StageId::Discover), |tx, _| async move { + let _ = tx.send(Ok(1u64)).await; + }); + let (mut out_rx, map_join) = spawn_map_stage( + rx, + 2, + fetch.clone(), + progress.stage(StageId::Parse), + "parsing", + |_n, _| async move { Err::(anyhow!("boom")) }, + ); + let err = out_rx.recv().await.unwrap().unwrap_err(); + assert!(format!("{err:#}").contains("boom")); + join_pipeline_stages(vec![join, map_join]).await.unwrap(); + } + + #[test] + fn buffer_constants_match_import_shape() { + assert_eq!(FETCH_TO_PARSE_BUFFER, 8); + assert_eq!(PARSE_TO_COMMIT_BUFFER, 8); + assert_eq!(FETCH_STAGE_CONCURRENCY, 4); + assert_eq!(PARSE_STAGE_CONCURRENCY, 4); + assert_eq!(StageId::Discover.noun(), "listing"); + assert_eq!(StageId::Fetch.verb(), "reading"); + assert_eq!(StageId::Parse.verb(), "parsing"); + assert_eq!(StageId::Commit.noun(), "commit"); + } +} diff --git a/crates/persisting-pchronicle-cli/src/exchange/progress.rs b/crates/persisting-pchronicle-cli/src/exchange/progress.rs new file mode 100644 index 00000000..728a55ef --- /dev/null +++ b/crates/persisting-pchronicle-cli/src/exchange/progress.rs @@ -0,0 +1,959 @@ +//! Pipeline-oriented CLI progress. +//! +//! Each pipeline stage owns one status line: +//! `listing ok=12 skipped=2 empty=1 error=0 queue=3/64 1.2GiB [listing] path.json` +//! The GiB column is attributed **source** bytes for that stage (not Lance/S3 +//! on-disk size). Commit attributes bytes when a trajectory enters its write +//! batch so the column stays aligned with reading/parsing under backpressure. +//! Bracket status: +//! - `waiting` — stalled on upstream (no item yet) +//! - `pending→X` — blocked because downstream buffer `X` is full +//! - AIMD detail always follows the word `aimd` on fetch/commit + +use super::super::*; +use std::io::Write; +use std::sync::Arc; + +/// Stable id for a progress line / pipeline stage. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub(crate) enum StageId { + Discover, + Fetch, + Parse, + Commit, + Delete, +} + +impl StageId { + pub(crate) fn noun(self) -> &'static str { + match self { + Self::Discover => "listing", + Self::Fetch => "reading", + Self::Parse => "parsing", + Self::Commit => "commit", + Self::Delete => "delete", + } + } + + pub(crate) fn verb(self) -> &'static str { + match self { + Self::Discover => "listing", + Self::Fetch => "reading", + Self::Parse => "parsing", + Self::Commit => "committing", + Self::Delete => "deleting", + } + } +} + +#[derive(Debug, Clone, Default)] +struct StageState { + /// Successfully processed items (files / trajectories). + ok: u64, + skipped: u64, + empty: u64, + error: u64, + bytes: u64, + /// Items accepted from the upstream stage (commit: trajectories ready). + inbound: u64, + /// Live inbound channel depth (pushed on enqueue, popped on dequeue). + queue_depth: u64, + /// Inbound channel capacity for `queue=depth/cap` display. + queue_cap: Option, + /// Optional known total (delete wipe, prelisted discover). + total_items: Option, + current: String, + last_error: Option, + /// Refcount of workers blocked on downstream backpressure. + flow_waiters: u32, + /// Refcount of workers blocked waiting for an upstream item. + upstream_waiters: u32, + /// Human-readable downstream pending reason (e.g. `pending→parsing`). + flow: Option, + /// Active object-store AIMD wait reason (`throttle` / `admit` / `backoff`), if any. + aimd_event: Option, + /// Remaining AIMD wait from the latest gate tick (ms); drives live `cd=`. + aimd_wait_ms: Option, +} + +impl StageState { + fn processed(&self) -> u64 { + self.ok + .saturating_add(self.skipped) + .saturating_add(self.empty) + .saturating_add(self.error) + } + + fn format_line(&self, id: StageId, queue: &str, status: &str) -> String { + let activity = if self.current.is_empty() { + "-".into() + } else { + truncate_middle(&self.current, 72) + }; + let bracket = stage_bracket(id, self.upstream_waiters > 0, self.flow_waiters > 0, status); + format!( + "{}\tok={} skipped={} empty={} error={}\tqueue={}\t{}\t[{bracket}] {}", + id.noun(), + self.ok, + self.skipped, + self.empty, + self.error, + queue, + format_byte_count(self.bytes), + activity, + ) + } +} + +/// Shared handle a running stage uses to report work / errors. +#[derive(Clone)] +pub(crate) struct StageHandle { + id: StageId, + state: Arc>, + painter: Arc>, +} + +impl StageHandle { + #[allow(dead_code)] + pub(crate) fn id(&self) -> StageId { + self.id + } + + pub(crate) fn set_current(&self, item: impl Into) { + if let Ok(mut state) = self.state.lock() { + state.current = item.into(); + } + let _ = self.repaint(); + } + + pub(crate) fn clear_current(&self) { + if let Ok(mut state) = self.state.lock() { + state.current.clear(); + } + let _ = self.repaint(); + } + + pub(crate) fn set_total_items(&self, total: u64) { + if let Ok(mut state) = self.state.lock() { + state.total_items = Some(total); + } + let _ = self.repaint(); + } + + pub(crate) fn set_queue_cap(&self, cap: u64) { + if let Ok(mut state) = self.state.lock() { + state.queue_cap = Some(cap); + } + let _ = self.repaint(); + } + + /// One item entered this stage's inbound channel. + pub(crate) fn queue_push(&self) { + if let Ok(mut state) = self.state.lock() { + state.queue_depth = state.queue_depth.saturating_add(1); + if let Some(cap) = state.queue_cap { + state.queue_depth = state.queue_depth.min(cap); + } + } + let _ = self.repaint(); + } + + /// One item left this stage's inbound channel. + pub(crate) fn queue_pop(&self) { + if let Ok(mut state) = self.state.lock() { + state.queue_depth = state.queue_depth.saturating_sub(1); + } + let _ = self.repaint(); + } + + /// Set absolute queue depth (e.g. commit batch fill). + pub(crate) fn set_queue(&self, queue: u64) { + if let Ok(mut state) = self.state.lock() { + state.queue_depth = match state.queue_cap { + Some(cap) => queue.min(cap), + None => queue, + }; + } + let _ = self.repaint(); + } + + /// Mark this stage blocked on downstream backpressure (`pending→…`). + pub(crate) fn enter_flow_wait(&self, reason: impl Into) { + if let Ok(mut state) = self.state.lock() { + state.flow_waiters = state.flow_waiters.saturating_add(1); + state.flow = Some(reason.into()); + } + let _ = self.repaint(); + } + + /// Clear one downstream-pending waiter; label drops when the last waiter leaves. + pub(crate) fn leave_flow_wait(&self) { + if let Ok(mut state) = self.state.lock() { + state.flow_waiters = state.flow_waiters.saturating_sub(1); + if state.flow_waiters == 0 { + state.flow = None; + } + } + let _ = self.repaint(); + } + + /// Mark this stage blocked waiting for an upstream item. + pub(crate) fn enter_upstream_wait(&self) { + if let Ok(mut state) = self.state.lock() { + state.upstream_waiters = state.upstream_waiters.saturating_add(1); + } + let _ = self.repaint(); + } + + /// Clear one upstream-wait waiter. + pub(crate) fn leave_upstream_wait(&self) { + if let Ok(mut state) = self.state.lock() { + state.upstream_waiters = state.upstream_waiters.saturating_sub(1); + } + let _ = self.repaint(); + } + + /// Overlay AIMD reason + optional remaining wait; always repaints this stage. + pub(crate) fn set_aimd_status(&self, event: Option, wait_ms: Option) { + if let Ok(mut state) = self.state.lock() { + state.aimd_event = event; + state.aimd_wait_ms = wait_ms; + } + let _ = self.repaint(); + } + + pub(crate) fn record(&self, items: u64, bytes: u64) { + if let Ok(mut state) = self.state.lock() { + state.ok = state.ok.saturating_add(items); + state.bytes = state.bytes.saturating_add(bytes); + } + let _ = self.repaint(); + } + + pub(crate) fn record_skipped(&self, items: u64, bytes: u64) { + if let Ok(mut state) = self.state.lock() { + state.skipped = state.skipped.saturating_add(items); + state.bytes = state.bytes.saturating_add(bytes); + } + let _ = self.repaint(); + } + + pub(crate) fn record_empty(&self, items: u64, bytes: u64) { + if let Ok(mut state) = self.state.lock() { + state.empty = state.empty.saturating_add(items); + state.bytes = state.bytes.saturating_add(bytes); + } + let _ = self.repaint(); + } + + pub(crate) fn record_inbound(&self, items: u64) { + if let Ok(mut state) = self.state.lock() { + state.inbound = state.inbound.saturating_add(items); + } + let _ = self.repaint(); + } + + #[allow(dead_code)] + pub(crate) fn set_items(&self, items: u64) { + if let Ok(mut state) = self.state.lock() { + state.ok = items; + } + let _ = self.repaint(); + } + + #[allow(dead_code)] + pub(crate) fn record_bytes(&self, bytes: u64) { + if let Ok(mut state) = self.state.lock() { + state.bytes = state.bytes.saturating_add(bytes); + } + let _ = self.repaint(); + } + + /// Replace the size column with an absolute value (e.g. measured on-disk). + pub(crate) fn set_bytes(&self, bytes: u64) { + if let Ok(mut state) = self.state.lock() { + state.bytes = bytes; + } + let _ = self.repaint(); + } + + pub(crate) fn record_error(&self, error: impl std::fmt::Display) { + if let Ok(mut state) = self.state.lock() { + state.last_error = Some(error.to_string()); + state.error = state.error.saturating_add(1); + } + let _ = self.repaint(); + } + + /// In-place activity override (e.g. index build note on the commit line). + pub(crate) fn set_activity_override(&self, activity: &str) { + if let Ok(mut painter) = self.painter.lock() { + painter.activity_override = Some((self.id, activity.to_owned())); + let _ = painter.paint(); + } + } + + #[allow(dead_code)] + pub(crate) fn clear_activity_override(&self) { + if let Ok(mut painter) = self.painter.lock() { + painter.activity_override = None; + let _ = painter.paint(); + } + } + + pub(crate) fn note_committed(&self, committed: u64, batch_bytes: u64) { + if let Ok(mut state) = self.state.lock() { + state.ok = committed; + state.bytes = state.bytes.saturating_add(batch_bytes); + state.current = format!("trajectories={committed}"); + } + if let Ok(mut painter) = self.painter.lock() { + painter.activity_override = None; + if let Ok(state) = self.state.lock() { + painter.stages.insert(self.id, state.clone()); + } + if !painter.tty { + let line = painter.line_for(self.id); + painter.log_lines.push(line); + } + let _ = painter.paint(); + } + } + + fn repaint(&self) -> Result<()> { + let snapshot = self + .state + .lock() + .map(|state| state.clone()) + .unwrap_or_default(); + if let Ok(mut painter) = self.painter.lock() { + painter.stages.insert(self.id, snapshot); + painter.paint()?; + } + Ok(()) + } +} + +#[derive(Debug, Default)] +struct PipelinePainter { + tty: bool, + painted_lines: usize, + /// When true, only the Delete stage line is shown (replace wipe). + delete_mode: bool, + stages: std::collections::HashMap, + order: Vec, + activity_override: Option<(StageId, String)>, + log_lines: Vec, + last_paint: Option, +} + +impl PipelinePainter { + fn stage_state(&self, id: StageId) -> StageState { + self.stages.get(&id).cloned().unwrap_or_default() + } + + /// Live inbound channel depth / capacity (not derived from ok counters). + /// Depth is clamped to capacity so concurrent push/pop races never paint + /// impossible values like `65/64`. + fn queue_for(&self, id: StageId) -> String { + let state = self.stage_state(id); + match id { + StageId::Discover | StageId::Delete => state + .total_items + .map(|total| { + let remaining = total.saturating_sub(state.processed()); + format!("{remaining}/{total}") + }) + .unwrap_or_else(|| "-".into()), + StageId::Fetch | StageId::Parse | StageId::Commit => match state.queue_cap { + Some(cap) => format!("{}/{}", state.queue_depth.min(cap), cap), + None => format!("{}", state.queue_depth), + }, + } + } + + fn line_for(&self, id: StageId) -> String { + let state = self.stage_state(id); + let queue = self.queue_for(id); + let status = enriched_status_label(id, &state); + if let Some((override_id, activity)) = &self.activity_override + && *override_id == id + { + let size = format_byte_count(state.bytes); + let bracket = if status.is_empty() { + "writing".to_owned() + } else { + format!("writing {status}") + }; + return format!( + "{}\tok={} skipped={} empty={} error={}\tqueue={queue}\t{size}\t[{bracket}] {}", + id.noun(), + state.ok, + state.skipped, + state.empty, + state.error, + truncate_middle(activity, 72), + ); + } + state.format_line(id, &queue, &status) + } + + fn visible_ids(&self) -> Vec { + if self.delete_mode { + vec![StageId::Delete] + } else { + self.order.clone() + } + } + + fn paint(&mut self) -> Result<()> { + let lines: Vec = self + .visible_ids() + .into_iter() + .map(|id| self.line_for(id)) + .collect(); + if self.tty { + let mut err = std::io::stderr(); + if self.painted_lines > 0 { + write!(err, "\x1b[{}A", self.painted_lines) + .context("move pipeline progress cursor")?; + } + for line in &lines { + write!(err, "\r\x1b[2K{line}\n").context("paint pipeline progress")?; + } + // Clear leftover lines if stage count shrank (e.g. leaving delete mode). + for _ in lines.len()..self.painted_lines { + write!(err, "\r\x1b[2K\n").context("clear stale progress line")?; + } + if lines.len() < self.painted_lines { + write!(err, "\x1b[{}A", self.painted_lines - lines.len()) + .context("rewind after clearing stale lines")?; + } + err.flush().context("flush pipeline progress")?; + self.painted_lines = lines.len(); + self.last_paint = Some(std::time::Instant::now()); + return Ok(()); + } + Ok(()) + } + + fn should_throttle(&self) -> bool { + self.tty + && self + .last_paint + .map(|at| at.elapsed() < std::time::Duration::from_millis(100)) + .unwrap_or(false) + } + + fn finish_tty(&mut self) -> Result<()> { + if self.tty && self.painted_lines > 0 { + let mut err = std::io::stderr(); + writeln!(err).context("finish pipeline progress")?; + err.flush().context("flush pipeline progress")?; + self.painted_lines = 0; + } + Ok(()) + } +} + +/// Multi-stage progress surface used by import (and reusable by export/sync). +pub(crate) struct CliProgress { + painter: Arc>, + handles: std::collections::HashMap, + /// Index-build callbacks paint onto the commit stage. + index_surface: Arc>, +} + +struct IndexActivityBridge { + commit: Option, +} + +impl CliProgress { + pub(crate) fn new(tty: bool) -> Self { + let order = vec![ + StageId::Discover, + StageId::Fetch, + StageId::Parse, + StageId::Commit, + ]; + let painter = Arc::new(std::sync::Mutex::new(PipelinePainter { + tty, + painted_lines: 0, + delete_mode: false, + stages: std::collections::HashMap::new(), + order: order.clone(), + activity_override: None, + log_lines: Vec::new(), + last_paint: None, + })); + let mut handles = std::collections::HashMap::new(); + for id in order { + let state = Arc::new(std::sync::Mutex::new(StageState::default())); + if let Ok(mut painter) = painter.lock() { + painter.stages.insert(id, StageState::default()); + } + handles.insert( + id, + StageHandle { + id, + state, + painter: Arc::clone(&painter), + }, + ); + } + // Delete stage exists but is only shown in delete_mode. + let delete_state = Arc::new(std::sync::Mutex::new(StageState::default())); + handles.insert( + StageId::Delete, + StageHandle { + id: StageId::Delete, + state: delete_state, + painter: Arc::clone(&painter), + }, + ); + let index_surface = Arc::new(std::sync::Mutex::new(IndexActivityBridge { + commit: handles.get(&StageId::Commit).cloned(), + })); + Self { + painter, + handles, + index_surface, + } + } + + pub(crate) fn stage(&self, id: StageId) -> StageHandle { + self.handles + .get(&id) + .cloned() + .expect("stage registered in CliProgress::new") + } + + pub(crate) fn attach_index_progress( + &self, + ) -> persisting_pchronicle::storage::IndexBuildProgressGuard { + let bridge = Arc::clone(&self.index_surface); + persisting_pchronicle::storage::install_index_build_progress(Arc::new(move |message| { + if let Ok(bridge) = bridge.lock() + && let Some(commit) = &bridge.commit + { + commit.set_activity_override(message); + } + })) + } + + /// Mirror object-store AIMD / admit waits onto fetch (read) and commit (write). + pub(crate) fn attach_object_store_throttle( + &self, + ) -> persisting_pchronicle::storage::ObjectStoreThrottleHookGuard { + let fetch = self.stage(StageId::Fetch); + let commit = self.stage(StageId::Commit); + persisting_pchronicle::storage::install_object_store_throttle_hook(Arc::new(move |event| { + let apply = |kind: persisting_pchronicle::storage::ObjectStoreIoKind, + reason: &str, + wait_ms: Option| { + let (primary, sibling) = match kind { + persisting_pchronicle::storage::ObjectStoreIoKind::Read => (&fetch, &commit), + persisting_pchronicle::storage::ObjectStoreIoKind::Write => (&commit, &fetch), + }; + let overlay = match reason { + "recover" | "ok" | "" => None, + other => Some(other.to_owned()), + }; + primary.set_aimd_status(overlay, wait_ms); + // Sibling line also re-reads the shared AIMD snapshot. + let _ = sibling.repaint(); + }; + match event { + persisting_pchronicle::storage::ObjectStoreThrottleEvent::Enter { + kind, + reason, + wait_ms, + .. + } + | persisting_pchronicle::storage::ObjectStoreThrottleEvent::Update { + kind, + reason, + wait_ms, + .. + } => { + let wait = (wait_ms > 0).then_some(wait_ms); + apply(kind, reason, wait); + } + persisting_pchronicle::storage::ObjectStoreThrottleEvent::Leave { kind } => { + apply(kind, "", None); + } + } + })) + } + + #[allow(dead_code)] + pub(crate) fn reset_import_counters(&mut self) { + for id in [ + StageId::Discover, + StageId::Fetch, + StageId::Parse, + StageId::Commit, + StageId::Delete, + ] { + let handle = self.stage(id); + if let Ok(mut state) = handle.state.lock() { + *state = StageState::default(); + } + let _ = handle.repaint(); + } + if let Ok(mut painter) = self.painter.lock() { + painter.delete_mode = false; + painter.activity_override = None; + for id in &painter.order.clone() { + painter.stages.insert(*id, StageState::default()); + } + } + } + + pub(crate) fn set_discovered(&mut self, files: u64, bytes: u64) -> Result<()> { + let discover = self.stage(StageId::Discover); + if let Ok(mut state) = discover.state.lock() { + state.ok = files; + state.bytes = bytes; + state.total_items = Some(files); + state.current.clear(); + } + for id in [StageId::Fetch, StageId::Parse] { + let handle = self.stage(id); + if let Ok(mut state) = handle.state.lock() { + state.total_items = Some(files); + } + } + discover.repaint() + } + + pub(crate) fn note_discovered(&mut self, file: &str, bytes: u64) -> Result<()> { + let discover = self.stage(StageId::Discover); + discover.set_current(file); + let throttle = self + .painter + .lock() + .map(|p| p.should_throttle()) + .unwrap_or(false); + if let Ok(mut state) = discover.state.lock() { + state.ok = state.ok.saturating_add(1); + state.bytes = state.bytes.saturating_add(bytes); + } + if throttle { + return Ok(()); + } + discover.repaint() + } + + pub(crate) fn note_fetched(&mut self, file: &str, bytes: u64) -> Result<()> { + let fetch = self.stage(StageId::Fetch); + fetch.set_current(file); + fetch.record(1, bytes); + Ok(()) + } + + pub(crate) fn note_parsed(&mut self, file: &str, bytes: u64) -> Result<()> { + let parse = self.stage(StageId::Parse); + parse.set_current(file); + parse.record(1, bytes); + // Non-TTY: emit a dense completed line when a source finishes parse. + if let Ok(mut painter) = self.painter.lock() + && !painter.tty + { + let line = format!( + "{}; {}; {}; {}", + painter.line_for(StageId::Discover), + painter.line_for(StageId::Fetch), + painter.line_for(StageId::Parse), + painter.line_for(StageId::Commit), + ); + painter.log_lines.push(line); + } + Ok(()) + } + + pub(crate) fn note_deleted(&mut self, deleted: u64, total: u64, path: &str) -> Result<()> { + if let Ok(mut painter) = self.painter.lock() { + painter.delete_mode = true; + } + let delete = self.stage(StageId::Delete); + delete.set_total_items(total); + if let Ok(mut state) = delete.state.lock() { + state.ok = deleted; + state.current = path.to_owned(); + } + if deleted == total { + delete.repaint()?; + if let Ok(mut painter) = self.painter.lock() { + if !painter.tty { + let line = painter.line_for(StageId::Delete); + painter.log_lines.push(line); + } + painter.delete_mode = false; + } + return Ok(()); + } + let throttle = self + .painter + .lock() + .map(|p| p.should_throttle()) + .unwrap_or(false); + if throttle && deleted > 1 && !deleted.is_multiple_of(64) { + return Ok(()); + } + delete.repaint() + } + + #[allow(dead_code)] + pub(crate) fn note_committed(&self, committed: u64, batch_bytes: u64) -> Result<()> { + self.stage(StageId::Commit) + .note_committed(committed, batch_bytes); + Ok(()) + } + + pub(crate) fn finish(&mut self) -> Result<()> { + if let Ok(mut painter) = self.painter.lock() { + painter.activity_override = None; + painter.finish_tty()?; + } + Ok(()) + } + + pub(crate) fn notice(&mut self, message: &str) -> Result<()> { + self.finish()?; + if let Ok(painter) = self.painter.lock() { + if painter.tty { + let mut err = std::io::stderr(); + writeln!(err, "{message}").context("write import notice")?; + err.flush().context("flush import notice")?; + } else { + drop(painter); + if let Ok(mut painter) = self.painter.lock() { + painter.log_lines.push(message.to_owned()); + } + } + } + Ok(()) + } + + pub(crate) fn flush_log(self, out: &mut dyn Write) -> Result<()> { + if let Ok(painter) = self.painter.lock() { + for line in &painter.log_lines { + writeln!(out, "{line}").context("flush import progress log")?; + } + } + Ok(()) + } +} + +pub(crate) fn format_byte_count(bytes: u64) -> String { + const KIB: f64 = 1024.0; + const MIB: f64 = 1024.0 * 1024.0; + const GIB: f64 = 1024.0 * 1024.0 * 1024.0; + let value = bytes as f64; + if value >= GIB { + format!("{:.1}GiB", value / GIB) + } else if value >= MIB { + format!("{:.1}MiB", value / MIB) + } else if value >= KIB { + format!("{:.1}KiB", value / KIB) + } else { + format!("{bytes}B") + } +} + +pub(crate) fn truncate_middle(value: &str, max_chars: usize) -> String { + let chars: Vec = value.chars().collect(); + if chars.len() <= max_chars { + return value.to_owned(); + } + if max_chars <= 3 { + return chars.into_iter().take(max_chars).collect(); + } + let head = (max_chars - 1) / 2; + let tail = max_chars - 1 - head; + let mut out: String = chars.iter().take(head).collect(); + out.push('…'); + out.extend(chars.iter().skip(chars.len() - tail)); + out +} + +fn stage_bracket( + id: StageId, + waiting_upstream: bool, + pending_downstream: bool, + status: &str, +) -> String { + let status = status.replace(',', " ").trim().to_owned(); + let verb = if waiting_upstream { + "waiting".to_owned() + } else if pending_downstream { + // Prefer an explicit `pending→…` token already in status. + if status + .split_whitespace() + .any(|part| part.starts_with("pending→")) + { + String::new() + } else { + "pending".to_owned() + } + } else { + id.verb().to_owned() + }; + + match (verb.is_empty(), status.is_empty()) { + (true, true) => id.verb().into(), + (true, false) => status, + (false, true) => verb, + (false, false) => format!("{verb} {status}"), + } +} + +fn enriched_status_label(id: StageId, state: &StageState) -> String { + let mut parts = Vec::new(); + if let Some(flow) = &state.flow { + parts.push(flow.clone()); + } + if matches!(id, StageId::Fetch | StageId::Commit) { + let snap = persisting_pchronicle::storage::object_store_gate_snapshot(); + // Prefer the live tick's remaining wait when present so `cd=` moves + // even if the paint lands between gate sleeps. + let mut snap = snap; + if let Some(wait_ms) = state.aimd_wait_ms { + snap.cooldown_remaining_ms = wait_ms; + } + parts.push( + persisting_pchronicle::storage::format_object_store_aimd_flow_label( + &snap, + state.aimd_event.as_deref(), + ), + ); + } + parts.join(" ") +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn format_byte_count_uses_binary_units() { + assert_eq!(format_byte_count(512), "512B"); + assert_eq!(format_byte_count(1536), "1.5KiB"); + assert_eq!(format_byte_count(2 * 1024 * 1024), "2.0MiB"); + } + + #[test] + fn stage_line_matches_pipeline_shape() { + let mut state = StageState { + ok: 12, + skipped: 2, + empty: 1, + error: 0, + bytes: 1536, + inbound: 0, + queue_depth: 0, + queue_cap: None, + total_items: None, + current: "a/long.json".into(), + last_error: None, + flow_waiters: 0, + upstream_waiters: 0, + flow: None, + aimd_event: None, + aimd_wait_ms: None, + }; + let line = state.format_line(StageId::Discover, "-", ""); + assert!( + line.starts_with("listing\tok=12 skipped=2 empty=1 error=0\tqueue=-\t1.5KiB\t"), + "{line}" + ); + assert!(!line.contains("flow="), "{line}"); + assert!(line.contains("[listing]"), "{line}"); + assert!(line.contains("a/long.json"), "{line}"); + + state.error = 3; + state.flow = Some("pending→parsing".into()); + state.flow_waiters = 1; + let wait_line = state.format_line(StageId::Fetch, "4/64", "pending→parsing"); + assert!(wait_line.contains("error=3"), "{wait_line}"); + assert!(wait_line.contains("queue=4/64"), "{wait_line}"); + assert!(!wait_line.contains("flow="), "{wait_line}"); + assert!(wait_line.contains("[pending→parsing]"), "{wait_line}"); + + state.flow_waiters = 0; + state.flow = None; + state.upstream_waiters = 1; + let upstream_line = state.format_line(StageId::Fetch, "0/64", "aimd ok s=0/4 p=1/1"); + assert!( + upstream_line.contains("[waiting aimd ok"), + "{upstream_line}" + ); + } + + #[test] + fn queue_tracks_inbound_channel_depth() { + let progress = CliProgress::new(false); + let fetch = progress.stage(StageId::Fetch); + let parse = progress.stage(StageId::Parse); + let commit = progress.stage(StageId::Commit); + fetch.set_queue_cap(64); + parse.set_queue_cap(8); + commit.set_queue_cap(4096); + fetch.queue_push(); + fetch.queue_push(); + parse.queue_push(); + commit.set_queue(128); + // Concurrent races must never paint above capacity. + for _ in 0..62 { + fetch.queue_push(); + } + + let painter = progress.painter.lock().unwrap(); + assert_eq!(painter.queue_for(StageId::Fetch), "64/64"); + assert_eq!(painter.queue_for(StageId::Parse), "1/8"); + assert_eq!(painter.queue_for(StageId::Commit), "128/4096"); + } + + #[test] + fn bracket_embeds_wait_pending_and_aimd_status() { + assert_eq!( + stage_bracket(StageId::Parse, true, false, "aimd ok s=0/4 p=1/1"), + "waiting aimd ok s=0/4 p=1/1" + ); + assert_eq!( + stage_bracket( + StageId::Fetch, + false, + true, + "pending→parsing aimd ok s=0/4 p=0/1" + ), + "pending→parsing aimd ok s=0/4 p=0/1" + ); + assert_eq!( + stage_bracket(StageId::Commit, false, false, "aimd ok s=0/4 p=1/1"), + "committing aimd ok s=0/4 p=1/1" + ); + } + + #[test] + fn non_tty_progress_logs_parse_and_commit_lines() { + let mut progress = CliProgress::new(false); + progress.set_discovered(2, 300).unwrap(); + progress.note_discovered("a.json", 100).unwrap(); + progress.note_discovered("b.json", 200).unwrap(); + progress.note_fetched("a.json", 100).unwrap(); + progress.note_parsed("a.json", 100).unwrap(); + progress.note_fetched("b.json", 200).unwrap(); + progress.note_parsed("b.json", 200).unwrap(); + progress.note_committed(3, 0).unwrap(); + let mut out = Vec::new(); + progress.flush_log(&mut out).unwrap(); + let text = String::from_utf8(out).unwrap(); + assert!(text.contains("listing\t"), "{text}"); + assert!(text.contains("reading\t"), "{text}"); + assert!(text.contains("parsing\t"), "{text}"); + assert!(text.contains("commit\t"), "{text}"); + assert!(!text.contains("status=fetching"), "{text}"); + } +} diff --git a/crates/persisting-pchronicle-cli/src/exchange/staging.rs b/crates/persisting-pchronicle-cli/src/exchange/staging.rs new file mode 100644 index 00000000..9abf6b16 --- /dev/null +++ b/crates/persisting-pchronicle-cli/src/exchange/staging.rs @@ -0,0 +1,153 @@ +//! Local staging and atomic publish helpers. + +use super::super::*; +use super::progress::CliProgress; +use anyhow::{Context, Result, anyhow}; +use std::ffi::CString; +use std::path::{Path, PathBuf}; + +pub(crate) struct StagingPathGuard { + path: Option, +} + +impl StagingPathGuard { + pub(crate) fn new(path: PathBuf) -> Self { + Self { path: Some(path) } + } + + pub(crate) fn disarm(&mut self) { + self.path = None; + } +} + +impl Drop for StagingPathGuard { + fn drop(&mut self) { + if let Some(path) = &self.path { + let _ = std::fs::remove_dir_all(path); + } + } +} + +pub(crate) async fn publish_staged_dataset( + staging: &Path, + output: &Path, + replace_existing: bool, + progress: Option<&mut CliProgress>, +) -> Result<()> { + let parent = output + .parent() + .context("Dataset output must have a parent directory")?; + if !replace_existing { + rename_noreplace(staging, output) + .with_context(|| format!("publish new Dataset {}", output.display()))?; + sync_dataset_parent(parent)?; + return Ok(()); + } + + let backup = parent.join(format!( + ".pchronicle-replace-{}-{}", + output + .file_name() + .map(|name| name.to_string_lossy()) + .unwrap_or_else(|| std::borrow::Cow::Borrowed("dataset")), + uuid::Uuid::new_v4().simple() + )); + rename_noreplace(output, &backup) + .with_context(|| format!("move existing Dataset to {}", backup.display()))?; + if let Err(error) = sync_dataset_parent(parent) { + return Err(rollback_replacement(output, &backup, error)); + } + if let Err(error) = rename_noreplace(staging, output) + .with_context(|| format!("publish replacement Dataset {}", output.display())) + { + return Err(rollback_replacement(output, &backup, error)); + } + sync_dataset_parent(parent).with_context(|| { + format!( + "sync replacement Dataset parent {}; old Dataset remains at {}", + parent.display(), + backup.display() + ) + })?; + let backup_location = DatasetLocation::parse( + backup + .to_str() + .context("replaced Dataset backup path is not valid UTF-8")?, + )?; + if let Some(progress) = progress { + backup_location + .remove_all_with_progress(|deleted, total, path| { + progress.note_deleted(deleted, total, path) + }) + .await + .with_context(|| format!("delete replaced Dataset backup {}", backup.display()))?; + progress.finish()?; + } else { + backup_location + .remove_all() + .await + .with_context(|| format!("delete replaced Dataset backup {}", backup.display()))?; + } + sync_dataset_parent(parent)?; + Ok(()) +} + +pub(crate) fn rollback_replacement( + output: &Path, + backup: &Path, + error: anyhow::Error, +) -> anyhow::Error { + match rename_noreplace(backup, output) { + Ok(()) => error, + Err(rollback_error) => anyhow!( + "{error}; failed to restore old Dataset from {} to {}: {rollback_error}", + backup.display(), + output.display() + ), + } +} + +pub(crate) fn sync_dataset_parent(parent: &Path) -> Result<()> { + std::fs::File::open(parent) + .and_then(|directory| directory.sync_all()) + .with_context(|| format!("sync Dataset parent {}", parent.display()))?; + Ok(()) +} + +#[cfg(any(target_os = "linux", target_os = "macos"))] +pub(crate) fn rename_noreplace(from: &Path, to: &Path) -> std::io::Result<()> { + use std::os::unix::ffi::OsStrExt; + + let from = CString::new(from.as_os_str().as_bytes())?; + let to = CString::new(to.as_os_str().as_bytes())?; + #[cfg(target_os = "linux")] + // SAFETY: both pointers come from live CString values and are NUL-terminated. + // Call SYS_renameat2 directly so the binary still links on manylinux2014 + // (glibc 2.17). The renameat2() wrapper only exists in glibc 2.28+. + let result = unsafe { + libc::syscall( + libc::SYS_renameat2, + libc::AT_FDCWD, + from.as_ptr(), + libc::AT_FDCWD, + to.as_ptr(), + libc::RENAME_NOREPLACE, + ) + }; + #[cfg(target_os = "macos")] + // SAFETY: both pointers come from live CString values and are NUL-terminated. + let result = unsafe { libc::renamex_np(from.as_ptr(), to.as_ptr(), libc::RENAME_EXCL) }; + if result == 0 { + Ok(()) + } else { + Err(std::io::Error::last_os_error()) + } +} + +#[cfg(not(any(target_os = "linux", target_os = "macos")))] +pub(crate) fn rename_noreplace(_from: &Path, _to: &Path) -> std::io::Result<()> { + Err(std::io::Error::new( + std::io::ErrorKind::Unsupported, + "atomic create-only Dataset publish is unsupported on this platform", + )) +} diff --git a/crates/persisting-pchronicle-cli/src/exchange/sync.rs b/crates/persisting-pchronicle-cli/src/exchange/sync.rs new file mode 100644 index 00000000..b3f14bbb --- /dev/null +++ b/crates/persisting-pchronicle-cli/src/exchange/sync.rs @@ -0,0 +1,102 @@ +//! Resident sync worker snapshot import. + +use super::super::*; +use super::import::{run_compact_jsonl_import, run_import}; +use anyhow::{Context, Result}; +use std::io::Write; + +/// Run one coalesced snapshot for the resident sync worker. +/// +/// - `--mirror` writes a Compact JSONL Lance Dataset (record-level ingest). +/// - `--to` writes a Storyline Lance Dataset (trajectory conversion). +/// +/// Either or both destinations may be set. Each reuses the import pipeline +/// (stage progress, replace semantics, publication) so sync and import share +/// the same listing → reading → parsing → commit surface. +#[allow(clippy::too_many_arguments)] +pub(crate) async fn sync_snapshot( + source: &str, + mirror: Option<&str>, + storyline: Option<&str>, + input_format: ExchangeFormat, + suggested_format: Option, + columns: &[String], + stderr: &mut dyn Write, + stderr_is_terminal: bool, +) -> Result<()> { + anyhow::ensure!( + mirror.is_some() || storyline.is_some(), + "sync requires --mirror and/or --to" + ); + + if let Some(mirror) = mirror { + let mut stdout = std::io::sink(); + run_compact_jsonl_import( + ImportArgs { + from: source.to_owned(), + output: Some(mirror.to_owned()), + format: ExchangeFormat::CompactJsonl, + suggested_format: None, + output_format: Some(ImportOutputFormat::CompactJsonl), + replace: true, + append: false, + on_duplicate: None, + yes: true, + stream: false, + max_input_bytes: Some(256 * 1024 * 1024), + commit_every: None, + resume: false, + wal_dir: None, + reset: false, + columns: columns.to_vec(), + }, + mirror, + &mut stdout, + stderr, + stderr_is_terminal, + ) + .await + .context("sync source into Compact JSONL mirror")?; + } + + if let Some(storyline) = storyline { + anyhow::ensure!( + input_format != ExchangeFormat::CompactJsonl, + "sync --to requires a trajectory input format; use --mirror for compact-jsonl sources" + ); + // ponytail: rebuild one atomic snapshot per coalesced batch; add affected-document + // mutation when profiling shows full-directory rebuilds are the bottleneck. + let mut stdout = std::io::sink(); + let mut stdin = std::io::empty(); + run_import( + ImportArgs { + from: source.to_owned(), + output: Some(storyline.to_owned()), + format: input_format, + suggested_format, + output_format: Some(ImportOutputFormat::Storyline), + replace: true, + append: false, + on_duplicate: None, + yes: true, + stream: false, + max_input_bytes: Some(256 * 1024 * 1024), + commit_every: None, + resume: false, + wal_dir: None, + reset: false, + columns: Vec::new(), + }, + None, + false, + stderr_is_terminal, + &mut stdin, + &mut stdout, + stderr, + ) + .await + .context("sync source into Storyline Lance")?; + } + + Ok(()) +} diff --git a/crates/persisting-pchronicle-cli/src/exchange/wal.rs b/crates/persisting-pchronicle-cli/src/exchange/wal.rs new file mode 100644 index 00000000..89c6b7d7 --- /dev/null +++ b/crates/persisting-pchronicle-cli/src/exchange/wal.rs @@ -0,0 +1,368 @@ +//! Local checkpoint WAL for resumable Storyline imports. +//! +//! Stores only source-path completion state (not payload bytes). The remote +//! progressive Storyline generation remains the source of truth for written +//! data; the WAL avoids re-fetching / re-parsing sources that already committed. + +use anyhow::{Context, Result}; +use serde::{Deserialize, Serialize}; +use std::collections::HashSet; +use std::fs::{self, File, OpenOptions}; +use std::io::{BufRead, BufReader, Write}; +use std::path::{Path, PathBuf}; +use std::time::{SystemTime, UNIX_EPOCH}; + +const WAL_ROOT_DIRNAME: &str = ".pchronicle-import-wal"; +const JOB_FILE: &str = "job.json"; +const DONE_FILE: &str = "done.jsonl"; +const FAILED_FILE: &str = "failed.jsonl"; +const CURSOR_FILE: &str = "cursor.json"; + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub(crate) struct ImportWalJob { + pub(crate) job_id: String, + pub(crate) from: String, + pub(crate) to: String, + pub(crate) output_format: String, + pub(crate) suggested_format: Option, + pub(crate) created_unix_secs: u64, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +struct DoneRecord { + path: String, + trajectories: u64, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +struct FailedRecord { + path: String, + error: String, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +struct CursorRecord { + path: String, + updated_unix_secs: u64, +} + +#[derive(Debug)] +pub(crate) struct ImportWal { + dir: PathBuf, + job: ImportWalJob, + done: HashSet, + failed: HashSet, +} + +impl ImportWal { + pub(crate) fn job_id( + from: &str, + to: &str, + output_format: &str, + suggested_format: Option<&str>, + ) -> String { + let mut hasher = blake3::Hasher::new(); + hasher.update(from.as_bytes()); + hasher.update(&[0]); + hasher.update(to.as_bytes()); + hasher.update(&[0]); + hasher.update(output_format.as_bytes()); + hasher.update(&[0]); + hasher.update(suggested_format.unwrap_or("").as_bytes()); + hasher.finalize().to_hex()[..32].to_string() + } + + pub(crate) fn default_root() -> PathBuf { + PathBuf::from(WAL_ROOT_DIRNAME) + } + + pub(crate) fn job_dir(root: &Path, job_id: &str) -> PathBuf { + root.join(job_id) + } + + pub(crate) fn open_or_create( + root: &Path, + from: &str, + to: &str, + output_format: &str, + suggested_format: Option<&str>, + resume: bool, + reset: bool, + ) -> Result { + let job_id = Self::job_id(from, to, output_format, suggested_format); + let dir = Self::job_dir(root, &job_id); + if reset && dir.exists() { + fs::remove_dir_all(&dir) + .with_context(|| format!("reset import WAL {}", dir.display()))?; + } + if resume { + anyhow::ensure!( + dir.join(JOB_FILE).is_file(), + "no import WAL at {} for --resume; omit --resume to start a new job or pass --reset", + dir.display() + ); + } + fs::create_dir_all(&dir).with_context(|| format!("create import WAL {}", dir.display()))?; + let job_path = dir.join(JOB_FILE); + let job = if job_path.is_file() { + let text = fs::read_to_string(&job_path) + .with_context(|| format!("read import WAL job {}", job_path.display()))?; + let existing: ImportWalJob = serde_json::from_str(&text) + .with_context(|| format!("parse import WAL job {}", job_path.display()))?; + anyhow::ensure!( + existing.from == from && existing.to == to, + "import WAL job fingerprint mismatch at {}", + dir.display() + ); + existing + } else { + let created = ImportWalJob { + job_id: job_id.clone(), + from: from.to_owned(), + to: to.to_owned(), + output_format: output_format.to_owned(), + suggested_format: suggested_format.map(str::to_owned), + created_unix_secs: unix_secs(), + }; + let encoded = serde_json::to_vec_pretty(&created).context("encode import WAL job")?; + fs::write(&job_path, encoded) + .with_context(|| format!("write import WAL job {}", job_path.display()))?; + created + }; + let done = load_done_paths(&dir.join(DONE_FILE))?; + let failed = load_failed_paths(&dir.join(FAILED_FILE))?; + Ok(Self { + dir, + job, + done, + failed, + }) + } + + pub(crate) fn dir(&self) -> &Path { + &self.dir + } + + pub(crate) fn job(&self) -> &ImportWalJob { + &self.job + } + + #[cfg(test)] + pub(crate) fn should_skip(&self, path: &str) -> bool { + self.done.contains(path) || self.failed.contains(path) + } + + pub(crate) fn done_count(&self) -> usize { + self.done.len() + } + + pub(crate) fn failed_count(&self) -> usize { + self.failed.len() + } + + pub(crate) fn skip_paths(&self) -> HashSet { + self.done.iter().chain(self.failed.iter()).cloned().collect() + } + + pub(crate) fn mark_done(&mut self, path: &str, trajectories: u64) -> Result<()> { + if !self.done.insert(path.to_owned()) { + return Ok(()); + } + self.failed.remove(path); + append_jsonl( + &self.dir.join(DONE_FILE), + &DoneRecord { + path: path.to_owned(), + trajectories, + }, + )?; + self.write_cursor(path)?; + Ok(()) + } + + pub(crate) fn mark_failed(&mut self, path: &str, error: &str) -> Result<()> { + if self.done.contains(path) { + return Ok(()); + } + let first = self.failed.insert(path.to_owned()); + if first { + append_jsonl( + &self.dir.join(FAILED_FILE), + &FailedRecord { + path: path.to_owned(), + error: truncate_error(error), + }, + )?; + } + self.write_cursor(path)?; + Ok(()) + } + + fn write_cursor(&self, path: &str) -> Result<()> { + let cursor = CursorRecord { + path: path.to_owned(), + updated_unix_secs: unix_secs(), + }; + let encoded = serde_json::to_vec_pretty(&cursor).context("encode import WAL cursor")?; + fs::write(self.dir.join(CURSOR_FILE), encoded) + .with_context(|| format!("write import WAL cursor in {}", self.dir.display()))?; + Ok(()) + } +} + +fn unix_secs() -> u64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|duration| duration.as_secs()) + .unwrap_or(0) +} + +fn truncate_error(error: &str) -> String { + const MAX: usize = 2_048; + if error.len() <= MAX { + error.to_owned() + } else { + format!("{}…", &error[..MAX]) + } +} + +fn append_jsonl(path: &Path, value: &impl Serialize) -> Result<()> { + let mut file = OpenOptions::new() + .create(true) + .append(true) + .open(path) + .with_context(|| format!("open import WAL {}", path.display()))?; + serde_json::to_writer(&mut file, value) + .with_context(|| format!("encode import WAL record for {}", path.display()))?; + file.write_all(b"\n") + .with_context(|| format!("append import WAL newline to {}", path.display()))?; + Ok(()) +} + +fn load_done_paths(path: &Path) -> Result> { + if !path.is_file() { + return Ok(HashSet::new()); + } + let file = File::open(path).with_context(|| format!("read import WAL {}", path.display()))?; + let mut out = HashSet::new(); + for (index, line) in BufReader::new(file).lines().enumerate() { + let line = line.with_context(|| format!("read import WAL {} line {}", path.display(), index + 1))?; + if line.trim().is_empty() { + continue; + } + let record: DoneRecord = serde_json::from_str(&line).with_context(|| { + format!("parse import WAL {} line {}", path.display(), index + 1) + })?; + out.insert(record.path); + } + Ok(out) +} + +fn load_failed_paths(path: &Path) -> Result> { + if !path.is_file() { + return Ok(HashSet::new()); + } + let file = File::open(path).with_context(|| format!("read import WAL {}", path.display()))?; + let mut out = HashSet::new(); + for (index, line) in BufReader::new(file).lines().enumerate() { + let line = line.with_context(|| format!("read import WAL {} line {}", path.display(), index + 1))?; + if line.trim().is_empty() { + continue; + } + let record: FailedRecord = serde_json::from_str(&line).with_context(|| { + format!("parse import WAL {} line {}", path.display(), index + 1) + })?; + out.insert(record.path); + } + Ok(out) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn job_id_is_stable_for_same_fingerprint() { + let left = ImportWal::job_id("@a", "@b", "storyline-lance", Some("actf")); + let right = ImportWal::job_id("@a", "@b", "storyline-lance", Some("actf")); + assert_eq!(left, right); + assert_ne!( + left, + ImportWal::job_id("@a", "@b", "storyline-lance", None) + ); + } + + #[test] + fn resume_requires_existing_wal_and_skip_sets_work() { + let root = tempfile::tempdir().unwrap(); + let err = ImportWal::open_or_create( + root.path(), + "from", + "to", + "storyline-lance", + None, + true, + false, + ) + .unwrap_err() + .to_string(); + assert!(err.contains("--resume"), "{err}"); + + let mut wal = ImportWal::open_or_create( + root.path(), + "from", + "to", + "storyline-lance", + None, + false, + false, + ) + .unwrap(); + wal.mark_done("a.json", 2).unwrap(); + wal.mark_failed("b.json", "boom").unwrap(); + + let resumed = ImportWal::open_or_create( + root.path(), + "from", + "to", + "storyline-lance", + None, + true, + false, + ) + .unwrap(); + assert!(resumed.should_skip("a.json")); + assert!(resumed.should_skip("b.json")); + assert!(!resumed.should_skip("c.json")); + assert_eq!(resumed.done_count(), 1); + assert_eq!(resumed.failed_count(), 1); + } + + #[test] + fn reset_clears_prior_state() { + let root = tempfile::tempdir().unwrap(); + let mut wal = ImportWal::open_or_create( + root.path(), + "from", + "to", + "storyline-lance", + None, + false, + false, + ) + .unwrap(); + wal.mark_done("a.json", 1).unwrap(); + let reset = ImportWal::open_or_create( + root.path(), + "from", + "to", + "storyline-lance", + None, + false, + true, + ) + .unwrap(); + assert!(!reset.should_skip("a.json")); + assert_eq!(reset.done_count(), 0); + } +} diff --git a/crates/persisting-pchronicle-cli/src/lib.rs b/crates/persisting-pchronicle-cli/src/lib.rs index 3c6d83e1..fff6121b 100644 --- a/crates/persisting-pchronicle-cli/src/lib.rs +++ b/crates/persisting-pchronicle-cli/src/lib.rs @@ -20,7 +20,6 @@ use output::*; use settings::*; use std::collections::{BTreeMap, HashMap, HashSet}; -use std::ffi::CString; use std::fmt::Write as _; use std::io::{Error as IoError, Read, Write}; use std::net::SocketAddr; @@ -33,23 +32,20 @@ use anyhow::{Context, Result, anyhow, bail}; use clap::{ArgGroup, Args, Parser, Subcommand, ValueEnum}; use futures::{StreamExt, stream, stream::FuturesUnordered}; use persisting_events::{CHRONICLE_SERVE_READY_VERSION, ChronicleServeReady}; -use persisting_pchronicle::document::{ - DocumentFormat, InputIssue, InputIssueKind, decode_json_storylines, detect_format, - encode_json_storylines, open_document, -}; -use persisting_pchronicle::model::StorylineDocument; use persisting_pchronicle::query::ChronicleQueryEngine; use persisting_pchronicle::search::{ FindExpr, FindJsonOperator, FindJsonPredicate, FindTextPredicate, combine_match_expressions, search_storyline_step_matches_fts_in_columns, }; +#[cfg(test)] +use persisting_pchronicle::storage::StorylineLanceStore; use persisting_pchronicle::storage::{ AutomaticProjectionInspection, AutomaticProjectionState, CatalogErrorPolicy, - CatalogSnapshotOptions, CatalogSourceKind, CatalogSourceStatus, CatalogStorylineKey, - DEFAULT_DATASET_NAME, DatasetCatalogSnapshot, DatasetLocation, DatasetMount, DiscoveredSource, - EventFactSnapshot, ObjectStoreManifestWriteMode, StorylineLanceStore, - StorylineProjectionBuildOutcome, automatic_projection_inventory, build_storyline_projection, - inspect_automatic_storyline_projection, probe_canonical_event_store, + CatalogSnapshotOptions, CatalogSourceKind, CatalogSourceStatus, DEFAULT_DATASET_NAME, + DatasetCatalogSnapshot, DatasetLocation, DatasetMount, DiscoveredSource, EventFactSnapshot, + ObjectStoreManifestWriteMode, StorylineProjectionBuildOutcome, automatic_projection_inventory, + build_storyline_projection, inspect_automatic_storyline_projection, + probe_canonical_event_store, }; use serde::{Deserialize, Serialize}; @@ -262,10 +258,10 @@ enum Command { Drop(DropArgs), /// Export complete Trajectories to an exchange format. Export(ExportArgs), - /// Mirror a changing directory into snapshot Datasets. + /// Mirror a changing directory into optional Compact and/or Storyline snapshots. /// - /// With --input-format compact-jsonl, each batch atomically replaces the - /// compact Lance Dataset at --convert; --to remains required but is not written. + /// `--mirror` replaces a Compact JSONL Lance Dataset; `--to` replaces a + /// Storyline Lance Dataset. Provide either or both. Sync(sync::SyncArgs), /// Run a deterministic local LLM upstream for Gateway testing. #[command(hide = true)] @@ -704,7 +700,7 @@ enum ImportOutputFormat { CompactJsonl, } -#[derive(Debug, Clone, Copy, PartialEq, Eq, ValueEnum)] +#[derive(Debug, Clone, Copy, PartialEq, Eq)] enum ImportMode { /// Require a new destination and publish it atomically. Create, @@ -747,6 +743,11 @@ struct ImportArgs { #[arg(short = 'i', long = "input-format", alias = "format", value_enum, default_value_t = ExchangeFormat::Auto)] format: ExchangeFormat, + /// When --format auto cannot decide, try this format if the file is weakly compatible. + /// Does not force decode; use --format to hard-pin. Only valid with --format auto. + #[arg(long = "suggested-format", value_enum, value_name = "FORMAT")] + suggested_format: Option, + /// Dataset layout: preserve, normalized Storyline (combine all inputs into one Storyline Lance Store at the Dataset root), or record-level compact JSONL. #[arg(short = 'o', long = "output-format", value_enum)] output_format: Option, @@ -759,10 +760,6 @@ struct ImportArgs { #[arg(long, conflicts_with = "replace")] append: bool, - /// Deprecated alias for --replace/--append/--create. Prefer --replace or --append. - #[arg(long, value_enum, hide = true)] - mode: Option, - /// How append handles an existing document ID. #[arg(long, value_enum, value_name = "suffix|skip")] on_duplicate: Option, @@ -785,6 +782,19 @@ struct ImportArgs { #[arg(long, value_name = "N")] commit_every: Option, + /// Resume a previous import using the local checkpoint WAL for the same + /// --from/--to fingerprint. Skips sources already recorded as done or failed. + #[arg(long)] + resume: bool, + + /// Root directory for import checkpoint WALs (default: ./.pchronicle-import-wal). + #[arg(long = "wal-dir", value_name = "DIR")] + wal_dir: Option, + + /// Delete the WAL for this --from/--to job before starting (implies a fresh checkpoint). + #[arg(long)] + reset: bool, + /// Compact JSONL mapping. id/timestamp override $.id/$.timestamp; missing or invalid id values /// use source_filename#line_number; other names add JSONB columns. /// Example: --column id=$.event.id --column model=$.payload.model. @@ -794,18 +804,11 @@ struct ImportArgs { impl ImportArgs { fn mode(&self) -> Result { - match (self.replace, self.append, self.mode) { - (true, true, _) => Err(anyhow!("--replace and --append cannot be combined")), - (true, false, Some(ImportMode::Append)) => Err(anyhow!( - "--replace conflicts with --mode append; omit --mode" - )), - (false, true, Some(ImportMode::Replace)) => Err(anyhow!( - "--append conflicts with --mode replace; omit --mode" - )), - (true, false, _) => Ok(ImportMode::Replace), - (false, true, _) => Ok(ImportMode::Append), - (false, false, Some(mode)) => Ok(mode), - (false, false, None) => Ok(ImportMode::Create), + match (self.replace, self.append) { + (true, true) => Err(anyhow!("--replace and --append cannot be combined")), + (true, false) => Ok(ImportMode::Replace), + (false, true) => Ok(ImportMode::Append), + (false, false) => Ok(ImportMode::Create), } } } @@ -1507,6 +1510,9 @@ struct ImportResponse { fact_rows: Option, #[serde(skip_serializing_if = "Option::is_none")] input_bytes: Option, + /// Physical Dataset size after import (Lance/object-store bytes). + #[serde(skip_serializing_if = "Option::is_none")] + on_disk_bytes: Option, } #[derive(Debug, Deserialize)] @@ -1635,7 +1641,9 @@ pub async fn run_with_stdio( .await } Command::Export(args) => run_export(args, config, stdout, &mut diagnostics).await, - Command::Sync(args) => sync::run(args, config, &mut diagnostics).await, + Command::Sync(args) => { + sync::run(args, config, &mut diagnostics, stderr_is_terminal).await + } Command::Echo(args) => run_echo(args, &mut diagnostics).await, Command::Dev(DevArgs { command: DevCommand::Echo(args), diff --git a/crates/persisting-pchronicle-cli/src/onboard.rs b/crates/persisting-pchronicle-cli/src/onboard.rs index fe488e7b..fd29b4ba 100644 --- a/crates/persisting-pchronicle-cli/src/onboard.rs +++ b/crates/persisting-pchronicle-cli/src/onboard.rs @@ -7,9 +7,9 @@ use clap::{Args, Subcommand}; use super::{ AnalysisOptions, DatasetArgs, DatasetCommand, ErrorMode, ExchangeFormat, ExportArgs, - ExportFormat, FindArgs, ImportArgs, ImportOutputFormat, ListArgs, OutputFormat, - QueryArgs, QueryOutputFormat, StatsReport, StatusArgs, run_dataset, run_export, run_find, - run_import, run_list, run_query, run_stats_report, run_status, + ExportFormat, FindArgs, ImportArgs, ImportOutputFormat, ListArgs, OutputFormat, QueryArgs, + QueryOutputFormat, StatsReport, StatusArgs, run_dataset, run_export, run_find, run_import, + run_list, run_query, run_stats_report, run_status, }; const DEMO_ATIF: &str = include_str!("../assets/onboard/support-ticket.json"); @@ -802,15 +802,18 @@ async fn capture_exchange(demo: &DemoWorkspace) -> Result { .into_owned(), ), format: ExchangeFormat::Atif, + suggested_format: None, output_format: Some(ImportOutputFormat::Preserve), replace: false, append: false, - mode: None, on_duplicate: None, yes: false, stream: false, max_input_bytes: None, commit_every: None, + resume: false, + wal_dir: None, + reset: false, columns: Vec::new(), }, Some(&settings), @@ -834,15 +837,18 @@ async fn capture_exchange(demo: &DemoWorkspace) -> Result { from: demo.atif_source().to_string_lossy().into_owned(), output: Some(storyline_output.to_string_lossy().into_owned()), format: ExchangeFormat::Atif, + suggested_format: None, output_format: Some(ImportOutputFormat::Storyline), replace: false, append: false, - mode: None, on_duplicate: None, yes: false, stream: false, max_input_bytes: None, commit_every: None, + resume: false, + wal_dir: None, + reset: false, columns: Vec::new(), }, Some(&settings), diff --git a/crates/persisting-pchronicle-cli/src/server/explorer.rs b/crates/persisting-pchronicle-cli/src/server/explorer.rs index 3d3496e0..7cde9e9b 100644 --- a/crates/persisting-pchronicle-cli/src/server/explorer.rs +++ b/crates/persisting-pchronicle-cli/src/server/explorer.rs @@ -64,6 +64,7 @@ pub(crate) struct CatalogTreeChild { pub(crate) entries: Vec, } +#[allow(dead_code)] pub(crate) fn catalog_tree( summaries: &[RunSummary], dataset: Option<&str>, @@ -101,24 +102,24 @@ pub(crate) fn catalog_tree_with_mounts( } }) .sum(); - let children = if dataset.is_none() { - fold_tree_children( + let children = match dataset { + None => fold_tree_children( merge_dataset_children(dataset_children(&scoped), datasets), max_children, prefix, - ) - } else { - let dataset_name = dataset.expect("dataset scope is some"); - let sources = datasets - .iter() - .find(|row| row.mount.name == dataset_name) - .map(|row| row.sources.as_slice()) - .unwrap_or(&[]); - fold_tree_children( - merge_file_children(file_children(&scoped, prefix), sources, prefix), - max_children, - prefix, - ) + ), + Some(dataset_name) => { + let sources = datasets + .iter() + .find(|row| row.mount.name == dataset_name) + .map(|row| row.sources.as_slice()) + .unwrap_or(&[]); + fold_tree_children( + merge_file_children(file_children(&scoped, prefix), sources, prefix), + max_children, + prefix, + ) + } }; CatalogTree { dataset: dataset.map(str::to_string), @@ -160,16 +161,13 @@ pub(crate) fn append_shallow_nav_children( } else { "file".into() }, - data_type: entry - .dataset_kind - .clone() - .unwrap_or_else(|| { - if entry.is_dir { - "directory".into() - } else { - "other".into() - } - }), + data_type: entry.dataset_kind.clone().unwrap_or_else(|| { + if entry.is_dir { + "directory".into() + } else { + "other".into() + } + }), path, run_count: 0, failed_count: 0, @@ -246,8 +244,7 @@ fn merge_file_children( Some((name, _)) => (name, true), None => ( rest, - source.kind == CatalogSourceKind::Directory - || source.file.contains('/'), + source.kind == CatalogSourceKind::Directory || source.file.contains('/'), ), }; // A Directory leaf under this prefix is always a folder to open. @@ -1755,7 +1752,11 @@ mod tests { assert_eq!( prod.children .iter() - .map(|child| (child.name.as_str(), child.kind.as_str(), child.data_type.as_str())) + .map(|child| ( + child.name.as_str(), + child.kind.as_str(), + child.data_type.as_str() + )) .collect::>(), vec![("infra", "dir", "directory")] ); diff --git a/crates/persisting-pchronicle-cli/src/server/mod.rs b/crates/persisting-pchronicle-cli/src/server/mod.rs index cd34b86c..0b56ee98 100644 --- a/crates/persisting-pchronicle-cli/src/server/mod.rs +++ b/crates/persisting-pchronicle-cli/src/server/mod.rs @@ -1248,9 +1248,7 @@ async fn tree_run_summaries( } if !runtime.snapshot.datasets().iter().any(|dataset| { dataset.sources.iter().any(|source| { - if source.kind - == persisting_pchronicle::storage::CatalogSourceKind::Directory - { + if source.kind == persisting_pchronicle::storage::CatalogSourceKind::Directory { return false; } match source.format.as_deref() { @@ -1436,9 +1434,7 @@ async fn resolve_run_summary( matches.retain(|run| run.root_session_id.as_ref() == Some(root)); } if matches.is_empty() { - if let Some(run) = - try_resolve_on_demand_storyline_run(state, query, request_id).await? - { + if let Some(run) = try_resolve_on_demand_storyline_run(state, query, request_id).await? { return Ok(run); } return Err(ApiError::not_found("run was not found")); @@ -1525,14 +1521,7 @@ async fn try_resolve_on_demand_storyline_run( if !ids.iter().any(|id| id == session_id) { return Ok(None); } - let path = explorer::explorer_run_path( - dataset_name, - file, - session_id, - session_id, - None, - None, - ); + let path = explorer::explorer_run_path(dataset_name, file, session_id, session_id, None, None); Ok(Some(RunSummary { dataset: dataset_name.to_string(), file: file.to_string(), @@ -1593,7 +1582,7 @@ async fn load_on_demand_storyline_bundle( run.document_id.clone() }; let stories = store - .get_storylines_by_document_ids(&[document_id.clone()]) + .get_storylines_by_document_ids(std::slice::from_ref(&document_id)) .await .map_err(|error| fail(request_id, op, error))?; let Some(Some(storyline)) = stories.into_iter().next() else { @@ -2140,70 +2129,70 @@ async fn explorer_turns( search_mode = "memory"; (loaded.turns.clone(), Some(needle)) } else { - let expression = crate::combine_match_expressions(&[needle.to_owned()]) - .map_err(|error| ApiError::invalid_request(error.to_string()))? - .ok_or_else(|| ApiError::invalid_request("search query must not be empty"))?; - let runtime = current_catalog(&state, &request_id).await?; - let (predicate, available, fts_errors) = crate::find_expression_predicate_for_dataset( - &runtime.snapshot, - &expression, - Some(&loaded.run.file), - Some(&loaded.run.dataset), - ) - .await - .map_err(|error| fail(&request_id, "explorer_turns", error))?; - fts.extend(fts_errors); - fts_available = fts_available || available; - let turns = if expression.has_text() || expression.has_step_json() { - let predicate = predicate.ok_or_else(|| { - fail( - &request_id, - "explorer_turns", - anyhow::anyhow!("turn search expression did not produce a predicate"), - ) - })?; - let sql = format!( - "SELECT DISTINCT step_id FROM {}.steps WHERE _file_ = {} AND document_id = {} AND session_id = {} AND ({predicate})", - loaded.run.dataset, - crate::sql_string(&loaded.run.file), - crate::sql_string(&loaded.run.document_id), - crate::sql_string(&loaded.run.session_id), - ); - let jsonl = runtime - .engine - .query_jsonl(&sql) - .await - .map_err(|error| fail(&request_id, "explorer_turns", error))?; - let step_ids = jsonl - .lines() - .filter(|line| !line.trim().is_empty()) - .filter_map(|line| { - serde_json::from_str::(line) - .ok() - .and_then(|row| row.get("step_id").and_then(Value::as_i64)) - }) - .collect::>(); - search_mode = if expression.has_text() && expression.has_json() { - "fts+json" - } else if expression.has_text() { - "fts" + let expression = crate::combine_match_expressions(&[needle.to_owned()]) + .map_err(|error| ApiError::invalid_request(error.to_string()))? + .ok_or_else(|| ApiError::invalid_request("search query must not be empty"))?; + let runtime = current_catalog(&state, &request_id).await?; + let (predicate, available, fts_errors) = crate::find_expression_predicate_for_dataset( + &runtime.snapshot, + &expression, + Some(&loaded.run.file), + Some(&loaded.run.dataset), + ) + .await + .map_err(|error| fail(&request_id, "explorer_turns", error))?; + fts.extend(fts_errors); + fts_available = fts_available || available; + let turns = if expression.has_text() || expression.has_step_json() { + let predicate = predicate.ok_or_else(|| { + fail( + &request_id, + "explorer_turns", + anyhow::anyhow!("turn search expression did not produce a predicate"), + ) + })?; + let sql = format!( + "SELECT DISTINCT step_id FROM {}.steps WHERE _file_ = {} AND document_id = {} AND session_id = {} AND ({predicate})", + loaded.run.dataset, + crate::sql_string(&loaded.run.file), + crate::sql_string(&loaded.run.document_id), + crate::sql_string(&loaded.run.session_id), + ); + let jsonl = runtime + .engine + .query_jsonl(&sql) + .await + .map_err(|error| fail(&request_id, "explorer_turns", error))?; + let step_ids = jsonl + .lines() + .filter(|line| !line.trim().is_empty()) + .filter_map(|line| { + serde_json::from_str::(line) + .ok() + .and_then(|row| row.get("step_id").and_then(Value::as_i64)) + }) + .collect::>(); + search_mode = if expression.has_text() && expression.has_json() { + "fts+json" + } else if expression.has_text() { + "fts" + } else { + "json" + }; + loaded + .turns + .iter() + .filter(|item| step_ids.contains(&item.turn.id)) + .cloned() + .collect::>() } else { - "json" + // Run-level JSON predicates have no step identity to display in + // this view. Keep the detail search scoped to Step expressions, + // matching the CLI find scope instead of applying an ad-hoc + // in-memory text filter. + Vec::new() }; - loaded - .turns - .iter() - .filter(|item| step_ids.contains(&item.turn.id)) - .cloned() - .collect::>() - } else { - // Run-level JSON predicates have no step identity to display in - // this view. Keep the detail search scoped to Step expressions, - // matching the CLI find scope instead of applying an ad-hoc - // in-memory text filter. - Vec::new() - }; - (turns, None) + (turns, None) } } else { (loaded.turns.clone(), query.q.as_deref()) diff --git a/crates/persisting-pchronicle-cli/src/settings.rs b/crates/persisting-pchronicle-cli/src/settings.rs index f2d2927c..a4bc258a 100644 --- a/crates/persisting-pchronicle-cli/src/settings.rs +++ b/crates/persisting-pchronicle-cli/src/settings.rs @@ -711,6 +711,7 @@ fn expand_catalog_pin( "catalog pin '@{name}' requires a dataset, for example '@{name}/prod'" ); let (dataset, path) = suffix.split_once('/').unwrap_or((suffix, "")); + let path = normalize_pin_suffix(path); if !path.is_empty() { validate_pin_suffix(path)?; } @@ -928,6 +929,9 @@ pub(super) fn expand_dataset_reference( } else { let rest = &input[1..]; let (name, suffix) = rest.split_once('/').unwrap_or((rest, "")); + // Directory-style refs often end with `/` (e.g. `@origin/foo/`); treat + // that as equivalent to the same path without trailing separators. + let suffix = normalize_pin_suffix(suffix); validate_pin_suffix(suffix)?; if name == DEFAULT_PIN_NAME { let root = resolve_default_pin(settings_override)?; @@ -967,7 +971,12 @@ pub(super) fn expand_dataset_reference( } } +fn normalize_pin_suffix(suffix: &str) -> &str { + suffix.trim_matches('/') +} + fn validate_pin_suffix(suffix: &str) -> Result<()> { + let suffix = normalize_pin_suffix(suffix); if suffix.is_empty() { return Ok(()); } @@ -1144,4 +1153,37 @@ secret_key = "sk" ); assert_eq!(settings.pins["testcata"].uri, "catalog://127.0.0.1:6001"); } + + #[test] + fn pin_suffix_allows_trailing_and_leading_slashes() { + assert!(validate_pin_suffix("SweEval/guoxu1/").is_ok()); + assert!(validate_pin_suffix("/SweEval/guoxu1///").is_ok()); + assert!(validate_pin_suffix("/").is_ok()); + assert!(validate_pin_suffix("").is_ok()); + } + + #[test] + fn pin_suffix_still_rejects_dot_and_empty_middle_segments() { + assert!(validate_pin_suffix("a/../b").is_err()); + assert!(validate_pin_suffix("a/./b").is_err()); + assert!(validate_pin_suffix("a//b").is_err()); + } + + #[test] + fn expand_dataset_reference_trims_trailing_slash() { + let temporary = tempfile::tempdir().expect("tempdir"); + let config = temporary.path().join("config.toml"); + std::fs::write( + &config, + r#" +[pins.origin] +uri = "s3://example-bucket/root" +"#, + ) + .expect("write config"); + let expanded = + expand_dataset_reference("@origin/SweEval/guoxu1/", Some(&config), false) + .expect("expand"); + assert_eq!(expanded, "s3://example-bucket/root/SweEval/guoxu1"); + } } diff --git a/crates/persisting-pchronicle-cli/src/sync.rs b/crates/persisting-pchronicle-cli/src/sync.rs index 4f3fc4f9..c7b9e2cd 100644 --- a/crates/persisting-pchronicle-cli/src/sync.rs +++ b/crates/persisting-pchronicle-cli/src/sync.rs @@ -13,19 +13,24 @@ pub(crate) struct SyncArgs { #[arg(long, value_name = "DATASET")] pub(crate) from: String, - /// Warehouse Dataset receiving source files; unused for compact-jsonl. - #[arg(long = "to", alias = "warehouse", value_name = "DATASET")] - pub(crate) to: String, + /// Compact JSONL Lance Dataset receiving each snapshot (record-level ingest). + #[arg(long, value_name = "DATASET")] + pub(crate) mirror: Option, - /// Storyline or compact JSONL Lance Dataset receiving each snapshot. - #[arg(long = "convert", alias = "storyline", value_name = "DATASET")] - pub(crate) convert: String, + /// Storyline Lance Dataset receiving each converted snapshot. + #[arg(long = "to", value_name = "DATASET")] + pub(crate) to: Option, - /// Input format. compact-jsonl requires a tree of .jsonl files. + /// Input format for --to trajectory conversion. Auto detects run data. + /// Compact-jsonl sources are only valid with --mirror (not --to). #[arg(long = "input-format", value_enum, default_value_t = ExchangeFormat::Auto)] pub(crate) input_format: ExchangeFormat, - /// Compact JSONL mapping; id/timestamp override $.id/$.timestamp defaults. + /// When --input-format auto cannot decide for --to, try this format if weakly compatible. + #[arg(long = "suggested-format", value_enum, value_name = "FORMAT")] + pub(crate) suggested_format: Option, + + /// Compact JSONL column mapping for --mirror. Same rules as import --column. #[arg(long = "column", value_name = "NAME=JSON_PATH", action = clap::ArgAction::Append)] pub(crate) columns: Vec, @@ -48,30 +53,83 @@ pub(crate) async fn run( args: SyncArgs, settings_override: Option<&Path>, stderr: &mut dyn Write, + stderr_is_terminal: bool, ) -> Result<()> { + anyhow::ensure!( + args.mirror.is_some() || args.to.is_some(), + "sync requires --mirror and/or --to" + ); + anyhow::ensure!( + args.columns.is_empty() || args.mirror.is_some(), + "--column is only valid with --mirror" + ); + if let Some(suggested) = args.suggested_format { + anyhow::ensure!( + args.to.is_some(), + "--suggested-format is only valid with --to" + ); + anyhow::ensure!( + args.input_format == ExchangeFormat::Auto, + "--suggested-format is only valid with --input-format auto" + ); + anyhow::ensure!( + suggested != ExchangeFormat::Auto, + "--suggested-format cannot be auto" + ); + anyhow::ensure!( + suggested != ExchangeFormat::CompactJsonl, + "--suggested-format cannot be compact-jsonl" + ); + } + if args.input_format == ExchangeFormat::CompactJsonl { + anyhow::ensure!( + args.to.is_none(), + "sync --input-format compact-jsonl cannot use --to; pass --mirror only" + ); + anyhow::ensure!( + args.mirror.is_some(), + "sync --input-format compact-jsonl requires --mirror" + ); + } + let source_uri = expand_dataset_reference(&args.from, settings_override, true) .with_context(|| format!("resolve sync source '{}'", args.from))?; - let warehouse_uri = expand_dataset_reference(&args.to, settings_override, false) - .with_context(|| format!("resolve sync Warehouse '{}'", args.to))?; - let convert_uri = expand_dataset_reference(&args.convert, settings_override, false) - .with_context(|| format!("resolve sync convert '{}'", args.convert))?; + let mirror_uri = match args.mirror.as_deref() { + Some(mirror) => Some(prepare_destination( + &expand_dataset_reference(mirror, settings_override, false) + .with_context(|| format!("resolve sync mirror '{mirror}'"))?, + "mirror", + )?), + None => None, + }; + let to_uri = match args.to.as_deref() { + Some(to) => Some(prepare_destination( + &expand_dataset_reference(to, settings_override, false) + .with_context(|| format!("resolve sync --to '{to}'"))?, + "to", + )?), + None => None, + }; - let warehouse_uri = prepare_destination(&warehouse_uri, "Warehouse")?; - let convert_uri = prepare_destination(&convert_uri, "conversion")?; - anyhow::ensure!( - warehouse_uri != convert_uri, - "sync targets must be different" - ); - ensure_targets_outside_source(&source_uri, &warehouse_uri, &convert_uri)?; + if let (Some(mirror), Some(to)) = (&mirror_uri, &to_uri) { + anyhow::ensure!(mirror != to, "sync --mirror and --to must be different"); + } + ensure_targets_outside_source(&source_uri, mirror_uri.as_deref(), to_uri.as_deref())?; - writeln!( - stderr, - "sync from={} to={} convert={}", - source_uri, warehouse_uri, convert_uri - ) - .context("write sync resolved targets")?; + let mut banner = format!("sync from={source_uri}"); + if let Some(mirror) = &mirror_uri { + banner.push_str(&format!(" mirror={mirror}")); + } + if let Some(to) = &to_uri { + banner.push_str(&format!(" to={to}")); + } + writeln!(stderr, "{banner}").context("write sync resolved targets")?; let interval = Duration::from_secs(args.interval_seconds.max(1)); + let input_format = args.input_format; + let suggested_format = args.suggested_format; + let columns = args.columns.clone(); + if args.once { let initial = scan_source(&source_uri).await?; anyhow::ensure!( @@ -80,10 +138,13 @@ pub(crate) async fn run( ); super::exchange::sync_snapshot( &source_uri, - &warehouse_uri, - &convert_uri, - args.input_format, - &args.columns, + mirror_uri.as_deref(), + to_uri.as_deref(), + input_format, + suggested_format, + &columns, + stderr, + stderr_is_terminal, ) .await?; writeln!(stderr, "sync batch={} status=ok", initial.len()) @@ -123,10 +184,13 @@ pub(crate) async fn run( match super::exchange::sync_snapshot( &source_uri, - &warehouse_uri, - &convert_uri, - args.input_format, - &args.columns, + mirror_uri.as_deref(), + to_uri.as_deref(), + input_format, + suggested_format, + &columns, + stderr, + stderr_is_terminal, ) .await { @@ -194,24 +258,32 @@ fn prepare_destination(uri: &str, name: &str) -> Result { Ok(parent.join(filename).to_string_lossy().into_owned()) } -fn ensure_targets_outside_source(source: &str, warehouse: &str, convert: &str) -> Result<()> { +fn ensure_targets_outside_source( + source: &str, + mirror: Option<&str>, + to: Option<&str>, +) -> Result<()> { let source = DatasetLocation::parse(source)?; - let warehouse = DatasetLocation::parse(warehouse)?; - let convert = DatasetLocation::parse(convert)?; let Some(source_path) = source.local_path() else { return Ok(()); }; - if let Some(warehouse_path) = warehouse.local_path() { - anyhow::ensure!( - !warehouse_path.starts_with(source_path), - "sync Warehouse target must be outside the source directory" - ); + if let Some(mirror) = mirror { + let mirror = DatasetLocation::parse(mirror)?; + if let Some(mirror_path) = mirror.local_path() { + anyhow::ensure!( + !mirror_path.starts_with(source_path), + "sync mirror target must be outside the source directory" + ); + } } - if let Some(convert_path) = convert.local_path() { - anyhow::ensure!( - !convert_path.starts_with(source_path), - "sync conversion target must be outside the source directory" - ); + if let Some(to) = to { + let to = DatasetLocation::parse(to)?; + if let Some(to_path) = to.local_path() { + anyhow::ensure!( + !to_path.starts_with(source_path), + "sync --to target must be outside the source directory" + ); + } } Ok(()) } @@ -279,7 +351,7 @@ mod tests { #[test] fn prepare_destination_preserves_object_store_uri() { assert_eq!( - prepare_destination("s3://bucket/prod/infra/agent/agentcompass", "Warehouse").unwrap(), + prepare_destination("s3://bucket/prod/infra/agent/agentcompass", "mirror").unwrap(), "s3://bucket/prod/infra/agent/agentcompass" ); } @@ -290,23 +362,22 @@ mod tests { let error = run( SyncArgs { from: "@origin/agentcompass".into(), - to: "/tmp/pchronicle-sync-warehouse".into(), - convert: "/tmp/pchronicle-sync-convert".into(), + mirror: None, + to: Some("/tmp/pchronicle-sync-convert".into()), input_format: ExchangeFormat::Auto, + suggested_format: None, columns: Vec::new(), interval_seconds: 1, once: true, }, None, &mut stderr, + false, ) .await .expect_err("pin must expand through settings, not local canonicalize"); let message = format!("{error:#}"); - assert!( - !message.contains("canonicalize sync source"), - "{message}" - ); + assert!(!message.contains("canonicalize sync source"), "{message}"); assert!( message.contains("unknown Dataset pin") || message.contains("resolve sync source"), "{message}" @@ -336,7 +407,33 @@ mod tests { } #[tokio::test] - async fn sync_once_rebuilds_warehouse_and_storyline() -> Result<()> { + async fn sync_once_requires_mirror_or_to() { + let mut stderr = Vec::new(); + let error = run( + SyncArgs { + from: "/tmp/unused".into(), + mirror: None, + to: None, + input_format: ExchangeFormat::Auto, + suggested_format: None, + columns: Vec::new(), + interval_seconds: 1, + once: true, + }, + None, + &mut stderr, + false, + ) + .await + .expect_err("at least one destination required"); + assert!( + format!("{error:#}").contains("requires --mirror and/or --to"), + "{error:#}" + ); + } + + #[tokio::test] + async fn sync_once_rebuilds_storyline() -> Result<()> { let temporary = tempfile::tempdir()?; let source = temporary.path().join("source"); fs::create_dir_all(&source)?; @@ -344,37 +441,66 @@ mod tests { Path::new(env!("CARGO_MANIFEST_DIR")).join("assets/onboard/support-ticket.json"), source.join("support-ticket.json"), )?; - let source_bytes = fs::read(source.join("support-ticket.json"))?; + let storyline = temporary.path().join("storyline"); let mut stderr = Vec::new(); run( SyncArgs { from: source.to_string_lossy().into_owned(), - to: temporary - .path() - .join("warehouse") - .to_string_lossy() - .into_owned(), - convert: temporary - .path() - .join("storyline") - .to_string_lossy() - .into_owned(), - input_format: ExchangeFormat::Auto, + mirror: None, + to: Some(storyline.to_string_lossy().into_owned()), + input_format: ExchangeFormat::Atif, + suggested_format: None, columns: Vec::new(), interval_seconds: 1, once: true, }, None, &mut stderr, + false, ) .await?; - assert_eq!( - fs::read(temporary.path().join("warehouse/support-ticket.json"))?, - source_bytes + assert!(storyline.join("CURRENT").is_file()); + Ok(()) + } + + #[tokio::test] + async fn sync_once_mirror_builds_compact_lance() -> Result<()> { + let temporary = tempfile::tempdir()?; + let source = temporary.path().join("source"); + fs::create_dir_all(&source)?; + fs::write( + source.join("events.jsonl"), + r#"{"id":"a","timestamp":"2026-01-01T00:00:00Z","payload":1} +{"id":"b","timestamp":"2026-01-01T00:00:01Z","payload":2} +"#, + )?; + + let mirror = temporary.path().join("mirror"); + let mut stderr = Vec::new(); + run( + SyncArgs { + from: source.to_string_lossy().into_owned(), + mirror: Some(mirror.to_string_lossy().into_owned()), + to: None, + input_format: ExchangeFormat::CompactJsonl, + suggested_format: None, + columns: Vec::new(), + interval_seconds: 1, + once: true, + }, + None, + &mut stderr, + false, + ) + .await?; + + assert!( + mirror.join("CURRENT").is_file() + || mirror.join("_versions").is_dir() + || mirror.exists() ); - assert!(temporary.path().join("storyline/CURRENT").is_file()); Ok(()) } } diff --git a/crates/persisting-pchronicle-cli/src/tests.rs b/crates/persisting-pchronicle-cli/src/tests.rs index e99bf80b..73fc5818 100644 --- a/crates/persisting-pchronicle-cli/src/tests.rs +++ b/crates/persisting-pchronicle-cli/src/tests.rs @@ -498,17 +498,19 @@ fn canonical_parser_surface_matches_the_cli_guide() -> Result<()> { assert!(!import.append); assert!(import.yes); - assert!(Cli::try_parse_from([ - "pchronicle", - "import", - "-f", - "input.json", - "-t", - "./imported", - "--replace", - "--append", - ]) - .is_err()); + assert!( + Cli::try_parse_from([ + "pchronicle", + "import", + "-f", + "input.json", + "-t", + "./imported", + "--replace", + "--append", + ]) + .is_err() + ); let cli = Cli::try_parse_from(["pchronicle", "drop", "./imported", "--yes"])?; let Command::Drop(drop) = cli.command else { @@ -2711,11 +2713,7 @@ async fn directory_import_auto_detects_each_file_and_skips_unknown_json() -> Res assert_eq!(response["trajectories"], 3, "{output_format:?}: {response}"); let warnings = String::from_utf8(stderr)?; assert!( - warnings.contains("import source=root.json status=processing"), - "{output_format:?}: {warnings}" - ); - assert!( - warnings.contains("import source=root.json status=completed"), + warnings.contains("root.json"), "{output_format:?}: {warnings}" ); assert!( @@ -2833,6 +2831,13 @@ async fn object_store_replace_clears_existing_prefix_before_import() -> Result<( stderr.contains("deleted:total =") || stderr.contains("[deleting]"), "replace should report delete progress, got: {stderr}" ); + assert!( + existing + .read_relative_bytes(".dataset-marker") + .await + .is_err(), + "replace should remove objects left by the previous Dataset" + ); let store = StorylineLanceStore::open_uri(&output).await?; let ids = store @@ -3063,16 +3068,20 @@ async fn object_store_directory_import_recurses_json_files() -> Result<()> { ) .await?; - let output = tempfile::tempdir()?; + let output_root = tempfile::tempdir()?; + let output = output_root.path().join("dataset"); + let wal_root = tempfile::tempdir()?; let cli = Cli::try_parse_from([ "pchronicle", "import", "--from", &source, "--to", - output.path().to_str().unwrap(), + output.to_str().unwrap(), "--output-format", "storyline", + "--wal-dir", + wal_root.path().to_str().unwrap(), ])?; let mut stdout = Vec::new(); let mut stderr = Vec::new(); @@ -3081,7 +3090,7 @@ async fn object_store_directory_import_recurses_json_files() -> Result<()> { assert!(stderr.contains("status=discovering")); assert!(stderr.contains("status=discovered files=2")); - let store = StorylineLanceStore::open(output.path()).await?; + let store = StorylineLanceStore::open(&output).await?; let ids = store .document_ids_snapshot() .await? @@ -3347,6 +3356,8 @@ async fn append_storyline_import_suffixes_or_skips_existing_document_ids() -> Re .map(|row| row["document_id"].as_str().unwrap().to_string()) .collect::>(); assert_eq!(ids, ["shared", "shared#1"]); + let manifest = persisting_pchronicle::storage::load_manifest(&output)?.context("manifest")?; + assert_eq!(manifest.stats.as_ref().unwrap().record_count, 2); Ok(()) } @@ -3519,12 +3530,13 @@ async fn directory_import_dedupes_unknown_warnings_across_sources() -> Result<() } #[tokio::test] -async fn directory_import_failure_does_not_publish_partial_output() -> Result<()> { +async fn directory_import_skips_invalid_json_and_publishes_valid_sources() -> Result<()> { let temp = tempfile::tempdir()?; let input = temp.path().join("input"); fs::create_dir_all(&input)?; fs::copy(example_source("atif"), input.join("a-valid.json"))?; fs::write(input.join("z-invalid.json"), "not json")?; + let wal_dir = temp.path().join("wal"); for output_format in [ImportOutputFormat::Preserve, ImportOutputFormat::Storyline] { let output = temp @@ -3537,18 +3549,28 @@ async fn directory_import_failure_does_not_publish_partial_output() -> Result<() input.to_string_lossy().into_owned(), "--output".to_owned(), output.to_string_lossy().into_owned(), + "--wal-dir".to_owned(), + wal_dir.to_string_lossy().into_owned(), + "--reset".to_owned(), ]; if output_format == ImportOutputFormat::Storyline { argv.extend(["--output-format".to_owned(), "storyline".to_owned()]); } let cli = Cli::try_parse_from(argv)?; - let error = run(cli, false, &mut Vec::new(), &mut Vec::new()) - .await - .unwrap_err(); - assert!(format!("{error:#}").contains("z-invalid.json"), "{error:#}"); - if output_format == ImportOutputFormat::Preserve { - assert!(!output.exists()); - } + let mut stdout = Vec::new(); + let mut stderr = Vec::new(); + run(cli, false, &mut stdout, &mut stderr).await?; + let response: Value = serde_json::from_slice(&stdout)?; + assert!( + response["trajectories"].as_u64().unwrap_or(0) >= 1, + "{output_format:?}: {response}" + ); + let stderr = String::from_utf8(stderr)?; + assert!( + stderr.contains("z-invalid.json") || stderr.to_lowercase().contains("skip"), + "{output_format:?}: expected skip warning for invalid JSON, got: {stderr}" + ); + assert!(output.exists(), "{output_format:?}: valid sources should publish"); } assert!(!fs::read_dir(temp.path())?.any(|entry| { entry diff --git a/crates/persisting-pchronicle/src/formats/actf/convert.rs b/crates/persisting-pchronicle/src/formats/actf/convert.rs index 2000d16b..cea2e082 100644 --- a/crates/persisting-pchronicle/src/formats/actf/convert.rs +++ b/crates/persisting-pchronicle/src/formats/actf/convert.rs @@ -1,6 +1,6 @@ //! ACTF ⇄ Storyline conversion. -use std::collections::BTreeMap; +use std::collections::{BTreeMap, HashSet}; use anyhow::Context as _; use serde_json::{Map, Value, json}; @@ -59,12 +59,14 @@ fn actf_tool_to_storyline( fn actf_observation_to_storyline_with_call_id( observation: &ActfObservation, - fallback_call_id: Option<&str>, + override_call_id: Option<&str>, ) -> Value { let mut result = serde_json::to_value(observation).unwrap_or_else(|_| Value::Object(Map::new())); if let Some(object) = result.as_object_mut() { - if let Some(source_call_id) = actf_observation_call_id(observation).or(fallback_call_id) { + if let Some(source_call_id) = + override_call_id.or_else(|| actf_observation_call_id(observation)) + { object.insert( "source_call_id".into(), Value::String(source_call_id.to_string()), @@ -81,12 +83,11 @@ fn actf_observation_to_storyline_with_call_id( result } -fn actf_observation_fallback_call_id( +fn actf_observation_fallback_call_index( observation: &ActfObservation, source_tools: &[ActfToolCall], - step_id: i64, assigned: &mut [bool], -) -> Option { +) -> Option { if actf_observation_call_id(observation).is_some() { return None; } @@ -120,7 +121,37 @@ fn actf_observation_fallback_call_id( }) .map(|(index, _)| index)?; assigned[position] = true; - Some(source_tools[position].effective_id(step_id, position)) + Some(position) +} + +fn actf_observation_tool_index( + observation: &ActfObservation, + source_tools: &[ActfToolCall], + step_id: i64, + assigned: &mut [bool], +) -> Option { + if let Some(call_id) = actf_observation_call_id(observation) { + return source_tools.iter().enumerate().find_map(|(index, call)| { + (call.effective_id(step_id, index) == call_id).then_some(index) + }); + } + actf_observation_fallback_call_index(observation, source_tools, assigned) +} + +/// Skillsbench / retry dumps often reuse the same tool call id across steps. +/// Storyline requires document-unique ids, so allocate a stable suffix here. +fn allocate_unique_tool_call_id(preferred: String, seen: &mut HashSet) -> String { + if seen.insert(preferred.clone()) { + return preferred; + } + let mut suffix = 2u32; + loop { + let candidate = format!("{preferred}#{suffix}"); + if seen.insert(candidate.clone()) { + return candidate; + } + suffix = suffix.saturating_add(1); + } } pub(crate) fn actf_to_storylines(document: &ActfDocument) -> Result> { @@ -171,6 +202,7 @@ fn attempt_to_storyline( .as_ref() .and_then(|(system, user)| StorylinePrompt::from_pair(system, user)); let mut turns = Vec::with_capacity(attempt.trajectory.steps.len()); + let mut seen_tool_call_ids = HashSet::new(); for (step, pair) in attempt.trajectory.steps.iter().zip(prompt_pairs) { let source_tools = step.effective_tools(); let mut assigned_observation_calls = vec![false; source_tools.len()]; @@ -183,13 +215,23 @@ fn attempt_to_storyline( assigned_observation_calls[position] = true; } } + let unique_ids = source_tools + .iter() + .enumerate() + .map(|(call_index, call)| { + allocate_unique_tool_call_id( + call.effective_id(step.step_id, call_index), + &mut seen_tool_call_ids, + ) + }) + .collect::>(); let tool_calls = (!source_tools.is_empty()) .then(|| { source_tools .iter() .enumerate() .map(|(call_index, call)| { - Ok(actf_tool_to_storyline( + let mut converted = actf_tool_to_storyline( call, if source_tools.len() == 1 { step.metric.env_action_ms.as_f64().map(|value| value as i64) @@ -198,7 +240,9 @@ fn attempt_to_storyline( }, step.step_id, call_index, - )) + ); + converted.tool_call_id = unique_ids[call_index].clone(); + Ok(converted) }) .collect::>>() }) @@ -208,16 +252,14 @@ fn attempt_to_storyline( .observation .iter() .map(|observation| { - let fallback_call_id = actf_observation_fallback_call_id( + let call_index = actf_observation_tool_index( observation, source_tools, step.step_id, &mut assigned_observation_calls, ); - actf_observation_to_storyline_with_call_id( - observation, - fallback_call_id.as_deref(), - ) + let unique_call_id = call_index.map(|index| unique_ids[index].as_str()); + actf_observation_to_storyline_with_call_id(observation, unique_call_id) }) .collect::>(); json!({"results": results}) @@ -425,8 +467,7 @@ fn openclaw_message_to_turn(event: &Value, id: i64) -> Result Ok(Some(StorylineTurn { @@ -1579,6 +1620,50 @@ mod tests { assert_eq!(storyline_to_actf(&story).unwrap(), document); } + #[test] + fn actf_reused_tool_call_ids_across_steps_are_uniquified() { + let document = parse_actf_document( + r#"{ + "task_id":"task-reuse","category":"software-engineering","k":1, + "correct":false,"attempts_tried":1,"solved_at":null, + "attempts":{"1":{"correct":false,"final_answer":null,"ground_truth":"expected", + "trajectory":{"schema_version":"ACTF_v1.0","steps":[{ + "step_id":1, + "assistant_content":{"content":"one","reasoning_content":"","tool_calls":[{"type":"tool_use","id":"call_ab31e377d3db4d3187f55bdc","name":"Bash","input":{"command":"pwd"}}]}, + "metric":{"prompt_tokens_len":1,"completion_tokens_len":2,"llm_infer_ms":3.5,"env_action_ms":4.5,"stop_reason":null}, + "system_prompt":"sys","user_content":"task", + "tools":[{"type":"tool_use","id":"call_ab31e377d3db4d3187f55bdc","name":"Bash","input":{"command":"pwd"}}], + "observation":[{"tool_use_id":"call_ab31e377d3db4d3187f55bdc","type":"tool_result","content":"/app","is_error":false}], + "started_at":"2026-01-01 00:00:00+00:00","finished_at":"2026-01-01 00:00:01+00:00" + },{ + "step_id":2, + "assistant_content":{"content":"two","reasoning_content":"","tool_calls":[{"type":"tool_use","id":"call_ab31e377d3db4d3187f55bdc","name":"Bash","input":{"command":"ls"}}]}, + "metric":{"prompt_tokens_len":1,"completion_tokens_len":2,"llm_infer_ms":3.5,"env_action_ms":4.5,"stop_reason":null}, + "system_prompt":"sys","user_content":"task", + "tools":[{"type":"tool_use","id":"call_ab31e377d3db4d3187f55bdc","name":"Bash","input":{"command":"ls"}}], + "observation":[{"tool_use_id":"call_ab31e377d3db4d3187f55bdc","type":"tool_result","content":"ok","is_error":false}], + "started_at":"2026-01-01 00:00:02+00:00","finished_at":"2026-01-01 00:00:03+00:00" + }],"started_at":"2026-01-01 00:00:00+00:00","finished_at":"2026-01-01 00:00:03+00:00"}, + "status":"completed","score":null,"error":"","artifacts":{},"extra":{},"analysis_result":{},"meta":{}}} + }"#, + ) + .unwrap(); + let story = actf_to_storyline(&document).unwrap(); + story.validate().unwrap(); + assert_eq!( + story.turns[0].tool_calls.as_ref().unwrap()[0].tool_call_id, + "call_ab31e377d3db4d3187f55bdc" + ); + assert_eq!( + story.turns[1].tool_calls.as_ref().unwrap()[0].tool_call_id, + "call_ab31e377d3db4d3187f55bdc#2" + ); + assert_eq!( + story.turns[1].observation.as_ref().unwrap()["results"][0]["source_call_id"], + "call_ab31e377d3db4d3187f55bdc#2" + ); + } + #[test] fn actf_noncanonical_source_fields_are_unknown_without_source_extra() { let document = parse_actf_document( diff --git a/crates/persisting-pchronicle/src/formats/actf/mod.rs b/crates/persisting-pchronicle/src/formats/actf/mod.rs index c77c6e4e..b0668978 100644 --- a/crates/persisting-pchronicle/src/formats/actf/mod.rs +++ b/crates/persisting-pchronicle/src/formats/actf/mod.rs @@ -91,14 +91,31 @@ fn path_has_actf_hint(path: Option<&Path>) -> bool { fn looks_like_actf_attempt(attempt: &Value) -> bool { match attempt.get("trajectory") { + // Error dumps: missing / null / empty placeholder trajectory. + None => true, + Some(trajectory) if trajectory.is_null() => true, + Some(trajectory) + if trajectory + .as_object() + .is_some_and(|object| object.is_empty()) => + { + true + } + // Pinchbench / harness dumps sometimes stringify a Python Trajectory repr + // instead of emitting a JSON object/array. + Some(trajectory) if trajectory.is_string() => trajectory.as_str().is_some_and(|text| { + let trimmed = text.trim_start(); + trimmed.starts_with("Trajectory(") || trimmed.contains("ACTF_") + }), + // skillsbench / pinchbench OpenClaw event-stream dumps Some(trajectory) if trajectory.is_array() => trajectory .as_array() .is_some_and(|events| events.iter().all(Value::is_object)), + // Canonical ACTF steps trajectory requires an ACTF_* schema_version. Some(trajectory) => trajectory .get("schema_version") .and_then(Value::as_str) .is_some_and(|version| version.starts_with("ACTF_")), - None => false, } } @@ -117,25 +134,32 @@ fn content_has_actf_fingerprint(content: &[u8]) -> bool { return false; }; let trimmed = text.trim_start(); - if trimmed.starts_with('{') || trimmed.starts_with('[') { - if let Ok(value) = serde_json::from_str::(trimmed) + if !(trimmed.starts_with('{') || trimmed.starts_with('[')) { + return false; + } + let sanitized = super::common::sanitize_json_nonfinite(trimmed); + if let Ok(value) = serde_json::from_str::(sanitized.as_ref()) + && looks_like_actf_value(&value) + { + return true; + } + for line in sanitized + .lines() + .filter(|line| !line.trim().is_empty()) + .take(32) + { + if let Ok(value) = serde_json::from_str::(line) && looks_like_actf_value(&value) { return true; } - for line in trimmed - .lines() - .filter(|line| !line.trim().is_empty()) - .take(32) - { - if let Ok(value) = serde_json::from_str::(line) - && looks_like_actf_value(&value) - { - return true; - } - } } - false + // Frontier-engineering dumps put a huge `final_answer` before + // `trajectory.schema_version`. When non-finite tokens still break parse, + // accept the structural markers that uniquely identify ACTF. + sanitized.contains("\"task_id\"") + && sanitized.contains("\"attempts\"") + && (sanitized.contains("\"ACTF_") || sanitized.contains("'ACTF_")) } fn decode_json( @@ -146,12 +170,14 @@ fn decode_json( reader .read_to_string(&mut input) .map_err(|error| InputIssue::invalid(error.to_string()))?; - let mut value: Value = - serde_json::from_str(&input).map_err(|error| InputIssue::invalid(error.to_string()))?; + let sanitized = super::common::sanitize_json_nonfinite(&input); + let mut value: Value = serde_json::from_str(sanitized.as_ref()) + .map_err(|error| InputIssue::invalid(error.to_string()))?; let envelope = take_unknown_fields_envelope(&mut value)?; let mut document: ActfDocument = serde_json::from_value(value).map_err(|error| InputIssue::invalid(error.to_string()))?; normalize_solved_at(&mut document.solved_at); + reconcile_document_tool_lists(&mut document); document.validate()?; let mut stories = actf_to_storylines(&document).map_err(|error| InputIssue::invalid(error.to_string()))?; @@ -191,6 +217,8 @@ fn decode_json( #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] pub struct ActfDocument { pub task_id: String, + /// Some error dumps emit numeric categories (`2`) instead of strings. + #[serde(deserialize_with = "stringish")] pub category: String, pub k: u64, pub correct: bool, @@ -263,6 +291,39 @@ impl ActfTrajectory { extra: Map::new(), } } + + fn normalize_timestamps(&mut self) { + const PLACEHOLDER: &str = "1970-01-01T00:00:00Z"; + if self.started_at.trim().is_empty() { + self.started_at = self + .steps + .iter() + .find_map(|step| { + let trimmed = step.started_at.trim(); + (!trimmed.is_empty()).then(|| trimmed.to_string()) + }) + .unwrap_or_else(|| PLACEHOLDER.into()); + } + if self.finished_at.trim().is_empty() { + self.finished_at = self + .steps + .iter() + .rev() + .find_map(|step| { + let trimmed = step.finished_at.trim(); + (!trimmed.is_empty()).then(|| trimmed.to_string()) + }) + .unwrap_or_else(|| self.started_at.clone()); + } + for step in &mut self.steps { + if step.started_at.trim().is_empty() { + step.started_at = self.started_at.clone(); + } + if step.finished_at.trim().is_empty() { + step.finished_at = self.finished_at.clone(); + } + } + } } #[derive(Deserialize)] @@ -274,8 +335,11 @@ enum ActfTrajectoryWire { Events(Vec), Canonical { schema_version: String, + #[serde(default)] steps: Vec, + #[serde(default)] started_at: String, + #[serde(default)] finished_at: String, #[serde(default)] events: Vec, @@ -289,7 +353,22 @@ impl<'de> Deserialize<'de> for ActfTrajectory { where D: Deserializer<'de>, { - match ActfTrajectoryWire::deserialize(deserializer)? { + let value = Value::deserialize(deserializer)?; + // Error dumps often ship `trajectory: null`, `trajectory: {}`, or a + // Python `Trajectory(...)` repr string instead of canonical JSON. + if value.is_null() || value.as_object().is_some_and(|object| object.is_empty()) { + return Ok(Self::from_event_log(Vec::new())); + } + if let Some(text) = value.as_str() { + let trimmed = text.trim_start(); + if trimmed.starts_with("Trajectory(") || trimmed.contains("ACTF_") { + return Ok(Self::from_event_log(Vec::new())); + } + return Err(serde::de::Error::custom( + "ACTF trajectory string is not a Trajectory(...) / ACTF dump", + )); + } + match ActfTrajectoryWire::deserialize(value).map_err(serde::de::Error::custom)? { ActfTrajectoryWire::Events(events) => Ok(Self::from_event_log(events)), ActfTrajectoryWire::Canonical { schema_version, @@ -298,14 +377,18 @@ impl<'de> Deserialize<'de> for ActfTrajectory { finished_at, events, extra, - } => Ok(Self { - schema_version, - steps, - started_at, - finished_at, - events, - extra, - }), + } => { + let mut trajectory = Self { + schema_version, + steps, + started_at, + finished_at, + events, + extra, + }; + trajectory.normalize_timestamps(); + Ok(trajectory) + } } } } @@ -323,7 +406,9 @@ pub struct ActfStep { pub tools: Vec, #[serde(default, deserialize_with = "null_as_default")] pub observation: Vec, + #[serde(default, deserialize_with = "null_as_empty_string")] pub started_at: String, + #[serde(default, deserialize_with = "null_as_empty_string")] pub finished_at: String, #[serde(flatten)] pub extra: Map, @@ -416,6 +501,23 @@ where Ok(Option::::deserialize(deserializer)?.unwrap_or_default()) } +/// Accept JSON string, number, or bool as a string field (common in error dumps). +fn stringish<'de, D>(deserializer: D) -> std::result::Result +where + D: Deserializer<'de>, +{ + let value = Value::deserialize(deserializer)?; + match value { + Value::Null => Ok(String::new()), + Value::String(text) => Ok(text), + Value::Number(number) => Ok(number.to_string()), + Value::Bool(flag) => Ok(flag.to_string()), + other => Err(serde::de::Error::custom(format!( + "expected string, number, bool, or null; got {other}" + ))), + } +} + fn null_as_default<'de, T, D>(deserializer: D) -> std::result::Result where T: Default + Deserialize<'de>, @@ -436,12 +538,34 @@ fn normalize_solved_at(value: &mut Value) { } } +/// Producers often ship both `tools` and `assistant_content.tool_calls` with +/// small drift (extra keys, id formatting). Prefer top-level `tools`, matching +/// [`ActfStep::effective_tools`]. +fn reconcile_step_tool_lists(step: &mut ActfStep) { + if step.tools.is_empty() || step.assistant_content.tool_calls.is_empty() { + return; + } + if step.tools != step.assistant_content.tool_calls { + step.assistant_content.tool_calls = step.tools.clone(); + } +} + +fn reconcile_document_tool_lists(document: &mut ActfDocument) { + for attempt in document.attempts.values_mut() { + for step in &mut attempt.trajectory.steps { + reconcile_step_tool_lists(step); + } + } +} + impl ActfDocument { #[cfg(any(test, feature = "lance-store"))] pub fn from_json_str(input: &str) -> InputResult { - let mut document: Self = - serde_json::from_str(input).map_err(|error| InputIssue::invalid(error.to_string()))?; + let sanitized = super::common::sanitize_json_nonfinite(input); + let mut document: Self = serde_json::from_str(sanitized.as_ref()) + .map_err(|error| InputIssue::invalid(error.to_string()))?; normalize_solved_at(&mut document.solved_at); + reconcile_document_tool_lists(&mut document); document.validate()?; Ok(document) } @@ -509,17 +633,14 @@ impl ActfTrajectory { "ACTF trajectory started_at and finished_at are required", )); } - if self.steps.is_empty() && self.events.is_empty() { - return Err(InputIssue::invalid( - "ACTF trajectory steps must not be empty", - )); - } + // Error dumps may ship null/`{}` trajectories (no steps, no events). + // Keep timestamps + schema; allow empty content. let mut previous_step = None; for step in &self.steps { - if step.step_id < 1 { + if step.step_id < 0 { return Err(InputIssue::invalid(format!( - "ACTF step_id must be positive, got {}", + "ACTF step_id must be non-negative, got {}", step.step_id ))); } @@ -536,15 +657,9 @@ impl ActfTrajectory { step.step_id ))); } - if !step.tools.is_empty() - && !step.assistant_content.tool_calls.is_empty() - && step.assistant_content.tool_calls != step.tools - { - return Err(InputIssue::invalid(format!( - "ACTF step {} assistant_content.tool_calls must equal tools", - step.step_id - ))); - } + // Divergent tools vs assistant_content.tool_calls is common in corpus + // dumps; import reconciles via reconcile_step_tool_lists, and convert + // already prefers top-level tools through effective_tools(). if !(step.metric.prompt_tokens_len.is_null() || step.metric.prompt_tokens_len.is_number()) || !(step.metric.completion_tokens_len.is_null() @@ -561,12 +676,9 @@ impl ActfTrajectory { let mut step_call_ids = HashSet::new(); for (call_index, call) in step.effective_tools().iter().enumerate() { let call_id = call.effective_id(step.step_id, call_index); - if !step_call_ids.insert(call_id) { - return Err(InputIssue::invalid(format!( - "duplicate ACTF tool call id '{}'", - call.effective_id(step.step_id, call_index) - ))); - } + // Duplicate ids within a step are reconciled at Storyline convert + // time; keep validating observation refs against the first insert. + let _ = step_call_ids.insert(call_id); } for observation in &step.observation { let referenced_id = observation @@ -650,6 +762,66 @@ mod tests { .unwrap() } + #[test] + fn accepts_null_or_empty_object_trajectory_as_empty_event_log() { + for trajectory in [json!(null), json!({})] { + let value = json!({ + "task_id": "frontierscience_research_0053", + "category": "research", + "correct": false, + "solved_at": null, + "attempts_tried": 1, + "k": 1, + "attempts": { + "1": { + "correct": false, + "final_answer": null, + "ground_truth": "rubric", + "trajectory": trajectory, + "meta": { + "status": "error", + "error": "TimeoutError: " + } + } + } + }); + let document: ActfDocument = serde_json::from_value(value).unwrap(); + let attempt = &document.attempts["1"]; + assert!(attempt.trajectory.steps.is_empty()); + assert!(attempt.trajectory.events.is_empty()); + document.validate().unwrap(); + let stories = super::convert::actf_to_storylines(&document).unwrap(); + assert_eq!(stories.len(), 1); + assert!(stories[0].turns.is_empty()); + assert_eq!(stories[0].session_id, "frontierscience_research_0053"); + } + } + + #[test] + fn accepts_python_trajectory_repr_string_as_empty_event_log() { + let value = json!({ + "task_id": "task_15_daily_summary", + "category": "synthesis", + "correct": false, + "solved_at": null, + "attempts_tried": 1, + "k": 1, + "attempts": { + "1": { + "correct": false, + "final_answer": "LLM request failed: network connection error.\n", + "trajectory": "Trajectory(schema_version='ACTF_v1.0', steps=[StepInfo(step_id=1)], started_at=datetime.datetime(2026, 6, 26, 7, 35, 16), finished_at=datetime.datetime(2026, 6, 26, 7, 35, 46))", + "status": null, + "score": 0.0 + } + } + }); + assert!(looks_like_actf_value(&value)); + let document: ActfDocument = serde_json::from_value(value).unwrap(); + assert!(document.attempts["1"].trajectory.steps.is_empty()); + document.validate().unwrap(); + } + #[test] fn accepts_numeric_solved_at_by_coercing_to_string() { let mut value = serde_json::to_value(fixture()).unwrap(); @@ -719,6 +891,34 @@ mod tests { ); } + #[test] + fn reconciles_divergent_tools_and_assistant_tool_calls_on_import() { + let mut value = serde_json::to_value(fixture()).unwrap(); + value["attempts"]["1"]["trajectory"]["steps"][0]["tools"] = json!([{ + "type": "tool_use", + "id": "call-tools", + "name": "Bash", + "input": {"command": "pwd"} + }]); + value["attempts"]["1"]["trajectory"]["steps"][0]["assistant_content"]["tool_calls"] = json!([{ + "type": "function", + "id": "call-assistant", + "function": {"name": "bash_command", "arguments": {"keystrokes": "pwd\n"}} + }]); + value["attempts"]["1"]["trajectory"]["steps"][0]["observation"] = json!([{ + "tool_use_id": "call-tools", + "type": "tool_result", + "content": "/app", + "is_error": false + }]); + let document = ActfDocument::from_json_str(&value.to_string()).unwrap(); + let step = &document.attempts["1"].trajectory.steps[0]; + assert_eq!(step.tools, step.assistant_content.tool_calls); + assert_eq!(step.effective_tools()[0].id, "call-tools"); + let stories = super::convert::actf_to_storylines(&document).unwrap(); + assert_eq!(stories.len(), 1); + } + #[cfg(feature = "proptest")] mod proptests { use proptest::prelude::*; @@ -766,8 +966,8 @@ mod tests { #[test] fn trajectory_validation_enforces_strictly_increasing_step_ids( - first in 1i64..10_000, - second in 1i64..10_000, + first in 0i64..10_000, + second in 0i64..10_000, ) { let mut document = fixture(); let trajectory = &mut document.attempts.get_mut("1").unwrap().trajectory; @@ -803,6 +1003,24 @@ mod tests { } } + #[test] + fn accepts_zero_based_step_ids() { + let mut document = fixture(); + let trajectory = &mut document.attempts.get_mut("1").unwrap().trajectory; + trajectory.steps[0].step_id = 0; + trajectory.steps[0].tools.clear(); + trajectory.steps[0].assistant_content.tool_calls.clear(); + trajectory.steps[0].observation.clear(); + if trajectory.steps.len() > 1 { + trajectory.steps[1].step_id = 1; + trajectory.steps[1].tools.clear(); + trajectory.steps[1].assistant_content.tool_calls.clear(); + trajectory.steps[1].observation.clear(); + } + trajectory.validate().unwrap(); + document.validate().unwrap(); + } + #[test] fn accepts_observation_without_type() { let mut value = serde_json::to_value(fixture()).unwrap(); @@ -819,6 +1037,48 @@ mod tests { document.validate().unwrap(); } + #[test] + fn parses_wireless_channel_dump_with_nan_and_numeric_solved_at() { + let document = parse_actf_document( + r#"{ + "task_id":"WirelessChannelSimulation/HighReliableSimulation", + "category":"WirelessChannelSimulation", + "correct":true, + "solved_at":1, + "attempts_tried":1, + "k":1, + "attempts":{"1":{ + "correct":true, + "final_answer":"print(1)", + "ground_truth":"", + "trajectory":{ + "schema_version":"ACTF_v1.0", + "steps":[{ + "step_id":1, + "assistant_content":{"content":"iteration=0","reasoning_content":"","tool_calls":[]}, + "metric":{"prompt_tokens_len":null,"completion_tokens_len":null,"llm_infer_ms":null,"env_action_ms":13653.41,"stop_reason":null}, + "system_prompt":"", + "user_content":"WirelessChannelSimulation/HighReliableSimulation", + "tools":[], + "observation":[{"combined_score": NaN}], + "started_at":"2026-01-01 00:00:00+00:00", + "finished_at":"2026-01-01 00:00:01+00:00" + }], + "started_at":"2026-01-01 00:00:00+00:00", + "finished_at":"2026-01-01 00:00:01+00:00" + }, + "status":"completed" + }} + }"#, + ) + .unwrap(); + assert_eq!(document.solved_at, Value::String("1".into())); + assert!( + document.attempts["1"].trajectory.steps[0].observation[0].extra["combined_score"] + .is_null() + ); + } + #[test] fn accepts_openclaw_event_log_as_trajectory() { let mut value = serde_json::to_value(fixture()).unwrap(); @@ -838,6 +1098,37 @@ mod tests { ); } + #[test] + fn accepts_numeric_category_and_empty_trajectory_object() { + let value = json!({ + "task_id": "f2feb6a4-363c-4c09-a804-0db564eafd68", + "category": 2, + "correct": false, + "solved_at": null, + "attempts_tried": 1, + "k": 1, + "attempts": { + "1": { + "correct": false, + "final_answer": null, + "ground_truth": "900000", + "trajectory": {}, + "meta": { + "status": "error", + "service_metrics": {}, + "service_task_id": null, + "error": "ClientConnectorError: Cannot connect to host" + } + } + } + }); + let document: ActfDocument = serde_json::from_value(value).unwrap(); + assert_eq!(document.category, "2"); + assert!(document.attempts["1"].trajectory.steps.is_empty()); + assert!(document.attempts["1"].trajectory.events.is_empty()); + document.validate().unwrap(); + } + #[test] fn treats_null_reasoning_content_as_empty_string() { let mut value = serde_json::to_value(fixture()).unwrap(); diff --git a/crates/persisting-pchronicle/src/formats/atif.rs b/crates/persisting-pchronicle/src/formats/atif.rs index 6d130f07..85e0d56f 100644 --- a/crates/persisting-pchronicle/src/formats/atif.rs +++ b/crates/persisting-pchronicle/src/formats/atif.rs @@ -408,8 +408,7 @@ fn atif_to_storyline_node( timestamp: step .timestamp .as_deref() - .map(StorylineTimestamp::from_rfc3339) - .transpose()?, + .and_then(StorylineTimestamp::from_rfc3339_lenient), source: step.source.clone(), message: step.message.clone(), reasoning_content: step.reasoning_content.clone(), diff --git a/crates/persisting-pchronicle/src/formats/common/json_sanitize.rs b/crates/persisting-pchronicle/src/formats/common/json_sanitize.rs new file mode 100644 index 00000000..841ba6b3 --- /dev/null +++ b/crates/persisting-pchronicle/src/formats/common/json_sanitize.rs @@ -0,0 +1,102 @@ +//! Repair non-standard JSON tokens that scientific / Python dumps emit. + +use std::borrow::Cow; + +/// Replace bare `NaN` / `Infinity` / `-Infinity` tokens with `null`. +/// +/// Python `json.dumps` allows these by default; `serde_json` rejects them, so +/// ACTF fingerprinting and decode both fail with "cannot detect import format" +/// even when the document is otherwise a clear ACTF dump. +pub(crate) fn sanitize_json_nonfinite(input: &str) -> Cow<'_, str> { + if !input.contains("NaN") && !input.contains("Infinity") { + return Cow::Borrowed(input); + } + let bytes = input.as_bytes(); + let mut out = Vec::with_capacity(bytes.len()); + let mut i = 0; + let mut in_string = false; + let mut escape = false; + while i < bytes.len() { + let b = bytes[i]; + if in_string { + out.push(b); + if escape { + escape = false; + } else if b == b'\\' { + escape = true; + } else if b == b'"' { + in_string = false; + } + i += 1; + continue; + } + if b == b'"' { + in_string = true; + out.push(b); + i += 1; + continue; + } + if match_bare_token(bytes, i, b"-Infinity") { + out.extend_from_slice(b"null"); + i += "-Infinity".len(); + continue; + } + if match_bare_token(bytes, i, b"Infinity") { + out.extend_from_slice(b"null"); + i += "Infinity".len(); + continue; + } + if match_bare_token(bytes, i, b"NaN") { + out.extend_from_slice(b"null"); + i += "NaN".len(); + continue; + } + out.push(b); + i += 1; + } + match String::from_utf8(out) { + Ok(text) => Cow::Owned(text), + Err(_) => Cow::Borrowed(input), + } +} + +fn match_bare_token(bytes: &[u8], index: usize, token: &[u8]) -> bool { + if !bytes[index..].starts_with(token) { + return false; + } + let before_ok = index == 0 + || matches!( + bytes[index - 1], + b':' | b'[' | b',' | b' ' | b'\t' | b'\n' | b'\r' + ); + let after = index + token.len(); + let after_ok = after >= bytes.len() + || matches!( + bytes[after], + b',' | b']' | b'}' | b' ' | b'\t' | b'\n' | b'\r' + ); + before_ok && after_ok +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::Value; + + #[test] + fn replaces_nonfinite_outside_strings_only() { + let input = r#"{"score": NaN, "note": "NaN", "hi": Infinity, "lo": -Infinity}"#; + let sanitized = sanitize_json_nonfinite(input); + let value: Value = serde_json::from_str(&sanitized).unwrap(); + assert!(value["score"].is_null()); + assert_eq!(value["note"], "NaN"); + assert!(value["hi"].is_null()); + assert!(value["lo"].is_null()); + } + + #[test] + fn leaves_standard_json_untouched() { + let input = r#"{"score": 1.5}"#; + assert!(matches!(sanitize_json_nonfinite(input), Cow::Borrowed(_))); + } +} diff --git a/crates/persisting-pchronicle/src/formats/common/mod.rs b/crates/persisting-pchronicle/src/formats/common/mod.rs index ecb5831e..38c0ec46 100644 --- a/crates/persisting-pchronicle/src/formats/common/mod.rs +++ b/crates/persisting-pchronicle/src/formats/common/mod.rs @@ -1,2 +1,5 @@ +pub(crate) mod json_sanitize; pub(crate) mod json_stream; pub(crate) mod jsonl; + +pub(crate) use json_sanitize::sanitize_json_nonfinite; diff --git a/crates/persisting-pchronicle/src/formats/detect.rs b/crates/persisting-pchronicle/src/formats/detect.rs index a2715a89..388d5036 100644 --- a/crates/persisting-pchronicle/src/formats/detect.rs +++ b/crates/persisting-pchronicle/src/formats/detect.rs @@ -118,6 +118,105 @@ mod tests { ); } + #[test] + fn does_not_guess_actf_from_steps_alone() { + let input = r#"{ + "task_id":"travel-planning", + "attempts":{"1":{ + "correct":false, + "trajectory":{ + "steps":[], + "started_at":"2026-06-17T07:26:27Z", + "finished_at":"2026-06-17T07:26:28Z" + } + }} + }"#; + assert_eq!(detect_format_from_content(input).unwrap(), None); + } + + #[test] + fn detects_actf_error_dump_with_empty_or_null_trajectory() { + for trajectory in [r#"{}"#, "null"] { + let input = format!( + r#"{{ + "task_id":"frontierscience_research_0053", + "category":"research", + "correct":false, + "attempts_tried":1, + "k":1, + "attempts":{{"1":{{ + "correct":false, + "trajectory":{trajectory}, + "meta":{{"status":"error","error":"TimeoutError: "}} + }}}} + }}"# + ); + assert_eq!( + detect_format_from_content(&input).unwrap(), + Some(DocumentFormat::Actf), + "trajectory={trajectory}" + ); + } + } + + #[test] + fn detects_actf_with_python_trajectory_repr_string() { + let input = r#"{ + "task_id":"task_15_daily_summary", + "category":"synthesis", + "correct":false, + "attempts_tried":1, + "k":1, + "attempts":{"1":{ + "correct":false, + "trajectory":"Trajectory(schema_version='ACTF_v1.0', steps=[])" + }} + }"#; + assert_eq!( + detect_format_from_content(input).unwrap(), + Some(DocumentFormat::Actf) + ); + } + + #[test] + fn detects_wireless_channel_actf_with_nan_observation_score() { + // Frontier-engineering dumps emit Python NaN and put schema_version + // after a large final_answer; full serde_json parse used to fail. + let input = r#"{ + "task_id":"WirelessChannelSimulation/HighReliableSimulation", + "category":"WirelessChannelSimulation", + "correct":true, + "solved_at":1, + "attempts_tried":1, + "k":1, + "attempts":{"1":{ + "correct":true, + "final_answer":"print(1)", + "ground_truth":"", + "trajectory":{ + "schema_version":"ACTF_v1.0", + "steps":[{ + "step_id":1, + "assistant_content":{"content":"iteration=0","reasoning_content":"","tool_calls":[]}, + "metric":{"prompt_tokens_len":null,"completion_tokens_len":null,"llm_infer_ms":null,"env_action_ms":1.0,"stop_reason":null}, + "system_prompt":"", + "user_content":"WirelessChannelSimulation/HighReliableSimulation", + "tools":[], + "observation":[{"combined_score": NaN}], + "started_at":"2026-01-01 00:00:00+00:00", + "finished_at":"2026-01-01 00:00:01+00:00" + }], + "started_at":"2026-01-01 00:00:00+00:00", + "finished_at":"2026-01-01 00:00:01+00:00" + } + }} + }"#; + assert_eq!( + detect_format_from_content(input).unwrap(), + Some(DocumentFormat::Actf) + ); + } + #[test] fn detects_atif_json_by_schema_and_agent_steps() { let versioned = r#"{"schema_version":"ATIF-v1.7","trajectory_id":"one","agent":{"name":"a","version":"1"},"steps":[]}"#; diff --git a/crates/persisting-pchronicle/src/formats/openai_corpus.rs b/crates/persisting-pchronicle/src/formats/openai_corpus.rs index ec97294c..4bc417ee 100644 --- a/crates/persisting-pchronicle/src/formats/openai_corpus.rs +++ b/crates/persisting-pchronicle/src/formats/openai_corpus.rs @@ -1474,9 +1474,7 @@ fn rows_to_storyline( .get("created_at") .filter(|value| !value.is_null()) .cloned() - .map(StorylineTimestamp::from_json) - .transpose() - .map_err(|issue| issue.at(format!("rows[{ordinal}].created_at")))?; + .and_then(StorylineTimestamp::from_json_lenient); let latency_ms = env_state .as_ref() .and_then(|state| state.get("total_latency_ms")) diff --git a/crates/persisting-pchronicle/src/formats/storyline.rs b/crates/persisting-pchronicle/src/formats/storyline.rs index 66c675ad..5763073e 100644 --- a/crates/persisting-pchronicle/src/formats/storyline.rs +++ b/crates/persisting-pchronicle/src/formats/storyline.rs @@ -14,7 +14,7 @@ use serde_json::{Map, Value}; use super::codec::{ DecodeContext, DecodeReport, FormatCapabilities, ProbeConfidence, TrajectoryFormat, }; -use super::timestamp::StorylineTimestamp; +use super::timestamp::{StorylineTimestamp, deserialize_optional_timestamp}; use super::unknown_fields::{StorylineUnknownFields, UnknownKeyCounts, compute_unknown_key_counts}; use crate::format::DocumentFormat; use crate::{InputIssue, InputResult, Result}; @@ -51,9 +51,17 @@ pub struct StorylineDocument { pub task: Option, #[serde(default, skip_serializing_if = "skip_optional_empty_prompt")] pub prompt: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde( + default, + deserialize_with = "deserialize_optional_timestamp", + skip_serializing_if = "Option::is_none" + )] pub started_at: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde( + default, + deserialize_with = "deserialize_optional_timestamp", + skip_serializing_if = "Option::is_none" + )] pub finished_at: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub final_metrics: Option, @@ -128,7 +136,12 @@ pub struct StorylineTurn { pub id: i64, #[serde(default, skip_serializing_if = "Option::is_none")] pub kind: Option, - #[serde(rename = "ts", default, skip_serializing_if = "Option::is_none")] + #[serde( + rename = "ts", + default, + deserialize_with = "deserialize_optional_timestamp", + skip_serializing_if = "Option::is_none" + )] pub timestamp: Option, #[serde(rename = "src")] pub source: String, @@ -162,7 +175,11 @@ pub struct StorylineTurn { pub env: Option, #[serde(default, skip_serializing_if = "skip_turn_prompt")] pub prompt: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] + #[serde( + default, + deserialize_with = "deserialize_optional_timestamp", + skip_serializing_if = "Option::is_none" + )] pub finished_at: Option, } @@ -1160,14 +1177,46 @@ mod tests { for value in [ serde_json::Value::Null, serde_json::json!(true), - serde_json::json!("2026/08/20 00:00:00"), + serde_json::json!("not-a-timestamp"), ] { assert!(crate::model::StorylineTimestamp::from_json(value).is_err()); } } #[test] - fn storyline_decode_rejects_non_rfc3339_timestamps() { + fn typed_timestamp_accepts_common_alternate_string_forms() { + for value in [ + serde_json::json!("2026/08/20 00:00:00"), + serde_json::json!("2026-08-20 12:00:00"), + serde_json::json!("2026-08-20T12:00:00"), + ] { + assert!( + crate::model::StorylineTimestamp::from_json(value.clone()).is_ok(), + "{value}" + ); + } + } + + #[test] + fn storyline_decode_keeps_unparseable_timestamps_empty() { + let input = serde_json::json!({ + "schema_version": STORYLINE_SCHEMA_VERSION, + "session": "session", + "agent": {"id": "agent"}, + "turns": [{ + "id": 1, + "ts": "definitely-not-a-time", + "src": "user", + "msg": "hello" + }] + }); + + let story = StorylineDocument::from_json_str(&input.to_string()).unwrap(); + assert!(story.turns[0].timestamp.is_none()); + } + + #[test] + fn storyline_decode_accepts_slash_separated_timestamps() { let input = serde_json::json!({ "schema_version": STORYLINE_SCHEMA_VERSION, "session": "session", @@ -1180,8 +1229,9 @@ mod tests { }] }); - let error = StorylineDocument::from_json_str(&input.to_string()).unwrap_err(); - assert!(error.to_string().contains("RFC3339"), "{error}"); + let story = StorylineDocument::from_json_str(&input.to_string()).unwrap(); + let ts = story.turns[0].timestamp.as_ref().expect("parsed timestamp"); + assert_eq!(ts.source_string(), Some("2026/08/20 12:00:00")); } #[test] diff --git a/crates/persisting-pchronicle/src/formats/timestamp.rs b/crates/persisting-pchronicle/src/formats/timestamp.rs index 15533c60..70aaa65a 100644 --- a/crates/persisting-pchronicle/src/formats/timestamp.rs +++ b/crates/persisting-pchronicle/src/formats/timestamp.rs @@ -1,4 +1,4 @@ -use chrono::{DateTime, SecondsFormat, Utc}; +use chrono::{DateTime, NaiveDateTime, SecondsFormat, Utc}; use serde::{Deserialize, Deserializer, Serialize, Serializer}; use serde_json::Value; @@ -19,9 +19,11 @@ pub struct StorylineTimestamp { impl StorylineTimestamp { pub fn from_json(source: Value) -> InputResult { let instant = match &source { - Value::String(value) => DateTime::parse_from_rfc3339(value) - .map_err(|_| InputIssue::invalid("timestamp string must be RFC3339"))? - .with_timezone(&Utc), + Value::String(value) => parse_timestamp_string(value).ok_or_else(|| { + InputIssue::invalid( + "timestamp string must be RFC3339 or a recognized date/time / Unix form", + ) + })?, Value::Number(value) => { let nanos = decimal_seconds_to_nanos(&value.to_string())?; DateTime::::from_timestamp_nanos(nanos) @@ -42,10 +44,20 @@ impl StorylineTimestamp { }) } + /// Best-effort parse for optional timestamps: try alternate forms, else `None`. + pub fn from_json_lenient(source: Value) -> Option { + Self::from_json(source).ok() + } + pub fn from_rfc3339(value: &str) -> InputResult { Self::from_json(Value::String(value.to_string())) } + /// Soft string parse used by converters: unrecognized values become `None`. + pub fn from_rfc3339_lenient(value: &str) -> Option { + Self::from_json_lenient(Value::String(value.to_string())) + } + pub fn from_utc(instant: DateTime) -> InputResult { let unix_nanos = instant .timestamp_nanos_opt() @@ -104,6 +116,131 @@ impl<'de> Deserialize<'de> for StorylineTimestamp { } } +/// Deserialize `Option`: null/missing → None; unparseable → None. +pub fn deserialize_optional_timestamp<'de, D>( + deserializer: D, +) -> Result, D::Error> +where + D: Deserializer<'de>, +{ + let value = Option::::deserialize(deserializer)?; + Ok(match value { + None | Some(Value::Null) => None, + Some(value) => StorylineTimestamp::from_json_lenient(value), + }) +} + +/// Parse common timestamp string forms into UTC. +/// +/// Order: RFC3339 → RFC3339-ish with assumed UTC → Naive local forms as UTC → +/// offset forms → Unix seconds/millis/micros encoded as decimal strings. +fn parse_timestamp_string(value: &str) -> Option> { + let value = value.trim(); + if value.is_empty() { + return None; + } + + if let Ok(dt) = DateTime::parse_from_rfc3339(value) { + return Some(dt.with_timezone(&Utc)); + } + + // Space separator / missing `Z`: normalize then retry RFC3339. + if let Some(normalized) = normalize_toward_rfc3339(value) + && let Ok(dt) = DateTime::parse_from_rfc3339(&normalized) + { + return Some(dt.with_timezone(&Utc)); + } + + const WITH_OFFSET: &[&str] = &[ + "%Y-%m-%d %H:%M:%S%.f%:z", + "%Y-%m-%d %H:%M:%S%:z", + "%Y-%m-%dT%H:%M:%S%.f%:z", + "%Y-%m-%dT%H:%M:%S%:z", + "%Y/%m/%d %H:%M:%S%.f%:z", + "%Y/%m/%d %H:%M:%S%:z", + "%Y-%m-%d %H:%M:%S%.f%z", + "%Y-%m-%d %H:%M:%S%z", + "%Y-%m-%dT%H:%M:%S%.f%z", + "%Y-%m-%dT%H:%M:%S%z", + ]; + for fmt in WITH_OFFSET { + if let Ok(dt) = DateTime::parse_from_str(value, fmt) { + return Some(dt.with_timezone(&Utc)); + } + } + + const NAIVE_UTC: &[&str] = &[ + "%Y-%m-%dT%H:%M:%S%.f", + "%Y-%m-%dT%H:%M:%S", + "%Y-%m-%d %H:%M:%S%.f", + "%Y-%m-%d %H:%M:%S", + "%Y/%m/%d %H:%M:%S%.f", + "%Y/%m/%d %H:%M:%S", + "%Y/%m/%dT%H:%M:%S%.f", + "%Y/%m/%dT%H:%M:%S", + ]; + for fmt in NAIVE_UTC { + if let Ok(naive) = NaiveDateTime::parse_from_str(value, fmt) { + return Some(naive.and_utc()); + } + } + + parse_unix_string(value) +} + +fn normalize_toward_rfc3339(value: &str) -> Option { + let trimmed = value.trim(); + if trimmed.len() < 11 { + return None; + } + // `2026-08-20 12:00:00` / `2026/08/20 12:00:00` → `T` + optional `Z` + let mut candidate = trimmed.replace('/', "-"); + if candidate.as_bytes().get(10) == Some(&b' ') { + candidate.replace_range(10..11, "T"); + } + let tail = &candidate[10..]; + let has_zone = candidate.ends_with('Z') + || candidate.ends_with('z') + || tail.contains('+') + || tail.rfind('-').is_some_and(|idx| idx > 0); + if !has_zone { + candidate.push('Z'); + } + if candidate == trimmed { + None + } else { + Some(candidate) + } +} + +fn parse_unix_string(value: &str) -> Option> { + if let Ok(n) = value.parse::() { + return unix_i64_to_utc(n); + } + // `"1710000000.25"` → seconds with fraction + if value.contains('.') + && let Ok(nanos) = decimal_seconds_to_nanos(value) + { + return Some(DateTime::::from_timestamp_nanos(nanos)); + } + None +} + +fn unix_i64_to_utc(n: i64) -> Option> { + let abs = n.unsigned_abs(); + // Heuristic by magnitude (absolute value): + // < 1e11 → seconds (year ~5138) + // < 1e14 → millis + // else → micros + if abs < 100_000_000_000 { + DateTime::from_timestamp(n, 0) + } else if abs < 100_000_000_000_000 { + DateTime::from_timestamp_millis(n) + } else { + DateTime::from_timestamp_micros(n) + } +} + fn decimal_seconds_to_nanos(input: &str) -> InputResult { let (negative, unsigned) = match input.strip_prefix('-') { Some(value) => (true, value), @@ -179,9 +316,37 @@ fn parse_digits(digits: &str) -> InputResult { #[cfg(test)] mod tests { - #[cfg(feature = "proptest")] use super::*; + #[test] + fn parses_common_non_rfc3339_strings() { + let cases = [ + "2026-08-20 12:00:00", + "2026/08/20 12:00:00", + "2026-08-20T12:00:00", + "2026-08-20 12:00:00.123456", + "2026/08/20T12:00:00.5", + ]; + for raw in cases { + let ts = StorylineTimestamp::from_rfc3339(raw) + .unwrap_or_else(|error| panic!("expected parse for {raw}: {error}")); + assert_eq!(ts.source_string(), Some(raw)); + assert!(ts.instant().timestamp() > 0, "{raw}"); + } + } + + #[test] + fn parses_unix_seconds_as_string() { + let ts = StorylineTimestamp::from_rfc3339("1710000000").unwrap(); + assert_eq!(ts.instant().timestamp(), 1710000000); + } + + #[test] + fn lenient_returns_none_for_garbage() { + assert!(StorylineTimestamp::from_rfc3339_lenient("not-a-time").is_none()); + assert!(StorylineTimestamp::from_rfc3339_lenient("").is_none()); + } + #[cfg(feature = "proptest")] mod proptests { use super::*; diff --git a/crates/persisting-pchronicle/src/formats/unknown_fields.rs b/crates/persisting-pchronicle/src/formats/unknown_fields.rs index 3115f88f..0fec0e74 100644 --- a/crates/persisting-pchronicle/src/formats/unknown_fields.rs +++ b/crates/persisting-pchronicle/src/formats/unknown_fields.rs @@ -63,6 +63,9 @@ pub struct UnknownFieldImportWarnings { } impl UnknownFieldImportWarnings { + pub fn merge(&mut self, other: &Self) { + self.observe(&other.counts); + } /// Observe all Storylines decoded from one physical input Source. /// /// Converters may attach a document-level unknown pointer to multiple diff --git a/crates/persisting-pchronicle/src/storage.rs b/crates/persisting-pchronicle/src/storage.rs index 0f63a95f..d00bde9b 100644 --- a/crates/persisting-pchronicle/src/storage.rs +++ b/crates/persisting-pchronicle/src/storage.rs @@ -29,6 +29,13 @@ pub use crate::discovery::{ pub use crate::store::index_build_progress::{ Guard as IndexBuildProgressGuard, install as install_index_build_progress, }; +#[cfg(feature = "lance-store")] +pub use crate::store::object_store_io_gate::{ + IoKind as ObjectStoreIoKind, ObjectStoreGateSnapshot, ObjectStoreThrottleEvent, + ObjectStoreThrottleHookGuard, format_aimd_flow_label as format_object_store_aimd_flow_label, + install_throttle_hook as install_object_store_throttle_hook, + snapshot as object_store_gate_snapshot, +}; #[cfg(feature = "lance-store")] pub use crate::store::{ @@ -36,25 +43,26 @@ pub use crate::store::{ CatalogErrorPolicy, CatalogEventProvenance, CatalogEventView, CatalogNamespace, CatalogPage, CatalogProjectionStatus, CatalogSnapshotOptions, CatalogSourceDescription, CatalogSourceKind, CatalogSourceRevision, CatalogSourceStatus, CatalogStorylineKey, CatalogTrajectoryBundle, - ChronicleManifest, CommitRunOutcome, CompactJsonlColumn, CompactJsonlOffload, - CompactJsonlOptions, CompactJsonlRecord, CompactJsonlStore, DEFAULT_CONTENT_OFFLOAD_THRESHOLD, - DEFAULT_CONTENT_PREVIEW_BYTES, DEFAULT_DATASET_NAME, DEFAULT_MAX_EVENT_FALLBACK_BYTES, - DEFAULT_MAX_EVENT_FALLBACK_ROWS, DEFAULT_PHYSICAL_PAGE_LIMIT, DatasetCatalogSnapshot, - DatasetLocation, DatasetLocationKind, DatasetMount, DiscoveredSource, EventFactSnapshot, - ImportableObjectEvent, ShallowNavEntry, - EventLogLayoutStats, EventWriterFence, ExportOutcome, LanceMaintenanceOptions, - LanceMaintenanceReport, LeaseAcquireOutcome, ManifestKind, ManifestStats, NamespacePath, - ObjectStoreManifestWriteMode, PhysicalColumn, PhysicalDataFile, PhysicalFileLayout, - PhysicalFragment, PhysicalLayout, PhysicalPage, PhysicalPagePreview, PhysicalPageQuery, - PhysicalSource, PhysicalTable, ProjectionSourceSnapshot, RawEventLanceAppender, - RawEventLanceStore, ReplayOutcome, RunControlStore, StorylineContentOptions, - StorylineContentReadMode, StorylineDataSource, StorylineDataSourceOptions, StorylineLanceStore, + ChronicleManifest, CommitRunOutcome, CompactJsonlBuildPhase, CompactJsonlColumn, + CompactJsonlImportEvent, CompactJsonlOffload, CompactJsonlOptions, CompactJsonlRecord, + CompactJsonlStore, DEFAULT_CONTENT_OFFLOAD_THRESHOLD, DEFAULT_CONTENT_PREVIEW_BYTES, + DEFAULT_MAX_CHUNK_BYTES, + DEFAULT_DATASET_NAME, DEFAULT_MAX_EVENT_FALLBACK_BYTES, DEFAULT_MAX_EVENT_FALLBACK_ROWS, + DEFAULT_PHYSICAL_PAGE_LIMIT, DatasetCatalogSnapshot, DatasetLocation, DatasetLocationKind, + DatasetMount, DiscoveredSource, EventFactSnapshot, EventLogLayoutStats, EventWriterFence, + ExportOutcome, ImportableObjectEvent, LanceMaintenanceOptions, LanceMaintenanceReport, + LeaseAcquireOutcome, ManifestKind, ManifestStats, NamespacePath, ObjectStoreManifestWriteMode, + PhysicalColumn, PhysicalDataFile, PhysicalFileLayout, PhysicalFragment, PhysicalLayout, + PhysicalPage, PhysicalPagePreview, PhysicalPageQuery, PhysicalSource, PhysicalTable, + ProjectionSourceSnapshot, RawEventLanceAppender, RawEventLanceStore, ReplayOutcome, + RunControlStore, ShallowNavEntry, StorylineContentOptions, StorylineContentReadMode, + StorylineDataSource, StorylineDataSourceOptions, StorylineLanceStore, StorylineMaintenanceReport, StorylineProjectionLineage, StorylineStreamImportReport, - StorylineStreamOptions, StorylineTablePaths, TrajectoryStats, attempt_registry_now_ms, distinct_session_ids_in_run, - export_source_dirs, export_story_bundle, inspect_physical_file, inspect_physical_layout, - inspect_physical_page, list_physical_sources, load_manifest, load_manifest_at_uri, - raw_event_lance_path, write_compact_jsonl_manifest, write_storyline_manifest, - write_storyline_manifest_at_uri, + StorylineStreamOptions, StorylineTablePaths, TrajectoryStats, attempt_registry_now_ms, + distinct_session_ids_in_run, export_source_dirs, export_story_bundle, inspect_physical_file, + inspect_physical_layout, inspect_physical_page, list_physical_sources, load_manifest, + load_manifest_at_uri, raw_event_lance_path, write_compact_jsonl_manifest, + write_storyline_manifest, write_storyline_manifest_at_uri, }; // Compatibility exports; new callers should use `crate::search`. diff --git a/crates/persisting-pchronicle/src/store/catalog/discovery.rs b/crates/persisting-pchronicle/src/store/catalog/discovery.rs index 73cfbf16..5c7974c4 100644 --- a/crates/persisting-pchronicle/src/store/catalog/discovery.rs +++ b/crates/persisting-pchronicle/src/store/catalog/discovery.rs @@ -929,7 +929,11 @@ async fn object_shallow_children( let mut child_dirs = BTreeSet::new(); let mut files = Vec::new(); for entry in entries { - let path = entry.path.trim_start_matches(&prefix).trim_matches('/'); + let path = entry + .path + .strip_prefix(&prefix) + .unwrap_or(&entry.path) + .trim_matches('/'); if path.is_empty() { continue; } @@ -1047,14 +1051,11 @@ async fn probe_object_prefix( }))); } if manifest.is_storyline_leaf() { - let current = store - .stat_file(&join("CURRENT")) - .await? - .ok_or_else(|| { - anyhow::anyhow!( - "storyline chronicle.manifest requires CURRENT under {relative}" - ) - })?; + let current = store.stat_file(&join("CURRENT")).await?.ok_or_else(|| { + anyhow::anyhow!( + "storyline chronicle.manifest requires CURRENT under {relative}" + ) + })?; let current_meta = RemoteObjectMeta::from(current); return Ok(Some(ObjectProbe::Source(Candidate::Storyline { file: source_file, diff --git a/crates/persisting-pchronicle/src/store/compact_jsonl.rs b/crates/persisting-pchronicle/src/store/compact_jsonl.rs index 787f5691..f9d81cb1 100644 --- a/crates/persisting-pchronicle/src/store/compact_jsonl.rs +++ b/crates/persisting-pchronicle/src/store/compact_jsonl.rs @@ -25,6 +25,8 @@ const RAW_COLUMN: &str = "_raw_"; const OFFLOAD_COLUMN: &str = "_offload_"; const FORMAT_KEY: &str = "pchronicle.format"; const FORMAT_NAME: &str = "compact-jsonl/v1"; +/// How often Building phases emit processed/total ticks. +const BUILD_PROGRESS_EVERY: u64 = 8192; #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub struct CompactJsonlColumn { @@ -112,6 +114,53 @@ pub struct CompactJsonlRecord { pub filename: String, } +/// Progress events emitted while building a Compact JSONL Lance snapshot. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum CompactJsonlImportEvent { + Listed { + files: u64, + bytes: u64, + }, + /// Periodic updates while reading one input file (`done` is true on completion). + Reading { + relative: String, + file_bytes: u64, + file_rows: u64, + total_rows: u64, + done: bool, + }, + Building { + phase: CompactJsonlBuildPhase, + rows: u64, + /// When set, UI shows `phase processed/rows` for long in-phase work. + processed: Option, + }, + Written { + rows: u64, + }, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum CompactJsonlBuildPhase { + Keys, + Offload, + Columns, + Lance, + Manifest, +} + +impl CompactJsonlBuildPhase { + pub fn as_str(self) -> &'static str { + match self { + Self::Keys => "keys", + Self::Offload => "offload", + Self::Columns => "columns", + Self::Lance => "lance", + Self::Manifest => "manifest", + } + } +} + pub struct CompactJsonlStore; impl CompactJsonlStore { @@ -316,6 +365,15 @@ impl CompactJsonlStore { input: impl AsRef, output: impl AsRef, options: &CompactJsonlOptions, + ) -> Result { + Self::import_path_with_progress(input, output, options, |_| Ok(())).await + } + + pub async fn import_path_with_progress( + input: impl AsRef, + output: impl AsRef, + options: &CompactJsonlOptions, + mut on_progress: impl FnMut(CompactJsonlImportEvent) -> Result<()>, ) -> Result { let input = input.as_ref(); let output = output.as_ref(); @@ -325,6 +383,16 @@ impl CompactJsonlStore { !files.is_empty(), "compact JSONL input contains no .json, .jsonl, or .ndjson files" ); + let listed_bytes = files.iter().try_fold(0u64, |total, path| { + let len = fs::metadata(path).map(|meta| meta.len()).unwrap_or(0); + total + .checked_add(len) + .context("compact JSONL listed byte count overflow") + })?; + on_progress(CompactJsonlImportEvent::Listed { + files: files.len() as u64, + bytes: listed_bytes, + })?; if output.exists() { fs::remove_dir_all(output) .with_context(|| format!("replace compact JSONL output {}", output.display()))?; @@ -342,6 +410,7 @@ impl CompactJsonlStore { .context("compact JSONL filename is not UTF-8")? .replace('\\', "/"); let first_row = rows.len(); + let file_bytes = fs::metadata(&file).map(|meta| meta.len()).unwrap_or(0); if is_json_document(&file) { let raw = fs::read(&file)?; let value: Value = serde_json::from_slice(&raw) @@ -374,10 +443,32 @@ impl CompactJsonlStore { "compact JSONL {relative}:{line_no} must be a JSON object" ); rows.push((value, raw, relative.clone(), line_no)); + let file_rows = rows.len() - first_row; + if file_rows % 8192 == 0 { + on_progress(CompactJsonlImportEvent::Reading { + relative: relative.clone(), + file_bytes, + file_rows: file_rows as u64, + total_rows: rows.len() as u64, + done: false, + })?; + } } } ensure!(rows.len() > first_row, "compact JSONL {relative} is empty"); + on_progress(CompactJsonlImportEvent::Reading { + relative, + file_bytes, + file_rows: (rows.len() - first_row) as u64, + total_rows: rows.len() as u64, + done: true, + })?; } + on_progress(CompactJsonlImportEvent::Building { + phase: CompactJsonlBuildPhase::Keys, + rows: rows.len() as u64, + processed: None, + })?; let schema = schema(options)?; let mut arrays: Vec> = Vec::new(); let (ids, timestamps): (Vec<_>, Vec<_>) = rows @@ -403,9 +494,15 @@ impl CompactJsonlStore { ensure!(unique_ids.insert(id), "duplicate compact JSONL id '{id}'"); } let filenames: Vec = rows.iter().map(|(_, _, file, _)| file.clone()).collect(); + let total_rows = rows.len() as u64; + on_progress(CompactJsonlImportEvent::Building { + phase: CompactJsonlBuildPhase::Offload, + rows: total_rows, + processed: Some(0), + })?; let mut offloads = Vec::with_capacity(rows.len()); let offload_dir = output.join("_offload"); - for (_, raw, _, _) in &rows { + for (idx, (_, raw, _, _)) in rows.iter().enumerate() { if options.offload_threshold > 0 && raw.len() >= options.offload_threshold { fs::create_dir_all(&offload_dir)?; let key = blake3::hash(raw).to_hex().to_string(); @@ -427,22 +524,40 @@ impl CompactJsonlStore { } else { offloads.push(None); } + let processed = (idx + 1) as u64; + if processed == total_rows || processed.is_multiple_of(BUILD_PROGRESS_EVERY) { + on_progress(CompactJsonlImportEvent::Building { + phase: CompactJsonlBuildPhase::Offload, + rows: total_rows, + processed: Some(processed), + })?; + } } arrays.push(Arc::new(StringArray::from(ids))); arrays.push(Arc::new(StringArray::from(timestamps))); arrays.push(Arc::new(StringArray::from(filenames))); - let data = rows - .iter() - .zip(&offloads) - .map(|((value, _, _, _), offload)| { - let value = if offload.is_none() { - serde_json::to_string(value)? - } else { - "null".into() - }; - encode_json_bytes(&value) - }) - .collect::>>()?; + on_progress(CompactJsonlImportEvent::Building { + phase: CompactJsonlBuildPhase::Columns, + rows: total_rows, + processed: Some(0), + })?; + let mut data = Vec::with_capacity(rows.len()); + for (idx, ((value, _, _, _), offload)) in rows.iter().zip(&offloads).enumerate() { + let value = if offload.is_none() { + serde_json::to_string(value)? + } else { + "null".into() + }; + data.push(encode_json_bytes(&value)?); + let processed = (idx + 1) as u64; + if processed == total_rows || processed.is_multiple_of(BUILD_PROGRESS_EVERY) { + on_progress(CompactJsonlImportEvent::Building { + phase: CompactJsonlBuildPhase::Columns, + rows: total_rows, + processed: Some(processed), + })?; + } + } arrays.push(Arc::new(LargeBinaryArray::from( data.iter().map(Vec::as_slice).collect::>(), ))); @@ -450,17 +565,25 @@ impl CompactJsonlStore { if matches!(column.name.as_str(), "id" | "timestamp") { continue; } - let values = rows - .iter() - .map(|(v, _, _, _)| -> Result>> { + let mut values = Vec::with_capacity(rows.len()); + for (idx, (v, _, _, _)) in rows.iter().enumerate() { + values.push( path_value(v, &column.path) .map(|x| { let json = serde_json::to_string(x)?; encode_json_bytes(&json) }) - .transpose() - }) - .collect::>>()?; + .transpose()?, + ); + let processed = (idx + 1) as u64; + if processed == total_rows || processed.is_multiple_of(BUILD_PROGRESS_EVERY) { + on_progress(CompactJsonlImportEvent::Building { + phase: CompactJsonlBuildPhase::Columns, + rows: total_rows, + processed: Some(processed), + })?; + } + } arrays.push(Arc::new(LargeBinaryArray::from( values.iter().map(|x| x.as_deref()).collect::>(), ))); @@ -490,17 +613,49 @@ impl CompactJsonlStore { .collect::>(), ))); let batch = RecordBatch::try_new(schema.clone(), arrays)?; - InsertBuilder::new(output.to_string_lossy().as_ref()) - .with_params(&WriteParams { + on_progress(CompactJsonlImportEvent::Building { + phase: CompactJsonlBuildPhase::Lance, + rows: total_rows, + processed: None, + })?; + let uri = output.to_string_lossy().into_owned(); + let mut write = Box::pin(async move { + let write_params = WriteParams { mode: WriteMode::Create, ..Default::default() - }) - .execute_stream(RecordBatchIterator::new(vec![Ok(batch)], schema)) - .await - .context("write compact JSONL Lance dataset")?; + }; + InsertBuilder::new(uri.as_str()) + .with_params(&write_params) + .execute_stream(RecordBatchIterator::new(vec![Ok(batch)], schema)) + .await + .context("write compact JSONL Lance dataset") + }); + loop { + tokio::select! { + result = &mut write => { + result?; + break; + } + _ = tokio::time::sleep(std::time::Duration::from_secs(2)) => { + on_progress(CompactJsonlImportEvent::Building { + phase: CompactJsonlBuildPhase::Lance, + rows: total_rows, + processed: None, + })?; + } + } + } // Store-layer contract: every published compact dataset carries // chronicle.manifest. import and sync both end here. + on_progress(CompactJsonlImportEvent::Building { + phase: CompactJsonlBuildPhase::Manifest, + rows: total_rows, + processed: None, + })?; Self::publish_manifest(output).await?; + on_progress(CompactJsonlImportEvent::Written { + rows: rows.len() as u64, + })?; Ok(rows.len()) } diff --git a/crates/persisting-pchronicle/src/store/location.rs b/crates/persisting-pchronicle/src/store/location.rs index 113c3381..99fd69be 100644 --- a/crates/persisting-pchronicle/src/store/location.rs +++ b/crates/persisting-pchronicle/src/store/location.rs @@ -286,23 +286,19 @@ impl DatasetLocation { if let Some(entry) = store .stat_file(&join(crate::store::CHRONICLE_MANIFEST_FILE)) .await? + && let Some((bytes, _)) = store.read(&entry.path).await? + && let Ok(text) = std::str::from_utf8(&bytes) + && let Ok(manifest) = toml::from_str::(text) + && manifest.validate().is_ok() { - if let Some((bytes, _)) = store.read(&entry.path).await? { - if let Ok(text) = std::str::from_utf8(&bytes) - && let Ok(manifest) = - toml::from_str::(text) - && manifest.validate().is_ok() - { - if manifest.is_storyline_leaf() { - return Ok(Some("storyline")); - } - if manifest.is_compact_jsonl_leaf() { - return Ok(Some("compact-jsonl")); - } - if matches!(manifest.kind, crate::store::ManifestKind::Leaf) { - return Ok(Some("other")); - } - } + if manifest.is_storyline_leaf() { + return Ok(Some("storyline")); + } + if manifest.is_compact_jsonl_leaf() { + return Ok(Some("compact-jsonl")); + } + if matches!(manifest.kind, crate::store::ManifestKind::Leaf) { + return Ok(Some("other")); } } if store.stat_file(&join("CURRENT")).await?.is_some() { @@ -410,7 +406,11 @@ impl DatasetLocation { let mut dirs = BTreeSet::new(); let mut files = BTreeSet::new(); for entry in entries { - let path = entry.path.trim_start_matches(&prefix).trim_matches('/'); + let path = entry + .path + .strip_prefix(&prefix) + .unwrap_or(&entry.path) + .trim_matches('/'); if path.is_empty() { continue; } @@ -570,7 +570,11 @@ impl DatasetLocation { })?; let mut child_dirs = BTreeSet::new(); for entry in entries { - let path = entry.path.trim_start_matches(&list_prefix).trim_matches('/'); + let path = entry + .path + .strip_prefix(&list_prefix) + .unwrap_or(&entry.path) + .trim_matches('/'); if path.is_empty() { continue; } @@ -595,7 +599,10 @@ impl DatasetLocation { on_event(ImportableObjectEvent::File { key: child_rel, size: entry.metadata.content_length(), - modified: entry.metadata.last_modified().map(|value| value.to_string()), + modified: entry + .metadata + .last_modified() + .map(|value| value.to_string()), }) .await?; continue; @@ -732,8 +739,7 @@ where .unwrap_or(file.as_path()) .to_string_lossy() .replace('\\', "/"); - std::fs::remove_file(&file) - .with_context(|| format!("delete file {}", file.display()))?; + std::fs::remove_file(&file).with_context(|| format!("delete file {}", file.display()))?; deleted = deleted.saturating_add(1); on_progress(deleted, total, &relative)?; } @@ -766,11 +772,7 @@ fn list_local_files_recursive(root: &Path) -> Result> { } fn is_nav_child_name(name: &str) -> bool { - !name.is_empty() - && name != "." - && name != ".." - && !name.starts_with('.') - && name != "_meta" + !name.is_empty() && name != "." && name != ".." && !name.starts_with('.') && name != "_meta" } fn is_storyline_interior_name(name: &str) -> bool { diff --git a/crates/persisting-pchronicle/src/store/mod.rs b/crates/persisting-pchronicle/src/store/mod.rs index 098e3372..598a4b77 100644 --- a/crates/persisting-pchronicle/src/store/mod.rs +++ b/crates/persisting-pchronicle/src/store/mod.rs @@ -35,8 +35,8 @@ mod events; mod files; #[cfg(feature = "lance-store")] pub(crate) mod index_build_gate; +#[cfg(feature = "lance-store")] pub(crate) mod index_build_progress; -pub(crate) mod object_store_io_gate; #[cfg(feature = "lance-store")] mod inspect; #[cfg(feature = "lance-store")] @@ -44,6 +44,8 @@ mod local_query_manifest; #[cfg(feature = "lance-store")] mod location; #[cfg(feature = "lance-store")] +pub(crate) mod object_store_io_gate; +#[cfg(feature = "lance-store")] pub(crate) mod opendal_store; #[cfg(feature = "lance-store")] mod query_engine; @@ -84,8 +86,8 @@ pub use chronicle_manifest::{ }; #[cfg(feature = "lance-store")] pub use compact_jsonl::{ - CompactJsonlColumn, CompactJsonlOffload, CompactJsonlOptions, CompactJsonlRecord, - CompactJsonlStore, + CompactJsonlBuildPhase, CompactJsonlColumn, CompactJsonlImportEvent, CompactJsonlOffload, + CompactJsonlOptions, CompactJsonlRecord, CompactJsonlStore, }; #[cfg(feature = "lance-store")] pub(crate) use document_source::{DocumentSourceImpl, open_document_source}; @@ -123,9 +125,7 @@ pub(crate) use local_query_manifest::{ LocalQueryInputFile, LocalQueryManifest, LocalQueryManifestOptions, }; #[cfg(feature = "lance-store")] -pub use location::{ - DatasetLocation, DatasetLocationKind, ImportableObjectEvent, ShallowNavEntry, -}; +pub use location::{DatasetLocation, DatasetLocationKind, ImportableObjectEvent, ShallowNavEntry}; #[cfg(feature = "lance-store")] pub use query_engine::{ ChronicleQueryEngine, ChronicleQueryExecutionOptions, DEFAULT_QUERY_MEMORY_LIMIT_BYTES, @@ -139,10 +139,11 @@ pub(crate) use storyline::StorylineProjectionPublicationOutcome; #[cfg(feature = "lance-store")] pub use storyline::{ DATAFUSION_RUNS_TABLE, DATAFUSION_STEPS_TABLE, DATAFUSION_TOOL_CALLS_TABLE, - DEFAULT_CONTENT_OFFLOAD_THRESHOLD, DEFAULT_CONTENT_PREVIEW_BYTES, ProjectionSourceSnapshot, - StorylineContentOptions, StorylineContentReadMode, StorylineDataFusionTableNames, - StorylineDataSource, StorylineDataSourceOptions, StorylineLanceStore, - StorylineMaintenanceReport, StorylineProjectionLineage, StorylineStreamImportReport, + DEFAULT_CONTENT_OFFLOAD_THRESHOLD, DEFAULT_CONTENT_PREVIEW_BYTES, DEFAULT_MAX_CHUNK_BYTES, + ProjectionSourceSnapshot, StorylineContentOptions, StorylineContentReadMode, + StorylineDataFusionTableNames, StorylineDataSource, StorylineDataSourceOptions, + StorylineLanceStore, StorylineMaintenanceReport, StorylineProjectionLineage, + StorylineStreamImportReport, StorylineStreamOptions, StorylineTableKind, StorylineTablePaths, story_runs_arrow_schema, story_runs_from_batch, story_runs_to_batch, story_steps_arrow_schema, story_steps_from_batch, story_steps_to_batch, story_tool_calls_arrow_schema, story_tool_calls_from_batch, diff --git a/crates/persisting-pchronicle/src/store/object_store_io_gate.rs b/crates/persisting-pchronicle/src/store/object_store_io_gate.rs index 5a515213..15c1a0c8 100644 --- a/crates/persisting-pchronicle/src/store/object_store_io_gate.rs +++ b/crates/persisting-pchronicle/src/store/object_store_io_gate.rs @@ -21,13 +21,13 @@ const SUCCESS_STREAK_TO_DECAY: u32 = 4; /// Whether the gated op is primarily reading metadata/objects or writing them. #[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub(crate) enum IoKind { +pub enum IoKind { Read, Write, } impl IoKind { - fn as_str(self) -> &'static str { + pub fn as_str(self) -> &'static str { match self { Self::Read => "read", Self::Write => "write", @@ -35,6 +35,88 @@ impl IoKind { } } +/// Live UI event while the process-wide object-store gate is throttling. +#[derive(Debug, Clone)] +pub enum ObjectStoreThrottleEvent { + Enter { + kind: IoKind, + /// Why the wait happened: `throttle` (AIMD sleep) or `admit` (semaphore). + reason: &'static str, + wait_ms: u64, + delay_ms: u64, + failures: u64, + }, + /// Cooldown tick / backoff / recovery — UI should refresh AIMD fields. + Update { + kind: IoKind, + /// `throttle` | `admit` | `backoff` | `recover` | `ok` + reason: &'static str, + wait_ms: u64, + delay_ms: u64, + failures: u64, + }, + Leave { + kind: IoKind, + }, +} + +/// Point-in-time gate status for progress painting between waits. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct ObjectStoreGateSnapshot { + pub kind: IoKind, + /// Current AIMD delay applied before the next remote acquire (0 = healthy). + pub delay_ms: u64, + pub cooldown_remaining_ms: u64, + pub failures: u64, + /// Successes toward the next multiplicative decay (`/` [`SUCCESS_STREAK_TO_DECAY`]). + pub success_streak: u32, + pub success_streak_target: u32, + pub active_waiters: u32, + /// Semaphore slots still free / configured remote concurrency. + pub available_permits: usize, + pub max_permits: usize, +} + +type ThrottleHook = Arc; + +fn throttle_hook_slot() -> &'static Mutex> { + static SLOT: OnceLock>> = OnceLock::new(); + SLOT.get_or_init(|| Mutex::new(None)) +} + +/// Restores the previous throttle UI hook when dropped. +pub struct ObjectStoreThrottleHookGuard { + previous: Option, +} + +impl Drop for ObjectStoreThrottleHookGuard { + fn drop(&mut self) { + if let Ok(mut slot) = throttle_hook_slot().lock() { + *slot = self.previous.take(); + } + } +} + +/// Install a process-wide S3/object-store throttle listener for the current scope. +pub fn install_throttle_hook( + hook: Arc, +) -> ObjectStoreThrottleHookGuard { + let previous = match throttle_hook_slot().lock() { + Ok(mut slot) => slot.replace(hook), + Err(_) => None, + }; + ObjectStoreThrottleHookGuard { previous } +} + +fn emit_throttle(event: ObjectStoreThrottleEvent) { + let Ok(slot) = throttle_hook_slot().lock() else { + return; + }; + if let Some(hook) = slot.as_ref() { + hook(event); + } +} + #[derive(Debug)] struct AimdState { /// Extra sleep applied before each remote acquire while degraded. @@ -45,6 +127,8 @@ struct AimdState { failures: u64, /// Last classified op that hit the gate (for progress UI). last_kind: IoKind, + /// Nested enter/leave count for active throttle waits. + active_waiters: u32, } impl Default for AimdState { @@ -55,12 +139,14 @@ impl Default for AimdState { successes_since_backoff: 0, failures: 0, last_kind: IoKind::Read, + active_waiters: 0, } } } struct Gate { semaphore: Arc, + concurrency: usize, state: Mutex, } @@ -74,6 +160,7 @@ fn gate() -> &'static Gate { .clamp(1, MAX_REMOTE_CONCURRENCY); Gate { semaphore: Arc::new(Semaphore::new(concurrency)), + concurrency, state: Mutex::new(AimdState::default()), } }) @@ -87,6 +174,73 @@ pub(crate) fn is_remote_uri(uri: &str) -> bool { !matches!(scheme, "file" | "file+uring" | "memory" | "shared-memory") } +/// Snapshot AIMD / cooldown state for progress UI. +pub fn snapshot() -> ObjectStoreGateSnapshot { + let g = gate(); + let available_permits = g.semaphore.available_permits(); + let max_permits = g.concurrency; + let Ok(state) = g.state.lock() else { + return ObjectStoreGateSnapshot { + kind: IoKind::Read, + delay_ms: 0, + cooldown_remaining_ms: 0, + failures: 0, + success_streak: 0, + success_streak_target: SUCCESS_STREAK_TO_DECAY, + active_waiters: 0, + available_permits, + max_permits, + }; + }; + let cooldown_remaining_ms = state + .cooldown_until + .and_then(|until| until.checked_duration_since(Instant::now())) + .map(|d| d.as_millis() as u64) + .unwrap_or(0); + ObjectStoreGateSnapshot { + kind: state.last_kind, + delay_ms: state.delay_ms, + cooldown_remaining_ms, + failures: state.failures, + success_streak: state.successes_since_backoff, + success_streak_target: SUCCESS_STREAK_TO_DECAY, + active_waiters: state.active_waiters, + available_permits, + max_permits, + } +} + +/// Compact AIMD label for progress brackets. All AIMD fields follow `aimd`. +pub fn format_aimd_flow_label(snap: &ObjectStoreGateSnapshot, event: Option<&str>) -> String { + let permits = format!("p={}/{}", snap.available_permits, snap.max_permits); + let streak = format!("s={}/{}", snap.success_streak, snap.success_streak_target); + let event = event + .map(str::trim) + .filter(|value| !value.is_empty()) + .unwrap_or(""); + let head = if event.is_empty() { + "aimd".to_owned() + } else { + format!("aimd {event}") + }; + if snap.cooldown_remaining_ms > 0 { + return format!( + "{head} cd={:.1}s d={}ms f={} {streak} w={} {permits}", + snap.cooldown_remaining_ms as f32 / 1000.0, + snap.delay_ms, + snap.failures, + snap.active_waiters, + ); + } + if snap.delay_ms > 0 || snap.failures > 0 || snap.active_waiters > 0 || !event.is_empty() { + return format!( + "{head} d={}ms f={} {streak} w={} {permits}", + snap.delay_ms, snap.failures, snap.active_waiters, + ); + } + format!("{head} ok {streak} {permits}") +} + pub(crate) struct Permit { _permit: Option, } @@ -100,18 +254,69 @@ pub(crate) async fn acquire(uri: &str, kind: IoKind) -> Permit { state.last_kind = kind; } wait_out_degradation(kind).await; - let permit = gate() - .semaphore - .clone() - .acquire_owned() - .await - .expect("object-store I/O semaphore is never closed"); + let permit = match gate().semaphore.clone().try_acquire_owned() { + Ok(permit) => permit, + Err(_) => { + enter_wait(kind, "admit", 0); + let permit = match gate().semaphore.clone().acquire_owned().await { + Ok(permit) => permit, + Err(error) => { + leave_wait(kind); + tracing::error!(?error, "object-store I/O semaphore closed unexpectedly"); + return Permit { _permit: None }; + } + }; + leave_wait(kind); + permit + } + }; wait_out_degradation(kind).await; Permit { _permit: Some(permit), } } +fn enter_wait(kind: IoKind, reason: &'static str, wait_ms: u64) { + let (delay_ms, failures) = { + let Ok(mut state) = gate().state.lock() else { + return; + }; + state.last_kind = kind; + state.active_waiters = state.active_waiters.saturating_add(1); + (state.delay_ms, state.failures) + }; + emit_throttle(ObjectStoreThrottleEvent::Enter { + kind, + reason, + wait_ms, + delay_ms, + failures, + }); +} + +fn leave_wait(kind: IoKind) { + if let Ok(mut state) = gate().state.lock() { + state.active_waiters = state.active_waiters.saturating_sub(1); + } + emit_throttle(ObjectStoreThrottleEvent::Leave { kind }); +} + +fn emit_update(kind: IoKind, reason: &'static str, wait_ms: u64) { + let (delay_ms, failures) = { + let Ok(state) = gate().state.lock() else { + return; + }; + (state.delay_ms, state.failures) + }; + emit_throttle(ObjectStoreThrottleEvent::Update { + kind, + reason, + wait_ms, + delay_ms, + failures, + }); +} + async fn wait_out_degradation(kind: IoKind) { let (sleep_for, delay_ms, failures) = { let Ok(state) = gate().state.lock() else { @@ -126,6 +331,8 @@ async fn wait_out_degradation(kind: IoKind) { if sleep_for.is_zero() { return; } + let wait_ms = sleep_for.as_millis() as u64; + enter_wait(kind, "throttle", wait_ms); crate::store::index_build_progress::note(format!( "s3 {} throttle wait {:.1}s (failures={failures}, delay={delay_ms}ms)", kind.as_str(), @@ -134,12 +341,24 @@ async fn wait_out_degradation(kind: IoKind) { tracing::warn!( target: "pchronicle.object_store_gate", kind = kind.as_str(), - wait_ms = sleep_for.as_millis() as u64, + wait_ms, delay_ms, failures, "object-store I/O gate cooling down before next remote op" ); - tokio::time::sleep(sleep_for).await; + // Tick the progress UI while cooling down so `cd=` counts down live. + let deadline = Instant::now() + sleep_for; + const TICK: Duration = Duration::from_millis(250); + loop { + let now = Instant::now(); + if now >= deadline { + break; + } + let remaining = deadline - now; + emit_update(kind, "throttle", remaining.as_millis() as u64); + tokio::time::sleep(remaining.min(TICK)).await; + } + leave_wait(kind); } /// Publish the current I/O phase for progress UI without taking a permit. @@ -155,26 +374,45 @@ pub(crate) fn note_success(uri: &str) { if !is_remote_uri(uri) { return; } - let Ok(mut state) = gate().state.lock() else { - return; + let (kind, changed, delay_ms, failures) = { + let Ok(mut state) = gate().state.lock() else { + return; + }; + let kind = state.last_kind; + state.successes_since_backoff = state.successes_since_backoff.saturating_add(1); + if state.delay_ms == 0 { + return; + } + let mut changed = false; + if state.successes_since_backoff >= SUCCESS_STREAK_TO_DECAY { + state.delay_ms /= 2; + state.successes_since_backoff = 0; + if state.delay_ms < 100 { + state.delay_ms = 0; + state.cooldown_until = None; + } + changed = true; + tracing::info!( + target: "pchronicle.object_store_gate", + delay_ms = state.delay_ms, + "object-store I/O gate recovered toward steady state" + ); + } + (kind, changed, state.delay_ms, state.failures) }; - state.successes_since_backoff = state.successes_since_backoff.saturating_add(1); - if state.delay_ms == 0 { + // Always publish streak / delay movement so the progress line can refresh. + if !changed && delay_ms == 0 { + // Healthy path: skip per-op UI spam; paints from commit/fetch cover s=. return; } - if state.successes_since_backoff >= SUCCESS_STREAK_TO_DECAY { - state.delay_ms /= 2; - state.successes_since_backoff = 0; - if state.delay_ms < 100 { - state.delay_ms = 0; - state.cooldown_until = None; - } - tracing::info!( - target: "pchronicle.object_store_gate", - delay_ms = state.delay_ms, - "object-store I/O gate recovered toward steady state" - ); - } + let reason = if delay_ms == 0 { "recover" } else { "ok" }; + emit_throttle(ObjectStoreThrottleEvent::Update { + kind, + reason, + wait_ms: 0, + delay_ms, + failures, + }); } /// Record a transient remote failure: grow shared delay and set a cooldown. @@ -182,39 +420,40 @@ pub(crate) fn note_failure(uri: &str, kind: IoKind) { if !is_remote_uri(uri) { return; } - let Ok(mut state) = gate().state.lock() else { - return; - }; - state.last_kind = kind; - state.failures = state.failures.saturating_add(1); - state.successes_since_backoff = 0; - state.delay_ms = if state.delay_ms == 0 { - 500 - } else { - state.delay_ms.saturating_mul(2).min(MAX_DELAY_MS) + let (delay_ms, failures, wait_ms) = { + let Ok(mut state) = gate().state.lock() else { + return; + }; + state.last_kind = kind; + state.failures = state.failures.saturating_add(1); + state.successes_since_backoff = 0; + state.delay_ms = if state.delay_ms == 0 { + 500 + } else { + state.delay_ms.saturating_mul(2).min(MAX_DELAY_MS) + }; + state.cooldown_until = Some(Instant::now() + Duration::from_millis(state.delay_ms)); + tracing::warn!( + target: "pchronicle.object_store_gate", + kind = kind.as_str(), + delay_ms = state.delay_ms, + failures = state.failures, + "object-store I/O gate backing off after transient failure" + ); + crate::store::index_build_progress::note(format!( + "s3 {} throttle backoff {}ms", + kind.as_str(), + state.delay_ms + )); + (state.delay_ms, state.failures, state.delay_ms) }; - state.cooldown_until = Some(Instant::now() + Duration::from_millis(state.delay_ms)); - tracing::warn!( - target: "pchronicle.object_store_gate", - kind = kind.as_str(), - delay_ms = state.delay_ms, - failures = state.failures, - "object-store I/O gate backing off after transient failure" - ); - crate::store::index_build_progress::note(format!( - "s3 {} throttle backoff {}ms", - kind.as_str(), - state.delay_ms - )); -} - -#[cfg(test)] -pub(crate) fn debug_delay_ms() -> u64 { - gate() - .state - .lock() - .map(|state| state.delay_ms) - .unwrap_or(0) + emit_throttle(ObjectStoreThrottleEvent::Update { + kind, + reason: "backoff", + wait_ms, + delay_ms, + failures, + }); } #[cfg(test)] @@ -222,11 +461,38 @@ mod tests { use super::*; #[test] - fn classifies_remote_uris() { - assert!(is_remote_uri("s3://bucket/prefix")); - assert!(is_remote_uri("gs://bucket/prefix")); - assert!(!is_remote_uri("/tmp/local")); - assert!(!is_remote_uri("file:///tmp/local")); - assert!(!is_remote_uri("shared-memory://x")); + fn aimd_flow_label_healthy_and_degraded() { + let healthy = ObjectStoreGateSnapshot { + kind: IoKind::Write, + delay_ms: 0, + cooldown_remaining_ms: 0, + failures: 0, + success_streak: 2, + success_streak_target: SUCCESS_STREAK_TO_DECAY, + active_waiters: 0, + available_permits: 1, + max_permits: 1, + }; + assert_eq!( + format_aimd_flow_label(&healthy, None), + "aimd ok s=2/4 p=1/1" + ); + + let cooling = ObjectStoreGateSnapshot { + delay_ms: 2000, + cooldown_remaining_ms: 1500, + failures: 3, + success_streak: 0, + active_waiters: 1, + available_permits: 0, + ..healthy + }; + let label = format_aimd_flow_label(&cooling, Some("throttle")); + assert!(label.starts_with("aimd throttle "), "{label}"); + assert!(label.contains("cd=1.5s"), "{label}"); + assert!(label.contains("d=2000ms"), "{label}"); + assert!(label.contains("f=3"), "{label}"); + assert!(label.contains("w=1"), "{label}"); + assert!(label.contains("p=0/1"), "{label}"); } } diff --git a/crates/persisting-pchronicle/src/store/storyline/content.rs b/crates/persisting-pchronicle/src/store/storyline/content.rs index 00b33cf0..ff3b4866 100644 --- a/crates/persisting-pchronicle/src/store/storyline/content.rs +++ b/crates/persisting-pchronicle/src/store/storyline/content.rs @@ -34,6 +34,9 @@ use crate::formats::unknown_fields::{ pub const STORYLINE_OBJECTS_DATASET: &str = "objects.lance"; pub const DEFAULT_CONTENT_OFFLOAD_THRESHOLD: usize = 64 * 1024; pub const DEFAULT_CONTENT_PREVIEW_BYTES: usize = 256; +/// Soft ceiling for one stream write chunk. Keeps Arrow UTF8/Binary builders +/// under the ~2GiB i32 offset limit when many medium-sized cells accumulate. +pub const DEFAULT_MAX_CHUNK_BYTES: usize = 256 * 1024 * 1024; pub(crate) const CONTENT_REF_MAGIC: &str = "\u{001e}PCHRONICLE-CONTENT:"; const CONTENT_INDEX_NAME: &str = "pchronicle_content_id_idx"; const CONTENT_ID_COLUMN: &str = "content_id"; @@ -93,7 +96,7 @@ impl Default for StorylineContentOptions { max_document_rows: None, max_document_bytes: None, max_chunk_rows: None, - max_chunk_bytes: None, + max_chunk_bytes: Some(DEFAULT_MAX_CHUNK_BYTES), max_import_documents: None, max_unknown_fields: DEFAULT_MAX_UNKNOWN_FIELDS, max_unknown_bytes: DEFAULT_MAX_UNKNOWN_BYTES, @@ -422,6 +425,13 @@ fn externalize_batch( continue; } let value = values.value(row); + // Already-published content refs must not be wrapped again. User + // payloads that only look like the magic prefix still offload. + let already_ref = matches!(ContentRef::parse(value), Ok(Some(_))); + if already_ref { + encoded.push(Some(value.to_string())); + continue; + } let should_offload = value.len() >= options.offload_threshold || value.starts_with(CONTENT_REF_MAGIC); if !should_offload { @@ -526,6 +536,51 @@ fn build_object( }) } +/// Encode a JSON content cell, offloading to `objects.lance` before Arrow Utf8 +/// materialization so large import batches cannot hit the 2GiB StringArray limit. +pub(crate) fn encode_json_content_cell( + value: &T, + options: StorylineContentOptions, + pending: &mut PendingContent, +) -> Result { + let encoded = serde_json::to_vec(value).context("serialize Storyline content JSON cell")?; + let collides = match serde_json::from_slice::(&encoded) { + Ok(serde_json::Value::String(text)) => text.starts_with(CONTENT_REF_MAGIC), + _ => false, + }; + if encoded.len() < options.offload_threshold && !collides { + return String::from_utf8(encoded).context("Storyline JSON cell is not UTF-8"); + } + if let Ok(serde_json::Value::String(text)) = + serde_json::from_slice::(&encoded) + && matches!(ContentRef::parse(&text), Ok(Some(_))) + { + return Ok(text); + } + let object = build_object(&encoded, LogicalType::Json, options)?; + let descriptor = object.reference.encode(); + pending.insert(object)?; + Ok(descriptor) +} + +/// Encode a UTF-8 content cell with the same pre-Arrow offload policy. +pub(crate) fn encode_utf8_content_cell( + value: &str, + options: StorylineContentOptions, + pending: &mut PendingContent, +) -> Result { + if matches!(ContentRef::parse(value), Ok(Some(_))) { + return Ok(value.to_owned()); + } + if value.len() < options.offload_threshold && !value.starts_with(CONTENT_REF_MAGIC) { + return Ok(value.to_owned()); + } + let object = build_object(value.as_bytes(), LogicalType::Utf8, options)?; + let descriptor = object.reference.encode(); + pending.insert(object)?; + Ok(descriptor) +} + fn utf8_preview(bytes: &[u8], maximum: usize) -> Result { let value = std::str::from_utf8(bytes).context("UTF-8 content column contains invalid bytes")?; diff --git a/crates/persisting-pchronicle/src/store/storyline/mod.rs b/crates/persisting-pchronicle/src/store/storyline/mod.rs index 6244afcd..85ec57fd 100644 --- a/crates/persisting-pchronicle/src/store/storyline/mod.rs +++ b/crates/persisting-pchronicle/src/store/storyline/mod.rs @@ -28,7 +28,8 @@ use mutation::{ }; pub use content::{ - DEFAULT_CONTENT_OFFLOAD_THRESHOLD, DEFAULT_CONTENT_PREVIEW_BYTES, StorylineContentOptions, + DEFAULT_CONTENT_OFFLOAD_THRESHOLD, DEFAULT_CONTENT_PREVIEW_BYTES, DEFAULT_MAX_CHUNK_BYTES, + StorylineContentOptions, }; pub use datafusion::{ DATAFUSION_RUNS_TABLE, DATAFUSION_STEPS_TABLE, DATAFUSION_TOOL_CALLS_TABLE, @@ -561,11 +562,28 @@ impl StorylineLanceStore { &self.root } - /// The exact local path or object-store URI used for Lance datasets. + /// Exact local path or object-store URI used for Lance datasets. pub fn root_uri(&self) -> &str { &self.root_uri } + /// Sum of object/file sizes currently under this Dataset root. + /// + /// This is physical on-disk (or object-store) size, not attributed input + /// bytes. Listing large prefixes can be slow; call at import completion. + pub async fn on_disk_bytes(&self) -> Result { + let objects = self + .control_store + .list("") + .await + .with_context(|| format!("list Storyline Dataset objects at {}", self.root_uri))?; + let mut total = 0u64; + for object in objects { + total = total.saturating_add(object.metadata.content_length()); + } + Ok(total) + } + pub fn storage_scheme(&self) -> &str { self.root_uri .split_once("://") @@ -1013,32 +1031,31 @@ impl StorylineLanceStore { #[cfg(test)] release_waiting_content_create(&self.root_uri, first_content_create); let objects_version = objects_result?; - let (runs_version, steps_version, tool_calls_version) = - join3_remote_aware( - self.is_remote_object_store(), - write_batches( - &created.runs, - run_batches, - story_runs_arrow_schema(), - &RUN_INDEXES, - stream_options.optimize_indices, - ), - write_batches( - &created.steps, - step_batches, - story_steps_arrow_schema(), - &STEP_INDEXES, - stream_options.optimize_indices, - ), - write_batches( - &created.tool_calls, - tool_call_batches, - story_tool_calls_arrow_schema(), - &TOOL_CALL_INDEXES, - stream_options.optimize_indices, - ), - ) - .await?; + let (runs_version, steps_version, tool_calls_version) = join3_remote_aware( + self.is_remote_object_store(), + write_batches( + &created.runs, + run_batches, + story_runs_arrow_schema(), + &RUN_INDEXES, + stream_options.optimize_indices, + ), + write_batches( + &created.steps, + step_batches, + story_steps_arrow_schema(), + &STEP_INDEXES, + stream_options.optimize_indices, + ), + write_batches( + &created.tool_calls, + tool_call_batches, + story_tool_calls_arrow_schema(), + &TOOL_CALL_INDEXES, + stream_options.optimize_indices, + ), + ) + .await?; created.runs_version = runs_version; created.steps_version = steps_version; created.tool_calls_version = tool_calls_version; @@ -1056,35 +1073,34 @@ impl StorylineLanceStore { stream_options.optimize_indices, ) .await?; - let (runs_version, steps_version, tool_calls_version) = - join3_remote_aware( - self.is_remote_object_store(), - replace_table_batches( - ¤t.runs, - current.runs_version, - &predicate, - &["document_id"], - run_batches, - story_runs_arrow_schema(), - ), - replace_table_batches( - ¤t.steps, - current.steps_version, - &predicate, - &["document_id", "step_id"], - step_batches, - story_steps_arrow_schema(), - ), - replace_table_batches( - ¤t.tool_calls, - current.tool_calls_version, - &predicate, - &["document_id", "step_id", "call_index"], - tool_call_batches, - story_tool_calls_arrow_schema(), - ), - ) - .await?; + let (runs_version, steps_version, tool_calls_version) = join3_remote_aware( + self.is_remote_object_store(), + replace_table_batches( + ¤t.runs, + current.runs_version, + &predicate, + &["document_id"], + run_batches, + story_runs_arrow_schema(), + ), + replace_table_batches( + ¤t.steps, + current.steps_version, + &predicate, + &["document_id", "step_id"], + step_batches, + story_steps_arrow_schema(), + ), + replace_table_batches( + ¤t.tool_calls, + current.tool_calls_version, + &predicate, + &["document_id", "step_id", "call_index"], + tool_call_batches, + story_tool_calls_arrow_schema(), + ), + ) + .await?; current.runs_version = runs_version; current.steps_version = steps_version; current.tool_calls_version = tool_calls_version; diff --git a/crates/persisting-pchronicle/src/store/storyline/mutation.rs b/crates/persisting-pchronicle/src/store/storyline/mutation.rs index 199d1fa2..6d9f3cfc 100644 --- a/crates/persisting-pchronicle/src/store/storyline/mutation.rs +++ b/crates/persisting-pchronicle/src/store/storyline/mutation.rs @@ -145,15 +145,18 @@ fn serialized_document_bytes(story: &StorylineDocument) -> Result { Ok(writer.0) } -struct EncodedBatchIterator { +struct EncodedBatchIterator { rows: std::sync::Arc<[T]>, offset: usize, emitted_empty: bool, - encode: fn(&[T]) -> Result, + encode: F, } -impl EncodedBatchIterator { - fn new(rows: Vec, encode: fn(&[T]) -> Result) -> Self { +impl EncodedBatchIterator +where + F: FnMut(&[T]) -> Result, +{ + fn new(rows: Vec, encode: F) -> Self { Self { rows: rows.into(), offset: 0, @@ -163,7 +166,10 @@ impl EncodedBatchIterator { } } -impl Iterator for EncodedBatchIterator { +impl Iterator for EncodedBatchIterator +where + F: FnMut(&[T]) -> Result, +{ type Item = std::result::Result; fn next(&mut self) -> Option { @@ -172,25 +178,48 @@ impl Iterator for EncodedBatchIterator { return None; } self.emitted_empty = true; - return Some( - (self.encode)(&[]).map_err(|error| ArrowError::ComputeError(error.to_string())), - ); + return Some(catch_encode_panic(|| (self.encode)(&[]))); } if self.offset >= self.rows.len() { return None; } let end = (self.offset + WRITE_BATCH_ROWS).min(self.rows.len()); - let result = (self.encode)(&self.rows[self.offset..end]) - .map_err(|error| ArrowError::ComputeError(error.to_string())); + let slice = &self.rows[self.offset..end]; + let result = catch_encode_panic(|| (self.encode)(slice)); self.offset = end; Some(result) } } -fn encode_rows( - rows: Vec, - encode: fn(&[T]) -> Result, -) -> Result> { +fn catch_encode_panic( + encode: impl FnOnce() -> Result, +) -> std::result::Result { + match std::panic::catch_unwind(std::panic::AssertUnwindSafe(encode)) { + Ok(Ok(batch)) => Ok(batch), + Ok(Err(error)) => Err(ArrowError::ComputeError(error.to_string())), + Err(panic) => { + let message = panic_message(&panic); + Err(ArrowError::ComputeError(format!( + "Storyline Arrow encode panicked ({message}); reduce commit batch size or skip oversized sources" + ))) + } + } +} + +fn panic_message(panic: &Box) -> String { + if let Some(message) = panic.downcast_ref::<&str>() { + (*message).to_string() + } else if let Some(message) = panic.downcast_ref::() { + message.clone() + } else { + "unknown panic".into() + } +} + +fn encode_rows(rows: Vec, encode: F) -> Result> +where + F: FnMut(&[T]) -> Result, +{ EncodedBatchIterator::new(rows, encode) .map(|batch| batch.map_err(anyhow::Error::from)) .collect() @@ -213,20 +242,31 @@ pub(super) fn externalize_rows( for run in &mut runs { externalize_unknown_field_values(&mut run.unknown_fields, options, &mut pending)?; } + // Offload large Utf8 content cells while encoding so Arrow StringArray + // construction never sees multi-GiB payloads (i32 offset overflow). let runs = externalize_batches( - encode_rows(runs, story_runs_to_batch)?, + encode_rows(runs, |chunk| { + super::rows::story_runs_to_batch_with_content(chunk, Some((options, &mut pending))) + })?, StorylineTableKind::Runs, options, &mut pending, )?; let steps = externalize_batches( - encode_rows(steps, story_steps_to_batch)?, + encode_rows(steps, |chunk| { + super::rows::story_steps_to_batch_with_content(chunk, Some((options, &mut pending))) + })?, StorylineTableKind::Steps, options, &mut pending, )?; let tool_calls = externalize_batches( - encode_rows(tool_calls, story_tool_calls_to_batch)?, + encode_rows(tool_calls, |chunk| { + super::rows::story_tool_calls_to_batch_with_content( + chunk, + Some((options, &mut pending)), + ) + })?, StorylineTableKind::ToolCalls, options, &mut pending, @@ -318,7 +358,9 @@ async fn write_record_batch_reader( build_indexes: bool, ) -> Result { let uri = path.to_string_lossy().into_owned(); - crate::store::object_store_io_gate::mark_kind(crate::store::object_store_io_gate::IoKind::Write); + crate::store::object_store_io_gate::mark_kind( + crate::store::object_store_io_gate::IoKind::Write, + ); let mut dataset = InsertBuilder::new(&uri) .with_params(&WriteParams { mode: WriteMode::Create, diff --git a/crates/persisting-pchronicle/src/store/storyline/rows.rs b/crates/persisting-pchronicle/src/store/storyline/rows.rs index 0663e175..4f83ebb0 100644 --- a/crates/persisting-pchronicle/src/store/storyline/rows.rs +++ b/crates/persisting-pchronicle/src/store/storyline/rows.rs @@ -266,14 +266,62 @@ fn json_array_owned(values: Vec>) -> Result { .context("encode Lance JSON column") } +/// Optional pre-Arrow content offload into `objects.lance`. +pub(crate) type ContentEncode<'a> = Option<( + super::content::StorylineContentOptions, + &'a mut super::content::PendingContent, +)>; + +fn json_content(value: &T, content: &mut ContentEncode<'_>) -> Result { + match content { + Some((options, pending)) => { + super::content::encode_json_content_cell(value, *options, pending) + } + None => json(value), + } +} + +fn opt_json_content( + value: &Option, + content: &mut ContentEncode<'_>, +) -> Result> { + value + .as_ref() + .map(|value| json_content(value, content)) + .transpose() +} + +fn utf8_content(value: &str, content: &mut ContentEncode<'_>) -> Result { + match content { + Some((options, pending)) => { + super::content::encode_utf8_content_cell(value, *options, pending) + } + None => Ok(value.to_owned()), + } +} + +fn opt_utf8_content( + value: Option<&str>, + content: &mut ContentEncode<'_>, +) -> Result> { + value.map(|value| utf8_content(value, content)).transpose() +} + pub fn story_runs_to_batch(rows: &[StoryRunRow]) -> Result { + story_runs_to_batch_with_content(rows, None) +} + +pub(crate) fn story_runs_to_batch_with_content( + rows: &[StoryRunRow], + mut content: ContentEncode<'_>, +) -> Result { RecordBatch::try_new( story_runs_arrow_schema(), vec![ Arc::new(req_utf8(rows.iter().map(|r| r.schema_version.as_str()))), Arc::new(opt_utf8_owned( rows.iter() - .map(|r| opt_json(&r.origin)) + .map(|r| opt_json_content(&r.origin, &mut content)) .collect::>>()?, )), Arc::new(req_utf8(rows.iter().map(|r| r.document_id.as_str()))), @@ -294,7 +342,7 @@ pub fn story_runs_to_batch(rows: &[StoryRunRow]) -> Result { Arc::new(opt_utf8(rows.iter().map(|r| r.agent_model_name.as_deref()))), Arc::new(opt_utf8_owned( rows.iter() - .map(|r| opt_json(&r.agent_tool_definitions)) + .map(|r| opt_json_content(&r.agent_tool_definitions, &mut content)) .collect::>>()?, )), Arc::new(json_array_owned( @@ -304,22 +352,28 @@ pub fn story_runs_to_batch(rows: &[StoryRunRow]) -> Result { )?), Arc::new(opt_utf8_owned( rows.iter() - .map(|r| opt_json(&r.parent)) + .map(|r| opt_json_content(&r.parent, &mut content)) .collect::>>()?, )), Arc::new(opt_utf8_owned( rows.iter() - .map(|r| opt_json(&r.child_session_ids)) + .map(|r| opt_json_content(&r.child_session_ids, &mut content)) + .collect::>>()?, + )), + Arc::new(opt_utf8_owned( + rows.iter() + .map(|r| opt_utf8_content(r.notes.as_deref(), &mut content)) .collect::>>()?, )), - Arc::new(opt_utf8(rows.iter().map(|r| r.notes.as_deref()))), Arc::new(json_array_owned( rows.iter() .map(|r| opt_json(&r.final_metrics)) .collect::>>()?, )?), - Arc::new(opt_utf8( - rows.iter().map(|r| r.continued_trajectory_ref.as_deref()), + Arc::new(opt_utf8_owned( + rows.iter() + .map(|r| opt_utf8_content(r.continued_trajectory_ref.as_deref(), &mut content)) + .collect::>>()?, )), Arc::new(json_array_owned( rows.iter() @@ -343,22 +397,24 @@ pub fn story_runs_to_batch(rows: &[StoryRunRow]) -> Result { Arc::new(opt_utf8_owned( rows.iter() .map(|r| { - (!r.unknown_key_counts.is_empty()) - .then(|| json(&r.unknown_key_counts)) - .transpose() + if r.unknown_key_counts.is_empty() { + Ok(None) + } else { + Ok(Some(json_content(&r.unknown_key_counts, &mut content)?)) + } }) .collect::>>()?, )), Arc::new(opt_utf8_owned( rows.iter() - .map(|r| opt_json(&r.task)) + .map(|r| opt_json_content(&r.task, &mut content)) .collect::>>()?, )), Arc::new(timestamp_array(rows.iter().map(|r| r.started_at.as_ref()))), Arc::new(timestamp_array(rows.iter().map(|r| r.finished_at.as_ref()))), Arc::new(opt_utf8_owned( rows.iter() - .map(|r| opt_json(&r.prompt)) + .map(|r| opt_json_content(&r.prompt, &mut content)) .collect::>>()?, )), ], @@ -367,6 +423,13 @@ pub fn story_runs_to_batch(rows: &[StoryRunRow]) -> Result { } pub fn story_steps_to_batch(rows: &[StoryStepRow]) -> Result { + story_steps_to_batch_with_content(rows, None) +} + +pub(crate) fn story_steps_to_batch_with_content( + rows: &[StoryStepRow], + mut content: ContentEncode<'_>, +) -> Result { RecordBatch::try_new( story_steps_arrow_schema(), vec![ @@ -389,11 +452,13 @@ pub fn story_steps_to_batch(rows: &[StoryStepRow]) -> Result { )), Arc::new(req_utf8_owned( rows.iter() - .map(|r| json(&r.message)) + .map(|r| json_content(&r.message, &mut content)) .collect::>()?, )), - Arc::new(opt_utf8( - rows.iter().map(|r| r.reasoning_content.as_deref()), + Arc::new(opt_utf8_owned( + rows.iter() + .map(|r| opt_utf8_content(r.reasoning_content.as_deref(), &mut content)) + .collect::>()?, )), Arc::new(opt_utf8(rows.iter().map(|r| { r.reasoning_effort @@ -402,7 +467,7 @@ pub fn story_steps_to_batch(rows: &[StoryStepRow]) -> Result { }))), Arc::new(opt_utf8_owned( rows.iter() - .map(|r| opt_json(&r.reasoning_effort)) + .map(|r| opt_json_content(&r.reasoning_effort, &mut content)) .collect::>()?, )), Arc::new(json_array_owned( @@ -431,7 +496,7 @@ pub fn story_steps_to_batch(rows: &[StoryStepRow]) -> Result { )), Arc::new(opt_utf8_owned( rows.iter() - .map(|r| opt_json(&r.observation)) + .map(|r| opt_json_content(&r.observation, &mut content)) .collect::>()?, )), Arc::new(json_array_owned( @@ -441,13 +506,13 @@ pub fn story_steps_to_batch(rows: &[StoryStepRow]) -> Result { )?), Arc::new(opt_utf8_owned( rows.iter() - .map(|r| opt_json(&r.env)) + .map(|r| opt_json_content(&r.env, &mut content)) .collect::>()?, )), Arc::new(timestamp_array(rows.iter().map(|r| r.finished_at.as_ref()))), Arc::new(opt_utf8_owned( rows.iter() - .map(|r| opt_json(&r.prompt)) + .map(|r| opt_json_content(&r.prompt, &mut content)) .collect::>()?, )), ], @@ -456,6 +521,13 @@ pub fn story_steps_to_batch(rows: &[StoryStepRow]) -> Result { } pub fn story_tool_calls_to_batch(rows: &[StoryToolCallRow]) -> Result { + story_tool_calls_to_batch_with_content(rows, None) +} + +pub(crate) fn story_tool_calls_to_batch_with_content( + rows: &[StoryToolCallRow], + mut content: ContentEncode<'_>, +) -> Result { RecordBatch::try_new( story_tool_calls_arrow_schema(), vec![ @@ -472,17 +544,17 @@ pub fn story_tool_calls_to_batch(rows: &[StoryToolCallRow]) -> Result>()?, )), Arc::new(opt_utf8_owned( rows.iter() - .map(|r| r.result.as_ref().map(json).transpose()) + .map(|r| opt_json_content(&r.result, &mut content)) .collect::>>()?, )), Arc::new(req_utf8_owned( rows.iter() - .map(|r| json(&r.results)) + .map(|r| json_content(&r.results, &mut content)) .collect::>()?, )), Arc::new(Int64Array::from( diff --git a/crates/persisting-pchronicle/src/store/storyline/tests.rs b/crates/persisting-pchronicle/src/store/storyline/tests.rs index 5b472e24..ee747852 100644 --- a/crates/persisting-pchronicle/src/store/storyline/tests.rs +++ b/crates/persisting-pchronicle/src/store/storyline/tests.rs @@ -344,6 +344,8 @@ async fn repeated_unknown_value_is_stored_once() { .await .unwrap(); assert_eq!(objects.count_rows(None).await.unwrap(), 1); + let on_disk = store.on_disk_bytes().await.unwrap(); + assert!(on_disk > 0, "committed Storyline Dataset should occupy disk"); let hydrated = store .get_storyline_full("unknown-first") .await diff --git a/crates/persisting-pchronicle/src/store/storyline/writer_control.rs b/crates/persisting-pchronicle/src/store/storyline/writer_control.rs index 4a5ec50a..600f223f 100644 --- a/crates/persisting-pchronicle/src/store/storyline/writer_control.rs +++ b/crates/persisting-pchronicle/src/store/storyline/writer_control.rs @@ -390,9 +390,12 @@ impl StorylineLanceStore { .current_if_match_unreliable .load(std::sync::atomic::Ordering::Relaxed); if !skip_if_match { + let expected = expected + .as_ref() + .context("missing expected version for conditional Storyline CURRENT write")?; match self .control_store - .write_match(CURRENT_FILE, contents.clone(), expected.as_ref().unwrap()) + .write_match(CURRENT_FILE, contents.clone(), expected) .await { Ok(()) => return Ok(true), @@ -415,10 +418,9 @@ impl StorylineLanceStore { return Ok(false); } // Remember for this store handle: avoid 412 spam on every commit. - let first = !self.current_if_match_unreliable.swap( - true, - std::sync::atomic::Ordering::Relaxed, - ); + let first = !self + .current_if_match_unreliable + .swap(true, std::sync::atomic::Ordering::Relaxed); if first { tracing::warn!( root_uri = %self.root_uri, @@ -470,11 +472,7 @@ impl StorylineLanceStore { }; let expected_version = current.version.clone(); let wrote = self - .try_write_current_control( - &next, - expected_version, - Some(¤t.control), - ) + .try_write_current_control(&next, expected_version, Some(¤t.control)) .await?; if wrote { return Ok(outcome); @@ -573,8 +571,9 @@ impl StorylineLanceStore { .await { Ok(true) => Err(conflict), - Ok(false) => Err(conflict - .context("mismatched writer lease was lost before release")), + Ok(false) => { + Err(conflict.context("mismatched writer lease was lost before release")) + } Err(error) => Err(conflict.context(format!( "failed to release mismatched writer lease: {error:#}" ))), diff --git a/docs/src/en/pchronicle/reference/cli.md b/docs/src/en/pchronicle/reference/cli.md index 464c7ac9..10ab3670 100644 --- a/docs/src/en/pchronicle/reference/cli.md +++ b/docs/src/en/pchronicle/reference/cli.md @@ -186,6 +186,7 @@ pchronicle import -f|--from SOURCE -t|--to NEW_DATASET [-i|--input-format auto|atif|actf|openai-messages|storyline|codex|claude-code|compact-jsonl] [-o|--output-format preserve|storyline|compact-jsonl] [|--replace] [--append] [--on-duplicate suffix|skip] [--yes] + [--resume] [--wal-dir DIR] [--reset] [--column NAME=JSON_PATH]... [OPTIONS] ``` @@ -194,6 +195,7 @@ pchronicle import -f input.json -t ./imported -i atif cat input.json | pchronicle import -f - -t ./imported -i openai-messages pchronicle import -f more.json -t ./normalized --append --on-duplicate skip pchronicle import -f rebuilt.json -t ./normalized --replace --yes +pchronicle import -f s3://bucket/corpus -t s3://bucket/out -o storyline --resume pchronicle import -f ./jsonl-root -t ./records.lance \ -o compact-jsonl \ --column id=$.event.id --column timestamp=$.event.time \ @@ -208,6 +210,13 @@ transaction, and only then removes the old data. It requires interactive confirmation or `--yes`; an existing object-store Dataset cannot currently be replaced in place. +Long Storyline imports write a local checkpoint WAL under +`./.pchronicle-import-wal//` (`job.json`, `done.jsonl`, `failed.jsonl`). +Use `--resume` with the same `--from`/`--to` fingerprint to skip sources already +recorded as done or failed; `--wal-dir` overrides the WAL root; `--reset` +deletes that job's WAL before starting. Decode and skippable commit failures are +recorded in the WAL and `import.log` so the job can continue. + Compact JSONL is a record store, not a trajectory conversion. Either `--input-format compact-jsonl` or `--output-format compact-jsonl` selects it. It recursively reads local `.json`, `.jsonl`, and `.ndjson` files. JSON objects @@ -223,27 +232,29 @@ local `create` and confirmed `replace`, but not stdin, object-store targets, or ### Sync ```text -pchronicle sync --from DIRECTORY --to DIRECTORY --convert DIRECTORY - [--input-format FORMAT] [--column NAME=JSON_PATH]... +pchronicle sync --from DIRECTORY + [--mirror DIRECTORY] [--to DIRECTORY] + [--input-format FORMAT] [--suggested-format FORMAT] + [--column NAME=JSON_PATH]... [--interval DURATION] [--once] ``` `sync` is a resident polling worker for `.json`, `.jsonl`, and `.ndjson` files. -For run-data formats it coalesces changes into a pending set and, on each -interval, atomically mirrors the source files byte-for-byte into a local -Warehouse Dataset and writes a Storyline Lance Dataset to `--convert`. -Pending changes are cleared only after both outputs succeed; failures retain -the set and retry with bounded exponential backoff. Use `--once` for one -initial batch and exit. The two destinations must be local directories outside -the source directory. - -With `--input-format compact-jsonl`, the source must be a local `.json`, `.jsonl`, -or `.ndjson` tree -and the same `--column` rules as compact import apply. Every successful batch -rescans the whole tree and atomically replaces the compact Lance snapshot at -`--convert`, so additions, changes, and deletions are reflected without -row-level incremental updates. In this mode `--to` is retained as a required -compatibility argument but is not written. +It coalesces changes into a pending set and, on each interval, rebuilds full +snapshots for the destinations you enable. Provide `--mirror`, `--to`, or both: + +- `--mirror` writes a Compact JSONL Lance Dataset (record-level ingest; optional + `--column` mapping). Each successful batch atomically replaces that target. +- `--to` converts trajectories into a Storyline Lance Dataset (`--input-format` / + `--suggested-format`). + +Pending changes clear only after every enabled destination succeeds; failures +retain the set and retry with bounded exponential backoff. Use `--once` for one +initial batch and exit. Local destinations must sit outside the source tree. + +With `--input-format compact-jsonl`, only `--mirror` is allowed (not `--to`). +Each successful batch rescans the tree and atomically replaces the compact Lance +snapshot at `--mirror`. ### Drop diff --git a/docs/src/en/rfcs/0014-compact-jsonl.md b/docs/src/en/rfcs/0014-compact-jsonl.md index 9925d063..31ba6513 100644 --- a/docs/src/en/rfcs/0014-compact-jsonl.md +++ b/docs/src/en/rfcs/0014-compact-jsonl.md @@ -176,16 +176,16 @@ pchronicle import \ ```text pchronicle sync \ --from ./jsonl-root \ - --to ./warehouse-copy \ - --convert ./records.lance \ + --mirror ./records.lance \ --input-format compact-jsonl \ --column id=$.event.id \ --column timestamp=$.event.time ``` -v1 sync 是 snapshot sync。每批变化 MUST 重新扫描完整 input root，并用一个完整的新 compact -snapshot 替换 `--convert`。创建、修改和删除源文件都必须反映到下一快照。v1 不承诺行级增量 -更新。 +v1 sync is snapshot sync. Each changed batch MUST rescan the full input root and +atomically replace `--mirror` with a complete compact snapshot. Adds, edits, and +deletes MUST appear in the next snapshot. v1 does not promise row-level +incremental updates. ## Export diff --git a/docs/src/zh/pchronicle/reference/cli.md b/docs/src/zh/pchronicle/reference/cli.md index 8f01003e..46b14ad8 100644 --- a/docs/src/zh/pchronicle/reference/cli.md +++ b/docs/src/zh/pchronicle/reference/cli.md @@ -71,7 +71,7 @@ pchronicle ├── find [DATASET] ├── query [DATASET] ├── import --from SOURCE --to DATASET -├── sync --from DIRECTORY --to DIRECTORY --convert DIRECTORY +├── sync --from DIRECTORY [--mirror DIRECTORY] [--to DIRECTORY] ├── export --from DATASET --to TARGET ├── agent codex|claude [DATASET] └── serve DATASET... @@ -272,6 +272,7 @@ pchronicle query \ pchronicle import -f|--from SOURCE -t|--to NEW_DATASET [-i|--input-format FORMAT] [-o|--output-format preserve|storyline|compact-jsonl] [--replace] [--append] [--on-duplicate suffix|skip] [--yes] + [--resume] [--wal-dir DIR] [--reset] [--column NAME=JSON_PATH]... [--max-input-bytes BYTES] ``` @@ -286,6 +287,8 @@ pchronicle import \ -f more.json -t ./normalized --append --on-duplicate skip pchronicle import \ -f rebuilt.json -t ./normalized --replace --yes +pchronicle import \ + -f s3://bucket/corpus -t s3://bucket/out -o storyline --resume pchronicle import \ -f ./jsonl-root -t ./records.lance \ -o compact-jsonl \ @@ -298,6 +301,11 @@ pchronicle import \ 必须显式指定 `-i`。`preserve` 保留文件边界和相对路径，`storyline` 合并为 normalized Store； 对象存储目标必须使用 `storyline`。 +长时间 Storyline import 会在 `./.pchronicle-import-wal//` 写入本地 checkpoint WAL +（`job.json`、`done.jsonl`、`failed.jsonl`）。同一 `--from`/`--to` 指纹下使用 `--resume` 可跳过 +已标记 done/failed 的源；`--wal-dir` 覆盖 WAL 根目录；`--reset` 会先删除该 job 的 WAL。 +decode 与可跳过的 commit 失败会写入 WAL 与 `import.log`，进程继续处理其余源。 + | Format | Import | Export | |---|---:|---:| | `atif` | 是 | 是 | @@ -326,21 +334,26 @@ Compact JSONL 是记录存储，不会转换或推断轨迹语义。指定 ### 2.8 `sync` ```text -pchronicle sync --from DIRECTORY --to DIRECTORY --convert DIRECTORY - [--input-format FORMAT] [--column NAME=JSON_PATH]... +pchronicle sync --from DIRECTORY + [--mirror DIRECTORY] [--to DIRECTORY] + [--input-format FORMAT] [--suggested-format FORMAT] + [--column NAME=JSON_PATH]... [--interval DURATION] [--once] ``` -`sync` 是常驻轮询器：监听源目录下的 `.json`、`.jsonl` 和 `.ndjson`。对于运行数据格式，它会将 -变更合并到 pending 池，并按 `--interval` 将源文件逐字节批量镜像到本地 Warehouse 目录，同时将 -数据转换为 Storyline Lance 写入 `--convert` 目标。一个批次成功后才清理 pending；失败会保留 -变更并指数退避重试。`--once` 只执行一次初始批次后退出。当前目标必须是本地目录，两个目标 -必须位于源目录之外。 +`sync` 是常驻轮询器：监听源目录下的 `.json`、`.jsonl` 和 `.ndjson`，将变更合并到 pending +池，并按 `--interval` 做整树 snapshot 重建。`--mirror` 与 `--to` 至少提供一个，也可同时提供： + +- `--mirror`：把源树按 Compact JSONL 规则写入 Compact Lance Dataset（record-level；可用 + `--column`）。每个成功批次原子替换该目标。 +- `--to`：把源轨迹转换为 Storyline Lance Dataset（可用 `--input-format` / + `--suggested-format`）。 + +一个批次内启用的目标全部成功后才清理 pending；失败会保留变更并指数退避重试。`--once` +只执行一次初始批次后退出。本地目标必须位于源目录之外。 -指定 `--input-format compact-jsonl` 时，源目录必须是本地 `.json`、`.jsonl` 或 `.ndjson` 目录树，列映射规则与 Compact -import 相同。每个成功批次都会重新扫描整个目录，并原子替换 `--convert` 指向的 Compact Lance -快照，因此新增、修改和删除都会反映在下一快照中，但不提供行级增量更新。此模式仍要求传入 -`--to` 作为兼容参数，但不会写入该路径。 +若 `--input-format compact-jsonl`，只能配合 `--mirror`（不能与 `--to` 同用）：每个成功批次 +重新扫描整个目录，并原子替换 `--mirror` 指向的 Compact Lance 快照。 ### 2.9 `drop` diff --git a/docs/src/zh/rfcs/0014-compact-jsonl.md b/docs/src/zh/rfcs/0014-compact-jsonl.md index 7bf073bb..7471995e 100644 --- a/docs/src/zh/rfcs/0014-compact-jsonl.md +++ b/docs/src/zh/rfcs/0014-compact-jsonl.md @@ -176,15 +176,14 @@ pchronicle import \ ```text pchronicle sync \ --from ./jsonl-root \ - --to ./warehouse-copy \ - --convert ./records.lance \ + --mirror ./records.lance \ --input-format compact-jsonl \ --column id=$.event.id \ --column timestamp=$.event.time ``` v1 sync 是 snapshot sync。每批变化 MUST 重新扫描完整 input root，并用一个完整的新 compact -snapshot 替换 `--convert`。创建、修改和删除源文件都必须反映到下一快照。v1 不承诺行级增量 +snapshot 替换 `--mirror`。创建、修改和删除源文件都必须反映到下一快照。v1 不承诺行级增量 更新。 ## Export From 4010c62923b9c25a0fb79ba0e26fafc320e52857 Mon Sep 17 00:00:00 2001 From: Reiase Date: Thu, 10 Sep 2026 20:06:56 +0800 Subject: [PATCH 8/8] feat(tests): introduce a custom reqwest client for echo tests Added a `test_client` function to create a reqwest client that ignores proxy settings, enhancing the reliability of echo tests. Updated existing test cases to utilize this new client, ensuring consistent behavior across different test scenarios. --- crates/persisting-gateway/src/echo.rs | 17 +++++++++++++---- 1 file changed, 13 insertions(+), 4 deletions(-) diff --git a/crates/persisting-gateway/src/echo.rs b/crates/persisting-gateway/src/echo.rs index 36458a53..9e1968bd 100644 --- a/crates/persisting-gateway/src/echo.rs +++ b/crates/persisting-gateway/src/echo.rs @@ -605,6 +605,15 @@ mod tests { use super::*; use std::sync::{Arc, Mutex}; + fn test_client() -> reqwest::Client { + // Local echo binds 127.0.0.1; ignore ambient HTTP(S)_PROXY / ALL_PROXY + // (e.g. socks5h) which reqwest may not support without extra features. + reqwest::Client::builder() + .no_proxy() + .build() + .expect("reqwest client") + } + async fn spawn_echo() -> (String, tokio::sync::oneshot::Sender<()>) { let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); let address = listener.local_addr().unwrap(); @@ -622,7 +631,7 @@ mod tests { #[tokio::test] async fn raw_echo_supports_plain_and_base64() { let (base, stop) = spawn_echo().await; - let client = reqwest::Client::new(); + let client = test_client(); let plain = client .post(format!("{base}/echo")) .body("hello") @@ -645,7 +654,7 @@ mod tests { #[tokio::test] async fn chat_echo_uses_last_user_message_and_streams() { let (base, stop) = spawn_echo().await; - let client = reqwest::Client::new(); + let client = test_client(); let response: Value = client .post(format!("{base}/v1/chat/completions")) .header(ECHO_ENCODING_HEADER, "base64") @@ -687,7 +696,7 @@ mod tests { #[tokio::test] async fn native_protocol_endpoints_return_their_wire_shapes() { let (base, stop) = spawn_echo().await; - let client = reqwest::Client::new(); + let client = test_client(); let messages: Value = client .post(format!("{base}/v1/messages")) @@ -774,7 +783,7 @@ forward = "echo-upstream" }, )); - let response = reqwest::Client::new() + let response = test_client() .post(format!("http://{gateway_address}/v1/messages")) .header(ECHO_ENCODING_HEADER, "base64") .json(&json!({