Skip to main content

rivet/load/
reconcile.rs

1//! End-to-end load integrity — reconcile **source → file → warehouse** row
2//! counts before (and after) a load.
3//!
4//! The OSS engine records, in each run's `manifest.json`, two of the three legs
5//! of the chain: how many rows the source held at extraction time
6//! ([`ExtractionMetadata::source_row_count`](crate::manifest::ExtractionMetadata),
7//! when cheaply probed) and how many rows were actually written to files
8//! ([`RunManifest::row_count`]). The warehouse leg — how many rows the load
9//! landed — is the loaders' post-load `COUNT(*)`, enforced by their
10//! `expected_rows` gate.
11//!
12//! Until now that gate was dead in the production path: nothing read the
13//! manifest, so `rivet load` trusted whatever Parquet happened to sit under
14//! the prefix ("file in a bucket"). This module closes the loop:
15//!
16//! 1. read every `manifest.json` under the export's GCS prefix,
17//! 2. **refuse** to load unless each is a self-consistent `Success` run whose
18//!    source count (when known) matches what it extracted, and
19//! 3. return the summed, authoritative `file_rows` the loader's count gate then
20//!    checks against the warehouse.
21//!
22//! It is the value the OSS core deliberately stops short of — file-level
23//! integrity is free (`rivet validate`); reconciling that the rows *arrived in
24//! the warehouse* is the paid last mile.
25
26use crate::destination::gcs::GcsStore;
27use crate::manifest::{MANIFEST_FILENAME, ManifestStatus, RunManifest};
28use anyhow::{Context, Result, bail};
29
30/// The reconciled row-count chain for one export's load, derived from the run
31/// manifests under its GCS prefix. `file_rows` is what the warehouse must end
32/// up holding; the loader's `expected_rows` gate enforces `warehouse == file`.
33#[derive(Debug, Clone, PartialEq, Eq)]
34pub struct LoadIntegrity {
35    /// Rows the source held at extraction time, summed over the manifests that
36    /// probed it. `None` when *no* contributing manifest carried a
37    /// `source_row_count` — the source→file leg is then unverifiable from the
38    /// manifest alone (e.g. a full snapshot that did not count the source).
39    pub source_rows: Option<u64>,
40    /// Rows written to files — the sum of the trustworthy manifests'
41    /// `row_count`. This is the authoritative expected warehouse row count.
42    pub file_rows: u64,
43    /// How many run manifests contributed to the totals.
44    pub manifests: usize,
45}
46
47impl LoadIntegrity {
48    /// A one-line human summary of the chain, e.g.
49    /// `source 1000 → files 1000 → (warehouse pending)`. The warehouse leg is
50    /// filled in by the caller once the load's `COUNT(*)` is known.
51    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/// Fetch and parse every `manifest.json` under `gcs_prefix` (recursive), keeping
60/// each manifest's bucket-relative storage key — needed to resolve a manifest's
61/// (relative) part paths back to full object keys for per-run loading (see
62/// [`select_load_uris`]).
63///
64/// A rivet export writes one manifest per `run_id`; a prefix that has
65/// accumulated several incremental / CDC runs holds several. Transport is the
66/// native opendal client the extraction destination already uses (`store`), so
67/// auth is the export's own GCS credentials and this is offline-testable over a
68/// filesystem-backed store.
69///
70/// `pub` (a public-API root the lib keeps alive) even though its only caller is
71/// the binary-only `cli::dispatch`; `#[allow(private_interfaces)]` because the
72/// injected `GcsStore` is deliberately an internal (`pub(crate)`) type — the
73/// `destination` module stays crate-private. Same rationale as the `preflight`
74/// module note in `lib.rs`.
75#[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/// Full `gs://` URIs of the parquet to load for `new` (the not-yet-loaded run
93/// manifests), preferring each manifest's own parts over a blanket listing.
94/// See [`select_load_keys`] for the selection rule.
95#[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
113/// The bucket-relative part keys a manifest declares, each resolved against the
114/// manifest's own directory: `<dir(manifest_key)>/<part.path>`. Shared by
115/// [`select_load_keys`] (which intersects them with what's present) and
116/// [`gc_orphans`] (which treats them as the keep-set), so the two can't drift on
117/// how a manifest maps to its files.
118fn 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
132/// Pure selection: which bucket-relative parquet keys to load for the given
133/// (not-yet-loaded) run manifests.
134///
135/// Prefers each manifest's own parts (via [`resolve_parts`]) intersected with
136/// `all_parquet` — so a load pulls exactly the new runs' files, not every object
137/// under the prefix (the key to incremental loads once `cleanup_source` no
138/// longer wipes the bucket). Falls back to the whole `all_parquet` listing when
139/// ANY new manifest resolves to no present part (legacy/part-less manifests
140/// still load, at the cost of not pruning); the row-count gate then still guards
141/// correctness.
142pub 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            // Can't resolve this run to files — don't risk a partial selection.
156            return all_parquet.to_vec();
157        }
158    }
159    selected.into_iter().collect()
160}
161
162/// Delete every `.parquet` under `gcs_prefix` that no **`Success`** manifest
163/// references — crash leftovers from an interrupted extract (a run killed before
164/// it wrote its manifest leaves orphan parts the load already ignores, but which
165/// accumulate). Keeps every manifest, `_SUCCESS`, and every manifested part
166/// (including a `snapshot/` sub-prefix's, since `keyed` is fetched recursively).
167/// Strictly gentler than `cleanup_source`, which wipes the whole prefix.
168///
169/// ⚠️ INVARIANT: it cannot distinguish a crash orphan from a *live* extract's
170/// committed-but-not-yet-manifested parts (both are unmanifested `.parquet`), and
171/// there is NO age/lease guard. Only run a load with `gc_orphans` when no extract
172/// is writing the same prefix — the normal pipeline (a load AFTER a completed
173/// extract) satisfies this; a load fired while a `rivet run` streams into the same
174/// prefix would delete its in-flight parts. Returns `(removed_count, removed_bytes)`.
175#[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
199/// Full/chunked loads care only about the LATEST snapshot: from `keyed` (all run
200/// manifests under the prefix) pick the newest by `finished_at`. Full loads
201/// OVERWRITE, so exactly ONE snapshot may be loaded (loading every accumulated
202/// run would duplicate rows). Deliberately **not** ledger-gated: full
203/// re-materializes the latest snapshot on *every* load, so a re-load self-heals a
204/// drifted target and stays resilient to hidden in-place source updates. An empty
205/// input (no staged run — e.g. the staging was cleaned) yields an empty
206/// selection, so the caller no-ops WITHOUT truncating the target.
207///
208/// Ordering parses `finished_at` as an instant — a lexical byte compare mis-picks
209/// on mixed RFC3339 precision (`…00.5Z` sorts before `…00Z`) — and falls back to
210/// lexical only if a timestamp fails to parse, so a malformed manifest can't panic.
211pub 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
227/// Which run manifests to load for `mode`, given the ledger's already-`loaded`
228/// run_ids. The single mode→selection decision, pure and testable so the load's
229/// central invariant isn't buried in dispatch I/O:
230/// - `Full` → the single LATEST run ([`latest_full`]) — ledger-INDEPENDENT, so a
231///   re-load re-materializes the snapshot (self-heal) and, crucially, the
232///   STATELESS path (empty `loaded`) still picks the latest instead of blanket-
233///   loading every accumulated snapshot (the duplicate-rows bug).
234/// - `Incremental`/`Cdc` → every run not yet in `loaded` (append). Stateless →
235///   `loaded` empty → load all; at-least-once, absorbed by the dedup view.
236pub 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
250/// Bucket-relative keys of every run manifest under `base` (recursive).
251fn 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    // A run into a shared prefix leaves BOTH the canonical `manifest.json`
258    // (last-writer-wins — a pointer to the LATEST run) and an immutable
259    // `manifest-<run_id>.json` copy (one per run). Sum the per-run copies so a
260    // prefix that accumulated several CDC/incremental cycles counts EVERY run;
261    // counting the canonical pointer too would double-count the latest. Fall
262    // back to the canonical name only when no per-run copy exists (a single
263    // batch run, or a legacy prefix predating the run-unique copy).
264    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
276/// A listed key is a run manifest iff its final path segment is the canonical
277/// [`MANIFEST_FILENAME`] (`manifest.json`) or a run-unique copy
278/// (`manifest-<run_id>.json`) — so a data file merely *named* like it (an
279/// unlikely `…/x_manifest.json`) is not mistaken for one.
280fn 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
285/// A per-run manifest copy: `manifest-<token>.json` (the sidecar the OSS sink
286/// writes alongside the canonical pointer so cross-run reconcile can sum it).
287fn is_run_unique_manifest(base: &str) -> bool {
288    base.starts_with("manifest-") && base.ends_with(".json")
289}
290
291/// Reconcile a run's manifests into the authoritative expected warehouse row
292/// count, refusing to load anything that is not provably complete.
293///
294/// Every manifest must be a **`Success`** run (a `Failed` / `Interrupted`
295/// manifest describes a partial, untrustworthy export) and **self-consistent**
296/// (`row_count` == sum of committed parts — a mismatch is a writer bug). When a
297/// manifest recorded a `source_row_count`, the **source→file** leg must
298/// reconcile too: `source_row_count == row_count`, or the extract silently
299/// dropped rows. `allow_source_drift` downgrades only that last check to a
300/// warning (e.g. an incremental cursor window whose source moved under it);
301/// the completeness and self-consistency gates never yield.
302///
303/// Returns [`LoadIntegrity`] with the summed `file_rows` the loader's
304/// `expected_rows` gate then checks against the warehouse's `COUNT(*)`.
305pub 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        // Completeness: only a `Success` run may be loaded.
320        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        // Self-consistency: the recorded aggregates must match the committed
330        // parts (a divergence is a writer bug, per OSS `validate_self_consistency`).
331        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        // Source→file: reconcile the extract against the source when the run
349        // probed it. `None` = not probed (unverifiable, not a failure).
350        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    /// A minimal `Success` manifest with one committed part of `rows` rows and,
400    /// optionally, a probed `source_row_count`.
401    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        // row_count claims 100 but the only committed part holds 100 → make the
490        // aggregate lie by bumping the recorded row_count only.
491        let mut m = manifest("r1", 100, Some(100));
492        m.row_count = 999; // no longer equals committed parts' sum (100)
493        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        // Source had 120, only 100 extracted → 20 rows silently dropped.
500        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    /// A `(manifest_key, manifest)` pair with a single part named `part`.
519    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        // Two runs' files sit under the prefix; only r2 is "new".
528        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        // select_load_keys (the pure selection) is unit-tested throughout; this
543        // pins the store WRAPPER around it — list `.parquet` under the base, then
544        // reconstruct each `gs://<bucket>/<key>` the loader COPYs from. A mangled
545        // wrapper (empty/garbage URIs) would send the warehouse load at nothing,
546        // or the wrong object — invisible to the pure-selection tests.
547        let (store, _g) = fs_store(&[
548            ("base/r1-000.parquet", b"a".to_vec()),
549            ("base/manifest-r1.json", b"{}".to_vec()), // not .parquet — must be skipped
550        ]);
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        // A snapshot manifest lives under `base/snapshot/`; its part is relative
562        // to that dir. Resolution must reconstruct the full key.
563        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        // A manifest whose part isn't in the listing (legacy / renamed) → don't
578        // risk a partial selection; load everything under the prefix.
579        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        // When every run is already in the ledger the "new" set is empty. That
591        // must resolve to ZERO uris — NOT the blanket fallback, which would
592        // re-load the whole prefix on every up-to-date run (double-load). The
593        // fallback fires only for an unresolvable NON-empty manifest.
594        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()),   // manifested
608            ("base/orphan.parquet", b"junk".to_vec()), // crash leftover — no manifest
609            ("base/manifest.json", b"{}".to_vec()),    // kept — not a .parquet
610            ("base/_SUCCESS", b"".to_vec()),           // kept — not a .parquet
611        ]);
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()), // manifested under snapshot/
640            ("base/orphan.parquet", b"x".to_vec()),         // top-level orphan
641        ]);
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        // Only a Success manifest keeps its parts — a Failed/Interrupted run's
655        // files are themselves crash leftovers.
656        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        // Full loads OVERWRITE, so only the LATEST snapshot may be loaded —
671        // selecting all accumulated runs would load duplicate snapshots.
672        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        // Full is NOT ledger-skipped: a re-load re-OVERWRITEs from the latest
685        // snapshot, self-healing a drifted target and staying resilient to hidden
686        // in-place updates. The ledger must NOT suppress the re-materialization —
687        // full has to stay full (the old `latest_unloaded_full` skip was the bug).
688        let keyed = vec![
689            keyed_at("r1", "2026-01-01T00:00:00Z"),
690            keyed_at("r2", "2026-01-02T00:00:00Z"),
691        ];
692        // r2 is "already loaded" in the caller's ledger, yet full still selects it.
693        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        // Empty staging (e.g. cleaned, no fresh extract) → empty selection → the
701        // caller returns None and never truncates the target to empty.
702        assert!(latest_full(Vec::new()).is_empty());
703    }
704
705    #[test]
706    fn latest_full_orders_by_parsed_instant_not_lexical_bytes() {
707        // `…00.500Z` is LATER than `…00Z` but sorts EARLIER lexically ('.' < 'Z').
708        // Parsing to an instant picks the truly-newest run, not the lexical max.
709        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        // Stateful, r2 already loaded → still selects r2 (re-materialize/self-heal).
730        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        // STATELESS (empty loaded) → the latest, NEVER a blanket load of both
735        // snapshots (the duplicate-rows bug this fix closes).
736        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            // Stateless → empty loaded → load all (at-least-once, dedup absorbs).
755            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")); // no `-` after manifest
767        assert!(!is_run_unique_manifest("manifest-abc.txt")); // not .json
768        assert!(!is_run_unique_manifest("data.json")); // wrong prefix
769    }
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    // ---- offline (filesystem-backed store) transport tests ----
796
797    /// Build an fs-backed [`GcsStore`] over a fresh tempdir seeded with
798    /// `(bucket-relative-key, bytes)` objects. Returns the store and the guard —
799    /// hold the `TempDir` for the store's lifetime.
800    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        // A shared prefix that accumulated two runs holds the last-writer-wins
818        // `manifest.json` pointer AND one immutable per-run copy each. Summing
819        // must count the two run copies, never the pointer (double-count guard).
820        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()), // data file: not a manifest
825        ]);
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        // A single batch run (or a legacy prefix) has only the canonical name —
840        // with no per-run copy, it must still be found.
841        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        // Two runs' copies plus a canonical pointer to the latest: fetch returns
851        // exactly the two run copies, parsed — so reconcile sums 100 + 40 = 140
852        // without double-counting the pointer.
853        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    // ── fake-gcs-server: the storage contract over the REAL opendal-GCS transport
885    //
886    // Every test above runs the storage-contract fns over an Fs-backed GcsStore —
887    // proving the LOGIC, but not that opendal's GCS client speaks the same
888    // list/read/remove semantics a real bucket answers with. This cell closes that
889    // gap against fsouza/fake-gcs-server (the GCS JSON API emulator the extract
890    // side already uses), so a GCS-only regression — listing pagination, prefix
891    // handling, `remove_all` recursion — fails HERE, not in production. The
892    // warehouse half (bq/snow COPY) can't be emulated; it stays the live BigQuery
893    // matrix (`smoke_batch_mysql` / `StagingWiped`). This owns the BUCKET half.
894    const FAKE_GCS_ENDPOINT: &str = "http://127.0.0.1:4443";
895    const FAKE_GCS_BUCKET: &str = "rivet-load-emulator";
896
897    /// A real GCS-transport [`GcsStore`] against the local emulator, scoped to
898    /// `FAKE_GCS_BUCKET`. Ensures the bucket first via the emulator's JSON API
899    /// (fake-gcs does not auto-create on write) — an idempotent POST, so a re-run
900    /// needs no teardown. curl, not a new HTTP dep: this is a dev-only cell.
901    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    /// Per-object drain of `prefix` — list + single `remove` each. fake-gcs
929    /// implements single-object DELETE but NOT the GCS batch-delete endpoint that
930    /// opendal's `remove_all` issues for 2+ objects (that 400s with `deleted: 0`),
931    /// so the emulator can't run the `cleanup_source` batch wipe. That wipe
932    /// (`delete_under` → `remove_all`) is verified against a REAL bucket by the
933    /// live BigQuery `StagingWiped` matrix, and over the Fs seam above; here we
934    /// drain per-object for idempotent setup/teardown.
935    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        // Test-owned prefix, drained first so a re-run starts clean even though the
946        // emulator bucket persists (no cross-run bleed).
947        let prefix = "load-contract/orders";
948        drain(&store, prefix);
949        let gs = format!("gs://{FAKE_GCS_BUCKET}/{prefix}");
950
951        // Seed one Success run (its manifest + committed part) plus a crash orphan
952        // — an unmanifested `.parquet` an interrupted extract would leave behind.
953        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        // 1. fetch + parse the manifest back over real GCS list+read.
967        let keyed = fetch_manifests_keyed(&store, &gs).unwrap();
968        assert_eq!(keyed.len(), 1, "the run's manifest, read back over GCS");
969
970        // 2. reconcile → file_rows, the value the warehouse COUNT(*) gate enforces.
971        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        // 3. select_load_uris resolves the MANIFESTED part only — never the orphan.
979        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        // 4. gc_orphans deletes the orphan over real GCS — a REAL single-object
988        // DELETE against the emulator — keeping the manifested part + manifest.
989        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        // cleanup_source's batch wipe (`remove_all`) is NOT emulatable here — see
1006        // `drain`. Teardown per-object; the empty listing confirms the deletes
1007        // landed over real GCS transport.
1008        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}