1use crate::destination::gcs::GcsStore;
27use crate::manifest::{MANIFEST_FILENAME, ManifestStatus, RunManifest};
28use crate::pipeline::validate_manifest::MANIFEST_MAX_BYTES;
29use anyhow::{Context, Result, bail};
30
31#[derive(Debug, Clone, PartialEq, Eq)]
35pub struct LoadIntegrity {
36 pub source_rows: Option<u64>,
41 pub file_rows: u64,
44 pub manifests: usize,
46}
47
48impl LoadIntegrity {
49 pub fn chain_prefix(&self) -> String {
53 let src = self
54 .source_rows
55 .map_or_else(|| "?".to_string(), |n| n.to_string());
56 format!("source {src} → files {}", self.file_rows)
57 }
58}
59
60#[allow(private_interfaces)]
77pub fn fetch_manifests_keyed(
78 store: &GcsStore,
79 gcs_prefix: &str,
80) -> Result<Vec<(String, RunManifest)>> {
81 let (_, base) = crate::load::split_gs_uri(gcs_prefix)?;
82 let keys = list_manifest_keys(store, base)?;
83 keys.into_iter()
84 .map(|key| {
85 let sz = store.stat_size(&key)?;
90 if sz > MANIFEST_MAX_BYTES {
91 bail!(
92 "manifest {key} is {sz} bytes, over the {MANIFEST_MAX_BYTES}-byte cap — \
93 refusing to read a possibly-hostile manifest into memory (CWE-400)"
94 );
95 }
96 let bytes = store.read(&key)?;
97 let m = serde_json::from_slice::<RunManifest>(&bytes)
98 .with_context(|| format!("parsing manifest {key}"))?;
99 Ok((key, m))
100 })
101 .collect()
102}
103
104#[allow(private_interfaces)]
108pub fn select_load_uris(
109 store: &GcsStore,
110 gcs_prefix: &str,
111 new: &[(String, RunManifest)],
112) -> Result<Vec<String>> {
113 let (bucket, base) = crate::load::split_gs_uri(gcs_prefix)?;
114 let all_parquet: Vec<String> = store
115 .list_files(base)?
116 .into_iter()
117 .filter(|k| k.ends_with(".parquet"))
118 .collect();
119 Ok(select_load_keys(new, &all_parquet)
120 .into_iter()
121 .map(|k| format!("gs://{bucket}/{k}"))
122 .collect())
123}
124
125fn resolve_parts<'a>(
131 manifest_key: &'a str,
132 m: &'a RunManifest,
133) -> impl Iterator<Item = String> + 'a {
134 let dir = manifest_key.rsplit_once('/').map(|(d, _)| d).unwrap_or("");
135 m.parts.iter().map(move |p| {
136 if dir.is_empty() {
137 p.path.clone()
138 } else {
139 format!("{dir}/{}", p.path)
140 }
141 })
142}
143
144pub fn select_load_keys(new: &[(String, RunManifest)], all_parquet: &[String]) -> Vec<String> {
155 use std::collections::BTreeSet;
156 let present: std::collections::HashSet<&str> = all_parquet.iter().map(String::as_str).collect();
157 let mut selected: BTreeSet<String> = BTreeSet::new();
158 for (key, m) in new {
159 let mut resolved_any = false;
160 for full in resolve_parts(key, m) {
161 if present.contains(full.as_str()) {
162 selected.insert(full);
163 resolved_any = true;
164 }
165 }
166 if !resolved_any {
167 return all_parquet.to_vec();
169 }
170 }
171 selected.into_iter().collect()
172}
173
174#[allow(private_interfaces)]
199pub fn gc_orphans(
200 store: &GcsStore,
201 gcs_prefix: &str,
202 keyed: &[(String, RunManifest)],
203 active: bool,
204) -> Result<(usize, u64)> {
205 let (_bucket, base) = crate::load::split_gs_uri(gcs_prefix)?;
206 let mut keep: std::collections::HashSet<String> = std::collections::HashSet::new();
210 let mut terminal: std::collections::HashSet<String> = std::collections::HashSet::new();
211 for (key, m) in keyed {
212 match m.status {
213 ManifestStatus::Success => keep.extend(resolve_parts(key, m)),
214 ManifestStatus::Running => {}
219 ManifestStatus::Failed | ManifestStatus::Interrupted => {
220 terminal.extend(resolve_parts(key, m))
221 }
222 }
223 }
224 let mut removed = 0usize;
225 let mut removed_bytes = 0u64;
226 for key in store.list_files(base)? {
227 if !key.ends_with(".parquet") || keep.contains(&key) {
228 continue;
229 }
230 if active && !terminal.contains(&key) {
234 log::warn!(
235 "gc_orphans: sparing unmanifested `{key}` — a run is active on this prefix \
236 (run-status ledger); it is GC'd once no run is active or it gets a manifest"
237 );
238 continue;
239 }
240 removed_bytes += store.stat_size(&key).unwrap_or(0);
241 store.remove(&key)?;
242 removed += 1;
243 }
244 for (key, m) in keyed {
252 if m.status == ManifestStatus::Running && is_superseded(m, keyed) {
253 removed_bytes += store.stat_size(key).unwrap_or(0);
254 store.remove(key)?;
255 removed += 1;
256 }
257 }
258 Ok((removed, removed_bytes))
259}
260
261pub fn has_active_running_manifest(keyed: &[(String, RunManifest)]) -> bool {
271 keyed
272 .iter()
273 .any(|(_, m)| m.status == ManifestStatus::Running && !is_superseded(m, keyed))
274}
275
276fn is_superseded(m: &RunManifest, keyed: &[(String, RunManifest)]) -> bool {
283 keyed
284 .iter()
285 .any(|(_, o)| o.export_name == m.export_name && o.started_at > m.started_at)
286}
287
288pub fn latest_full(keyed: Vec<(String, RunManifest)>) -> Vec<(String, RunManifest)> {
301 keyed
302 .into_iter()
303 .max_by(|a, b| {
304 match (
305 chrono::DateTime::parse_from_rfc3339(&a.1.finished_at).ok(),
306 chrono::DateTime::parse_from_rfc3339(&b.1.finished_at).ok(),
307 ) {
308 (Some(x), Some(y)) => x.cmp(&y),
309 _ => a.1.finished_at.cmp(&b.1.finished_at),
310 }
311 })
312 .into_iter()
313 .collect()
314}
315
316pub fn select_runs(
326 keyed: Vec<(String, RunManifest)>,
327 loaded: &std::collections::HashSet<String>,
328 mode: crate::load::plan::LoadMode,
329) -> Vec<(String, RunManifest)> {
330 let keyed: Vec<(String, RunManifest)> = keyed
335 .into_iter()
336 .filter(|(_, m)| m.status != ManifestStatus::Running)
337 .collect();
338 match mode {
339 crate::load::plan::LoadMode::Full => latest_full(keyed),
340 _ => keyed
341 .into_iter()
342 .filter(|(_, m)| !loaded.contains(&m.run_id))
343 .collect(),
344 }
345}
346
347fn list_manifest_keys(store: &GcsStore, base: &str) -> Result<Vec<String>> {
349 let all: Vec<String> = store
350 .list_files(base)?
351 .into_iter()
352 .filter(|k| is_manifest_key(k))
353 .collect();
354 let run_unique: Vec<String> = all
362 .iter()
363 .filter(|k| is_run_unique_manifest(k.rsplit('/').next().unwrap_or("")))
364 .cloned()
365 .collect();
366 Ok(if run_unique.is_empty() {
367 all
368 } else {
369 run_unique
370 })
371}
372
373fn is_manifest_key(key: &str) -> bool {
378 let base = key.rsplit('/').next().unwrap_or("");
379 base == MANIFEST_FILENAME || is_run_unique_manifest(base)
380}
381
382fn is_run_unique_manifest(base: &str) -> bool {
385 crate::manifest::is_run_unique_manifest_name(base)
387}
388
389pub fn ensure_single_export(keyed: &[(String, RunManifest)]) -> Result<()> {
400 let mut names: std::collections::BTreeSet<&str> = keyed
401 .iter()
402 .map(|(_, m)| m.export_name.as_str())
403 .filter(|n| !n.is_empty())
404 .collect();
405 if names.len() > 1 {
406 let listed: Vec<&str> = std::mem::take(&mut names).into_iter().collect();
407 bail!(
408 "the load prefix holds manifests from {} distinct exports ({}) — the load sums every \
409 manifest under the prefix and cleanup wipes it recursively, so loading a shared \
410 prefix would cross-contaminate the row count and could delete a sibling export's \
411 parts. Give each export a DISTINCT destination prefix (or scope the load prefix to a \
412 single export) and re-run.",
413 listed.len(),
414 listed.join(", ")
415 );
416 }
417 Ok(())
418}
419
420pub fn reconcile(manifests: &[RunManifest], allow_source_drift: bool) -> Result<LoadIntegrity> {
435 if manifests.is_empty() {
436 bail!(
437 "no `{MANIFEST_FILENAME}` found under the export prefix — refusing to load \
438 unverified files. A rivet export writes a manifest on success; its absence \
439 means the run never completed (or points at the wrong prefix)."
440 );
441 }
442
443 let mut file_rows: u64 = 0;
444 let mut source_rows: u64 = 0;
445 let mut any_source = false;
446
447 for m in manifests {
448 if m.status != ManifestStatus::Success {
450 bail!(
451 "manifest for run `{}` (export `{}`) is {:?}, not Success — refusing to load a \
452 partial export",
453 m.run_id,
454 m.export_name,
455 m.status
456 );
457 }
458 m.validate_self_consistency().map_err(|e| {
461 anyhow::anyhow!(
462 "manifest for run `{}` (export `{}`) is internally inconsistent: {e} — refusing \
463 to load",
464 m.run_id,
465 m.export_name
466 )
467 })?;
468
469 let rows = u64::try_from(m.row_count).with_context(|| {
470 format!(
471 "manifest for run `{}` has a negative row_count ({})",
472 m.run_id, m.row_count
473 )
474 })?;
475 file_rows += rows;
476
477 if let Some(src) = m
480 .source
481 .extraction
482 .as_ref()
483 .and_then(|x| x.source_row_count)
484 {
485 let src = u64::try_from(src).with_context(|| {
486 format!(
487 "manifest for run `{}` has a negative source_row_count ({src})",
488 m.run_id
489 )
490 })?;
491 any_source = true;
492 source_rows += src;
493 if src != rows {
494 if allow_source_drift {
495 eprintln!(
496 "warning: source→file drift for run `{}` (export `{}`): source had {src} \
497 rows, extracted {rows} (--allow-source-drift)",
498 m.run_id, m.export_name
499 );
500 } else {
501 bail!(
502 "source→file mismatch for run `{}` (export `{}`): source had {src} rows \
503 but {rows} were extracted — the extract dropped {} row(s). Investigate \
504 before loading, or pass --allow-source-drift to override.",
505 m.run_id,
506 m.export_name,
507 src.abs_diff(rows)
508 );
509 }
510 }
511 }
512 }
513
514 Ok(LoadIntegrity {
515 source_rows: any_source.then_some(source_rows),
516 file_rows,
517 manifests: manifests.len(),
518 })
519}
520
521#[cfg(test)]
522mod tests {
523 use super::*;
524 use crate::manifest::{
525 ExtractionMetadata, ManifestDestination, ManifestPart, ManifestSource, PartStatus,
526 };
527
528 fn manifest(run: &str, rows: i64, source: Option<i64>) -> RunManifest {
531 RunManifest {
532 manifest_version: crate::manifest::MANIFEST_VERSION,
533 run_id: run.into(),
534 export_name: "orders".into(),
535 mode: "batch".into(),
536 started_at: "t".into(),
537 finished_at: "t".into(),
538 status: ManifestStatus::Success,
539 source: ManifestSource {
540 engine: "pg".into(),
541 schema: None,
542 table: None,
543 extraction: source.map(|n| ExtractionMetadata {
544 strategy: "full".into(),
545 cursor_column: None,
546 cursor_type: None,
547 cursor_low: None,
548 cursor_high: None,
549 source_row_count: Some(n),
550 }),
551 },
552 destination: ManifestDestination {
553 kind: "gcs".into(),
554 uri: "gs://b/p".into(),
555 },
556 format: "parquet".into(),
557 compression: "zstd".into(),
558 schema_fingerprint: "xxh3:0".into(),
559 row_count: rows,
560 part_count: 1,
561 parts: vec![ManifestPart {
562 part_id: 0,
563 path: "part-000000.parquet".into(),
564 rows,
565 size_bytes: 1,
566 content_fingerprint: "xxh3:0".into(),
567 content_md5: String::new(),
568 status: PartStatus::Committed,
569 }],
570 column_checksums: None,
571 checksum_key_column: None,
572 }
573 }
574
575 #[test]
576 fn sums_file_and_source_rows_across_manifests() {
577 let ms = vec![manifest("r1", 100, Some(100)), manifest("r2", 40, Some(40))];
578 let got = reconcile(&ms, false).unwrap();
579 assert_eq!(got.file_rows, 140);
580 assert_eq!(got.source_rows, Some(140));
581 assert_eq!(got.manifests, 2);
582 }
583
584 #[test]
585 fn ensure_single_export_refuses_a_prefix_shared_by_two_exports() {
586 let keyed = |m: RunManifest| ("gs://b/p/manifest-x.json".to_string(), m);
587 let orders = keyed(manifest("r1", 10, None)); let mut cust_m = manifest("r2", 10, None);
589 cust_m.export_name = "customers".into();
590 assert!(ensure_single_export(&[orders.clone(), keyed(cust_m)]).is_err());
593 assert!(ensure_single_export(&[orders.clone(), keyed(manifest("r3", 5, None))]).is_ok());
595 let mut legacy = manifest("r0", 3, None);
598 legacy.export_name = String::new();
599 assert!(ensure_single_export(&[orders, keyed(legacy)]).is_ok());
600 }
601
602 #[test]
603 fn source_rows_is_none_when_no_manifest_probed_the_source() {
604 let ms = vec![manifest("r1", 100, None), manifest("r2", 40, None)];
605 let got = reconcile(&ms, false).unwrap();
606 assert_eq!(got.file_rows, 140);
607 assert_eq!(
608 got.source_rows, None,
609 "unprobed source is unknown, not zero"
610 );
611 }
612
613 #[test]
614 fn source_rows_present_even_if_only_some_manifests_probed() {
615 let ms = vec![manifest("r1", 100, Some(100)), manifest("r2", 40, None)];
616 let got = reconcile(&ms, false).unwrap();
617 assert_eq!(got.source_rows, Some(100));
618 }
619
620 #[test]
621 fn empty_manifests_refuses_to_load() {
622 let err = reconcile(&[], false).unwrap_err().to_string();
623 assert!(err.contains("refusing to load"), "{err}");
624 }
625
626 #[test]
627 fn non_success_manifest_refuses_to_load() {
628 let mut m = manifest("r1", 100, Some(100));
629 m.status = ManifestStatus::Interrupted;
630 let err = reconcile(&[m], false).unwrap_err().to_string();
631 assert!(err.contains("not Success"), "{err}");
632 }
633
634 #[test]
635 fn self_inconsistent_manifest_refuses_to_load() {
636 let mut m = manifest("r1", 100, Some(100));
639 m.row_count = 999; let err = reconcile(&[m], false).unwrap_err().to_string();
641 assert!(err.contains("inconsistent"), "{err}");
642 }
643
644 #[test]
645 fn source_file_mismatch_hard_fails_by_default() {
646 let m = manifest("r1", 100, Some(120));
648 let err = reconcile(&[m], false).unwrap_err().to_string();
649 assert!(err.contains("source→file mismatch"), "{err}");
650 assert!(err.contains("dropped 20"), "{err}");
651 }
652
653 #[test]
654 fn source_file_mismatch_is_allowed_under_the_override() {
655 let m = manifest("r1", 100, Some(120));
656 let got = reconcile(&[m], true).expect("--allow-source-drift proceeds");
657 assert_eq!(got.file_rows, 100);
658 assert_eq!(
659 got.source_rows,
660 Some(120),
661 "the probed source count is still surfaced"
662 );
663 }
664
665 fn keyed(key: &str, run: &str, part: &str) -> (String, RunManifest) {
667 let mut m = manifest(run, 10, Some(10));
668 m.parts[0].path = part.into();
669 (key.to_string(), m)
670 }
671
672 #[test]
673 fn select_load_keys_picks_only_the_new_runs_parts() {
674 let all = vec![
676 "base/r1-000.parquet".to_string(),
677 "base/r2-000.parquet".to_string(),
678 ];
679 let new = vec![keyed("base/manifest-r2.json", "r2", "r2-000.parquet")];
680 assert_eq!(
681 select_load_keys(&new, &all),
682 vec!["base/r2-000.parquet".to_string()],
683 "loads r2's part only — not r1's already-loaded file"
684 );
685 }
686
687 #[test]
688 fn select_load_uris_lists_and_re_prefixes_selected_keys_with_the_source_bucket() {
689 let (store, _g) = fs_store(&[
695 ("base/r1-000.parquet", b"a".to_vec()),
696 ("base/manifest-r1.json", b"{}".to_vec()), ]);
698 let new = vec![keyed("base/manifest-r1.json", "r1", "r1-000.parquet")];
699 assert_eq!(
700 select_load_uris(&store, "gs://my-bucket/base", &new).unwrap(),
701 vec!["gs://my-bucket/base/r1-000.parquet".to_string()],
702 "the selected key, re-prefixed with the source bucket — never the manifest"
703 );
704 }
705
706 #[test]
707 fn select_load_keys_resolves_a_snapshot_subprefix_manifest() {
708 let all = vec!["base/snapshot/snap-000.parquet".to_string()];
711 let new = vec![keyed(
712 "base/snapshot/manifest-r1.json",
713 "r1",
714 "snap-000.parquet",
715 )];
716 assert_eq!(
717 select_load_keys(&new, &all),
718 vec!["base/snapshot/snap-000.parquet".to_string()]
719 );
720 }
721
722 #[test]
723 fn select_load_keys_falls_back_to_full_listing_when_a_manifest_has_no_present_part() {
724 let all = vec!["base/a.parquet".to_string(), "base/b.parquet".to_string()];
727 let new = vec![keyed("base/manifest-r1.json", "r1", "missing.parquet")];
728 assert_eq!(
729 select_load_keys(&new, &all),
730 all,
731 "unresolvable part → blanket fallback"
732 );
733 }
734
735 #[test]
736 fn select_load_keys_empty_new_set_selects_nothing_never_the_full_listing() {
737 let all = vec![
742 "base/r1-000.parquet".to_string(),
743 "base/r2-000.parquet".to_string(),
744 ];
745 assert!(
746 select_load_keys(&[], &all).is_empty(),
747 "no new runs ⇒ load nothing, not everything"
748 );
749 }
750
751 #[test]
752 fn gc_orphans_removes_unmanifested_parquet_only() {
753 let (store, _g) = fs_store(&[
754 ("base/r1-000.parquet", b"aa".to_vec()), ("base/orphan.parquet", b"junk".to_vec()), ("base/manifest.json", b"{}".to_vec()), ("base/_SUCCESS", b"".to_vec()), ]);
759 let keyed = vec![keyed("base/manifest-r1.json", "r1", "r1-000.parquet")];
760 let (removed, bytes) = gc_orphans(&store, "gs://b/base", &keyed, false).unwrap();
762 assert_eq!(removed, 1, "only the unmanifested part is removed");
763 assert_eq!(bytes, 4, "'junk' is 4 bytes");
764 let mut left = store.list_files("base").unwrap();
765 left.sort();
766 assert_eq!(
767 left,
768 vec![
769 "base/_SUCCESS".to_string(),
770 "base/manifest.json".to_string(),
771 "base/r1-000.parquet".to_string(),
772 ],
773 "the manifested part, the manifest, and _SUCCESS all survive"
774 );
775 }
776
777 #[test]
778 fn gc_orphans_of_an_all_manifested_prefix_removes_nothing() {
779 let (store, _g) = fs_store(&[("base/r1-000.parquet", b"a".to_vec())]);
780 let keyed = vec![keyed("base/manifest-r1.json", "r1", "r1-000.parquet")];
781 assert_eq!(
782 gc_orphans(&store, "gs://b/base", &keyed, false).unwrap().0,
783 0
784 );
785 }
786
787 #[test]
788 fn gc_orphans_keeps_a_snapshot_subprefix_manifested_part() {
789 let (store, _g) = fs_store(&[
790 ("base/snapshot/s-000.parquet", b"a".to_vec()), ("base/orphan.parquet", b"x".to_vec()), ]);
793 let keyed = vec![keyed("base/snapshot/manifest-s.json", "s", "s-000.parquet")];
794 let (removed, _) = gc_orphans(&store, "gs://b/base", &keyed, false).unwrap();
795 assert_eq!(removed, 1, "the top-level orphan goes");
796 assert_eq!(
797 store.list_files("base/snapshot").unwrap(),
798 vec!["base/snapshot/s-000.parquet".to_string()],
799 "the snapshot-subprefix manifested part is kept"
800 );
801 }
802
803 #[test]
804 fn gc_orphans_deletes_a_terminal_runs_parts_even_while_a_run_is_active() {
805 let (store, _g) = fs_store(&[("base/f-000.parquet", b"x".to_vec())]);
809 let mut kv = keyed("base/manifest-f.json", "f", "f-000.parquet");
810 kv.1.status = ManifestStatus::Failed;
811 assert_eq!(gc_orphans(&store, "gs://b/base", &[kv], true).unwrap().0, 1);
812 }
813
814 #[test]
815 fn gc_orphans_spares_an_unmanifested_part_while_a_run_is_active() {
816 let (store, _g) = fs_store(&[
822 ("base/r1-000.parquet", b"aa".to_vec()), ("base/inflight.parquet", b"x".to_vec()), ]);
825 let keyed = vec![keyed("base/manifest-r1.json", "r1", "r1-000.parquet")];
826 let (removed, _) = gc_orphans(&store, "gs://b/base", &keyed, true).unwrap();
827 assert_eq!(
828 removed, 0,
829 "an unmanifested part is spared while a run is active"
830 );
831 assert!(
832 store
833 .list_files("base")
834 .unwrap()
835 .iter()
836 .any(|k| k.ends_with("inflight.parquet")),
837 "the live run's in-flight part must survive gc_orphans"
838 );
839 }
840
841 #[test]
842 fn gc_orphans_collects_an_unmanifested_part_when_no_run_is_active() {
843 let (store, _g) = fs_store(&[
846 ("base/r1-000.parquet", b"aa".to_vec()),
847 ("base/dead-orphan.parquet", b"x".to_vec()),
848 ]);
849 let keyed = vec![keyed("base/manifest-r1.json", "r1", "r1-000.parquet")];
850 assert_eq!(
851 gc_orphans(&store, "gs://b/base", &keyed, false).unwrap().0,
852 1,
853 "with no active run, an unmanifested orphan is collected"
854 );
855 }
856
857 fn running(run: &str, started_at: &str) -> (String, RunManifest) {
859 let mut m = manifest(run, 0, None);
860 m.status = ManifestStatus::Running;
861 m.started_at = started_at.into();
862 m.finished_at = String::new();
863 m.parts.clear();
864 (format!("base/manifest-{run}.json"), m)
865 }
866
867 #[test]
868 fn has_active_running_manifest_true_for_a_lone_running_marker() {
869 assert!(has_active_running_manifest(&[running(
870 "r1",
871 "2026-01-01T00:00:00Z"
872 )]));
873 }
874
875 #[test]
876 fn has_active_running_manifest_false_when_none_is_running() {
877 assert!(!has_active_running_manifest(&[keyed(
879 "k",
880 "r1",
881 "p.parquet"
882 )]));
883 }
884
885 #[test]
886 fn has_active_running_manifest_false_for_a_superseded_running_marker() {
887 let r1 = running("r1", "2026-01-01T00:00:00Z");
891 let mut r2 = manifest("r2", 10, Some(10)); r2.started_at = "2026-01-02T00:00:00Z".into();
893 assert!(!has_active_running_manifest(&[
894 r1,
895 ("base/manifest-r2.json".into(), r2)
896 ]));
897 }
898
899 #[test]
900 fn select_runs_drops_a_running_manifest() {
901 let run_marker = running("r_running", "2026-01-02T00:00:00Z");
905 let ok = keyed("base/manifest-r_ok.json", "r_ok", "r_ok-000.parquet"); let sel = select_runs(
907 vec![run_marker, ok],
908 &std::collections::HashSet::new(),
909 crate::load::plan::LoadMode::Incremental,
910 );
911 assert_eq!(sel.len(), 1, "only the Success run is selected");
912 assert_eq!(sel[0].1.run_id, "r_ok");
913 assert!(sel.iter().all(|(_, m)| m.status != ManifestStatus::Running));
914 }
915
916 #[test]
917 fn gc_orphans_removes_a_superseded_running_marker_but_spares_a_live_one() {
918 let (store, _g) = fs_store(&[
923 ("base/manifest-r1.json", b"{}".to_vec()), ("base/manifest-r2.json", b"{}".to_vec()), ]);
926 let r1 = running("r1", "2026-01-01T00:00:01Z");
927 let r2 = running("r2", "2026-01-01T00:00:02Z");
928 let (removed, _) = gc_orphans(&store, "gs://b/base", &[r1, r2], true).unwrap();
929 assert_eq!(removed, 1, "only the superseded running marker is removed");
930 let left = store.list_files("base").unwrap();
931 assert!(
932 !left.iter().any(|k| k.ends_with("manifest-r1.json")),
933 "the superseded (dead) running marker is deleted"
934 );
935 assert!(
936 left.iter().any(|k| k.ends_with("manifest-r2.json")),
937 "the live (non-superseded) running marker survives — it is the active signal"
938 );
939 }
940
941 fn keyed_at(run: &str, finished_at: &str) -> (String, RunManifest) {
942 let mut m = manifest(run, 100, None);
943 m.finished_at = finished_at.into();
944 (format!("base/manifest-{run}.json"), m)
945 }
946
947 #[test]
948 fn latest_full_picks_the_newest_snapshot_not_all() {
949 let keyed = vec![
952 keyed_at("r1", "2026-01-01T00:00:00Z"),
953 keyed_at("r3", "2026-01-03T00:00:00Z"),
954 keyed_at("r2", "2026-01-02T00:00:00Z"),
955 ];
956 let sel = latest_full(keyed);
957 assert_eq!(sel.len(), 1, "exactly one snapshot, never all");
958 assert_eq!(sel[0].1.run_id, "r3", "the newest by finished_at");
959 }
960
961 #[test]
962 fn latest_full_re_materializes_even_when_the_latest_is_already_loaded() {
963 let keyed = vec![
968 keyed_at("r1", "2026-01-01T00:00:00Z"),
969 keyed_at("r2", "2026-01-02T00:00:00Z"),
970 ];
971 let sel = latest_full(keyed);
973 assert_eq!(sel.len(), 1);
974 assert_eq!(sel[0].1.run_id, "r2", "always the latest, loaded or not");
975 }
976
977 #[test]
978 fn latest_full_of_no_staged_runs_is_empty_so_the_caller_no_ops_without_truncating() {
979 assert!(latest_full(Vec::new()).is_empty());
982 }
983
984 #[test]
985 fn latest_full_orders_by_parsed_instant_not_lexical_bytes() {
986 let keyed = vec![
989 keyed_at("older", "2026-01-01T00:00:00Z"),
990 keyed_at("newer", "2026-01-01T00:00:00.500Z"),
991 ];
992 let sel = latest_full(keyed);
993 assert_eq!(sel.len(), 1);
994 assert_eq!(
995 sel[0].1.run_id, "newer",
996 "the fractional-second run is the newer instant"
997 );
998 }
999
1000 #[test]
1001 fn select_runs_full_picks_the_latest_even_when_loaded_and_even_when_stateless() {
1002 use crate::load::plan::LoadMode;
1003 use std::collections::HashSet;
1004 let keyed = vec![
1005 keyed_at("r1", "2026-01-01T00:00:00Z"),
1006 keyed_at("r2", "2026-01-02T00:00:00Z"),
1007 ];
1008 let loaded = HashSet::from(["r2".to_string()]);
1010 let sel = select_runs(keyed.clone(), &loaded, LoadMode::Full);
1011 assert_eq!(sel.len(), 1);
1012 assert_eq!(sel[0].1.run_id, "r2", "Full picks latest, loaded or not");
1013 let sel = select_runs(keyed, &HashSet::new(), LoadMode::Full);
1016 assert_eq!(sel.len(), 1, "stateless Full is not a blanket load");
1017 assert_eq!(sel[0].1.run_id, "r2");
1018 }
1019
1020 #[test]
1021 fn select_runs_append_modes_filter_loaded_and_load_all_when_stateless() {
1022 use crate::load::plan::LoadMode;
1023 use std::collections::HashSet;
1024 let keyed = vec![
1025 keyed_at("r1", "2026-01-01T00:00:00Z"),
1026 keyed_at("r2", "2026-01-02T00:00:00Z"),
1027 ];
1028 let loaded = HashSet::from(["r1".to_string()]);
1029 for mode in [LoadMode::Incremental, LoadMode::Cdc] {
1030 let sel = select_runs(keyed.clone(), &loaded, mode);
1031 assert_eq!(sel.len(), 1, "{mode:?}: only the unloaded run");
1032 assert_eq!(sel[0].1.run_id, "r2");
1033 assert_eq!(
1035 select_runs(keyed.clone(), &HashSet::new(), mode).len(),
1036 2,
1037 "{mode:?}: stateless loads every run"
1038 );
1039 }
1040 }
1041
1042 #[test]
1043 fn is_run_unique_manifest_needs_both_prefix_and_json() {
1044 assert!(is_run_unique_manifest("manifest-20260101T000000.json"));
1045 assert!(!is_run_unique_manifest("manifest.json")); assert!(!is_run_unique_manifest("manifest-abc.txt")); assert!(!is_run_unique_manifest("data.json")); }
1049
1050 #[test]
1051 fn chain_prefix_renders_source_and_files() {
1052 let known = LoadIntegrity {
1053 source_rows: Some(100),
1054 file_rows: 100,
1055 manifests: 1,
1056 };
1057 assert_eq!(known.chain_prefix(), "source 100 → files 100");
1058 let unknown = LoadIntegrity {
1059 source_rows: None,
1060 file_rows: 40,
1061 manifests: 1,
1062 };
1063 assert_eq!(unknown.chain_prefix(), "source ? → files 40");
1064 }
1065
1066 #[test]
1067 fn is_manifest_key_matches_only_the_final_segment() {
1068 assert!(is_manifest_key("gs://b/p/manifest.json"));
1069 assert!(is_manifest_key("manifest.json"));
1070 assert!(!is_manifest_key("gs://b/p/part-0.parquet"));
1071 assert!(!is_manifest_key("gs://b/p/x_manifest.json"));
1072 }
1073
1074 fn fs_store(files: &[(&str, Vec<u8>)]) -> (GcsStore, tempfile::TempDir) {
1080 let dir = tempfile::tempdir().unwrap();
1081 for (rel, bytes) in files {
1082 let p = dir.path().join(rel);
1083 std::fs::create_dir_all(p.parent().unwrap()).unwrap();
1084 std::fs::write(p, bytes).unwrap();
1085 }
1086 let store = GcsStore::open_fs(dir.path().to_str().unwrap()).unwrap();
1087 (store, dir)
1088 }
1089
1090 fn manifest_bytes(run: &str, rows: i64, source: Option<i64>) -> Vec<u8> {
1091 serde_json::to_vec(&manifest(run, rows, source)).unwrap()
1092 }
1093
1094 #[test]
1095 fn list_manifest_keys_prefers_run_unique_copies_over_the_canonical_pointer() {
1096 let (store, _g) = fs_store(&[
1100 ("base/manifest.json", b"{}".to_vec()),
1101 ("base/manifest-r1.json", b"{}".to_vec()),
1102 ("base/manifest-r2.json", b"{}".to_vec()),
1103 ("base/part-0.parquet", b"x".to_vec()), ]);
1105 let mut keys = list_manifest_keys(&store, "base").unwrap();
1106 keys.sort();
1107 assert_eq!(
1108 keys,
1109 vec![
1110 "base/manifest-r1.json".to_string(),
1111 "base/manifest-r2.json".to_string(),
1112 ]
1113 );
1114 }
1115
1116 #[test]
1117 fn list_manifest_keys_falls_back_to_the_canonical_name_for_a_single_run() {
1118 let (store, _g) = fs_store(&[("base/manifest.json", b"{}".to_vec())]);
1121 assert_eq!(
1122 list_manifest_keys(&store, "base").unwrap(),
1123 vec!["base/manifest.json".to_string()]
1124 );
1125 }
1126
1127 #[test]
1128 fn fetch_manifests_keyed_reads_and_parses_every_run_copy_under_the_prefix() {
1129 let (store, _g) = fs_store(&[
1133 ("base/manifest.json", manifest_bytes("r2", 40, Some(40))),
1134 (
1135 "base/manifest-r1.json",
1136 manifest_bytes("r1", 100, Some(100)),
1137 ),
1138 ("base/manifest-r2.json", manifest_bytes("r2", 40, Some(40))),
1139 ]);
1140 let manifests: Vec<_> = fetch_manifests_keyed(&store, "gs://my-bucket/base")
1141 .unwrap()
1142 .into_iter()
1143 .map(|(_, m)| m)
1144 .collect();
1145 assert_eq!(manifests.len(), 2);
1146 let integrity = reconcile(&manifests, false).unwrap();
1147 assert_eq!(integrity.file_rows, 140);
1148 assert_eq!(integrity.manifests, 2);
1149 }
1150
1151 #[test]
1152 fn fetch_manifests_keyed_names_the_key_when_a_manifest_is_unparseable() {
1153 let (store, _g) = fs_store(&[("base/manifest.json", b"{ not json".to_vec())]);
1154 let err = fetch_manifests_keyed(&store, "gs://my-bucket/base")
1155 .unwrap_err()
1156 .to_string();
1157 assert!(
1158 err.contains("parsing manifest") && err.contains("base/manifest.json"),
1159 "error should name the offending key: {err}"
1160 );
1161 }
1162
1163 const FAKE_GCS_ENDPOINT: &str = "http://127.0.0.1:4443";
1174 const FAKE_GCS_BUCKET: &str = "rivet-load-emulator";
1175
1176 fn fake_gcs_store() -> GcsStore {
1181 let created = std::process::Command::new("curl")
1182 .args([
1183 "-s",
1184 "-X",
1185 "POST",
1186 &format!("{FAKE_GCS_ENDPOINT}/storage/v1/b?project=rivet-test"),
1187 "-H",
1188 "Content-Type: application/json",
1189 "-d",
1190 &format!("{{\"name\":\"{FAKE_GCS_BUCKET}\"}}"),
1191 ])
1192 .output();
1193 assert!(
1194 created.is_ok_and(|o| o.status.success()),
1195 "could not reach fake-gcs to create the bucket — is `docker compose up -d fake-gcs` running on :4443?"
1196 );
1197 let cfg = crate::config::DestinationConfig {
1198 destination_type: crate::config::DestinationType::Gcs,
1199 bucket: Some(FAKE_GCS_BUCKET.into()),
1200 endpoint: Some(FAKE_GCS_ENDPOINT.into()),
1201 allow_anonymous: true,
1202 ..Default::default()
1203 };
1204 GcsStore::new(&cfg).expect("build GcsStore against fake-gcs")
1205 }
1206
1207 fn drain(store: &GcsStore, prefix: &str) {
1215 for key in store.list_files(prefix).unwrap() {
1216 store.remove(&key).unwrap();
1217 }
1218 }
1219
1220 #[test]
1221 #[ignore = "emulator: needs `docker compose up -d fake-gcs` (fsouza/fake-gcs-server :4443)"]
1222 fn storage_contract_over_fake_gcs() {
1223 let store = fake_gcs_store();
1224 let prefix = "load-contract/orders";
1227 drain(&store, prefix);
1228 let gs = format!("gs://{FAKE_GCS_BUCKET}/{prefix}");
1229
1230 store
1233 .put(
1234 &format!("{prefix}/manifest-r1.json"),
1235 &manifest_bytes("r1", 100, Some(100)),
1236 )
1237 .unwrap();
1238 store
1239 .put(&format!("{prefix}/part-000000.parquet"), b"rows-of-r1")
1240 .unwrap();
1241 store
1242 .put(&format!("{prefix}/orphan.parquet"), b"crash-leftover")
1243 .unwrap();
1244
1245 let keyed = fetch_manifests_keyed(&store, &gs).unwrap();
1247 assert_eq!(keyed.len(), 1, "the run's manifest, read back over GCS");
1248
1249 let manifests: Vec<_> = keyed.iter().map(|(_, m)| m.clone()).collect();
1251 assert_eq!(
1252 reconcile(&manifests, false).unwrap().file_rows,
1253 100,
1254 "file_rows drives the count-gate; a bad GCS read would corrupt it"
1255 );
1256
1257 assert_eq!(
1259 select_load_uris(&store, &gs, &keyed).unwrap(),
1260 vec![format!(
1261 "gs://{FAKE_GCS_BUCKET}/{prefix}/part-000000.parquet"
1262 )],
1263 "load pulls the manifested part, not the unmanifested crash orphan"
1264 );
1265
1266 let (removed, _bytes) = gc_orphans(&store, &gs, &keyed, false).unwrap();
1269 assert_eq!(
1270 removed, 1,
1271 "exactly the orphan parquet is GC'd over real GCS"
1272 );
1273 let mut left = store.list_files(prefix).unwrap();
1274 left.sort();
1275 assert_eq!(
1276 left,
1277 vec![
1278 format!("{prefix}/manifest-r1.json"),
1279 format!("{prefix}/part-000000.parquet"),
1280 ],
1281 "the manifested part + its manifest survive the orphan GC"
1282 );
1283
1284 drain(&store, prefix);
1288 assert!(
1289 store.list_files(prefix).unwrap().is_empty(),
1290 "teardown left the prefix clean over real GCS"
1291 );
1292 }
1293}