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 crate::pipeline::validate_manifest::MANIFEST_MAX_BYTES;
29use anyhow::{Context, Result, bail};
30
31/// The reconciled row-count chain for one export's load, derived from the run
32/// manifests under its GCS prefix. `file_rows` is what the warehouse must end
33/// up holding; the loader's `expected_rows` gate enforces `warehouse == file`.
34#[derive(Debug, Clone, PartialEq, Eq)]
35pub struct LoadIntegrity {
36    /// Rows the source held at extraction time, summed over the manifests that
37    /// probed it. `None` when *no* contributing manifest carried a
38    /// `source_row_count` — the source→file leg is then unverifiable from the
39    /// manifest alone (e.g. a full snapshot that did not count the source).
40    pub source_rows: Option<u64>,
41    /// Rows written to files — the sum of the trustworthy manifests'
42    /// `row_count`. This is the authoritative expected warehouse row count.
43    pub file_rows: u64,
44    /// How many run manifests contributed to the totals.
45    pub manifests: usize,
46}
47
48impl LoadIntegrity {
49    /// A one-line human summary of the chain, e.g.
50    /// `source 1000 → files 1000 → (warehouse pending)`. The warehouse leg is
51    /// filled in by the caller once the load's `COUNT(*)` is known.
52    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/// Fetch and parse every `manifest.json` under `gcs_prefix` (recursive), keeping
61/// each manifest's bucket-relative storage key — needed to resolve a manifest's
62/// (relative) part paths back to full object keys for per-run loading (see
63/// [`select_load_uris`]).
64///
65/// A rivet export writes one manifest per `run_id`; a prefix that has
66/// accumulated several incremental / CDC runs holds several. Transport is the
67/// native opendal client the extraction destination already uses (`store`), so
68/// auth is the export's own GCS credentials and this is offline-testable over a
69/// filesystem-backed store.
70///
71/// `pub` (a public-API root the lib keeps alive) even though its only caller is
72/// the binary-only `cli::dispatch`; `#[allow(private_interfaces)]` because the
73/// injected `GcsStore` is deliberately an internal (`pub(crate)`) type — the
74/// `destination` module stays crate-private. Same rationale as the `preflight`
75/// module note in `lib.rs`.
76#[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            // Round-5 (CWE-400): cap the manifest body before reading it whole into
86            // memory — a planted multi-GB manifest.json under the load prefix would
87            // otherwise OOM the loader. Mirrors the V21 cap the export-side manifest
88            // readers already enforce; this was the 4th, uncapped, read path.
89            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/// Full `gs://` URIs of the parquet to load for `new` (the not-yet-loaded run
105/// manifests), preferring each manifest's own parts over a blanket listing.
106/// See [`select_load_keys`] for the selection rule.
107#[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
125/// The bucket-relative part keys a manifest declares, each resolved against the
126/// manifest's own directory: `<dir(manifest_key)>/<part.path>`. Shared by
127/// [`select_load_keys`] (which intersects them with what's present) and
128/// [`gc_orphans`] (which treats them as the keep-set), so the two can't drift on
129/// how a manifest maps to its files.
130fn 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
144/// Pure selection: which bucket-relative parquet keys to load for the given
145/// (not-yet-loaded) run manifests.
146///
147/// Prefers each manifest's own parts (via [`resolve_parts`]) intersected with
148/// `all_parquet` — so a load pulls exactly the new runs' files, not every object
149/// under the prefix (the key to incremental loads once `cleanup_source` no
150/// longer wipes the bucket). Falls back to the whole `all_parquet` listing when
151/// ANY new manifest resolves to no present part (legacy/part-less manifests
152/// still load, at the cost of not pruning); the row-count gate then still guards
153/// correctness.
154pub 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            // Can't resolve this run to files — don't risk a partial selection.
168            return all_parquet.to_vec();
169        }
170    }
171    selected.into_iter().collect()
172}
173
174/// Delete every `.parquet` under `gcs_prefix` that no **`Success`** manifest
175/// references — crash leftovers from an interrupted extract that accumulate.
176/// Keeps every manifest, `_SUCCESS`, and every manifested part (including a
177/// `snapshot/` sub-prefix's, since `keyed` is fetched recursively). Strictly
178/// gentler than `cleanup_source`, which wipes the whole prefix.
179///
180/// A candidate falls into three classes:
181/// - referenced by a **`Success`** manifest → KEEP (live data).
182/// - referenced by a **`Failed`/`Interrupted`** manifest → DELETE unconditionally.
183///   A run that WROTE a manifest is terminal — no live extract is still streaming
184///   it — so its leftovers are unambiguous crash debris.
185/// - referenced by **NO manifest at all** → AMBIGUOUS: a crash-BEFORE-manifest
186///   orphan (delete) OR a concurrent extract's committed-but-not-yet-manifested
187///   part (must NOT delete). `active` decides — the caller's answer to "is a run
188///   currently WRITING this prefix?", read from the central run-status ledger
189///   (`StateStore::has_active_run_on_prefix`): authoritative and CLOCK-FREE.
190///   `active` → spare every unmanifested part; `!active` → no run is live, so an
191///   unmanifested part is dead crash debris → delete.
192///
193/// The ledger read is the SEAM. A co-located / shared-Postgres load gets a
194/// precise `active`; a stateless or foreign-host load passes `active = true`
195/// (conservative — spare rather than risk a live cross-host extract's parts),
196/// which the bucket-manifest `running`-status projection later refines.
197/// Returns `(removed, bytes)`.
198#[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    // Success parts → keep. Failed/Interrupted parts → terminal (their run is
207    // done, so deletable regardless of `active`). A part in NEITHER set has no
208    // manifest at all — the ambiguous case `active` gates.
209    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            // A `running` manifest is a LIVE run's marker (schema-less, no parts).
215            // Its parts (if any) must NOT be treated as terminal debris — its
216            // activeness is what makes `active` true (see the callsite), which
217            // spares the run's unmanifested in-flight parts below.
218            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        // A live run on this prefix + a part with NO manifest → maybe its
231        // in-flight write → spare. A terminal-manifested part is deleted even
232        // while a (different) run is active.
233        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    // Second pass: a SUPERSEDED `running` MARKER manifest (its run was overtaken by
245    // a newer run of the same export) is a dead crash marker — never the live
246    // signal — so its lingering `.json` is safe to remove. gc otherwise never
247    // touches it (the parquet pass above only lists `.parquet`), so a hard-crashed
248    // run's marker would accumulate forever. A NON-superseded running manifest is
249    // the ACTIVE signal and MUST survive — deletion here is gated on SUPERSESSION,
250    // never on `active`.
251    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
261/// Is a LIVE run's `running` MARKER manifest present under the prefix — the
262/// bucket-side projection of the run-status ledger, for a cross-boundary load
263/// (Airflow / a foreign-host `rivet load`) that cannot read the extract's state
264/// DB? True iff some `running` manifest is NOT superseded by a NEWER manifest
265/// (any status) of the SAME export — the same clock-free supersession the ledger
266/// uses (a newer `started_at` means the old running run crashed and its successor
267/// already re-ran). `gc_orphans`'s `active` is `ledger_active OR this`, so the two
268/// signals are belt-and-suspenders: the ledger is precise co-located, the marker
269/// covers cross-host.
270pub 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
276/// A `running` manifest is SUPERSEDED when a NEWER run of the SAME export exists
277/// (a higher `started_at`) — it crashed and its successor already re-ran, so it
278/// no longer protects anything. The ONE clock-free staleness predicate, shared by
279/// [`has_active_running_manifest`] (spare the non-superseded) and `gc_orphans`'s
280/// marker-GC sweep (delete the superseded). The ledger enforces the same rule in
281/// SQL — that copy cannot share this Rust.
282fn 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
288/// Full/chunked loads care only about the LATEST snapshot: from `keyed` (all run
289/// manifests under the prefix) pick the newest by `finished_at`. Full loads
290/// OVERWRITE, so exactly ONE snapshot may be loaded (loading every accumulated
291/// run would duplicate rows). Deliberately **not** ledger-gated: full
292/// re-materializes the latest snapshot on *every* load, so a re-load self-heals a
293/// drifted target and stays resilient to hidden in-place source updates. An empty
294/// input (no staged run — e.g. the staging was cleaned) yields an empty
295/// selection, so the caller no-ops WITHOUT truncating the target.
296///
297/// Ordering parses `finished_at` as an instant — a lexical byte compare mis-picks
298/// on mixed RFC3339 precision (`…00.5Z` sorts before `…00Z`) — and falls back to
299/// lexical only if a timestamp fails to parse, so a malformed manifest can't panic.
300pub 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
316/// Which run manifests to load for `mode`, given the ledger's already-`loaded`
317/// run_ids. The single mode→selection decision, pure and testable so the load's
318/// central invariant isn't buried in dispatch I/O:
319/// - `Full` → the single LATEST run ([`latest_full`]) — ledger-INDEPENDENT, so a
320///   re-load re-materializes the snapshot (self-heal) and, crucially, the
321///   STATELESS path (empty `loaded`) still picks the latest instead of blanket-
322///   loading every accumulated snapshot (the duplicate-rows bug).
323/// - `Incremental`/`Cdc` → every run not yet in `loaded` (append). Stateless →
324///   `loaded` empty → load all; at-least-once, absorbed by the dedup view.
325pub 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    // A `running` manifest is a LIVE run's in-flight marker — never loadable (no
331    // committed parts). Drop it BEFORE selection so it neither becomes the
332    // "latest" full snapshot nor an incremental delta, either of which would
333    // make `reconcile` refuse the whole load on a non-Success run.
334    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
347/// Bucket-relative keys of every run manifest under `base` (recursive).
348fn 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    // A run into a shared prefix leaves BOTH the canonical `manifest.json`
355    // (last-writer-wins — a pointer to the LATEST run) and an immutable
356    // `manifest-<run_id>.json` copy (one per run). Sum the per-run copies so a
357    // prefix that accumulated several CDC/incremental cycles counts EVERY run;
358    // counting the canonical pointer too would double-count the latest. Fall
359    // back to the canonical name only when no per-run copy exists (a single
360    // batch run, or a legacy prefix predating the run-unique copy).
361    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
373/// A listed key is a run manifest iff its final path segment is the canonical
374/// [`MANIFEST_FILENAME`] (`manifest.json`) or a run-unique copy
375/// (`manifest-<run_id>.json`) — so a data file merely *named* like it (an
376/// unlikely `…/x_manifest.json`) is not mistaken for one.
377fn 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
382/// A per-run manifest copy: `manifest-<token>.json` (the sidecar the OSS sink
383/// writes alongside the canonical pointer so cross-run reconcile can sum it).
384fn is_run_unique_manifest(base: &str) -> bool {
385    // The sidecar naming scheme lives once, in `manifest.rs` (the writer's home).
386    crate::manifest::is_run_unique_manifest_name(base)
387}
388
389/// Refuse a load whose prefix holds MORE THAN ONE export's manifests. The load
390/// sums EVERY manifest under `gcs_prefix` (fetched recursively) and cleanup
391/// removes the prefix recursively, so a shared base prefix — two `partition_by`
392/// exports whose prefix-before-`{partition}` coincides — would silently cross-
393/// contaminate the reconciled row count AND delete the sibling export's un-loaded
394/// parts. The `LoadPlan` carries no source `export_name` to disambiguate which
395/// export the operator meant, so fail LOUDLY rather than load-and-delete the
396/// wrong data. Legacy manifests with no recorded `export_name` (pre-0.x) are
397/// tolerated — an empty name matches any single named export, so a prefix that
398/// mixes old and new runs of the SAME export still loads.
399///
400/// The CDC `initial: snapshot` leg is NOT a sibling: `cdc_job` synthesizes it
401/// as `{export}__snapshot_{table}` and points it at the SAME prefix by design —
402/// the baseline and the drain are one table's rows, summed and cleaned
403/// together. The guard therefore compares export FAMILIES, folding a
404/// snapshot-leg name to its parent (see [`crate::manifest::snapshot_family`]).
405pub fn ensure_single_export(keyed: &[(String, RunManifest)]) -> Result<()> {
406    let mut names: std::collections::BTreeSet<&str> = keyed
407        .iter()
408        .map(|(_, m)| crate::manifest::snapshot_family(m.export_name.as_str()))
409        .filter(|n| !n.is_empty())
410        .collect();
411    if names.len() > 1 {
412        let listed: Vec<&str> = std::mem::take(&mut names).into_iter().collect();
413        bail!(
414            "the load prefix holds manifests from {} distinct exports ({}) — the load sums every \
415             manifest under the prefix and cleanup wipes it recursively, so loading a shared \
416             prefix would cross-contaminate the row count and could delete a sibling export's \
417             parts. Give each export a DISTINCT destination prefix (or scope the load prefix to a \
418             single export) and re-run.",
419            listed.len(),
420            listed.join(", ")
421        );
422    }
423    Ok(())
424}
425
426/// Reconcile a run's manifests into the authoritative expected warehouse row
427/// count, refusing to load anything that is not provably complete.
428///
429/// Every manifest must be a **`Success`** run (a `Failed` / `Interrupted`
430/// manifest describes a partial, untrustworthy export) and **self-consistent**
431/// (`row_count` == sum of committed parts — a mismatch is a writer bug). When a
432/// manifest recorded a `source_row_count`, the **source→file** leg must
433/// reconcile too: `source_row_count == row_count`, or the extract silently
434/// dropped rows. `allow_source_drift` downgrades only that last check to a
435/// warning (e.g. an incremental cursor window whose source moved under it);
436/// the completeness and self-consistency gates never yield.
437///
438/// Returns [`LoadIntegrity`] with the summed `file_rows` the loader's
439/// `expected_rows` gate then checks against the warehouse's `COUNT(*)`.
440pub fn reconcile(manifests: &[RunManifest], allow_source_drift: bool) -> Result<LoadIntegrity> {
441    if manifests.is_empty() {
442        bail!(
443            "no `{MANIFEST_FILENAME}` found under the export prefix — refusing to load \
444             unverified files. A rivet export writes a manifest on success; its absence \
445             means the run never completed (or points at the wrong prefix)."
446        );
447    }
448
449    let mut file_rows: u64 = 0;
450    let mut source_rows: u64 = 0;
451    let mut any_source = false;
452
453    for m in manifests {
454        // Completeness: only a `Success` run may be loaded.
455        if m.status != ManifestStatus::Success {
456            bail!(
457                "manifest for run `{}` (export `{}`) is {:?}, not Success — refusing to load a \
458                 partial export",
459                m.run_id,
460                m.export_name,
461                m.status
462            );
463        }
464        // Self-consistency: the recorded aggregates must match the committed
465        // parts (a divergence is a writer bug, per OSS `validate_self_consistency`).
466        m.validate_self_consistency().map_err(|e| {
467            anyhow::anyhow!(
468                "manifest for run `{}` (export `{}`) is internally inconsistent: {e} — refusing \
469                 to load",
470                m.run_id,
471                m.export_name
472            )
473        })?;
474
475        let rows = u64::try_from(m.row_count).with_context(|| {
476            format!(
477                "manifest for run `{}` has a negative row_count ({})",
478                m.run_id, m.row_count
479            )
480        })?;
481        file_rows += rows;
482
483        // Source→file: reconcile the extract against the source when the run
484        // probed it. `None` = not probed (unverifiable, not a failure).
485        if let Some(src) = m
486            .source
487            .extraction
488            .as_ref()
489            .and_then(|x| x.source_row_count)
490        {
491            let src = u64::try_from(src).with_context(|| {
492                format!(
493                    "manifest for run `{}` has a negative source_row_count ({src})",
494                    m.run_id
495                )
496            })?;
497            any_source = true;
498            source_rows += src;
499            if src != rows {
500                if allow_source_drift {
501                    eprintln!(
502                        "warning: source→file drift for run `{}` (export `{}`): source had {src} \
503                         rows, extracted {rows} (--allow-source-drift)",
504                        m.run_id, m.export_name
505                    );
506                } else {
507                    bail!(
508                        "source→file mismatch for run `{}` (export `{}`): source had {src} rows \
509                         but {rows} were extracted — the extract dropped {} row(s). Investigate \
510                         before loading, or pass --allow-source-drift to override.",
511                        m.run_id,
512                        m.export_name,
513                        src.abs_diff(rows)
514                    );
515                }
516            }
517        }
518    }
519
520    Ok(LoadIntegrity {
521        source_rows: any_source.then_some(source_rows),
522        file_rows,
523        manifests: manifests.len(),
524    })
525}
526
527#[cfg(test)]
528mod tests {
529    use super::*;
530    use crate::manifest::{
531        ExtractionMetadata, ManifestDestination, ManifestPart, ManifestSource, PartStatus,
532    };
533
534    /// A minimal `Success` manifest with one committed part of `rows` rows and,
535    /// optionally, a probed `source_row_count`.
536    fn manifest(run: &str, rows: i64, source: Option<i64>) -> RunManifest {
537        RunManifest {
538            row_hash: None,
539            manifest_version: crate::manifest::MANIFEST_VERSION,
540            run_id: run.into(),
541            export_name: "orders".into(),
542            mode: "batch".into(),
543            started_at: "t".into(),
544            finished_at: "t".into(),
545            status: ManifestStatus::Success,
546            source: ManifestSource {
547                engine: "pg".into(),
548                schema: None,
549                table: None,
550                extraction: source.map(|n| ExtractionMetadata {
551                    strategy: "full".into(),
552                    cursor_column: None,
553                    cursor_type: None,
554                    cursor_low: None,
555                    cursor_high: None,
556                    source_row_count: Some(n),
557                }),
558            },
559            destination: ManifestDestination {
560                kind: "gcs".into(),
561                uri: "gs://b/p".into(),
562            },
563            format: "parquet".into(),
564            compression: "zstd".into(),
565            schema_fingerprint: "xxh3:0".into(),
566            row_count: rows,
567            part_count: 1,
568            parts: vec![ManifestPart {
569                part_id: 0,
570                path: "part-000000.parquet".into(),
571                rows,
572                size_bytes: 1,
573                content_fingerprint: "xxh3:0".into(),
574                content_md5: String::new(),
575                status: PartStatus::Committed,
576            }],
577            column_checksums: None,
578            checksum_key_column: None,
579        }
580    }
581
582    #[test]
583    fn sums_file_and_source_rows_across_manifests() {
584        let ms = vec![manifest("r1", 100, Some(100)), manifest("r2", 40, Some(40))];
585        let got = reconcile(&ms, false).unwrap();
586        assert_eq!(got.file_rows, 140);
587        assert_eq!(got.source_rows, Some(140));
588        assert_eq!(got.manifests, 2);
589    }
590
591    #[test]
592    fn ensure_single_export_refuses_a_prefix_shared_by_two_exports() {
593        let keyed = |m: RunManifest| ("gs://b/p/manifest-x.json".to_string(), m);
594        let orders = keyed(manifest("r1", 10, None)); // export_name "orders"
595        let mut cust_m = manifest("r2", 10, None);
596        cust_m.export_name = "customers".into();
597        // Two DISTINCT exports under one load prefix → loud refusal (else the load
598        // sums both and cleanup deletes the sibling's parts).
599        assert!(ensure_single_export(&[orders.clone(), keyed(cust_m)]).is_err());
600        // Many runs of ONE export → fine.
601        assert!(ensure_single_export(&[orders.clone(), keyed(manifest("r3", 5, None))]).is_ok());
602        // A legacy (no export_name) run mixed with the SAME named export → still
603        // one logical export, tolerated (an upgrade mid-prefix must not brick load).
604        let mut legacy = manifest("r0", 3, None);
605        legacy.export_name = String::new();
606        assert!(ensure_single_export(&[orders, keyed(legacy)]).is_ok());
607    }
608
609    /// Regression: the CDC `initial: snapshot` leg shares the drain's prefix BY
610    /// DESIGN (`cdc_job` synthesizes `{export}__snapshot_{table}` pointed at the
611    /// same destination). 0.24.0's cross-run manifest summing put both names in
612    /// front of this guard for the first time and the ordinary snapshot → drain
613    /// → `rivet load` flow refused itself. The pair is ONE export family.
614    #[test]
615    fn ensure_single_export_admits_the_cdc_snapshot_leg_of_the_same_export() {
616        let keyed = |m: RunManifest| ("gs://b/p/manifest-x.json".to_string(), m);
617        let drain = keyed(manifest("r1", 10, None)); // export_name "orders"
618        let mut snap = manifest("r2", 10, None);
619        snap.export_name = "orders__snapshot_orders".into();
620        // Snapshot baseline + drain of the SAME export → one family, loads.
621        assert!(ensure_single_export(&[drain.clone(), keyed(snap)]).is_ok());
622        // A snapshot leg of a DIFFERENT export is still a sibling → refuse.
623        let mut foreign = manifest("r3", 10, None);
624        foreign.export_name = "customers__snapshot_customers".into();
625        assert!(ensure_single_export(&[drain, keyed(foreign)]).is_err());
626    }
627
628    #[test]
629    fn source_rows_is_none_when_no_manifest_probed_the_source() {
630        let ms = vec![manifest("r1", 100, None), manifest("r2", 40, None)];
631        let got = reconcile(&ms, false).unwrap();
632        assert_eq!(got.file_rows, 140);
633        assert_eq!(
634            got.source_rows, None,
635            "unprobed source is unknown, not zero"
636        );
637    }
638
639    #[test]
640    fn source_rows_present_even_if_only_some_manifests_probed() {
641        let ms = vec![manifest("r1", 100, Some(100)), manifest("r2", 40, None)];
642        let got = reconcile(&ms, false).unwrap();
643        assert_eq!(got.source_rows, Some(100));
644    }
645
646    #[test]
647    fn empty_manifests_refuses_to_load() {
648        let err = reconcile(&[], false).unwrap_err().to_string();
649        assert!(err.contains("refusing to load"), "{err}");
650    }
651
652    #[test]
653    fn non_success_manifest_refuses_to_load() {
654        let mut m = manifest("r1", 100, Some(100));
655        m.status = ManifestStatus::Interrupted;
656        let err = reconcile(&[m], false).unwrap_err().to_string();
657        assert!(err.contains("not Success"), "{err}");
658    }
659
660    #[test]
661    fn self_inconsistent_manifest_refuses_to_load() {
662        // row_count claims 100 but the only committed part holds 100 → make the
663        // aggregate lie by bumping the recorded row_count only.
664        let mut m = manifest("r1", 100, Some(100));
665        m.row_count = 999; // no longer equals committed parts' sum (100)
666        let err = reconcile(&[m], false).unwrap_err().to_string();
667        assert!(err.contains("inconsistent"), "{err}");
668    }
669
670    #[test]
671    fn source_file_mismatch_hard_fails_by_default() {
672        // Source had 120, only 100 extracted → 20 rows silently dropped.
673        let m = manifest("r1", 100, Some(120));
674        let err = reconcile(&[m], false).unwrap_err().to_string();
675        assert!(err.contains("source→file mismatch"), "{err}");
676        assert!(err.contains("dropped 20"), "{err}");
677    }
678
679    #[test]
680    fn source_file_mismatch_is_allowed_under_the_override() {
681        let m = manifest("r1", 100, Some(120));
682        let got = reconcile(&[m], true).expect("--allow-source-drift proceeds");
683        assert_eq!(got.file_rows, 100);
684        assert_eq!(
685            got.source_rows,
686            Some(120),
687            "the probed source count is still surfaced"
688        );
689    }
690
691    /// A `(manifest_key, manifest)` pair with a single part named `part`.
692    fn keyed(key: &str, run: &str, part: &str) -> (String, RunManifest) {
693        let mut m = manifest(run, 10, Some(10));
694        m.parts[0].path = part.into();
695        (key.to_string(), m)
696    }
697
698    #[test]
699    fn select_load_keys_picks_only_the_new_runs_parts() {
700        // Two runs' files sit under the prefix; only r2 is "new".
701        let all = vec![
702            "base/r1-000.parquet".to_string(),
703            "base/r2-000.parquet".to_string(),
704        ];
705        let new = vec![keyed("base/manifest-r2.json", "r2", "r2-000.parquet")];
706        assert_eq!(
707            select_load_keys(&new, &all),
708            vec!["base/r2-000.parquet".to_string()],
709            "loads r2's part only — not r1's already-loaded file"
710        );
711    }
712
713    #[test]
714    fn select_load_uris_lists_and_re_prefixes_selected_keys_with_the_source_bucket() {
715        // select_load_keys (the pure selection) is unit-tested throughout; this
716        // pins the store WRAPPER around it — list `.parquet` under the base, then
717        // reconstruct each `gs://<bucket>/<key>` the loader COPYs from. A mangled
718        // wrapper (empty/garbage URIs) would send the warehouse load at nothing,
719        // or the wrong object — invisible to the pure-selection tests.
720        let (store, _g) = fs_store(&[
721            ("base/r1-000.parquet", b"a".to_vec()),
722            ("base/manifest-r1.json", b"{}".to_vec()), // not .parquet — must be skipped
723        ]);
724        let new = vec![keyed("base/manifest-r1.json", "r1", "r1-000.parquet")];
725        assert_eq!(
726            select_load_uris(&store, "gs://my-bucket/base", &new).unwrap(),
727            vec!["gs://my-bucket/base/r1-000.parquet".to_string()],
728            "the selected key, re-prefixed with the source bucket — never the manifest"
729        );
730    }
731
732    #[test]
733    fn select_load_keys_resolves_a_snapshot_subprefix_manifest() {
734        // A snapshot manifest lives under `base/snapshot/`; its part is relative
735        // to that dir. Resolution must reconstruct the full key.
736        let all = vec!["base/snapshot/snap-000.parquet".to_string()];
737        let new = vec![keyed(
738            "base/snapshot/manifest-r1.json",
739            "r1",
740            "snap-000.parquet",
741        )];
742        assert_eq!(
743            select_load_keys(&new, &all),
744            vec!["base/snapshot/snap-000.parquet".to_string()]
745        );
746    }
747
748    #[test]
749    fn select_load_keys_falls_back_to_full_listing_when_a_manifest_has_no_present_part() {
750        // A manifest whose part isn't in the listing (legacy / renamed) → don't
751        // risk a partial selection; load everything under the prefix.
752        let all = vec!["base/a.parquet".to_string(), "base/b.parquet".to_string()];
753        let new = vec![keyed("base/manifest-r1.json", "r1", "missing.parquet")];
754        assert_eq!(
755            select_load_keys(&new, &all),
756            all,
757            "unresolvable part → blanket fallback"
758        );
759    }
760
761    #[test]
762    fn select_load_keys_empty_new_set_selects_nothing_never_the_full_listing() {
763        // When every run is already in the ledger the "new" set is empty. That
764        // must resolve to ZERO uris — NOT the blanket fallback, which would
765        // re-load the whole prefix on every up-to-date run (double-load). The
766        // fallback fires only for an unresolvable NON-empty manifest.
767        let all = vec![
768            "base/r1-000.parquet".to_string(),
769            "base/r2-000.parquet".to_string(),
770        ];
771        assert!(
772            select_load_keys(&[], &all).is_empty(),
773            "no new runs ⇒ load nothing, not everything"
774        );
775    }
776
777    #[test]
778    fn gc_orphans_removes_unmanifested_parquet_only() {
779        let (store, _g) = fs_store(&[
780            ("base/r1-000.parquet", b"aa".to_vec()),   // manifested
781            ("base/orphan.parquet", b"junk".to_vec()), // crash leftover — no manifest
782            ("base/manifest.json", b"{}".to_vec()),    // kept — not a .parquet
783            ("base/_SUCCESS", b"".to_vec()),           // kept — not a .parquet
784        ]);
785        let keyed = vec![keyed("base/manifest-r1.json", "r1", "r1-000.parquet")];
786        // No run is active on this prefix → the unmanifested part is dead debris.
787        let (removed, bytes) = gc_orphans(&store, "gs://b/base", &keyed, false).unwrap();
788        assert_eq!(removed, 1, "only the unmanifested part is removed");
789        assert_eq!(bytes, 4, "'junk' is 4 bytes");
790        let mut left = store.list_files("base").unwrap();
791        left.sort();
792        assert_eq!(
793            left,
794            vec![
795                "base/_SUCCESS".to_string(),
796                "base/manifest.json".to_string(),
797                "base/r1-000.parquet".to_string(),
798            ],
799            "the manifested part, the manifest, and _SUCCESS all survive"
800        );
801    }
802
803    #[test]
804    fn gc_orphans_of_an_all_manifested_prefix_removes_nothing() {
805        let (store, _g) = fs_store(&[("base/r1-000.parquet", b"a".to_vec())]);
806        let keyed = vec![keyed("base/manifest-r1.json", "r1", "r1-000.parquet")];
807        assert_eq!(
808            gc_orphans(&store, "gs://b/base", &keyed, false).unwrap().0,
809            0
810        );
811    }
812
813    #[test]
814    fn gc_orphans_keeps_a_snapshot_subprefix_manifested_part() {
815        let (store, _g) = fs_store(&[
816            ("base/snapshot/s-000.parquet", b"a".to_vec()), // manifested under snapshot/
817            ("base/orphan.parquet", b"x".to_vec()),         // top-level orphan
818        ]);
819        let keyed = vec![keyed("base/snapshot/manifest-s.json", "s", "s-000.parquet")];
820        let (removed, _) = gc_orphans(&store, "gs://b/base", &keyed, false).unwrap();
821        assert_eq!(removed, 1, "the top-level orphan goes");
822        assert_eq!(
823            store.list_files("base/snapshot").unwrap(),
824            vec!["base/snapshot/s-000.parquet".to_string()],
825            "the snapshot-subprefix manifested part is kept"
826        );
827    }
828
829    #[test]
830    fn gc_orphans_deletes_a_terminal_runs_parts_even_while_a_run_is_active() {
831        // A Failed/Interrupted run's parts are terminal crash debris — deleted
832        // regardless of `active`, since no LIVE run is streaming THEM. Passing
833        // active=true proves the `active` gate applies only to UNmanifested parts.
834        let (store, _g) = fs_store(&[("base/f-000.parquet", b"x".to_vec())]);
835        let mut kv = keyed("base/manifest-f.json", "f", "f-000.parquet");
836        kv.1.status = ManifestStatus::Failed;
837        assert_eq!(gc_orphans(&store, "gs://b/base", &[kv], true).unwrap().0, 1);
838    }
839
840    #[test]
841    fn gc_orphans_spares_an_unmanifested_part_while_a_run_is_active() {
842        // The concurrent-extract guard. A part with NO manifest, while a run is
843        // ACTIVE on this prefix (the run-status ledger says so), is probably that
844        // run's committed-but-not-yet-manifested write — deleting it would be
845        // silent data loss on the live extract. RED against a mutant that ignores
846        // `active` (deletes every non-Success-manifested `.parquet`).
847        let (store, _g) = fs_store(&[
848            ("base/r1-000.parquet", b"aa".to_vec()),  // manifested (kept)
849            ("base/inflight.parquet", b"x".to_vec()), // unmanifested, a live run's part
850        ]);
851        let keyed = vec![keyed("base/manifest-r1.json", "r1", "r1-000.parquet")];
852        let (removed, _) = gc_orphans(&store, "gs://b/base", &keyed, true).unwrap();
853        assert_eq!(
854            removed, 0,
855            "an unmanifested part is spared while a run is active"
856        );
857        assert!(
858            store
859                .list_files("base")
860                .unwrap()
861                .iter()
862                .any(|k| k.ends_with("inflight.parquet")),
863            "the live run's in-flight part must survive gc_orphans"
864        );
865    }
866
867    #[test]
868    fn gc_orphans_collects_an_unmanifested_part_when_no_run_is_active() {
869        // The other half: with NO run active, an unmanifested part is a dead
870        // crash orphan and IS collected — the gate DEFERS cleanup, never abandons.
871        let (store, _g) = fs_store(&[
872            ("base/r1-000.parquet", b"aa".to_vec()),
873            ("base/dead-orphan.parquet", b"x".to_vec()),
874        ]);
875        let keyed = vec![keyed("base/manifest-r1.json", "r1", "r1-000.parquet")];
876        assert_eq!(
877            gc_orphans(&store, "gs://b/base", &keyed, false).unwrap().0,
878            1,
879            "with no active run, an unmanifested orphan is collected"
880        );
881    }
882
883    /// A `running` MARKER manifest (schema-less, no parts) started at `started_at`.
884    fn running(run: &str, started_at: &str) -> (String, RunManifest) {
885        let mut m = manifest(run, 0, None);
886        m.status = ManifestStatus::Running;
887        m.started_at = started_at.into();
888        m.finished_at = String::new();
889        m.parts.clear();
890        (format!("base/manifest-{run}.json"), m)
891    }
892
893    #[test]
894    fn has_active_running_manifest_true_for_a_lone_running_marker() {
895        assert!(has_active_running_manifest(&[running(
896            "r1",
897            "2026-01-01T00:00:00Z"
898        )]));
899    }
900
901    #[test]
902    fn has_active_running_manifest_false_when_none_is_running() {
903        // A Success manifest is not a live-run marker.
904        assert!(!has_active_running_manifest(&[keyed(
905            "k",
906            "r1",
907            "p.parquet"
908        )]));
909    }
910
911    #[test]
912    fn has_active_running_manifest_false_for_a_superseded_running_marker() {
913        // r1 crashed leaving a `running` marker; r2 (newer started_at, SAME
914        // export) already ran → r1 is stale and must NOT count as active. Same
915        // clock-free supersession the ledger uses.
916        let r1 = running("r1", "2026-01-01T00:00:00Z");
917        let mut r2 = manifest("r2", 10, Some(10)); // Success, same export ("orders")
918        r2.started_at = "2026-01-02T00:00:00Z".into();
919        assert!(!has_active_running_manifest(&[
920            r1,
921            ("base/manifest-r2.json".into(), r2)
922        ]));
923    }
924
925    #[test]
926    fn select_runs_drops_a_running_manifest() {
927        // A `running` (in-flight) manifest must NEVER be selected for load — else
928        // reconcile would refuse the whole load on a non-Success run. Incremental
929        // mode would otherwise include it (not yet loaded).
930        let run_marker = running("r_running", "2026-01-02T00:00:00Z");
931        let ok = keyed("base/manifest-r_ok.json", "r_ok", "r_ok-000.parquet"); // Success
932        let sel = select_runs(
933            vec![run_marker, ok],
934            &std::collections::HashSet::new(),
935            crate::load::plan::LoadMode::Incremental,
936        );
937        assert_eq!(sel.len(), 1, "only the Success run is selected");
938        assert_eq!(sel[0].1.run_id, "r_ok");
939        assert!(sel.iter().all(|(_, m)| m.status != ManifestStatus::Running));
940    }
941
942    #[test]
943    fn gc_orphans_removes_a_superseded_running_marker_but_spares_a_live_one() {
944        // r1 crashed leaving a `running` marker; r2 (newer, same export) is the
945        // live run. The DEAD superseded marker must be GC'd (else it accumulates
946        // forever — gc otherwise only lists `.parquet`); the LIVE one must survive
947        // (it is the active signal). Gated on supersession, so `active=true`.
948        let (store, _g) = fs_store(&[
949            ("base/manifest-r1.json", b"{}".to_vec()), // superseded running marker
950            ("base/manifest-r2.json", b"{}".to_vec()), // live running marker (newest)
951        ]);
952        let r1 = running("r1", "2026-01-01T00:00:01Z");
953        let r2 = running("r2", "2026-01-01T00:00:02Z");
954        let (removed, _) = gc_orphans(&store, "gs://b/base", &[r1, r2], true).unwrap();
955        assert_eq!(removed, 1, "only the superseded running marker is removed");
956        let left = store.list_files("base").unwrap();
957        assert!(
958            !left.iter().any(|k| k.ends_with("manifest-r1.json")),
959            "the superseded (dead) running marker is deleted"
960        );
961        assert!(
962            left.iter().any(|k| k.ends_with("manifest-r2.json")),
963            "the live (non-superseded) running marker survives — it is the active signal"
964        );
965    }
966
967    fn keyed_at(run: &str, finished_at: &str) -> (String, RunManifest) {
968        let mut m = manifest(run, 100, None);
969        m.finished_at = finished_at.into();
970        (format!("base/manifest-{run}.json"), m)
971    }
972
973    #[test]
974    fn latest_full_picks_the_newest_snapshot_not_all() {
975        // Full loads OVERWRITE, so only the LATEST snapshot may be loaded —
976        // selecting all accumulated runs would load duplicate snapshots.
977        let keyed = vec![
978            keyed_at("r1", "2026-01-01T00:00:00Z"),
979            keyed_at("r3", "2026-01-03T00:00:00Z"),
980            keyed_at("r2", "2026-01-02T00:00:00Z"),
981        ];
982        let sel = latest_full(keyed);
983        assert_eq!(sel.len(), 1, "exactly one snapshot, never all");
984        assert_eq!(sel[0].1.run_id, "r3", "the newest by finished_at");
985    }
986
987    #[test]
988    fn latest_full_re_materializes_even_when_the_latest_is_already_loaded() {
989        // Full is NOT ledger-skipped: a re-load re-OVERWRITEs from the latest
990        // snapshot, self-healing a drifted target and staying resilient to hidden
991        // in-place updates. The ledger must NOT suppress the re-materialization —
992        // full has to stay full (the old `latest_unloaded_full` skip was the bug).
993        let keyed = vec![
994            keyed_at("r1", "2026-01-01T00:00:00Z"),
995            keyed_at("r2", "2026-01-02T00:00:00Z"),
996        ];
997        // r2 is "already loaded" in the caller's ledger, yet full still selects it.
998        let sel = latest_full(keyed);
999        assert_eq!(sel.len(), 1);
1000        assert_eq!(sel[0].1.run_id, "r2", "always the latest, loaded or not");
1001    }
1002
1003    #[test]
1004    fn latest_full_of_no_staged_runs_is_empty_so_the_caller_no_ops_without_truncating() {
1005        // Empty staging (e.g. cleaned, no fresh extract) → empty selection → the
1006        // caller returns None and never truncates the target to empty.
1007        assert!(latest_full(Vec::new()).is_empty());
1008    }
1009
1010    #[test]
1011    fn latest_full_orders_by_parsed_instant_not_lexical_bytes() {
1012        // `…00.500Z` is LATER than `…00Z` but sorts EARLIER lexically ('.' < 'Z').
1013        // Parsing to an instant picks the truly-newest run, not the lexical max.
1014        let keyed = vec![
1015            keyed_at("older", "2026-01-01T00:00:00Z"),
1016            keyed_at("newer", "2026-01-01T00:00:00.500Z"),
1017        ];
1018        let sel = latest_full(keyed);
1019        assert_eq!(sel.len(), 1);
1020        assert_eq!(
1021            sel[0].1.run_id, "newer",
1022            "the fractional-second run is the newer instant"
1023        );
1024    }
1025
1026    #[test]
1027    fn select_runs_full_picks_the_latest_even_when_loaded_and_even_when_stateless() {
1028        use crate::load::plan::LoadMode;
1029        use std::collections::HashSet;
1030        let keyed = vec![
1031            keyed_at("r1", "2026-01-01T00:00:00Z"),
1032            keyed_at("r2", "2026-01-02T00:00:00Z"),
1033        ];
1034        // Stateful, r2 already loaded → still selects r2 (re-materialize/self-heal).
1035        let loaded = HashSet::from(["r2".to_string()]);
1036        let sel = select_runs(keyed.clone(), &loaded, LoadMode::Full);
1037        assert_eq!(sel.len(), 1);
1038        assert_eq!(sel[0].1.run_id, "r2", "Full picks latest, loaded or not");
1039        // STATELESS (empty loaded) → the latest, NEVER a blanket load of both
1040        // snapshots (the duplicate-rows bug this fix closes).
1041        let sel = select_runs(keyed, &HashSet::new(), LoadMode::Full);
1042        assert_eq!(sel.len(), 1, "stateless Full is not a blanket load");
1043        assert_eq!(sel[0].1.run_id, "r2");
1044    }
1045
1046    #[test]
1047    fn select_runs_append_modes_filter_loaded_and_load_all_when_stateless() {
1048        use crate::load::plan::LoadMode;
1049        use std::collections::HashSet;
1050        let keyed = vec![
1051            keyed_at("r1", "2026-01-01T00:00:00Z"),
1052            keyed_at("r2", "2026-01-02T00:00:00Z"),
1053        ];
1054        let loaded = HashSet::from(["r1".to_string()]);
1055        for mode in [LoadMode::Incremental, LoadMode::Cdc] {
1056            let sel = select_runs(keyed.clone(), &loaded, mode);
1057            assert_eq!(sel.len(), 1, "{mode:?}: only the unloaded run");
1058            assert_eq!(sel[0].1.run_id, "r2");
1059            // Stateless → empty loaded → load all (at-least-once, dedup absorbs).
1060            assert_eq!(
1061                select_runs(keyed.clone(), &HashSet::new(), mode).len(),
1062                2,
1063                "{mode:?}: stateless loads every run"
1064            );
1065        }
1066    }
1067
1068    #[test]
1069    fn is_run_unique_manifest_needs_both_prefix_and_json() {
1070        assert!(is_run_unique_manifest("manifest-20260101T000000.json"));
1071        assert!(!is_run_unique_manifest("manifest.json")); // no `-` after manifest
1072        assert!(!is_run_unique_manifest("manifest-abc.txt")); // not .json
1073        assert!(!is_run_unique_manifest("data.json")); // wrong prefix
1074    }
1075
1076    #[test]
1077    fn chain_prefix_renders_source_and_files() {
1078        let known = LoadIntegrity {
1079            source_rows: Some(100),
1080            file_rows: 100,
1081            manifests: 1,
1082        };
1083        assert_eq!(known.chain_prefix(), "source 100 → files 100");
1084        let unknown = LoadIntegrity {
1085            source_rows: None,
1086            file_rows: 40,
1087            manifests: 1,
1088        };
1089        assert_eq!(unknown.chain_prefix(), "source ? → files 40");
1090    }
1091
1092    #[test]
1093    fn is_manifest_key_matches_only_the_final_segment() {
1094        assert!(is_manifest_key("gs://b/p/manifest.json"));
1095        assert!(is_manifest_key("manifest.json"));
1096        assert!(!is_manifest_key("gs://b/p/part-0.parquet"));
1097        assert!(!is_manifest_key("gs://b/p/x_manifest.json"));
1098    }
1099
1100    // ---- offline (filesystem-backed store) transport tests ----
1101
1102    /// Build an fs-backed [`GcsStore`] over a fresh tempdir seeded with
1103    /// `(bucket-relative-key, bytes)` objects. Returns the store and the guard —
1104    /// hold the `TempDir` for the store's lifetime.
1105    fn fs_store(files: &[(&str, Vec<u8>)]) -> (GcsStore, tempfile::TempDir) {
1106        let dir = tempfile::tempdir().unwrap();
1107        for (rel, bytes) in files {
1108            let p = dir.path().join(rel);
1109            std::fs::create_dir_all(p.parent().unwrap()).unwrap();
1110            std::fs::write(p, bytes).unwrap();
1111        }
1112        let store = GcsStore::open_fs(dir.path().to_str().unwrap()).unwrap();
1113        (store, dir)
1114    }
1115
1116    fn manifest_bytes(run: &str, rows: i64, source: Option<i64>) -> Vec<u8> {
1117        serde_json::to_vec(&manifest(run, rows, source)).unwrap()
1118    }
1119
1120    #[test]
1121    fn list_manifest_keys_prefers_run_unique_copies_over_the_canonical_pointer() {
1122        // A shared prefix that accumulated two runs holds the last-writer-wins
1123        // `manifest.json` pointer AND one immutable per-run copy each. Summing
1124        // must count the two run copies, never the pointer (double-count guard).
1125        let (store, _g) = fs_store(&[
1126            ("base/manifest.json", b"{}".to_vec()),
1127            ("base/manifest-r1.json", b"{}".to_vec()),
1128            ("base/manifest-r2.json", b"{}".to_vec()),
1129            ("base/part-0.parquet", b"x".to_vec()), // data file: not a manifest
1130        ]);
1131        let mut keys = list_manifest_keys(&store, "base").unwrap();
1132        keys.sort();
1133        assert_eq!(
1134            keys,
1135            vec![
1136                "base/manifest-r1.json".to_string(),
1137                "base/manifest-r2.json".to_string(),
1138            ]
1139        );
1140    }
1141
1142    #[test]
1143    fn list_manifest_keys_falls_back_to_the_canonical_name_for_a_single_run() {
1144        // A single batch run (or a legacy prefix) has only the canonical name —
1145        // with no per-run copy, it must still be found.
1146        let (store, _g) = fs_store(&[("base/manifest.json", b"{}".to_vec())]);
1147        assert_eq!(
1148            list_manifest_keys(&store, "base").unwrap(),
1149            vec!["base/manifest.json".to_string()]
1150        );
1151    }
1152
1153    #[test]
1154    fn fetch_manifests_keyed_reads_and_parses_every_run_copy_under_the_prefix() {
1155        // Two runs' copies plus a canonical pointer to the latest: fetch returns
1156        // exactly the two run copies, parsed — so reconcile sums 100 + 40 = 140
1157        // without double-counting the pointer.
1158        let (store, _g) = fs_store(&[
1159            ("base/manifest.json", manifest_bytes("r2", 40, Some(40))),
1160            (
1161                "base/manifest-r1.json",
1162                manifest_bytes("r1", 100, Some(100)),
1163            ),
1164            ("base/manifest-r2.json", manifest_bytes("r2", 40, Some(40))),
1165        ]);
1166        let manifests: Vec<_> = fetch_manifests_keyed(&store, "gs://my-bucket/base")
1167            .unwrap()
1168            .into_iter()
1169            .map(|(_, m)| m)
1170            .collect();
1171        assert_eq!(manifests.len(), 2);
1172        let integrity = reconcile(&manifests, false).unwrap();
1173        assert_eq!(integrity.file_rows, 140);
1174        assert_eq!(integrity.manifests, 2);
1175    }
1176
1177    #[test]
1178    fn fetch_manifests_keyed_names_the_key_when_a_manifest_is_unparseable() {
1179        let (store, _g) = fs_store(&[("base/manifest.json", b"{ not json".to_vec())]);
1180        let err = fetch_manifests_keyed(&store, "gs://my-bucket/base")
1181            .unwrap_err()
1182            .to_string();
1183        assert!(
1184            err.contains("parsing manifest") && err.contains("base/manifest.json"),
1185            "error should name the offending key: {err}"
1186        );
1187    }
1188
1189    // ── fake-gcs-server: the storage contract over the REAL opendal-GCS transport
1190    //
1191    // Every test above runs the storage-contract fns over an Fs-backed GcsStore —
1192    // proving the LOGIC, but not that opendal's GCS client speaks the same
1193    // list/read/remove semantics a real bucket answers with. This cell closes that
1194    // gap against fsouza/fake-gcs-server (the GCS JSON API emulator the extract
1195    // side already uses), so a GCS-only regression — listing pagination, prefix
1196    // handling, `remove_all` recursion — fails HERE, not in production. The
1197    // warehouse half (bq/snow COPY) can't be emulated; it stays the live BigQuery
1198    // matrix (`smoke_batch_mysql` / `StagingWiped`). This owns the BUCKET half.
1199    const FAKE_GCS_ENDPOINT: &str = "http://127.0.0.1:4443";
1200    const FAKE_GCS_BUCKET: &str = "rivet-load-emulator";
1201
1202    /// A real GCS-transport [`GcsStore`] against the local emulator, scoped to
1203    /// `FAKE_GCS_BUCKET`. Ensures the bucket first via the emulator's JSON API
1204    /// (fake-gcs does not auto-create on write) — an idempotent POST, so a re-run
1205    /// needs no teardown. curl, not a new HTTP dep: this is a dev-only cell.
1206    fn fake_gcs_store() -> GcsStore {
1207        let created = std::process::Command::new("curl")
1208            .args([
1209                "-s",
1210                "-X",
1211                "POST",
1212                &format!("{FAKE_GCS_ENDPOINT}/storage/v1/b?project=rivet-test"),
1213                "-H",
1214                "Content-Type: application/json",
1215                "-d",
1216                &format!("{{\"name\":\"{FAKE_GCS_BUCKET}\"}}"),
1217            ])
1218            .output();
1219        assert!(
1220            created.is_ok_and(|o| o.status.success()),
1221            "could not reach fake-gcs to create the bucket — is `docker compose up -d fake-gcs` running on :4443?"
1222        );
1223        let cfg = crate::config::DestinationConfig {
1224            destination_type: crate::config::DestinationType::Gcs,
1225            bucket: Some(FAKE_GCS_BUCKET.into()),
1226            endpoint: Some(FAKE_GCS_ENDPOINT.into()),
1227            allow_anonymous: true,
1228            ..Default::default()
1229        };
1230        GcsStore::new(&cfg).expect("build GcsStore against fake-gcs")
1231    }
1232
1233    /// Per-object drain of `prefix` — list + single `remove` each. fake-gcs
1234    /// implements single-object DELETE but NOT the GCS batch-delete endpoint that
1235    /// opendal's `remove_all` issues for 2+ objects (that 400s with `deleted: 0`),
1236    /// so the emulator can't run the `cleanup_source` batch wipe. That wipe
1237    /// (`delete_under` → `remove_all`) is verified against a REAL bucket by the
1238    /// live BigQuery `StagingWiped` matrix, and over the Fs seam above; here we
1239    /// drain per-object for idempotent setup/teardown.
1240    fn drain(store: &GcsStore, prefix: &str) {
1241        for key in store.list_files(prefix).unwrap() {
1242            store.remove(&key).unwrap();
1243        }
1244    }
1245
1246    #[test]
1247    #[ignore = "emulator: needs `docker compose up -d fake-gcs` (fsouza/fake-gcs-server :4443)"]
1248    fn storage_contract_over_fake_gcs() {
1249        let store = fake_gcs_store();
1250        // Test-owned prefix, drained first so a re-run starts clean even though the
1251        // emulator bucket persists (no cross-run bleed).
1252        let prefix = "load-contract/orders";
1253        drain(&store, prefix);
1254        let gs = format!("gs://{FAKE_GCS_BUCKET}/{prefix}");
1255
1256        // Seed one Success run (its manifest + committed part) plus a crash orphan
1257        // — an unmanifested `.parquet` an interrupted extract would leave behind.
1258        store
1259            .put(
1260                &format!("{prefix}/manifest-r1.json"),
1261                &manifest_bytes("r1", 100, Some(100)),
1262            )
1263            .unwrap();
1264        store
1265            .put(&format!("{prefix}/part-000000.parquet"), b"rows-of-r1")
1266            .unwrap();
1267        store
1268            .put(&format!("{prefix}/orphan.parquet"), b"crash-leftover")
1269            .unwrap();
1270
1271        // 1. fetch + parse the manifest back over real GCS list+read.
1272        let keyed = fetch_manifests_keyed(&store, &gs).unwrap();
1273        assert_eq!(keyed.len(), 1, "the run's manifest, read back over GCS");
1274
1275        // 2. reconcile → file_rows, the value the warehouse COUNT(*) gate enforces.
1276        let manifests: Vec<_> = keyed.iter().map(|(_, m)| m.clone()).collect();
1277        assert_eq!(
1278            reconcile(&manifests, false).unwrap().file_rows,
1279            100,
1280            "file_rows drives the count-gate; a bad GCS read would corrupt it"
1281        );
1282
1283        // 3. select_load_uris resolves the MANIFESTED part only — never the orphan.
1284        assert_eq!(
1285            select_load_uris(&store, &gs, &keyed).unwrap(),
1286            vec![format!(
1287                "gs://{FAKE_GCS_BUCKET}/{prefix}/part-000000.parquet"
1288            )],
1289            "load pulls the manifested part, not the unmanifested crash orphan"
1290        );
1291
1292        // 4. gc_orphans deletes the orphan over real GCS — a REAL single-object
1293        // DELETE against the emulator — keeping the manifested part + manifest.
1294        let (removed, _bytes) = gc_orphans(&store, &gs, &keyed, false).unwrap();
1295        assert_eq!(
1296            removed, 1,
1297            "exactly the orphan parquet is GC'd over real GCS"
1298        );
1299        let mut left = store.list_files(prefix).unwrap();
1300        left.sort();
1301        assert_eq!(
1302            left,
1303            vec![
1304                format!("{prefix}/manifest-r1.json"),
1305                format!("{prefix}/part-000000.parquet"),
1306            ],
1307            "the manifested part + its manifest survive the orphan GC"
1308        );
1309
1310        // cleanup_source's batch wipe (`remove_all`) is NOT emulatable here — see
1311        // `drain`. Teardown per-object; the empty listing confirms the deletes
1312        // landed over real GCS transport.
1313        drain(&store, prefix);
1314        assert!(
1315            store.list_files(prefix).unwrap().is_empty(),
1316            "teardown left the prefix clean over real GCS"
1317        );
1318    }
1319}