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 row_hash: None,
533 manifest_version: crate::manifest::MANIFEST_VERSION,
534 run_id: run.into(),
535 export_name: "orders".into(),
536 mode: "batch".into(),
537 started_at: "t".into(),
538 finished_at: "t".into(),
539 status: ManifestStatus::Success,
540 source: ManifestSource {
541 engine: "pg".into(),
542 schema: None,
543 table: None,
544 extraction: source.map(|n| ExtractionMetadata {
545 strategy: "full".into(),
546 cursor_column: None,
547 cursor_type: None,
548 cursor_low: None,
549 cursor_high: None,
550 source_row_count: Some(n),
551 }),
552 },
553 destination: ManifestDestination {
554 kind: "gcs".into(),
555 uri: "gs://b/p".into(),
556 },
557 format: "parquet".into(),
558 compression: "zstd".into(),
559 schema_fingerprint: "xxh3:0".into(),
560 row_count: rows,
561 part_count: 1,
562 parts: vec![ManifestPart {
563 part_id: 0,
564 path: "part-000000.parquet".into(),
565 rows,
566 size_bytes: 1,
567 content_fingerprint: "xxh3:0".into(),
568 content_md5: String::new(),
569 status: PartStatus::Committed,
570 }],
571 column_checksums: None,
572 checksum_key_column: None,
573 }
574 }
575
576 #[test]
577 fn sums_file_and_source_rows_across_manifests() {
578 let ms = vec![manifest("r1", 100, Some(100)), manifest("r2", 40, Some(40))];
579 let got = reconcile(&ms, false).unwrap();
580 assert_eq!(got.file_rows, 140);
581 assert_eq!(got.source_rows, Some(140));
582 assert_eq!(got.manifests, 2);
583 }
584
585 #[test]
586 fn ensure_single_export_refuses_a_prefix_shared_by_two_exports() {
587 let keyed = |m: RunManifest| ("gs://b/p/manifest-x.json".to_string(), m);
588 let orders = keyed(manifest("r1", 10, None)); let mut cust_m = manifest("r2", 10, None);
590 cust_m.export_name = "customers".into();
591 assert!(ensure_single_export(&[orders.clone(), keyed(cust_m)]).is_err());
594 assert!(ensure_single_export(&[orders.clone(), keyed(manifest("r3", 5, None))]).is_ok());
596 let mut legacy = manifest("r0", 3, None);
599 legacy.export_name = String::new();
600 assert!(ensure_single_export(&[orders, keyed(legacy)]).is_ok());
601 }
602
603 #[test]
604 fn source_rows_is_none_when_no_manifest_probed_the_source() {
605 let ms = vec![manifest("r1", 100, None), manifest("r2", 40, None)];
606 let got = reconcile(&ms, false).unwrap();
607 assert_eq!(got.file_rows, 140);
608 assert_eq!(
609 got.source_rows, None,
610 "unprobed source is unknown, not zero"
611 );
612 }
613
614 #[test]
615 fn source_rows_present_even_if_only_some_manifests_probed() {
616 let ms = vec![manifest("r1", 100, Some(100)), manifest("r2", 40, None)];
617 let got = reconcile(&ms, false).unwrap();
618 assert_eq!(got.source_rows, Some(100));
619 }
620
621 #[test]
622 fn empty_manifests_refuses_to_load() {
623 let err = reconcile(&[], false).unwrap_err().to_string();
624 assert!(err.contains("refusing to load"), "{err}");
625 }
626
627 #[test]
628 fn non_success_manifest_refuses_to_load() {
629 let mut m = manifest("r1", 100, Some(100));
630 m.status = ManifestStatus::Interrupted;
631 let err = reconcile(&[m], false).unwrap_err().to_string();
632 assert!(err.contains("not Success"), "{err}");
633 }
634
635 #[test]
636 fn self_inconsistent_manifest_refuses_to_load() {
637 let mut m = manifest("r1", 100, Some(100));
640 m.row_count = 999; let err = reconcile(&[m], false).unwrap_err().to_string();
642 assert!(err.contains("inconsistent"), "{err}");
643 }
644
645 #[test]
646 fn source_file_mismatch_hard_fails_by_default() {
647 let m = manifest("r1", 100, Some(120));
649 let err = reconcile(&[m], false).unwrap_err().to_string();
650 assert!(err.contains("source→file mismatch"), "{err}");
651 assert!(err.contains("dropped 20"), "{err}");
652 }
653
654 #[test]
655 fn source_file_mismatch_is_allowed_under_the_override() {
656 let m = manifest("r1", 100, Some(120));
657 let got = reconcile(&[m], true).expect("--allow-source-drift proceeds");
658 assert_eq!(got.file_rows, 100);
659 assert_eq!(
660 got.source_rows,
661 Some(120),
662 "the probed source count is still surfaced"
663 );
664 }
665
666 fn keyed(key: &str, run: &str, part: &str) -> (String, RunManifest) {
668 let mut m = manifest(run, 10, Some(10));
669 m.parts[0].path = part.into();
670 (key.to_string(), m)
671 }
672
673 #[test]
674 fn select_load_keys_picks_only_the_new_runs_parts() {
675 let all = vec![
677 "base/r1-000.parquet".to_string(),
678 "base/r2-000.parquet".to_string(),
679 ];
680 let new = vec![keyed("base/manifest-r2.json", "r2", "r2-000.parquet")];
681 assert_eq!(
682 select_load_keys(&new, &all),
683 vec!["base/r2-000.parquet".to_string()],
684 "loads r2's part only — not r1's already-loaded file"
685 );
686 }
687
688 #[test]
689 fn select_load_uris_lists_and_re_prefixes_selected_keys_with_the_source_bucket() {
690 let (store, _g) = fs_store(&[
696 ("base/r1-000.parquet", b"a".to_vec()),
697 ("base/manifest-r1.json", b"{}".to_vec()), ]);
699 let new = vec![keyed("base/manifest-r1.json", "r1", "r1-000.parquet")];
700 assert_eq!(
701 select_load_uris(&store, "gs://my-bucket/base", &new).unwrap(),
702 vec!["gs://my-bucket/base/r1-000.parquet".to_string()],
703 "the selected key, re-prefixed with the source bucket — never the manifest"
704 );
705 }
706
707 #[test]
708 fn select_load_keys_resolves_a_snapshot_subprefix_manifest() {
709 let all = vec!["base/snapshot/snap-000.parquet".to_string()];
712 let new = vec![keyed(
713 "base/snapshot/manifest-r1.json",
714 "r1",
715 "snap-000.parquet",
716 )];
717 assert_eq!(
718 select_load_keys(&new, &all),
719 vec!["base/snapshot/snap-000.parquet".to_string()]
720 );
721 }
722
723 #[test]
724 fn select_load_keys_falls_back_to_full_listing_when_a_manifest_has_no_present_part() {
725 let all = vec!["base/a.parquet".to_string(), "base/b.parquet".to_string()];
728 let new = vec![keyed("base/manifest-r1.json", "r1", "missing.parquet")];
729 assert_eq!(
730 select_load_keys(&new, &all),
731 all,
732 "unresolvable part → blanket fallback"
733 );
734 }
735
736 #[test]
737 fn select_load_keys_empty_new_set_selects_nothing_never_the_full_listing() {
738 let all = vec![
743 "base/r1-000.parquet".to_string(),
744 "base/r2-000.parquet".to_string(),
745 ];
746 assert!(
747 select_load_keys(&[], &all).is_empty(),
748 "no new runs ⇒ load nothing, not everything"
749 );
750 }
751
752 #[test]
753 fn gc_orphans_removes_unmanifested_parquet_only() {
754 let (store, _g) = fs_store(&[
755 ("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()), ]);
760 let keyed = vec![keyed("base/manifest-r1.json", "r1", "r1-000.parquet")];
761 let (removed, bytes) = gc_orphans(&store, "gs://b/base", &keyed, false).unwrap();
763 assert_eq!(removed, 1, "only the unmanifested part is removed");
764 assert_eq!(bytes, 4, "'junk' is 4 bytes");
765 let mut left = store.list_files("base").unwrap();
766 left.sort();
767 assert_eq!(
768 left,
769 vec![
770 "base/_SUCCESS".to_string(),
771 "base/manifest.json".to_string(),
772 "base/r1-000.parquet".to_string(),
773 ],
774 "the manifested part, the manifest, and _SUCCESS all survive"
775 );
776 }
777
778 #[test]
779 fn gc_orphans_of_an_all_manifested_prefix_removes_nothing() {
780 let (store, _g) = fs_store(&[("base/r1-000.parquet", b"a".to_vec())]);
781 let keyed = vec![keyed("base/manifest-r1.json", "r1", "r1-000.parquet")];
782 assert_eq!(
783 gc_orphans(&store, "gs://b/base", &keyed, false).unwrap().0,
784 0
785 );
786 }
787
788 #[test]
789 fn gc_orphans_keeps_a_snapshot_subprefix_manifested_part() {
790 let (store, _g) = fs_store(&[
791 ("base/snapshot/s-000.parquet", b"a".to_vec()), ("base/orphan.parquet", b"x".to_vec()), ]);
794 let keyed = vec![keyed("base/snapshot/manifest-s.json", "s", "s-000.parquet")];
795 let (removed, _) = gc_orphans(&store, "gs://b/base", &keyed, false).unwrap();
796 assert_eq!(removed, 1, "the top-level orphan goes");
797 assert_eq!(
798 store.list_files("base/snapshot").unwrap(),
799 vec!["base/snapshot/s-000.parquet".to_string()],
800 "the snapshot-subprefix manifested part is kept"
801 );
802 }
803
804 #[test]
805 fn gc_orphans_deletes_a_terminal_runs_parts_even_while_a_run_is_active() {
806 let (store, _g) = fs_store(&[("base/f-000.parquet", b"x".to_vec())]);
810 let mut kv = keyed("base/manifest-f.json", "f", "f-000.parquet");
811 kv.1.status = ManifestStatus::Failed;
812 assert_eq!(gc_orphans(&store, "gs://b/base", &[kv], true).unwrap().0, 1);
813 }
814
815 #[test]
816 fn gc_orphans_spares_an_unmanifested_part_while_a_run_is_active() {
817 let (store, _g) = fs_store(&[
823 ("base/r1-000.parquet", b"aa".to_vec()), ("base/inflight.parquet", b"x".to_vec()), ]);
826 let keyed = vec![keyed("base/manifest-r1.json", "r1", "r1-000.parquet")];
827 let (removed, _) = gc_orphans(&store, "gs://b/base", &keyed, true).unwrap();
828 assert_eq!(
829 removed, 0,
830 "an unmanifested part is spared while a run is active"
831 );
832 assert!(
833 store
834 .list_files("base")
835 .unwrap()
836 .iter()
837 .any(|k| k.ends_with("inflight.parquet")),
838 "the live run's in-flight part must survive gc_orphans"
839 );
840 }
841
842 #[test]
843 fn gc_orphans_collects_an_unmanifested_part_when_no_run_is_active() {
844 let (store, _g) = fs_store(&[
847 ("base/r1-000.parquet", b"aa".to_vec()),
848 ("base/dead-orphan.parquet", b"x".to_vec()),
849 ]);
850 let keyed = vec![keyed("base/manifest-r1.json", "r1", "r1-000.parquet")];
851 assert_eq!(
852 gc_orphans(&store, "gs://b/base", &keyed, false).unwrap().0,
853 1,
854 "with no active run, an unmanifested orphan is collected"
855 );
856 }
857
858 fn running(run: &str, started_at: &str) -> (String, RunManifest) {
860 let mut m = manifest(run, 0, None);
861 m.status = ManifestStatus::Running;
862 m.started_at = started_at.into();
863 m.finished_at = String::new();
864 m.parts.clear();
865 (format!("base/manifest-{run}.json"), m)
866 }
867
868 #[test]
869 fn has_active_running_manifest_true_for_a_lone_running_marker() {
870 assert!(has_active_running_manifest(&[running(
871 "r1",
872 "2026-01-01T00:00:00Z"
873 )]));
874 }
875
876 #[test]
877 fn has_active_running_manifest_false_when_none_is_running() {
878 assert!(!has_active_running_manifest(&[keyed(
880 "k",
881 "r1",
882 "p.parquet"
883 )]));
884 }
885
886 #[test]
887 fn has_active_running_manifest_false_for_a_superseded_running_marker() {
888 let r1 = running("r1", "2026-01-01T00:00:00Z");
892 let mut r2 = manifest("r2", 10, Some(10)); r2.started_at = "2026-01-02T00:00:00Z".into();
894 assert!(!has_active_running_manifest(&[
895 r1,
896 ("base/manifest-r2.json".into(), r2)
897 ]));
898 }
899
900 #[test]
901 fn select_runs_drops_a_running_manifest() {
902 let run_marker = running("r_running", "2026-01-02T00:00:00Z");
906 let ok = keyed("base/manifest-r_ok.json", "r_ok", "r_ok-000.parquet"); let sel = select_runs(
908 vec![run_marker, ok],
909 &std::collections::HashSet::new(),
910 crate::load::plan::LoadMode::Incremental,
911 );
912 assert_eq!(sel.len(), 1, "only the Success run is selected");
913 assert_eq!(sel[0].1.run_id, "r_ok");
914 assert!(sel.iter().all(|(_, m)| m.status != ManifestStatus::Running));
915 }
916
917 #[test]
918 fn gc_orphans_removes_a_superseded_running_marker_but_spares_a_live_one() {
919 let (store, _g) = fs_store(&[
924 ("base/manifest-r1.json", b"{}".to_vec()), ("base/manifest-r2.json", b"{}".to_vec()), ]);
927 let r1 = running("r1", "2026-01-01T00:00:01Z");
928 let r2 = running("r2", "2026-01-01T00:00:02Z");
929 let (removed, _) = gc_orphans(&store, "gs://b/base", &[r1, r2], true).unwrap();
930 assert_eq!(removed, 1, "only the superseded running marker is removed");
931 let left = store.list_files("base").unwrap();
932 assert!(
933 !left.iter().any(|k| k.ends_with("manifest-r1.json")),
934 "the superseded (dead) running marker is deleted"
935 );
936 assert!(
937 left.iter().any(|k| k.ends_with("manifest-r2.json")),
938 "the live (non-superseded) running marker survives — it is the active signal"
939 );
940 }
941
942 fn keyed_at(run: &str, finished_at: &str) -> (String, RunManifest) {
943 let mut m = manifest(run, 100, None);
944 m.finished_at = finished_at.into();
945 (format!("base/manifest-{run}.json"), m)
946 }
947
948 #[test]
949 fn latest_full_picks_the_newest_snapshot_not_all() {
950 let keyed = vec![
953 keyed_at("r1", "2026-01-01T00:00:00Z"),
954 keyed_at("r3", "2026-01-03T00:00:00Z"),
955 keyed_at("r2", "2026-01-02T00:00:00Z"),
956 ];
957 let sel = latest_full(keyed);
958 assert_eq!(sel.len(), 1, "exactly one snapshot, never all");
959 assert_eq!(sel[0].1.run_id, "r3", "the newest by finished_at");
960 }
961
962 #[test]
963 fn latest_full_re_materializes_even_when_the_latest_is_already_loaded() {
964 let keyed = vec![
969 keyed_at("r1", "2026-01-01T00:00:00Z"),
970 keyed_at("r2", "2026-01-02T00:00:00Z"),
971 ];
972 let sel = latest_full(keyed);
974 assert_eq!(sel.len(), 1);
975 assert_eq!(sel[0].1.run_id, "r2", "always the latest, loaded or not");
976 }
977
978 #[test]
979 fn latest_full_of_no_staged_runs_is_empty_so_the_caller_no_ops_without_truncating() {
980 assert!(latest_full(Vec::new()).is_empty());
983 }
984
985 #[test]
986 fn latest_full_orders_by_parsed_instant_not_lexical_bytes() {
987 let keyed = vec![
990 keyed_at("older", "2026-01-01T00:00:00Z"),
991 keyed_at("newer", "2026-01-01T00:00:00.500Z"),
992 ];
993 let sel = latest_full(keyed);
994 assert_eq!(sel.len(), 1);
995 assert_eq!(
996 sel[0].1.run_id, "newer",
997 "the fractional-second run is the newer instant"
998 );
999 }
1000
1001 #[test]
1002 fn select_runs_full_picks_the_latest_even_when_loaded_and_even_when_stateless() {
1003 use crate::load::plan::LoadMode;
1004 use std::collections::HashSet;
1005 let keyed = vec![
1006 keyed_at("r1", "2026-01-01T00:00:00Z"),
1007 keyed_at("r2", "2026-01-02T00:00:00Z"),
1008 ];
1009 let loaded = HashSet::from(["r2".to_string()]);
1011 let sel = select_runs(keyed.clone(), &loaded, LoadMode::Full);
1012 assert_eq!(sel.len(), 1);
1013 assert_eq!(sel[0].1.run_id, "r2", "Full picks latest, loaded or not");
1014 let sel = select_runs(keyed, &HashSet::new(), LoadMode::Full);
1017 assert_eq!(sel.len(), 1, "stateless Full is not a blanket load");
1018 assert_eq!(sel[0].1.run_id, "r2");
1019 }
1020
1021 #[test]
1022 fn select_runs_append_modes_filter_loaded_and_load_all_when_stateless() {
1023 use crate::load::plan::LoadMode;
1024 use std::collections::HashSet;
1025 let keyed = vec![
1026 keyed_at("r1", "2026-01-01T00:00:00Z"),
1027 keyed_at("r2", "2026-01-02T00:00:00Z"),
1028 ];
1029 let loaded = HashSet::from(["r1".to_string()]);
1030 for mode in [LoadMode::Incremental, LoadMode::Cdc] {
1031 let sel = select_runs(keyed.clone(), &loaded, mode);
1032 assert_eq!(sel.len(), 1, "{mode:?}: only the unloaded run");
1033 assert_eq!(sel[0].1.run_id, "r2");
1034 assert_eq!(
1036 select_runs(keyed.clone(), &HashSet::new(), mode).len(),
1037 2,
1038 "{mode:?}: stateless loads every run"
1039 );
1040 }
1041 }
1042
1043 #[test]
1044 fn is_run_unique_manifest_needs_both_prefix_and_json() {
1045 assert!(is_run_unique_manifest("manifest-20260101T000000.json"));
1046 assert!(!is_run_unique_manifest("manifest.json")); assert!(!is_run_unique_manifest("manifest-abc.txt")); assert!(!is_run_unique_manifest("data.json")); }
1050
1051 #[test]
1052 fn chain_prefix_renders_source_and_files() {
1053 let known = LoadIntegrity {
1054 source_rows: Some(100),
1055 file_rows: 100,
1056 manifests: 1,
1057 };
1058 assert_eq!(known.chain_prefix(), "source 100 → files 100");
1059 let unknown = LoadIntegrity {
1060 source_rows: None,
1061 file_rows: 40,
1062 manifests: 1,
1063 };
1064 assert_eq!(unknown.chain_prefix(), "source ? → files 40");
1065 }
1066
1067 #[test]
1068 fn is_manifest_key_matches_only_the_final_segment() {
1069 assert!(is_manifest_key("gs://b/p/manifest.json"));
1070 assert!(is_manifest_key("manifest.json"));
1071 assert!(!is_manifest_key("gs://b/p/part-0.parquet"));
1072 assert!(!is_manifest_key("gs://b/p/x_manifest.json"));
1073 }
1074
1075 fn fs_store(files: &[(&str, Vec<u8>)]) -> (GcsStore, tempfile::TempDir) {
1081 let dir = tempfile::tempdir().unwrap();
1082 for (rel, bytes) in files {
1083 let p = dir.path().join(rel);
1084 std::fs::create_dir_all(p.parent().unwrap()).unwrap();
1085 std::fs::write(p, bytes).unwrap();
1086 }
1087 let store = GcsStore::open_fs(dir.path().to_str().unwrap()).unwrap();
1088 (store, dir)
1089 }
1090
1091 fn manifest_bytes(run: &str, rows: i64, source: Option<i64>) -> Vec<u8> {
1092 serde_json::to_vec(&manifest(run, rows, source)).unwrap()
1093 }
1094
1095 #[test]
1096 fn list_manifest_keys_prefers_run_unique_copies_over_the_canonical_pointer() {
1097 let (store, _g) = fs_store(&[
1101 ("base/manifest.json", b"{}".to_vec()),
1102 ("base/manifest-r1.json", b"{}".to_vec()),
1103 ("base/manifest-r2.json", b"{}".to_vec()),
1104 ("base/part-0.parquet", b"x".to_vec()), ]);
1106 let mut keys = list_manifest_keys(&store, "base").unwrap();
1107 keys.sort();
1108 assert_eq!(
1109 keys,
1110 vec![
1111 "base/manifest-r1.json".to_string(),
1112 "base/manifest-r2.json".to_string(),
1113 ]
1114 );
1115 }
1116
1117 #[test]
1118 fn list_manifest_keys_falls_back_to_the_canonical_name_for_a_single_run() {
1119 let (store, _g) = fs_store(&[("base/manifest.json", b"{}".to_vec())]);
1122 assert_eq!(
1123 list_manifest_keys(&store, "base").unwrap(),
1124 vec!["base/manifest.json".to_string()]
1125 );
1126 }
1127
1128 #[test]
1129 fn fetch_manifests_keyed_reads_and_parses_every_run_copy_under_the_prefix() {
1130 let (store, _g) = fs_store(&[
1134 ("base/manifest.json", manifest_bytes("r2", 40, Some(40))),
1135 (
1136 "base/manifest-r1.json",
1137 manifest_bytes("r1", 100, Some(100)),
1138 ),
1139 ("base/manifest-r2.json", manifest_bytes("r2", 40, Some(40))),
1140 ]);
1141 let manifests: Vec<_> = fetch_manifests_keyed(&store, "gs://my-bucket/base")
1142 .unwrap()
1143 .into_iter()
1144 .map(|(_, m)| m)
1145 .collect();
1146 assert_eq!(manifests.len(), 2);
1147 let integrity = reconcile(&manifests, false).unwrap();
1148 assert_eq!(integrity.file_rows, 140);
1149 assert_eq!(integrity.manifests, 2);
1150 }
1151
1152 #[test]
1153 fn fetch_manifests_keyed_names_the_key_when_a_manifest_is_unparseable() {
1154 let (store, _g) = fs_store(&[("base/manifest.json", b"{ not json".to_vec())]);
1155 let err = fetch_manifests_keyed(&store, "gs://my-bucket/base")
1156 .unwrap_err()
1157 .to_string();
1158 assert!(
1159 err.contains("parsing manifest") && err.contains("base/manifest.json"),
1160 "error should name the offending key: {err}"
1161 );
1162 }
1163
1164 const FAKE_GCS_ENDPOINT: &str = "http://127.0.0.1:4443";
1175 const FAKE_GCS_BUCKET: &str = "rivet-load-emulator";
1176
1177 fn fake_gcs_store() -> GcsStore {
1182 let created = std::process::Command::new("curl")
1183 .args([
1184 "-s",
1185 "-X",
1186 "POST",
1187 &format!("{FAKE_GCS_ENDPOINT}/storage/v1/b?project=rivet-test"),
1188 "-H",
1189 "Content-Type: application/json",
1190 "-d",
1191 &format!("{{\"name\":\"{FAKE_GCS_BUCKET}\"}}"),
1192 ])
1193 .output();
1194 assert!(
1195 created.is_ok_and(|o| o.status.success()),
1196 "could not reach fake-gcs to create the bucket — is `docker compose up -d fake-gcs` running on :4443?"
1197 );
1198 let cfg = crate::config::DestinationConfig {
1199 destination_type: crate::config::DestinationType::Gcs,
1200 bucket: Some(FAKE_GCS_BUCKET.into()),
1201 endpoint: Some(FAKE_GCS_ENDPOINT.into()),
1202 allow_anonymous: true,
1203 ..Default::default()
1204 };
1205 GcsStore::new(&cfg).expect("build GcsStore against fake-gcs")
1206 }
1207
1208 fn drain(store: &GcsStore, prefix: &str) {
1216 for key in store.list_files(prefix).unwrap() {
1217 store.remove(&key).unwrap();
1218 }
1219 }
1220
1221 #[test]
1222 #[ignore = "emulator: needs `docker compose up -d fake-gcs` (fsouza/fake-gcs-server :4443)"]
1223 fn storage_contract_over_fake_gcs() {
1224 let store = fake_gcs_store();
1225 let prefix = "load-contract/orders";
1228 drain(&store, prefix);
1229 let gs = format!("gs://{FAKE_GCS_BUCKET}/{prefix}");
1230
1231 store
1234 .put(
1235 &format!("{prefix}/manifest-r1.json"),
1236 &manifest_bytes("r1", 100, Some(100)),
1237 )
1238 .unwrap();
1239 store
1240 .put(&format!("{prefix}/part-000000.parquet"), b"rows-of-r1")
1241 .unwrap();
1242 store
1243 .put(&format!("{prefix}/orphan.parquet"), b"crash-leftover")
1244 .unwrap();
1245
1246 let keyed = fetch_manifests_keyed(&store, &gs).unwrap();
1248 assert_eq!(keyed.len(), 1, "the run's manifest, read back over GCS");
1249
1250 let manifests: Vec<_> = keyed.iter().map(|(_, m)| m.clone()).collect();
1252 assert_eq!(
1253 reconcile(&manifests, false).unwrap().file_rows,
1254 100,
1255 "file_rows drives the count-gate; a bad GCS read would corrupt it"
1256 );
1257
1258 assert_eq!(
1260 select_load_uris(&store, &gs, &keyed).unwrap(),
1261 vec![format!(
1262 "gs://{FAKE_GCS_BUCKET}/{prefix}/part-000000.parquet"
1263 )],
1264 "load pulls the manifested part, not the unmanifested crash orphan"
1265 );
1266
1267 let (removed, _bytes) = gc_orphans(&store, &gs, &keyed, false).unwrap();
1270 assert_eq!(
1271 removed, 1,
1272 "exactly the orphan parquet is GC'd over real GCS"
1273 );
1274 let mut left = store.list_files(prefix).unwrap();
1275 left.sort();
1276 assert_eq!(
1277 left,
1278 vec![
1279 format!("{prefix}/manifest-r1.json"),
1280 format!("{prefix}/part-000000.parquet"),
1281 ],
1282 "the manifested part + its manifest survive the orphan GC"
1283 );
1284
1285 drain(&store, prefix);
1289 assert!(
1290 store.list_files(prefix).unwrap().is_empty(),
1291 "teardown left the prefix clean over real GCS"
1292 );
1293 }
1294}