1use crate::destination::gcs::GcsStore;
27use crate::manifest::{MANIFEST_FILENAME, ManifestStatus, RunManifest};
28use anyhow::{Context, Result, bail};
29
30#[derive(Debug, Clone, PartialEq, Eq)]
34pub struct LoadIntegrity {
35 pub source_rows: Option<u64>,
40 pub file_rows: u64,
43 pub manifests: usize,
45}
46
47impl LoadIntegrity {
48 pub fn chain_prefix(&self) -> String {
52 let src = self
53 .source_rows
54 .map_or_else(|| "?".to_string(), |n| n.to_string());
55 format!("source {src} → files {}", self.file_rows)
56 }
57}
58
59#[allow(private_interfaces)]
76pub fn fetch_manifests_keyed(
77 store: &GcsStore,
78 gcs_prefix: &str,
79) -> Result<Vec<(String, RunManifest)>> {
80 let (_, base) = crate::load::split_gs_uri(gcs_prefix)?;
81 let keys = list_manifest_keys(store, base)?;
82 keys.into_iter()
83 .map(|key| {
84 let bytes = store.read(&key)?;
85 let m = serde_json::from_slice::<RunManifest>(&bytes)
86 .with_context(|| format!("parsing manifest {key}"))?;
87 Ok((key, m))
88 })
89 .collect()
90}
91
92#[allow(private_interfaces)]
96pub fn select_load_uris(
97 store: &GcsStore,
98 gcs_prefix: &str,
99 new: &[(String, RunManifest)],
100) -> Result<Vec<String>> {
101 let (bucket, base) = crate::load::split_gs_uri(gcs_prefix)?;
102 let all_parquet: Vec<String> = store
103 .list_files(base)?
104 .into_iter()
105 .filter(|k| k.ends_with(".parquet"))
106 .collect();
107 Ok(select_load_keys(new, &all_parquet)
108 .into_iter()
109 .map(|k| format!("gs://{bucket}/{k}"))
110 .collect())
111}
112
113fn resolve_parts<'a>(
119 manifest_key: &'a str,
120 m: &'a RunManifest,
121) -> impl Iterator<Item = String> + 'a {
122 let dir = manifest_key.rsplit_once('/').map(|(d, _)| d).unwrap_or("");
123 m.parts.iter().map(move |p| {
124 if dir.is_empty() {
125 p.path.clone()
126 } else {
127 format!("{dir}/{}", p.path)
128 }
129 })
130}
131
132pub fn select_load_keys(new: &[(String, RunManifest)], all_parquet: &[String]) -> Vec<String> {
143 use std::collections::BTreeSet;
144 let present: std::collections::HashSet<&str> = all_parquet.iter().map(String::as_str).collect();
145 let mut selected: BTreeSet<String> = BTreeSet::new();
146 for (key, m) in new {
147 let mut resolved_any = false;
148 for full in resolve_parts(key, m) {
149 if present.contains(full.as_str()) {
150 selected.insert(full);
151 resolved_any = true;
152 }
153 }
154 if !resolved_any {
155 return all_parquet.to_vec();
157 }
158 }
159 selected.into_iter().collect()
160}
161
162#[allow(private_interfaces)]
176pub fn gc_orphans(
177 store: &GcsStore,
178 gcs_prefix: &str,
179 keyed: &[(String, RunManifest)],
180) -> Result<(usize, u64)> {
181 let (_bucket, base) = crate::load::split_gs_uri(gcs_prefix)?;
182 let keep: std::collections::HashSet<String> = keyed
183 .iter()
184 .filter(|(_, m)| m.status == ManifestStatus::Success)
185 .flat_map(|(key, m)| resolve_parts(key, m))
186 .collect();
187 let mut removed = 0usize;
188 let mut removed_bytes = 0u64;
189 for key in store.list_files(base)? {
190 if key.ends_with(".parquet") && !keep.contains(&key) {
191 removed_bytes += store.stat_size(&key).unwrap_or(0);
192 store.remove(&key)?;
193 removed += 1;
194 }
195 }
196 Ok((removed, removed_bytes))
197}
198
199pub fn latest_full(keyed: Vec<(String, RunManifest)>) -> Vec<(String, RunManifest)> {
212 keyed
213 .into_iter()
214 .max_by(|a, b| {
215 match (
216 chrono::DateTime::parse_from_rfc3339(&a.1.finished_at).ok(),
217 chrono::DateTime::parse_from_rfc3339(&b.1.finished_at).ok(),
218 ) {
219 (Some(x), Some(y)) => x.cmp(&y),
220 _ => a.1.finished_at.cmp(&b.1.finished_at),
221 }
222 })
223 .into_iter()
224 .collect()
225}
226
227pub fn select_runs(
237 keyed: Vec<(String, RunManifest)>,
238 loaded: &std::collections::HashSet<String>,
239 mode: crate::load::plan::LoadMode,
240) -> Vec<(String, RunManifest)> {
241 match mode {
242 crate::load::plan::LoadMode::Full => latest_full(keyed),
243 _ => keyed
244 .into_iter()
245 .filter(|(_, m)| !loaded.contains(&m.run_id))
246 .collect(),
247 }
248}
249
250fn list_manifest_keys(store: &GcsStore, base: &str) -> Result<Vec<String>> {
252 let all: Vec<String> = store
253 .list_files(base)?
254 .into_iter()
255 .filter(|k| is_manifest_key(k))
256 .collect();
257 let run_unique: Vec<String> = all
265 .iter()
266 .filter(|k| is_run_unique_manifest(k.rsplit('/').next().unwrap_or("")))
267 .cloned()
268 .collect();
269 Ok(if run_unique.is_empty() {
270 all
271 } else {
272 run_unique
273 })
274}
275
276fn is_manifest_key(key: &str) -> bool {
281 let base = key.rsplit('/').next().unwrap_or("");
282 base == MANIFEST_FILENAME || is_run_unique_manifest(base)
283}
284
285fn is_run_unique_manifest(base: &str) -> bool {
288 base.starts_with("manifest-") && base.ends_with(".json")
289}
290
291pub fn reconcile(manifests: &[RunManifest], allow_source_drift: bool) -> Result<LoadIntegrity> {
306 if manifests.is_empty() {
307 bail!(
308 "no `{MANIFEST_FILENAME}` found under the export prefix — refusing to load \
309 unverified files. A rivet export writes a manifest on success; its absence \
310 means the run never completed (or points at the wrong prefix)."
311 );
312 }
313
314 let mut file_rows: u64 = 0;
315 let mut source_rows: u64 = 0;
316 let mut any_source = false;
317
318 for m in manifests {
319 if m.status != ManifestStatus::Success {
321 bail!(
322 "manifest for run `{}` (export `{}`) is {:?}, not Success — refusing to load a \
323 partial export",
324 m.run_id,
325 m.export_name,
326 m.status
327 );
328 }
329 m.validate_self_consistency().map_err(|e| {
332 anyhow::anyhow!(
333 "manifest for run `{}` (export `{}`) is internally inconsistent: {e} — refusing \
334 to load",
335 m.run_id,
336 m.export_name
337 )
338 })?;
339
340 let rows = u64::try_from(m.row_count).with_context(|| {
341 format!(
342 "manifest for run `{}` has a negative row_count ({})",
343 m.run_id, m.row_count
344 )
345 })?;
346 file_rows += rows;
347
348 if let Some(src) = m
351 .source
352 .extraction
353 .as_ref()
354 .and_then(|x| x.source_row_count)
355 {
356 let src = u64::try_from(src).with_context(|| {
357 format!(
358 "manifest for run `{}` has a negative source_row_count ({src})",
359 m.run_id
360 )
361 })?;
362 any_source = true;
363 source_rows += src;
364 if src != rows {
365 if allow_source_drift {
366 eprintln!(
367 "warning: source→file drift for run `{}` (export `{}`): source had {src} \
368 rows, extracted {rows} (--allow-source-drift)",
369 m.run_id, m.export_name
370 );
371 } else {
372 bail!(
373 "source→file mismatch for run `{}` (export `{}`): source had {src} rows \
374 but {rows} were extracted — the extract dropped {} row(s). Investigate \
375 before loading, or pass --allow-source-drift to override.",
376 m.run_id,
377 m.export_name,
378 src.abs_diff(rows)
379 );
380 }
381 }
382 }
383 }
384
385 Ok(LoadIntegrity {
386 source_rows: any_source.then_some(source_rows),
387 file_rows,
388 manifests: manifests.len(),
389 })
390}
391
392#[cfg(test)]
393mod tests {
394 use super::*;
395 use crate::manifest::{
396 ExtractionMetadata, ManifestDestination, ManifestPart, ManifestSource, PartStatus,
397 };
398
399 fn manifest(run: &str, rows: i64, source: Option<i64>) -> RunManifest {
402 RunManifest {
403 manifest_version: crate::manifest::MANIFEST_VERSION,
404 run_id: run.into(),
405 export_name: "orders".into(),
406 mode: "batch".into(),
407 started_at: "t".into(),
408 finished_at: "t".into(),
409 status: ManifestStatus::Success,
410 source: ManifestSource {
411 engine: "pg".into(),
412 schema: None,
413 table: None,
414 extraction: source.map(|n| ExtractionMetadata {
415 strategy: "full".into(),
416 cursor_column: None,
417 cursor_type: None,
418 cursor_low: None,
419 cursor_high: None,
420 source_row_count: Some(n),
421 }),
422 },
423 destination: ManifestDestination {
424 kind: "gcs".into(),
425 uri: "gs://b/p".into(),
426 },
427 format: "parquet".into(),
428 compression: "zstd".into(),
429 schema_fingerprint: "xxh3:0".into(),
430 row_count: rows,
431 part_count: 1,
432 parts: vec![ManifestPart {
433 part_id: 0,
434 path: "part-000000.parquet".into(),
435 rows,
436 size_bytes: 1,
437 content_fingerprint: "xxh3:0".into(),
438 content_md5: String::new(),
439 status: PartStatus::Committed,
440 }],
441 column_checksums: None,
442 checksum_key_column: None,
443 }
444 }
445
446 #[test]
447 fn sums_file_and_source_rows_across_manifests() {
448 let ms = vec![manifest("r1", 100, Some(100)), manifest("r2", 40, Some(40))];
449 let got = reconcile(&ms, false).unwrap();
450 assert_eq!(got.file_rows, 140);
451 assert_eq!(got.source_rows, Some(140));
452 assert_eq!(got.manifests, 2);
453 }
454
455 #[test]
456 fn source_rows_is_none_when_no_manifest_probed_the_source() {
457 let ms = vec![manifest("r1", 100, None), manifest("r2", 40, None)];
458 let got = reconcile(&ms, false).unwrap();
459 assert_eq!(got.file_rows, 140);
460 assert_eq!(
461 got.source_rows, None,
462 "unprobed source is unknown, not zero"
463 );
464 }
465
466 #[test]
467 fn source_rows_present_even_if_only_some_manifests_probed() {
468 let ms = vec![manifest("r1", 100, Some(100)), manifest("r2", 40, None)];
469 let got = reconcile(&ms, false).unwrap();
470 assert_eq!(got.source_rows, Some(100));
471 }
472
473 #[test]
474 fn empty_manifests_refuses_to_load() {
475 let err = reconcile(&[], false).unwrap_err().to_string();
476 assert!(err.contains("refusing to load"), "{err}");
477 }
478
479 #[test]
480 fn non_success_manifest_refuses_to_load() {
481 let mut m = manifest("r1", 100, Some(100));
482 m.status = ManifestStatus::Interrupted;
483 let err = reconcile(&[m], false).unwrap_err().to_string();
484 assert!(err.contains("not Success"), "{err}");
485 }
486
487 #[test]
488 fn self_inconsistent_manifest_refuses_to_load() {
489 let mut m = manifest("r1", 100, Some(100));
492 m.row_count = 999; let err = reconcile(&[m], false).unwrap_err().to_string();
494 assert!(err.contains("inconsistent"), "{err}");
495 }
496
497 #[test]
498 fn source_file_mismatch_hard_fails_by_default() {
499 let m = manifest("r1", 100, Some(120));
501 let err = reconcile(&[m], false).unwrap_err().to_string();
502 assert!(err.contains("source→file mismatch"), "{err}");
503 assert!(err.contains("dropped 20"), "{err}");
504 }
505
506 #[test]
507 fn source_file_mismatch_is_allowed_under_the_override() {
508 let m = manifest("r1", 100, Some(120));
509 let got = reconcile(&[m], true).expect("--allow-source-drift proceeds");
510 assert_eq!(got.file_rows, 100);
511 assert_eq!(
512 got.source_rows,
513 Some(120),
514 "the probed source count is still surfaced"
515 );
516 }
517
518 fn keyed(key: &str, run: &str, part: &str) -> (String, RunManifest) {
520 let mut m = manifest(run, 10, Some(10));
521 m.parts[0].path = part.into();
522 (key.to_string(), m)
523 }
524
525 #[test]
526 fn select_load_keys_picks_only_the_new_runs_parts() {
527 let all = vec![
529 "base/r1-000.parquet".to_string(),
530 "base/r2-000.parquet".to_string(),
531 ];
532 let new = vec![keyed("base/manifest-r2.json", "r2", "r2-000.parquet")];
533 assert_eq!(
534 select_load_keys(&new, &all),
535 vec!["base/r2-000.parquet".to_string()],
536 "loads r2's part only — not r1's already-loaded file"
537 );
538 }
539
540 #[test]
541 fn select_load_uris_lists_and_re_prefixes_selected_keys_with_the_source_bucket() {
542 let (store, _g) = fs_store(&[
548 ("base/r1-000.parquet", b"a".to_vec()),
549 ("base/manifest-r1.json", b"{}".to_vec()), ]);
551 let new = vec![keyed("base/manifest-r1.json", "r1", "r1-000.parquet")];
552 assert_eq!(
553 select_load_uris(&store, "gs://my-bucket/base", &new).unwrap(),
554 vec!["gs://my-bucket/base/r1-000.parquet".to_string()],
555 "the selected key, re-prefixed with the source bucket — never the manifest"
556 );
557 }
558
559 #[test]
560 fn select_load_keys_resolves_a_snapshot_subprefix_manifest() {
561 let all = vec!["base/snapshot/snap-000.parquet".to_string()];
564 let new = vec![keyed(
565 "base/snapshot/manifest-r1.json",
566 "r1",
567 "snap-000.parquet",
568 )];
569 assert_eq!(
570 select_load_keys(&new, &all),
571 vec!["base/snapshot/snap-000.parquet".to_string()]
572 );
573 }
574
575 #[test]
576 fn select_load_keys_falls_back_to_full_listing_when_a_manifest_has_no_present_part() {
577 let all = vec!["base/a.parquet".to_string(), "base/b.parquet".to_string()];
580 let new = vec![keyed("base/manifest-r1.json", "r1", "missing.parquet")];
581 assert_eq!(
582 select_load_keys(&new, &all),
583 all,
584 "unresolvable part → blanket fallback"
585 );
586 }
587
588 #[test]
589 fn select_load_keys_empty_new_set_selects_nothing_never_the_full_listing() {
590 let all = vec![
595 "base/r1-000.parquet".to_string(),
596 "base/r2-000.parquet".to_string(),
597 ];
598 assert!(
599 select_load_keys(&[], &all).is_empty(),
600 "no new runs ⇒ load nothing, not everything"
601 );
602 }
603
604 #[test]
605 fn gc_orphans_removes_unmanifested_parquet_only() {
606 let (store, _g) = fs_store(&[
607 ("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()), ]);
612 let keyed = vec![keyed("base/manifest-r1.json", "r1", "r1-000.parquet")];
613 let (removed, bytes) = gc_orphans(&store, "gs://b/base", &keyed).unwrap();
614 assert_eq!(removed, 1, "only the unmanifested part is removed");
615 assert_eq!(bytes, 4, "'junk' is 4 bytes");
616 let mut left = store.list_files("base").unwrap();
617 left.sort();
618 assert_eq!(
619 left,
620 vec![
621 "base/_SUCCESS".to_string(),
622 "base/manifest.json".to_string(),
623 "base/r1-000.parquet".to_string(),
624 ],
625 "the manifested part, the manifest, and _SUCCESS all survive"
626 );
627 }
628
629 #[test]
630 fn gc_orphans_of_an_all_manifested_prefix_removes_nothing() {
631 let (store, _g) = fs_store(&[("base/r1-000.parquet", b"a".to_vec())]);
632 let keyed = vec![keyed("base/manifest-r1.json", "r1", "r1-000.parquet")];
633 assert_eq!(gc_orphans(&store, "gs://b/base", &keyed).unwrap().0, 0);
634 }
635
636 #[test]
637 fn gc_orphans_keeps_a_snapshot_subprefix_manifested_part() {
638 let (store, _g) = fs_store(&[
639 ("base/snapshot/s-000.parquet", b"a".to_vec()), ("base/orphan.parquet", b"x".to_vec()), ]);
642 let keyed = vec![keyed("base/snapshot/manifest-s.json", "s", "s-000.parquet")];
643 let (removed, _) = gc_orphans(&store, "gs://b/base", &keyed).unwrap();
644 assert_eq!(removed, 1, "the top-level orphan goes");
645 assert_eq!(
646 store.list_files("base/snapshot").unwrap(),
647 vec!["base/snapshot/s-000.parquet".to_string()],
648 "the snapshot-subprefix manifested part is kept"
649 );
650 }
651
652 #[test]
653 fn gc_orphans_does_not_protect_a_failed_runs_parts() {
654 let (store, _g) = fs_store(&[("base/f-000.parquet", b"x".to_vec())]);
657 let mut kv = keyed("base/manifest-f.json", "f", "f-000.parquet");
658 kv.1.status = ManifestStatus::Failed;
659 assert_eq!(gc_orphans(&store, "gs://b/base", &[kv]).unwrap().0, 1);
660 }
661
662 fn keyed_at(run: &str, finished_at: &str) -> (String, RunManifest) {
663 let mut m = manifest(run, 100, None);
664 m.finished_at = finished_at.into();
665 (format!("base/manifest-{run}.json"), m)
666 }
667
668 #[test]
669 fn latest_full_picks_the_newest_snapshot_not_all() {
670 let keyed = vec![
673 keyed_at("r1", "2026-01-01T00:00:00Z"),
674 keyed_at("r3", "2026-01-03T00:00:00Z"),
675 keyed_at("r2", "2026-01-02T00:00:00Z"),
676 ];
677 let sel = latest_full(keyed);
678 assert_eq!(sel.len(), 1, "exactly one snapshot, never all");
679 assert_eq!(sel[0].1.run_id, "r3", "the newest by finished_at");
680 }
681
682 #[test]
683 fn latest_full_re_materializes_even_when_the_latest_is_already_loaded() {
684 let keyed = vec![
689 keyed_at("r1", "2026-01-01T00:00:00Z"),
690 keyed_at("r2", "2026-01-02T00:00:00Z"),
691 ];
692 let sel = latest_full(keyed);
694 assert_eq!(sel.len(), 1);
695 assert_eq!(sel[0].1.run_id, "r2", "always the latest, loaded or not");
696 }
697
698 #[test]
699 fn latest_full_of_no_staged_runs_is_empty_so_the_caller_no_ops_without_truncating() {
700 assert!(latest_full(Vec::new()).is_empty());
703 }
704
705 #[test]
706 fn latest_full_orders_by_parsed_instant_not_lexical_bytes() {
707 let keyed = vec![
710 keyed_at("older", "2026-01-01T00:00:00Z"),
711 keyed_at("newer", "2026-01-01T00:00:00.500Z"),
712 ];
713 let sel = latest_full(keyed);
714 assert_eq!(sel.len(), 1);
715 assert_eq!(
716 sel[0].1.run_id, "newer",
717 "the fractional-second run is the newer instant"
718 );
719 }
720
721 #[test]
722 fn select_runs_full_picks_the_latest_even_when_loaded_and_even_when_stateless() {
723 use crate::load::plan::LoadMode;
724 use std::collections::HashSet;
725 let keyed = vec![
726 keyed_at("r1", "2026-01-01T00:00:00Z"),
727 keyed_at("r2", "2026-01-02T00:00:00Z"),
728 ];
729 let loaded = HashSet::from(["r2".to_string()]);
731 let sel = select_runs(keyed.clone(), &loaded, LoadMode::Full);
732 assert_eq!(sel.len(), 1);
733 assert_eq!(sel[0].1.run_id, "r2", "Full picks latest, loaded or not");
734 let sel = select_runs(keyed, &HashSet::new(), LoadMode::Full);
737 assert_eq!(sel.len(), 1, "stateless Full is not a blanket load");
738 assert_eq!(sel[0].1.run_id, "r2");
739 }
740
741 #[test]
742 fn select_runs_append_modes_filter_loaded_and_load_all_when_stateless() {
743 use crate::load::plan::LoadMode;
744 use std::collections::HashSet;
745 let keyed = vec![
746 keyed_at("r1", "2026-01-01T00:00:00Z"),
747 keyed_at("r2", "2026-01-02T00:00:00Z"),
748 ];
749 let loaded = HashSet::from(["r1".to_string()]);
750 for mode in [LoadMode::Incremental, LoadMode::Cdc] {
751 let sel = select_runs(keyed.clone(), &loaded, mode);
752 assert_eq!(sel.len(), 1, "{mode:?}: only the unloaded run");
753 assert_eq!(sel[0].1.run_id, "r2");
754 assert_eq!(
756 select_runs(keyed.clone(), &HashSet::new(), mode).len(),
757 2,
758 "{mode:?}: stateless loads every run"
759 );
760 }
761 }
762
763 #[test]
764 fn is_run_unique_manifest_needs_both_prefix_and_json() {
765 assert!(is_run_unique_manifest("manifest-20260101T000000.json"));
766 assert!(!is_run_unique_manifest("manifest.json")); assert!(!is_run_unique_manifest("manifest-abc.txt")); assert!(!is_run_unique_manifest("data.json")); }
770
771 #[test]
772 fn chain_prefix_renders_source_and_files() {
773 let known = LoadIntegrity {
774 source_rows: Some(100),
775 file_rows: 100,
776 manifests: 1,
777 };
778 assert_eq!(known.chain_prefix(), "source 100 → files 100");
779 let unknown = LoadIntegrity {
780 source_rows: None,
781 file_rows: 40,
782 manifests: 1,
783 };
784 assert_eq!(unknown.chain_prefix(), "source ? → files 40");
785 }
786
787 #[test]
788 fn is_manifest_key_matches_only_the_final_segment() {
789 assert!(is_manifest_key("gs://b/p/manifest.json"));
790 assert!(is_manifest_key("manifest.json"));
791 assert!(!is_manifest_key("gs://b/p/part-0.parquet"));
792 assert!(!is_manifest_key("gs://b/p/x_manifest.json"));
793 }
794
795 fn fs_store(files: &[(&str, Vec<u8>)]) -> (GcsStore, tempfile::TempDir) {
801 let dir = tempfile::tempdir().unwrap();
802 for (rel, bytes) in files {
803 let p = dir.path().join(rel);
804 std::fs::create_dir_all(p.parent().unwrap()).unwrap();
805 std::fs::write(p, bytes).unwrap();
806 }
807 let store = GcsStore::open_fs(dir.path().to_str().unwrap()).unwrap();
808 (store, dir)
809 }
810
811 fn manifest_bytes(run: &str, rows: i64, source: Option<i64>) -> Vec<u8> {
812 serde_json::to_vec(&manifest(run, rows, source)).unwrap()
813 }
814
815 #[test]
816 fn list_manifest_keys_prefers_run_unique_copies_over_the_canonical_pointer() {
817 let (store, _g) = fs_store(&[
821 ("base/manifest.json", b"{}".to_vec()),
822 ("base/manifest-r1.json", b"{}".to_vec()),
823 ("base/manifest-r2.json", b"{}".to_vec()),
824 ("base/part-0.parquet", b"x".to_vec()), ]);
826 let mut keys = list_manifest_keys(&store, "base").unwrap();
827 keys.sort();
828 assert_eq!(
829 keys,
830 vec![
831 "base/manifest-r1.json".to_string(),
832 "base/manifest-r2.json".to_string(),
833 ]
834 );
835 }
836
837 #[test]
838 fn list_manifest_keys_falls_back_to_the_canonical_name_for_a_single_run() {
839 let (store, _g) = fs_store(&[("base/manifest.json", b"{}".to_vec())]);
842 assert_eq!(
843 list_manifest_keys(&store, "base").unwrap(),
844 vec!["base/manifest.json".to_string()]
845 );
846 }
847
848 #[test]
849 fn fetch_manifests_keyed_reads_and_parses_every_run_copy_under_the_prefix() {
850 let (store, _g) = fs_store(&[
854 ("base/manifest.json", manifest_bytes("r2", 40, Some(40))),
855 (
856 "base/manifest-r1.json",
857 manifest_bytes("r1", 100, Some(100)),
858 ),
859 ("base/manifest-r2.json", manifest_bytes("r2", 40, Some(40))),
860 ]);
861 let manifests: Vec<_> = fetch_manifests_keyed(&store, "gs://my-bucket/base")
862 .unwrap()
863 .into_iter()
864 .map(|(_, m)| m)
865 .collect();
866 assert_eq!(manifests.len(), 2);
867 let integrity = reconcile(&manifests, false).unwrap();
868 assert_eq!(integrity.file_rows, 140);
869 assert_eq!(integrity.manifests, 2);
870 }
871
872 #[test]
873 fn fetch_manifests_keyed_names_the_key_when_a_manifest_is_unparseable() {
874 let (store, _g) = fs_store(&[("base/manifest.json", b"{ not json".to_vec())]);
875 let err = fetch_manifests_keyed(&store, "gs://my-bucket/base")
876 .unwrap_err()
877 .to_string();
878 assert!(
879 err.contains("parsing manifest") && err.contains("base/manifest.json"),
880 "error should name the offending key: {err}"
881 );
882 }
883
884 const FAKE_GCS_ENDPOINT: &str = "http://127.0.0.1:4443";
895 const FAKE_GCS_BUCKET: &str = "rivet-load-emulator";
896
897 fn fake_gcs_store() -> GcsStore {
902 let created = std::process::Command::new("curl")
903 .args([
904 "-s",
905 "-X",
906 "POST",
907 &format!("{FAKE_GCS_ENDPOINT}/storage/v1/b?project=rivet-test"),
908 "-H",
909 "Content-Type: application/json",
910 "-d",
911 &format!("{{\"name\":\"{FAKE_GCS_BUCKET}\"}}"),
912 ])
913 .output();
914 assert!(
915 created.is_ok_and(|o| o.status.success()),
916 "could not reach fake-gcs to create the bucket — is `docker compose up -d fake-gcs` running on :4443?"
917 );
918 let cfg = crate::config::DestinationConfig {
919 destination_type: crate::config::DestinationType::Gcs,
920 bucket: Some(FAKE_GCS_BUCKET.into()),
921 endpoint: Some(FAKE_GCS_ENDPOINT.into()),
922 allow_anonymous: true,
923 ..Default::default()
924 };
925 GcsStore::new(&cfg).expect("build GcsStore against fake-gcs")
926 }
927
928 fn drain(store: &GcsStore, prefix: &str) {
936 for key in store.list_files(prefix).unwrap() {
937 store.remove(&key).unwrap();
938 }
939 }
940
941 #[test]
942 #[ignore = "emulator: needs `docker compose up -d fake-gcs` (fsouza/fake-gcs-server :4443)"]
943 fn storage_contract_over_fake_gcs() {
944 let store = fake_gcs_store();
945 let prefix = "load-contract/orders";
948 drain(&store, prefix);
949 let gs = format!("gs://{FAKE_GCS_BUCKET}/{prefix}");
950
951 store
954 .put(
955 &format!("{prefix}/manifest-r1.json"),
956 &manifest_bytes("r1", 100, Some(100)),
957 )
958 .unwrap();
959 store
960 .put(&format!("{prefix}/part-000000.parquet"), b"rows-of-r1")
961 .unwrap();
962 store
963 .put(&format!("{prefix}/orphan.parquet"), b"crash-leftover")
964 .unwrap();
965
966 let keyed = fetch_manifests_keyed(&store, &gs).unwrap();
968 assert_eq!(keyed.len(), 1, "the run's manifest, read back over GCS");
969
970 let manifests: Vec<_> = keyed.iter().map(|(_, m)| m.clone()).collect();
972 assert_eq!(
973 reconcile(&manifests, false).unwrap().file_rows,
974 100,
975 "file_rows drives the count-gate; a bad GCS read would corrupt it"
976 );
977
978 assert_eq!(
980 select_load_uris(&store, &gs, &keyed).unwrap(),
981 vec![format!(
982 "gs://{FAKE_GCS_BUCKET}/{prefix}/part-000000.parquet"
983 )],
984 "load pulls the manifested part, not the unmanifested crash orphan"
985 );
986
987 let (removed, _bytes) = gc_orphans(&store, &gs, &keyed).unwrap();
990 assert_eq!(
991 removed, 1,
992 "exactly the orphan parquet is GC'd over real GCS"
993 );
994 let mut left = store.list_files(prefix).unwrap();
995 left.sort();
996 assert_eq!(
997 left,
998 vec![
999 format!("{prefix}/manifest-r1.json"),
1000 format!("{prefix}/part-000000.parquet"),
1001 ],
1002 "the manifested part + its manifest survive the orphan GC"
1003 );
1004
1005 drain(&store, prefix);
1009 assert!(
1010 store.list_files(prefix).unwrap().is_empty(),
1011 "teardown left the prefix clean over real GCS"
1012 );
1013 }
1014}