Skip to main content

rivet/load/
plan.rs

1//! Config-driven load planning — derive a BigQuery load (native schema, table,
2//! partition, source URIs) from a rivet export config, so a client never
3//! hand-types column types. The schema comes from rivet's own type resolver
4//! via `rivet check --target bigquery --json` (the argv/process boundary,
5//! ADR-0026); the table/partition/destination come from the parsed config.
6
7use crate::types::target::{TargetColumnSpec, TargetStatus};
8use anyhow::{Context, Result, bail};
9use serde::Deserialize;
10use std::process::Command;
11
12/// The warehouse load target — the config's **top-level `load:` block**,
13/// declared ONCE for all exports. OSS accepts and ignores this block (a
14/// reserved passthrough); the loader reads it here. One config file drives
15/// both the export and the load — no second file, no per-table repetition.
16///
17/// `cleanup_source`/`cluster_by` are target-agnostic; the warehouse and its
18/// connection config live in [`LoadTarget`], keyed on the `target:`
19/// discriminator. NOTE: `#[serde(flatten)]` on `target` DISABLES serde's
20/// `deny_unknown_fields`, so cross-warehouse fields do NOT fail to deserialize —
21/// a `target: snowflake` block silently accepted BigQuery's `project:`/`dataset:`
22/// (dogfood LOW). [`reject_foreign_target_fields`] is the runtime guard that
23/// closes that gap; the "invalid combos fail to deserialize" belief was false.
24#[derive(Debug, Clone, Deserialize)]
25pub struct LoadSection {
26    #[serde(flatten)]
27    pub target: LoadTarget,
28    #[serde(default)]
29    pub cleanup_source: bool,
30    /// Primary key column(s) for the incremental/CDC current-state dedup view —
31    /// the view's PARTITION BY. Required for `mode: incremental` / `mode: cdc`;
32    /// ignored for `full` (which overwrites, no view). Composite key = several
33    /// columns, e.g. `pk: [tenant, id]`.
34    #[serde(default)]
35    pub pk: Vec<String>,
36    /// Load even when a run manifest's source count disagrees with what it
37    /// extracted (source→file drift): warn instead of blocking. The
38    /// file→warehouse count gate and manifest gates still apply.
39    #[serde(default)]
40    pub allow_source_drift: bool,
41    /// After a successful load, delete staged Parquet under the export prefix
42    /// that no `Success` manifest references — crash leftovers from an
43    /// interrupted extract. Keeps the current run's files, manifests, and
44    /// `_SUCCESS`; strictly gentler than `cleanup_source`, which wipes the whole
45    /// prefix. Off by default. ⚠️ Only enable when no extract writes this prefix
46    /// concurrently — it can't tell a crash orphan from a live run's in-flight
47    /// parts (see `reconcile::gc_orphans`); the normal load-after-extract flow is
48    /// safe.
49    #[serde(default)]
50    pub gc_orphans: bool,
51    /// Clustering key column(s) — BigQuery `CLUSTER BY` / Snowflake `CLUSTER BY`.
52    /// Empty = none. Applies at table creation.
53    #[serde(default)]
54    pub cluster_by: Vec<String>,
55}
56
57/// Per-export overrides of the top-level [`LoadSection`] — every field optional,
58/// `None` inherits the top-level value. `target` is present ONLY to reject it:
59/// the warehouse is shared (`plan_loads` runs one `rivet check --target`), so it
60/// stays top-level.
61// `deny_unknown_fields` so a per-export `load:` typo (`gc_orphan`, `cleanupsrc`)
62// fails loudly instead of silently deserializing to the default and dropping the
63// override. (LoadOverride has no `#[serde(flatten)]`, so unlike LoadSection this
64// works directly.)
65#[derive(Debug, Deserialize)]
66#[serde(deny_unknown_fields)]
67struct LoadOverride {
68    #[serde(default)]
69    pk: Option<Vec<String>>,
70    #[serde(default)]
71    cleanup_source: Option<bool>,
72    #[serde(default)]
73    gc_orphans: Option<bool>,
74    #[serde(default)]
75    cluster_by: Option<Vec<String>>,
76    #[serde(default)]
77    allow_source_drift: Option<bool>,
78    /// Only to REJECT — a per-export `load:` cannot re-target the warehouse.
79    #[serde(default)]
80    target: Option<serde_json::Value>,
81}
82
83impl LoadSection {
84    /// The effective load config for one export: this top-level section with the
85    /// export's [`LoadOverride`] applied — each `Some` field replaces, each
86    /// `None` inherits. `target` is never overridden.
87    fn with_override(&self, o: &LoadOverride) -> LoadSection {
88        let mut eff = self.clone();
89        if let Some(pk) = &o.pk {
90            eff.pk = pk.clone();
91        }
92        if let Some(c) = o.cleanup_source {
93            eff.cleanup_source = c;
94        }
95        if let Some(g) = o.gc_orphans {
96            eff.gc_orphans = g;
97        }
98        if let Some(cb) = &o.cluster_by {
99            eff.cluster_by = cb.clone();
100        }
101        if let Some(d) = o.allow_source_drift {
102            eff.allow_source_drift = d;
103        }
104        eff
105    }
106}
107
108/// A warehouse and its connection config. `target:` is the serde discriminator.
109#[derive(Debug, Clone, Deserialize)]
110#[serde(tag = "target", rename_all = "lowercase")]
111pub enum LoadTarget {
112    Bigquery {
113        project: String,
114        dataset: String,
115    },
116    Snowflake {
117        connection: String,
118        warehouse: String,
119        database: String,
120        schema: String,
121        storage_integration: String,
122    },
123}
124
125impl LoadTarget {
126    /// The `--target` name to pass to `rivet check`.
127    pub fn name(&self) -> &'static str {
128        match self {
129            LoadTarget::Bigquery { .. } => "bigquery",
130            LoadTarget::Snowflake { .. } => "snowflake",
131        }
132    }
133}
134
135/// Which load strategy an export's `mode` maps to. Drives BOTH the ledger's
136/// file selection and the warehouse write path:
137/// - `Full` — the export is a complete snapshot; load the LATEST run only and
138///   OVERWRITE (chunked is a parallel full snapshot, same handling).
139/// - `Incremental` — the export is a delta since a cursor; APPEND it to
140///   `<table>__changes` and dedup to current state ordered by the cursor.
141/// - `Cdc` — a change stream; APPEND + dedup by `(__pos, __seq)` with tombstones.
142#[derive(Debug, Clone, Copy, PartialEq, Eq)]
143pub enum LoadMode {
144    Full,
145    Incremental,
146    Cdc,
147}
148
149impl LoadMode {
150    /// The ledger's `mode` discriminator (the `load_run.mode` column) — the single
151    /// source of truth for the string that names each strategy in the state DB, so
152    /// no call site hand-writes a stringly-typed `"full"`/`"cdc"` that can drift.
153    pub fn ledger_str(self) -> &'static str {
154        match self {
155            LoadMode::Full => "full",
156            LoadMode::Incremental => "incremental",
157            LoadMode::Cdc => "cdc",
158        }
159    }
160}
161
162/// What a rivet config resolves to for a BigQuery load.
163#[derive(Debug, Clone)]
164pub struct LoadPlan {
165    /// The declared export NAME (config `name:`), distinct from the warehouse
166    /// `table` — error messages address the export the operator wrote, not the
167    /// table it resolves to (dogfood LOW: require_pk labelled the table as the
168    /// export).
169    pub export_name: String,
170    pub table: String,
171    pub partition_by: Option<String>,
172    pub specs: Vec<TargetColumnSpec>,
173    /// `gs://bucket/base/` — the destination prefix up to the `{partition}`
174    /// token, i.e. the root to list source Parquet under.
175    pub gcs_prefix: String,
176    /// The export's GCS destination (bucket + auth) — the native opendal client
177    /// the load layer lists / reads / deletes through.
178    pub destination: crate::config::DestinationConfig,
179    /// The `load:` target from the same config.
180    pub load: LoadSection,
181    /// The export's mode → the load strategy (see [`LoadMode`]).
182    pub mode: LoadMode,
183    /// The incremental cursor column (from `cursor_column:`) — the dedup view's
184    /// latest-per-PK ordering key. `Some` only for [`LoadMode::Incremental`].
185    pub cursor_column: Option<String>,
186}
187
188/// One export's slice of `rivet check --target X --json`. The tool emits one
189/// such JSON document **per export** (concatenated), so a multi-table config
190/// yields a stream of these — parsed with a streaming deserializer.
191#[derive(Deserialize)]
192struct ExportReport {
193    export: String,
194    columns: Vec<ColReport>,
195}
196
197#[derive(Deserialize)]
198struct ColReport {
199    column: String,
200    target_type: String,
201    target_status: String,
202}
203
204/// Resolve a rivet config into **one [`LoadPlan`] per export** — the shared
205/// top-level `load:` target plus each export's own table / partition / GCS
206/// destination / native schema. `rivet check --json` emits one JSON document
207/// per export, so a multi-table config produces a plan per table, all pointed
208/// at the same warehouse target.
209/// Every key a `load:` block may carry — the [`LoadSection`] fields plus the
210/// flattened [`LoadTarget`] variant fields. The top-level block can't use serde
211/// `deny_unknown_fields` (incompatible with the flattened `target` enum), so we
212/// check its keys by hand — else a typo (`gc_orphan`, `cleanupsource`) silently
213/// deserializes to the default and the setting never applies.
214const LOAD_KEYS: &[&str] = &[
215    "target",
216    "cleanup_source",
217    "pk",
218    "allow_source_drift",
219    "gc_orphans",
220    "cluster_by",
221    // LoadTarget::{Bigquery, Snowflake} variant fields (flattened in).
222    "project",
223    "dataset",
224    "connection",
225    "warehouse",
226    "database",
227    "schema",
228    "storage_integration",
229];
230
231/// Reject any key in a `load:` block that isn't in [`LOAD_KEYS`] — turns a
232/// silently-ignored typo into a loud error naming the valid keys.
233fn check_load_keys(value: &serde_json::Value, whose: &str) -> Result<()> {
234    if let Some(obj) = value.as_object() {
235        for k in obj.keys() {
236            if !LOAD_KEYS.contains(&k.as_str()) {
237                bail!(
238                    "unknown key `{k}` in the {whose} `load:` block — valid keys are: {}",
239                    LOAD_KEYS.join(", ")
240                );
241            }
242        }
243    }
244    Ok(())
245}
246
247/// Warehouse fields that belong to exactly ONE target.
248const BIGQUERY_ONLY: &[&str] = &["project", "dataset"];
249const SNOWFLAKE_ONLY: &[&str] = &[
250    "connection",
251    "warehouse",
252    "database",
253    "schema",
254    "storage_integration",
255];
256
257/// Reject fields that belong to a DIFFERENT warehouse than the resolved
258/// `target`. `#[serde(flatten)]` on the target enum can't `deny_unknown_fields`,
259/// so `LOAD_KEYS` (the union of all warehouses' fields) let a `target: snowflake`
260/// block silently carry — and ignore — BigQuery's `project:`/`dataset:` (and
261/// vice versa) (dogfood LOW). Name the offending key and its warehouse.
262fn reject_foreign_target_fields(
263    value: &serde_json::Value,
264    target: &str,
265    whose: &str,
266) -> Result<()> {
267    let (foreign, other) = match target {
268        "bigquery" => (SNOWFLAKE_ONLY, "snowflake"),
269        "snowflake" => (BIGQUERY_ONLY, "bigquery"),
270        _ => return Ok(()),
271    };
272    if let Some(obj) = value.as_object() {
273        for k in obj.keys() {
274            if foreign.contains(&k.as_str()) {
275                bail!(
276                    "the {whose} `load:` block targets `{target}` but carries `{k}`, a `{other}` \
277                     field — remove it (it would be silently ignored, masking a mis-configured load)"
278                );
279            }
280        }
281    }
282    Ok(())
283}
284
285/// Resolve a load's staging prefix from an export destination.
286///
287/// Expands the same `{date}`/`{export}`/`{table}` placeholders the export wrote
288/// with (`PlaceholderContext::for_today`) so the load lists the ACTUAL prefix
289/// (`exports/orders/`) rather than the literal config token (`exports/{export}/`)
290/// — without this the load found no manifests under the unexpanded path and
291/// reported "up to date" having loaded nothing (#100). `{partition}` is stripped
292/// (its per-partition sub-prefixes live below).
293///
294/// A `{date}` in the load-listed BASE is refused up front: it expands to the
295/// LOAD day, so a nightly export + an after-midnight load list DIFFERENT prefixes
296/// and the load silently reports "up to date" (bughunt HIGH — the same #100
297/// silent-no-load class the expansion above closes for the static tokens, left
298/// open for the day-specific one). `{run_id}` (and any token still unresolved
299/// after expansion) fails loud the same way.
300fn resolve_load_prefix(
301    dest: &crate::config::DestinationConfig,
302    export_name: &str,
303    bucket: &str,
304) -> Result<String> {
305    // Refuse a day-specific `{date}` in the load base BEFORE expansion — once
306    // expanded to the load day, the `contains('{')` guard below can never see it,
307    // so an export written on a different UTC day is silently missed.
308    let raw_prefix = dest.prefix.as_deref().unwrap_or("");
309    let raw_base = raw_prefix.split("{partition}").next().unwrap_or(raw_prefix);
310    if raw_base.contains("{date}") {
311        bail!(
312            "export `{}`: destination.prefix `{}` puts a day-specific `{{date}}` in the load base. \
313             `rivet load` lists the LOAD-day prefix, so an export written on a different day (a \
314             nightly export + an after-midnight load) lands under a different, EMPTY prefix and is \
315             silently reported 'up to date'. Remove `{{date}}` from the load base, or place it \
316             BELOW `{{partition}}` so the load can list a stable prefix.",
317            export_name,
318            raw_base
319        );
320    }
321    let ctx = crate::destination::placeholder::PlaceholderContext::for_today(export_name);
322    let expanded = crate::destination::placeholder::expand_destination(dest.clone(), &ctx);
323    let prefix = expanded.prefix.as_deref().unwrap_or("");
324    let base = prefix.split("{partition}").next().unwrap_or(prefix);
325    if base.contains('{') {
326        bail!(
327            "export `{}`: load prefix `{}` still has an unresolved placeholder after expansion — \
328             `rivet load` cannot reconstruct which run's output to load (a `{{run_id}}` prefix is \
329             run-specific). Drop the run-specific token from `destination.prefix`, or run \
330             `rivet load` from the context that wrote the export.",
331            export_name,
332            base
333        );
334    }
335    Ok(format!("gs://{bucket}/{base}"))
336}
337
338pub fn plan_loads(config_path: &str, rivet_bin: &str) -> Result<Vec<LoadPlan>> {
339    let yaml = std::fs::read_to_string(config_path)
340        .with_context(|| format!("reading config {config_path}"))?;
341    let cfg = crate::config::Config::from_yaml(&yaml).context("parsing rivet config")?;
342    if cfg.exports.is_empty() {
343        bail!("config has no exports");
344    }
345
346    // The `load:` target from the same config (OSS accepts + ignores it),
347    // shared by every export.
348    let load_value = cfg.load.clone().context(
349        "config has no top-level `load:` block — add `load: { target, ... }` to load into a warehouse",
350    )?;
351    check_load_keys(&load_value, "top-level")?;
352    let load: LoadSection = serde_json::from_value(load_value.clone())
353        .context("parsing the top-level `load:` block")?;
354    reject_foreign_target_fields(&load_value, load.target.name(), "top-level")?;
355
356    // Native schema from rivet's own resolver, for the load target — no
357    // hand-typing. One JSON document per export, so parse a stream.
358    let out = Command::new(rivet_bin)
359        .args([
360            "check",
361            "-c",
362            config_path,
363            "--target",
364            load.target.name(),
365            "--json",
366        ])
367        .output()
368        .with_context(|| {
369            format!("running `{rivet_bin} check` — is rivet on PATH? pass --rivet-bin")
370        })?;
371    if !out.status.success() {
372        bail!(
373            "rivet check failed: {}",
374            String::from_utf8_lossy(&out.stderr).trim()
375        );
376    }
377    let reports: Vec<ExportReport> = serde_json::Deserializer::from_slice(&out.stdout)
378        .into_iter::<ExportReport>()
379        .collect::<Result<_, _>>()
380        .context("parsing `rivet check --json` (one document per export)")?;
381
382    build_plans(&cfg, &load, reports)
383}
384
385/// The **pure core** of [`plan_loads`]: map the parsed `rivet check` reports onto
386/// one [`LoadPlan`] per export, given the config and the shared `load:` section.
387///
388/// No I/O — the subprocess and filesystem work is done by [`plan_loads`], and
389/// everything they produced arrives in the args. That makes the per-export
390/// resolution unit-testable without a `rivet` binary: the export→report name
391/// match, `target_status`→[`TargetStatus`] mapping, [`ExportMode`]→[`LoadMode`]
392/// mapping, the `gs://` prefix, the per-export `load:` override, and the
393/// duplicate-target guard.
394fn build_plans(
395    cfg: &crate::config::Config,
396    load: &LoadSection,
397    reports: Vec<ExportReport>,
398) -> Result<Vec<LoadPlan>> {
399    let mut plans = Vec::with_capacity(reports.len());
400    for report in reports {
401        let export = cfg
402            .exports
403            .iter()
404            .find(|e| e.name == report.export)
405            .with_context(|| {
406                format!(
407                    "rivet check reported export `{}` not found in config",
408                    report.export
409                )
410            })?;
411        let table = export.table.clone().unwrap_or_else(|| export.name.clone());
412
413        let dest = &export.destination;
414        let bucket = dest.bucket.as_deref().with_context(|| {
415            format!(
416                "export `{}` has no destination `bucket` — a GCS destination is required",
417                export.name
418            )
419        })?;
420        let gcs_prefix = resolve_load_prefix(dest, &export.name, bucket)?;
421
422        let specs = report
423            .columns
424            .into_iter()
425            .map(|c| TargetColumnSpec {
426                column_name: c.column,
427                target_type: c.target_type,
428                autoload_type: String::new(),
429                status: match c.target_status.as_str() {
430                    "fail" => TargetStatus::Fail,
431                    "warn" => TargetStatus::Warn,
432                    _ => TargetStatus::Ok,
433                },
434                note: None,
435                cast_sql: None,
436            })
437            .collect();
438
439        // Complete-snapshot modes → overwrite the latest run; delta modes → their
440        // own append path. Exhaustive (no `_`) on purpose: a future delta-style
441        // ExportMode then fails to COMPILE here until someone picks its load
442        // semantics, instead of silently defaulting to OVERWRITE (the
443        // incremental-overwrite data-loss class).
444        let mode = match export.mode {
445            crate::config::ExportMode::Cdc => LoadMode::Cdc,
446            crate::config::ExportMode::Incremental => LoadMode::Incremental,
447            crate::config::ExportMode::Full => LoadMode::Full, // whole result set
448            crate::config::ExportMode::Chunked => LoadMode::Full, // parallel full snapshot
449            crate::config::ExportMode::TimeWindow => LoadMode::Full, // whole rolling window
450        };
451        // Effective load config: the shared top-level `load:`, with this export's
452        // own `load:` block overriding the table-specific fields (pk, cleanup, …).
453        // The warehouse `target` is shared and cannot be re-targeted per export.
454        let eff_load = match &export.load {
455            Some(v) => {
456                let o: LoadOverride = serde_json::from_value(v.clone()).with_context(|| {
457                    format!("parsing export `{}` `load:` override", export.name)
458                })?;
459                if o.target.is_some() {
460                    bail!(
461                        "export `{}`: a per-export `load:` cannot override `target:` — the \
462                         warehouse is shared; set `target:` in the top-level `load:` block only",
463                        export.name
464                    );
465                }
466                load.with_override(&o)
467            }
468            None => load.clone(),
469        };
470        plans.push(LoadPlan {
471            export_name: export.name.clone(),
472            table,
473            partition_by: export.partition_by.clone(),
474            specs,
475            gcs_prefix,
476            destination: export.destination.clone(),
477            load: eff_load,
478            mode,
479            cursor_column: export.cursor_column.clone(),
480        });
481    }
482    reject_duplicate_target_tables(&plans.iter().map(|p| p.table.as_str()).collect::<Vec<_>>())?;
483    Ok(plans)
484}
485
486/// Reject two exports that resolve to the SAME warehouse table. The `target:` is
487/// shared, so two exports whose `table:` (or `name:`) resolves alike land on one
488/// warehouse object — a full OVERWRITE would clobber what a cdc/incremental
489/// export appends a `<table>__changes` view over, and they'd share one ledger
490/// skip-set. Pure + unit-testable; caught here, not silently at load time.
491fn reject_duplicate_target_tables(tables: &[&str]) -> Result<()> {
492    let mut seen = std::collections::HashSet::new();
493    for t in tables {
494        if !seen.insert(*t) {
495            bail!(
496                "two exports resolve to the same load target table `{t}` — each would clobber \
497                 the other (a full OVERWRITE vs a cdc/incremental append share the table and \
498                 its ledger). Give each export its own `table:` or destination."
499            );
500        }
501    }
502    Ok(())
503}
504
505/// Resolve the config's source engine into the CDC [`SourceEngine`] the dedup
506/// view's `__pos` parse is keyed on. One config has one source, so this is a
507/// job-wide property. MongoDB is supported too: its change stream carries a
508/// document `_id` (the dedup partition key) and an order-preserving `_data`
509/// resume token in `__pos`, so the current-state view applies just as it does to
510/// the relational engines.
511pub fn source_engine(config_path: &str) -> Result<crate::load::cdc::SourceEngine> {
512    use crate::config::SourceType;
513    use crate::load::cdc::SourceEngine;
514
515    let yaml = std::fs::read_to_string(config_path)
516        .with_context(|| format!("reading config {config_path}"))?;
517    let cfg = crate::config::Config::from_yaml(&yaml).context("parsing rivet config")?;
518    match cfg.source.source_type {
519        SourceType::Postgres => Ok(SourceEngine::Postgres),
520        SourceType::Mysql => Ok(SourceEngine::MySql),
521        SourceType::Mssql => Ok(SourceEngine::SqlServer),
522        SourceType::Mongo => Ok(SourceEngine::Mongo),
523    }
524}
525
526#[cfg(test)]
527mod tests {
528    use super::*;
529
530    #[test]
531    fn ledger_str_names_each_mode_stably() {
532        // The state DB's `load_run.mode` discriminator — every mode must map to
533        // its exact stable string, since retry/skip logic keys off it. A drifted
534        // value would mislabel loads in the ledger.
535        assert_eq!(LoadMode::Full.ledger_str(), "full");
536        assert_eq!(LoadMode::Incremental.ledger_str(), "incremental");
537        assert_eq!(LoadMode::Cdc.ledger_str(), "cdc");
538    }
539
540    #[test]
541    fn resolve_load_prefix_expands_deterministic_tokens_and_refuses_run_specific() {
542        use crate::config::{DestinationConfig, DestinationType};
543        let dest = |prefix: &str| DestinationConfig {
544            destination_type: DestinationType::Gcs,
545            bucket: Some("BKT".into()),
546            prefix: Some(prefix.into()),
547            ..Default::default()
548        };
549
550        // {export}/{table} are deterministic → the load must list the ACTUAL
551        // prefix, not the literal token. (#100: the load listed `exports/{export}/`
552        // verbatim, found no manifests, and reported "up to date" — loaded nothing.)
553        assert_eq!(
554            resolve_load_prefix(&dest("exports/{export}/"), "orders", "BKT").unwrap(),
555            "gs://BKT/exports/orders/"
556        );
557        // {partition} is stripped; {table} is an alias for {export}.
558        assert_eq!(
559            resolve_load_prefix(&dest("e/{table}/{partition}/"), "orders", "BKT").unwrap(),
560            "gs://BKT/e/orders/"
561        );
562        // {date} in the load base is DAY-specific → refused (bughunt HIGH: it
563        // expanded to the LOAD day, so a cross-midnight load silently listed an
564        // empty prefix). See resolve_load_prefix_refuses_day_specific_date_in_the_base.
565        assert!(resolve_load_prefix(&dest("d/{date}/{export}/"), "orders", "BKT").is_err());
566
567        // {run_id} is run-specific and unknowable here → refuse LOUD, never a
568        // literal-token listing that silently loads nothing.
569        let err = resolve_load_prefix(&dest("e/{run_id}/"), "orders", "BKT").unwrap_err();
570        assert!(
571            err.to_string().contains("unresolved placeholder"),
572            "a run-specific token must be refused, not silently listed: {err}"
573        );
574    }
575
576    /// A `ColReport` with an explicit `target_status`.
577    fn col(name: &str, status: &str) -> ColReport {
578        ColReport {
579            column: name.into(),
580            target_type: "STRING".into(),
581            target_status: status.into(),
582        }
583    }
584
585    /// Drive the PURE `build_plans` (no `rivet` subprocess) — the deepened core
586    /// of `plan_loads`. Kills the mutation survivors that live in the per-export
587    /// resolution: the export→report name match (`==`→`!=`) and the `fail`/`warn`
588    /// `target_status` arms. Also pins mode mapping, the `gs://` prefix, table
589    /// resolution, and the cursor column.
590    #[test]
591    fn build_plans_matches_by_name_maps_statuses_and_mode() {
592        let cfg = crate::config::Config::from_yaml(
593            r#"
594source:
595  type: postgres
596  url: "postgresql://localhost/test"
597exports:
598  - name: alpha
599    table: alpha_tbl
600    mode: full
601    format: parquet
602    destination:
603      type: gcs
604      bucket: b1
605      prefix: exports/alpha/
606  - name: beta
607    table: beta_tbl
608    mode: incremental
609    cursor_column: updated_at
610    format: parquet
611    destination:
612      type: gcs
613      bucket: b2
614      prefix: exports/beta/
615load:
616  target: bigquery
617  project: p
618  dataset: d
619"#,
620        )
621        .unwrap();
622        let load: LoadSection = serde_json::from_value(cfg.load.clone().unwrap()).unwrap();
623
624        // Reports arrive in the OPPOSITE order to the exports, so a plan only
625        // lands on the right table if its export is found by NAME, not position.
626        let reports = vec![
627            ExportReport {
628                export: "beta".into(),
629                columns: vec![col("id", "ok"), col("f", "fail"), col("w", "warn")],
630            },
631            ExportReport {
632                export: "alpha".into(),
633                columns: vec![col("id", "ok")],
634            },
635        ];
636
637        let plans = build_plans(&cfg, &load, reports).unwrap();
638        assert_eq!(plans.len(), 2);
639
640        // reports[0] = beta → matched by name to the 2nd export (kills `==`→`!=`,
641        // which would resolve the first NON-matching export instead).
642        assert_eq!(
643            plans[0].table, "beta_tbl",
644            "found beta by name, not position"
645        );
646        assert_eq!(plans[0].mode, LoadMode::Incremental);
647        assert_eq!(plans[0].cursor_column.as_deref(), Some("updated_at"));
648        assert_eq!(plans[0].gcs_prefix, "gs://b2/exports/beta/");
649        // target_status → spec.status (kills the `fail`/`warn` arm deletions,
650        // which would collapse those columns to Ok and load an unmappable column).
651        let statuses: Vec<_> = plans[0].specs.iter().map(|s| s.status).collect();
652        assert_eq!(
653            statuses,
654            vec![TargetStatus::Ok, TargetStatus::Fail, TargetStatus::Warn]
655        );
656
657        // reports[1] = alpha → the full-snapshot export.
658        assert_eq!(plans[1].table, "alpha_tbl");
659        assert_eq!(plans[1].mode, LoadMode::Full);
660        assert_eq!(plans[1].gcs_prefix, "gs://b1/exports/alpha/");
661    }
662
663    #[test]
664    fn build_plans_bails_on_a_report_for_an_unknown_export() {
665        let cfg = crate::config::Config::from_yaml(
666            "source:\n  type: postgres\n  url: \"postgresql://localhost/test\"\n\
667             exports:\n  - name: a\n    query: \"SELECT 1\"\n    format: parquet\n    \
668             destination:\n      type: gcs\n      bucket: b\n      prefix: p/\nload:\n  \
669             target: bigquery\n  project: p\n  dataset: d\n",
670        )
671        .unwrap();
672        let load: LoadSection = serde_json::from_value(cfg.load.clone().unwrap()).unwrap();
673        let reports = vec![ExportReport {
674            export: "ghost".into(),
675            columns: vec![],
676        }];
677        let err = build_plans(&cfg, &load, reports).unwrap_err().to_string();
678        assert!(err.contains("ghost") && err.contains("not found"), "{err}");
679    }
680
681    #[test]
682    fn reject_duplicate_target_tables_catches_a_collision() {
683        // Two exports resolving to the same warehouse table would clobber each
684        // other — caught at plan time, not silently at load time.
685        assert!(reject_duplicate_target_tables(&["orders", "events", "orders"]).is_err());
686        assert!(reject_duplicate_target_tables(&["orders", "events"]).is_ok());
687        assert!(reject_duplicate_target_tables(&[]).is_ok());
688    }
689
690    #[test]
691    fn multi_export_check_json_parses_as_a_stream() {
692        // `rivet check --json` emits ONE document per export (no wrapping array).
693        // A single-object parse fails "trailing characters"; the stream parser
694        // yields one ExportReport per table. Guards the multi-table regression.
695        let raw = concat!(
696            "{\"export\":\"orders\",\"columns\":[{\"column\":\"id\",\"target_type\":\"NUMBER\",\"target_status\":\"ok\"}]}\n",
697            "{\"export\":\"customers\",\"columns\":[{\"column\":\"cid\",\"target_type\":\"NUMBER\",\"target_status\":\"ok\"}]}\n"
698        );
699        let reports: Vec<ExportReport> = serde_json::Deserializer::from_slice(raw.as_bytes())
700            .into_iter::<ExportReport>()
701            .collect::<Result<_, _>>()
702            .unwrap();
703        assert_eq!(reports.len(), 2);
704        assert_eq!(reports[0].export, "orders");
705        assert_eq!(reports[1].export, "customers");
706        assert_eq!(reports[1].columns[0].column, "cid");
707    }
708
709    #[test]
710    fn bigquery_load_section_deserializes_into_its_variant() {
711        let value = serde_json::json!({
712            "target": "bigquery", "project": "p", "dataset": "d",
713            "cleanup_source": true, "cluster_by": ["customer"]
714        });
715        let load: LoadSection = serde_json::from_value(value).unwrap();
716        assert_eq!(load.target.name(), "bigquery");
717        assert!(load.cleanup_source);
718        assert_eq!(load.cluster_by, vec!["customer"]);
719        match load.target {
720            LoadTarget::Bigquery { project, dataset } => {
721                assert_eq!((project.as_str(), dataset.as_str()), ("p", "d"));
722            }
723            _ => panic!("expected Bigquery variant"),
724        }
725    }
726
727    #[test]
728    fn snowflake_missing_field_is_unrepresentable_no_runtime_validate() {
729        let full = serde_json::json!({
730            "target": "snowflake", "connection": "rivet", "warehouse": "wh",
731            "database": "db", "schema": "sc", "storage_integration": "si"
732        });
733        let load: LoadSection = serde_json::from_value(full).unwrap();
734        assert_eq!(load.target.name(), "snowflake");
735
736        // A snowflake block missing storage_integration doesn't deserialize —
737        // the type makes it unrepresentable, so there is no runtime validate().
738        let partial = serde_json::json!({
739            "target": "snowflake", "connection": "rivet", "warehouse": "wh",
740            "database": "db", "schema": "sc"
741        });
742        let err = serde_json::from_value::<LoadSection>(partial).unwrap_err();
743        assert!(
744            err.to_string().contains("storage_integration"),
745            "error should name the missing field: {err}"
746        );
747    }
748
749    #[test]
750    fn unknown_target_is_rejected_at_deserialize() {
751        let value = serde_json::json!({ "target": "redshift", "project": "p" });
752        assert!(serde_json::from_value::<LoadSection>(value).is_err());
753    }
754
755    fn top_level_load() -> LoadSection {
756        serde_json::from_value(serde_json::json!({
757            "target": "bigquery", "project": "p", "dataset": "d",
758            "pk": ["top"], "cleanup_source": true, "gc_orphans": false,
759            "cluster_by": ["c0"], "allow_source_drift": false,
760        }))
761        .unwrap()
762    }
763
764    #[test]
765    fn with_override_replaces_some_fields_and_inherits_the_rest() {
766        let top = top_level_load();
767        // Override ONLY pk + gc_orphans; the rest must inherit the top-level.
768        let o: LoadOverride =
769            serde_json::from_value(serde_json::json!({ "pk": ["id"], "gc_orphans": true }))
770                .unwrap();
771        let eff = top.with_override(&o);
772        assert_eq!(eff.pk, vec!["id"], "pk replaced");
773        assert!(eff.gc_orphans, "gc_orphans replaced");
774        assert!(
775            eff.cleanup_source,
776            "cleanup_source inherited (top-level true)"
777        );
778        assert_eq!(eff.cluster_by, vec!["c0"], "cluster_by inherited");
779        assert!(!eff.allow_source_drift, "allow_source_drift inherited");
780    }
781
782    #[test]
783    fn override_parsing_leaves_omitted_fields_none() {
784        let o: LoadOverride = serde_json::from_value(serde_json::json!({ "pk": ["id"] })).unwrap();
785        assert_eq!(o.pk.as_deref(), Some(&["id".to_string()][..]));
786        assert!(o.cleanup_source.is_none());
787        assert!(o.gc_orphans.is_none());
788        assert!(o.cluster_by.is_none());
789        assert!(o.allow_source_drift.is_none());
790        assert!(o.target.is_none());
791    }
792
793    #[test]
794    fn empty_override_is_distinct_from_inherit() {
795        let top = top_level_load();
796        // An EXPLICIT empty pk clears the inherited one; a missing pk keeps it.
797        let cleared: LoadOverride =
798            serde_json::from_value(serde_json::json!({ "pk": [] })).unwrap();
799        assert!(
800            top.with_override(&cleared).pk.is_empty(),
801            "explicit [] clears"
802        );
803        let inherit: LoadOverride = serde_json::from_value(serde_json::json!({})).unwrap();
804        assert_eq!(
805            top.with_override(&inherit).pk,
806            vec!["top"],
807            "omitted inherits"
808        );
809    }
810
811    #[test]
812    fn override_carrying_target_is_detected() {
813        // The plan_loads guard rejects a per-export `load:` that re-targets the
814        // warehouse; the override captures `target:` so the guard can see it.
815        let o: LoadOverride =
816            serde_json::from_value(serde_json::json!({ "target": "snowflake" })).unwrap();
817        assert!(
818            o.target.is_some(),
819            "target captured for the plan_loads guard"
820        );
821    }
822
823    #[test]
824    fn unknown_top_level_load_key_is_rejected() {
825        // A typo (`gc_orphan` for `gc_orphans`) must fail loudly, not silently
826        // deserialize to the default so the setting never applies.
827        let typo = serde_json::json!({
828            "target": "bigquery", "project": "p", "dataset": "d", "gc_orphan": true
829        });
830        let err = check_load_keys(&typo, "top-level").unwrap_err().to_string();
831        assert!(err.contains("gc_orphan"), "{err}");
832        // Every valid LoadSection + LoadTarget key passes.
833        let ok = serde_json::json!({
834            "target": "bigquery", "project": "p", "dataset": "d",
835            "gc_orphans": true, "cleanup_source": false, "pk": ["id"],
836            "allow_source_drift": true, "cluster_by": ["a"]
837        });
838        assert!(check_load_keys(&ok, "top-level").is_ok());
839    }
840
841    #[test]
842    fn unknown_per_export_override_key_is_rejected() {
843        // `deny_unknown_fields` on LoadOverride catches per-export typos.
844        let typo = serde_json::json!({ "pk": ["id"], "cluster_bye": ["x"] });
845        assert!(
846            serde_json::from_value::<LoadOverride>(typo).is_err(),
847            "a typo'd override key must fail to parse"
848        );
849        let ok = serde_json::json!({ "pk": ["id"], "cleanup_source": true });
850        assert!(serde_json::from_value::<LoadOverride>(ok).is_ok());
851    }
852
853    #[test]
854    fn resolve_load_prefix_refuses_day_specific_date_in_the_base() {
855        // #bughunt HIGH: {date} expands to the LOAD day, so a nightly export + an
856        // after-midnight load list DIFFERENT prefixes → silent "up to date". Refuse
857        // {date} in the load base; a static base (or {date} below {partition}) is ok.
858        let dest = |p: &str| crate::config::DestinationConfig {
859            prefix: Some(p.to_string()),
860            ..Default::default()
861        };
862        let err = resolve_load_prefix(&dest("exports/{date}/{export}/"), "orders", "bkt")
863            .unwrap_err()
864            .to_string();
865        assert!(
866            err.contains("{date}") && err.contains("load base"),
867            "must refuse a day-specific date base: {err}"
868        );
869        assert!(resolve_load_prefix(&dest("exports/{export}/"), "orders", "bkt").is_ok());
870        // {date} BELOW {partition} is not in the listed base — allowed.
871        assert!(
872            resolve_load_prefix(
873                &dest("exports/{export}/{partition}/{date}/"),
874                "orders",
875                "bkt"
876            )
877            .is_ok()
878        );
879    }
880
881    #[test]
882    fn cross_warehouse_load_fields_are_rejected() {
883        // #dogfood LOW: `#[serde(flatten)]` on the target enum disables
884        // deny_unknown_fields, so a `target: snowflake` block silently accepted
885        // (and ignored) BigQuery's `project:`/`dataset:`. Now a loud error.
886        let snow_with_bq = serde_json::json!({
887            "target": "snowflake", "connection": "c", "warehouse": "w",
888            "database": "db", "schema": "s", "storage_integration": "si",
889            "project": "STALE", "dataset": "STALE"
890        });
891        let err = reject_foreign_target_fields(&snow_with_bq, "snowflake", "top-level")
892            .unwrap_err()
893            .to_string();
894        assert!(
895            err.contains("project") && err.contains("bigquery"),
896            "snowflake+project must name the foreign field and warehouse: {err}"
897        );
898        // The reverse: bigquery target carrying a snowflake-only field.
899        let bq_with_snow = serde_json::json!({
900            "target": "bigquery", "project": "p", "dataset": "d", "warehouse": "WH"
901        });
902        assert!(reject_foreign_target_fields(&bq_with_snow, "bigquery", "top-level").is_err());
903        // A clean, target-matching block passes.
904        let clean = serde_json::json!({ "target": "bigquery", "project": "p", "dataset": "d" });
905        assert!(reject_foreign_target_fields(&clean, "bigquery", "top-level").is_ok());
906    }
907}