Skip to main content

rivet/config/
export.rs

1//! Per-export configuration: query/table/mode, chunking, format, destination link.
2//!
3//! `SchemaDriftPolicy` lives here because it is only ever read via
4//! [`ExportConfig::on_schema_drift`].
5
6use std::path::Path;
7
8use schemars::JsonSchema;
9use serde::{Deserialize, Serialize};
10
11use super::IncrementalCursorMode;
12use super::destination::DestinationConfig;
13use super::format::{CompressionProfile, CompressionType, FormatType, ParquetConfig};
14use super::resolve::{parse_file_size, resolve_vars};
15use crate::tuning::TuningConfig;
16
17/// What to do when structural schema drift is detected (column added, removed, or retyped).
18///
19/// ```yaml
20/// exports:
21///   - name: orders
22///     on_schema_drift: fail   # warn (default), continue, fail
23/// ```
24/// How deep `--validate` must verify each part's integrity.
25#[derive(Debug, Deserialize, Serialize, JsonSchema, Clone, Copy, PartialEq, Eq, Default)]
26#[serde(rename_all = "snake_case")]
27pub enum VerifyMode {
28    /// Accept size-only verification when no content checksum is available.
29    #[default]
30    Size,
31    /// Require every part's content to be MD5-verified against the store's
32    /// listing; fail validation for any part that is only size-verified.
33    Content,
34}
35
36impl VerifyMode {
37    /// Whether content (not just size) verification is required.
38    pub fn requires_content(self) -> bool {
39        matches!(self, VerifyMode::Content)
40    }
41}
42
43#[derive(Debug, Deserialize, Serialize, JsonSchema, Clone, Copy, PartialEq, Eq, Default)]
44#[serde(rename_all = "snake_case")]
45pub enum SchemaDriftPolicy {
46    /// Log a warning and continue. The new schema fingerprint is stored. (Default.)
47    #[default]
48    Warn,
49    /// Silently accept schema changes — store the new schema, no log output.
50    Continue,
51    /// Abort the run with a non-zero exit. The schema store is NOT updated so the
52    /// next run will detect the same change again.
53    Fail,
54}
55#[derive(Debug, Deserialize, JsonSchema, Clone)]
56#[serde(deny_unknown_fields)]
57pub struct ExportConfig {
58    pub name: String,
59    #[serde(default)]
60    pub query: Option<String>,
61    pub query_file: Option<String>,
62    /// Shortcut for `query: "SELECT * FROM <schema>.<table>"`.
63    ///
64    /// Accepts `table` or `schema.table` with ASCII-only identifiers
65    /// (`[A-Za-z_][A-Za-z0-9_]*`). Generates an unquoted single-table
66    /// query so the Postgres NUMERIC catalog-hint resolver recognises it
67    /// and auto-types `numeric(p,s)` columns without manual overrides.
68    ///
69    /// Mutually exclusive with `query` and `query_file`.
70    #[serde(default)]
71    pub table: Option<String>,
72    /// CDC only: capture **several** tables through ONE change stream (one
73    /// PostgreSQL slot / one MySQL binlog connection) instead of one export —
74    /// and one slot — per table. Each table's parts land under
75    /// `<destination>/<table>/` with their own `manifest.json` + `_SUCCESS`;
76    /// the checkpoint (stream position) is shared. Mutually exclusive with
77    /// `table:`. Not yet supported for SQL Server (capture instances are
78    /// per-table).
79    #[serde(default)]
80    pub tables: Option<Vec<String>>,
81    #[serde(default = "default_mode")]
82    pub mode: ExportMode,
83    /// Change-data-capture settings, required when `mode: cdc`. Reuses the
84    /// export's `table`, `destination`, and `format`; carries only the
85    /// CDC-specific knobs (resume checkpoint, per-engine stream params).
86    #[serde(default)]
87    pub cdc: Option<CdcExportConfig>,
88    pub cursor_column: Option<String>,
89    /// Secondary column for [`IncrementalCursorMode::Coalesce`] only (see ADR-0007).
90    #[serde(default)]
91    pub cursor_fallback_column: Option<String>,
92    /// How primary (and optional fallback) columns drive incremental progression.
93    #[serde(default)]
94    pub incremental_cursor_mode: IncrementalCursorMode,
95    pub chunk_column: Option<String>,
96    #[serde(default)]
97    pub chunk_dense: bool,
98    #[serde(default = "default_chunk_size")]
99    pub chunk_size: usize,
100    /// Target memory budget per chunk in MB. When set, `chunk_size` is derived
101    /// from this budget at plan-build time using a `pg_class` row-size estimate
102    /// (`pg_relation_size / reltuples`), clamped to `[10_000, 5_000_000]` rows.
103    ///
104    /// Mutually exclusive with an explicit non-default `chunk_size:`. Only
105    /// applies to `mode: chunked` on a Postgres source using the `table:`
106    /// shortcut (the row-size probe needs a known relation).
107    ///
108    /// ```yaml
109    /// exports:
110    ///   - name: page_views
111    ///     table: public.page_views
112    ///     mode: chunked
113    ///     chunk_size_memory_mb: 256
114    /// ```
115    #[serde(default)]
116    pub chunk_size_memory_mb: Option<u64>,
117    /// Divide the column range into exactly this many equal chunks.
118    /// Mutually exclusive with `chunk_dense` and `chunk_by_days`.
119    /// When set, `chunk_size` is computed dynamically from min/max.
120    pub chunk_count: Option<usize>,
121    pub chunk_by_days: Option<u32>,
122    /// Keyset (seek) pagination on this single index-backed unique key — the
123    /// source-safe shape for tables without a single-integer PK (OPT-4). The
124    /// column MUST be backed by a usable index (PK or unique); the planner
125    /// refuses a non-indexed key rather than emit a full-scan + filesort query.
126    pub chunk_by_key: Option<String>,
127    #[serde(default = "default_parallel")]
128    pub parallel: usize,
129
130    /// Advisory execution wave (1 = highest priority, run first). Written by
131    /// `rivet plan` from the source-aware prioritization score (see ADR-0006)
132    /// and consumed by `rivet apply`, which runs exports wave-by-wave in
133    /// ascending order. `None` = unscheduled (apply treats it as the last wave).
134    /// Operators may hand-edit it; a later `rivet plan` refreshes it in place.
135    #[serde(default, skip_serializing_if = "Option::is_none")]
136    pub wave: Option<u32>,
137
138    /// Whether this export is cheap enough to run concurrently with its
139    /// wave-mates under `rivet apply --parallel-export-processes`. Written by
140    /// `rivet plan` (true when the source-aware cost class is `Low`, i.e.
141    /// < ~100K rows); a heavier table already chunk-parallelizes internally, so
142    /// two of them at once would overload the source. `None`/`false` → the
143    /// export runs alone within its wave. Operators may hand-edit it; a later
144    /// `rivet plan` refreshes it in place.
145    #[serde(default, skip_serializing_if = "Option::is_none")]
146    pub parallel_safe: Option<bool>,
147    pub time_column: Option<String>,
148    #[serde(default = "default_time_column_type")]
149    pub time_column_type: TimeColumnType,
150    pub days_window: Option<u32>,
151
152    /// Date/time output partitioning: split this export's rows into one
153    /// destination sub-prefix per calendar bucket of this **DATE or TIMESTAMP**
154    /// column, bucketed by `partition_granularity`
155    /// (`day` / `month` / `year`), in a Hive-style `col=value/` layout
156    /// (`created_at=2023-01-01/`, `created_at=2023-01/`, `created_at=2023/`).
157    /// Requires a `{partition}` token in `destination.path` /
158    /// `destination.prefix`.
159    ///
160    /// This is **not** arbitrary value partitioning: the column's min/max is
161    /// read and parsed as a date to generate contiguous calendar buckets, so a
162    /// non-temporal column (e.g. `partition_by: status`) fails at run time with
163    /// "could not parse partition min `<value>` from column `<col>` as a date".
164    /// To split by a categorical column, write one export per value with a
165    /// `WHERE` filter instead.
166    ///
167    /// Orthogonal to `mode`: each partition runs the export's own mode, so
168    /// `mode: chunked` chunks *within* a day. Rows whose partition column is
169    /// NULL land in `col=__HIVE_DEFAULT_PARTITION__/` (Hive default partition)
170    /// so no row is silently dropped. Not compatible with `mode: time_window`.
171    ///
172    /// ```yaml
173    /// exports:
174    ///   - name: events
175    ///     table: events
176    ///     partition_by: created_at        # must be a DATE or TIMESTAMP column
177    ///     partition_granularity: day
178    ///     destination:
179    ///       type: s3
180    ///       bucket: my-bucket
181    ///       prefix: "events/{partition}/"   # → events/created_at=2023-01-01/
182    /// ```
183    #[serde(default)]
184    pub partition_by: Option<String>,
185
186    /// Calendar bucket width for `partition_by`:
187    /// `day` (default), `month`, or `year`. Determines how the partition
188    /// column's date/timestamp range is split into contiguous Hive buckets
189    /// (`col=2023-01-01/` / `col=2023-01/` / `col=2023/`). Has no effect
190    /// unless `partition_by` is set.
191    #[serde(default)]
192    pub partition_granularity: PartitionGranularity,
193    pub format: FormatType,
194    #[serde(default)]
195    pub compression: CompressionType,
196    pub compression_level: Option<u32>,
197    pub compression_profile: Option<CompressionProfile>,
198    #[serde(default)]
199    pub skip_empty: bool,
200    pub destination: DestinationConfig,
201    /// Integrity depth required of `--validate` for this export's parts.
202    /// `size` (default) accepts size-only verification; `content` requires every
203    /// part's content MD5 to be checked against the store's listing (no
204    /// download) and **fails** validation for any part that could only be
205    /// size-verified — e.g. a part too large to upload as a single PUT (raise
206    /// `max_file_size` down so it fits), or a backend that exposes no checksum.
207    #[serde(default)]
208    pub verify: VerifyMode,
209    #[serde(default)]
210    pub meta_columns: MetaColumns,
211    #[serde(default)]
212    pub quality: Option<QualityConfig>,
213    /// Rotate to a new part when the current file reaches this size.
214    /// Accepts `B`/`KB`/`MB`/`GB` (case-insensitive) or a bare byte count;
215    /// a fractional value is allowed (`1.5GB`). Units are binary (IEC-style):
216    /// `KB` = 1024 bytes, `MB` = 1024 KB, `GB` = 1024 MB. Example: `256MB`.
217    pub max_file_size: Option<String>,
218    #[serde(default)]
219    pub chunk_checkpoint: bool,
220    pub chunk_max_attempts: Option<u32>,
221    #[serde(default)]
222    pub tuning: Option<TuningConfig>,
223    /// Optional logical group for shared source capacity (replica, host). Advisory prioritization only.
224    #[serde(default)]
225    pub source_group: Option<String>,
226    /// Hint (Epic C / ADR-0006) that this export should always be treated as reconcile-heavy
227    /// by planning, independent of the `--reconcile` CLI flag. Advisory only.
228    #[serde(default)]
229    pub reconcile_required: bool,
230
231    /// Per-column type overrides (roadmap §8). Keys are column names; values
232    /// are short type strings such as `decimal(18,2)`, `timestamp_tz`, `json`.
233    ///
234    /// ```yaml
235    /// exports:
236    ///   - name: payments
237    ///     columns:
238    ///       amount: decimal(18,2)
239    ///       fee: decimal(18,6)
240    ///       created_at: timestamp_tz
241    /// ```
242    ///
243    /// Overrides take priority over autodetection and are validated at
244    /// plan time — an invalid type string fails before the export runs.
245    #[serde(default)]
246    pub columns: std::collections::HashMap<String, String>,
247
248    /// Downstream warehouse this export targets (`bigquery` / `bq`,
249    /// `duckdb`). When set, `rivet check --type-report` resolves each column
250    /// against it (native type, honest autoload type, recovery hint) without
251    /// needing `--target` on the CLI — the CLI flag still wins when both are
252    /// present. The Parquet interchange stays target-neutral (ADR-0014 T2);
253    /// `target:` only drives guidance and the future load-schema artifact.
254    ///
255    /// ```yaml
256    /// exports:
257    ///   - name: payments
258    ///     target: bigquery
259    /// ```
260    #[serde(default)]
261    pub target: Option<String>,
262
263    /// Per-export overrides for the top-level `load:` block (`pk`,
264    /// `cleanup_source`, `gc_orphans`, `cluster_by`, `allow_source_drift`); any
265    /// field omitted here inherits the top-level value. The warehouse `target`
266    /// is shared and stays in the top-level `load:` — it cannot be overridden
267    /// per export.
268    ///
269    /// ```yaml
270    /// load: { target: bigquery, project: p, dataset: d }   # shared default
271    /// exports:
272    ///   - name: orders
273    ///     table: orders
274    ///     mode: cdc
275    ///     load: { pk: [id] }                                # this table's pk
276    /// ```
277    ///
278    /// Raw JSON (parsed by the load module) so `config` carries no load types —
279    /// mirrors the top-level [`crate::config::Config::load`].
280    #[serde(default)]
281    pub load: Option<serde_json::Value>,
282
283    /// Policy applied when structural schema drift is detected (column added, removed, or retyped).
284    /// Defaults to `warn`: log a warning and continue.
285    #[serde(default)]
286    pub on_schema_drift: SchemaDriftPolicy,
287
288    /// Growth-factor threshold for data shape drift warnings (Epic 8).
289    /// When a string/binary column's max observed byte length in the current run
290    /// exceeds `stored_max * shape_drift_warn_factor`, Rivet logs a warning.
291    /// `None` uses the default of 2.0. Set to `0.0` to disable shape tracking.
292    #[serde(default)]
293    pub shape_drift_warn_factor: Option<f64>,
294
295    /// Parquet row group tuning. Only meaningful when `format: parquet`.
296    /// When absent, the parquet library default (1,048,576 rows/group) is used.
297    #[serde(default)]
298    pub parquet: Option<ParquetConfig>,
299}
300
301impl ExportConfig {
302    /// Resolve the effective `(CompressionType, level)` for this export.
303    /// `compression_profile` takes precedence over `compression` + `compression_level`.
304    ///
305    /// L24: when a profile is set *and* a conflicting explicit codec/level was
306    /// written, warn once that the profile wins rather than silently dropping the
307    /// explicit choice. An explicit codec is only detectable when it differs from
308    /// the `#[serde(default)]` (Zstd) — a literal `compression: zstd` alongside a
309    /// profile is indistinguishable from an omitted field and stays silent.
310    pub fn effective_compression(&self) -> (CompressionType, Option<u32>) {
311        if let Some(profile) = self.compression_profile {
312            let explicit_codec =
313                (self.compression != CompressionType::default()).then_some(self.compression);
314            if let Some(msg) = super::format::compression_profile_override_warning(
315                profile,
316                explicit_codec,
317                self.compression_level,
318            ) {
319                log::warn!("export '{}': {}", self.name, msg);
320            }
321            profile.to_codec()
322        } else {
323            (self.compression, self.compression_level)
324        }
325    }
326
327    pub fn max_file_size_bytes(&self) -> Option<u64> {
328        self.max_file_size
329            .as_ref()
330            .and_then(|s| parse_file_size(s).ok())
331    }
332
333    pub fn resolve_query(
334        &self,
335        config_dir: &Path,
336        params: Option<&std::collections::HashMap<String, String>>,
337    ) -> crate::error::Result<String> {
338        // table: shortcut takes precedence — already validated by
339        // `validate_business_rules` to be mutually exclusive with query/query_file.
340        if let Some(tbl) = &self.table {
341            validate_table_shortcut_ident(&self.name, tbl)?;
342            return Ok(format!("SELECT * FROM {tbl}"));
343        }
344        match (&self.query, &self.query_file) {
345            (Some(q), None) => {
346                if params.is_some() {
347                    resolve_vars(q, params)
348                } else {
349                    Ok(q.clone())
350                }
351            }
352            (None, Some(file)) => {
353                let file_path = std::path::Path::new(file);
354                // SecOps: block absolute paths and `..` traversal components.
355                if file_path.is_absolute() {
356                    anyhow::bail!(
357                        "export '{}': query_file must be a relative path: '{}'",
358                        self.name,
359                        file
360                    );
361                }
362                if file_path
363                    .components()
364                    .any(|c| c == std::path::Component::ParentDir)
365                {
366                    anyhow::bail!(
367                        "export '{}': query_file path must not contain '..': '{}'",
368                        self.name,
369                        file
370                    );
371                }
372                let joined = config_dir.join(file);
373                // Canonicalize-based check catches symlink-based evasion for files
374                // that already exist on disk.
375                if let Ok(canonical) = joined.canonicalize() {
376                    let base = config_dir
377                        .canonicalize()
378                        .unwrap_or_else(|_| config_dir.to_path_buf());
379                    if !canonical.starts_with(&base) {
380                        anyhow::bail!(
381                            "export '{}': query_file '{}' resolves outside the config directory",
382                            self.name,
383                            file
384                        );
385                    }
386                }
387                let raw = std::fs::read_to_string(&joined)?;
388                resolve_vars(&raw, params)
389            }
390            (Some(_), Some(_)) => {
391                anyhow::bail!(
392                    "export '{}': specify either 'query' or 'query_file', not both",
393                    self.name
394                )
395            }
396            (None, None) => {
397                anyhow::bail!(
398                    "export '{}': must specify exactly one of 'query', 'query_file', or 'table'",
399                    self.name
400                )
401            }
402        }
403    }
404}
405
406/// Validate the value of the `table:` YAML shortcut.
407///
408/// Accepts ASCII identifiers in the form `<table>` or `<schema>.<table>`. Each
409/// segment must match `[A-Za-z_][A-Za-z0-9_]*`. Anything else (quoted
410/// identifiers, exotic chars, three-part names, SQL injection attempts) is
411/// rejected — the user should fall back to `query:` for those cases.
412///
413/// The bound on identifier shape keeps generated SQL safe to interpolate
414/// without quoting and ensures the generated `SELECT * FROM <ident>` form is
415/// recognised by the PG catalog-hint parser ([src/source/postgres.rs]).
416fn validate_table_shortcut_ident(export_name: &str, raw: &str) -> crate::error::Result<()> {
417    let trimmed = raw.trim();
418    if trimmed.is_empty() {
419        anyhow::bail!("export '{export_name}': 'table' is empty");
420    }
421    let parts: Vec<&str> = trimmed.split('.').collect();
422    if parts.len() > 2 {
423        anyhow::bail!(
424            "export '{export_name}': 'table' must be `<name>` or `<schema>.<name>` (got '{raw}')"
425        );
426    }
427    for part in &parts {
428        if part.is_empty() {
429            anyhow::bail!("export '{export_name}': 'table' has an empty segment in '{raw}'");
430        }
431        let mut chars = part.chars();
432        let first = chars.next().unwrap();
433        if !(first.is_ascii_alphabetic() || first == '_') {
434            anyhow::bail!(
435                "export '{export_name}': 'table' segment '{part}' must start with a letter or underscore (use 'query:' for quoted identifiers)"
436            );
437        }
438        if !chars.all(|c| c.is_ascii_alphanumeric() || c == '_') {
439            anyhow::bail!(
440                "export '{export_name}': 'table' segment '{part}' contains non-identifier characters (use 'query:' for quoted identifiers)"
441            );
442        }
443    }
444    Ok(())
445}
446
447#[derive(Debug, Deserialize, Serialize, JsonSchema, Clone)]
448#[serde(deny_unknown_fields)]
449pub struct QualityConfig {
450    pub row_count_min: Option<usize>,
451    pub row_count_max: Option<usize>,
452    #[serde(default)]
453    pub null_ratio_max: std::collections::HashMap<String, f64>,
454    #[serde(default)]
455    pub unique_columns: Vec<String>,
456    /// Cap on the number of distinct values tracked per column during uniqueness checks.
457    /// When the limit is hit, a Warn issue is emitted and tracking stops for that column.
458    /// Prevents unbounded HashSet growth on high-cardinality columns.
459    pub unique_max_entries: Option<usize>,
460}
461
462#[derive(Debug, Deserialize, Serialize, JsonSchema, Clone, Default)]
463#[serde(deny_unknown_fields)]
464pub struct MetaColumns {
465    #[serde(default)]
466    pub exported_at: bool,
467    #[serde(default)]
468    pub row_hash: bool,
469}
470
471fn default_mode() -> ExportMode {
472    ExportMode::Full
473}
474
475fn default_chunk_size() -> usize {
476    100_000
477}
478
479fn default_parallel() -> usize {
480    1
481}
482
483fn default_time_column_type() -> TimeColumnType {
484    TimeColumnType::Timestamp
485}
486
487/// `until_current` defaults to `true` — the OSS model is the BOUNDED, scheduler-
488/// driven drain ("read to the log end and exit"). Continuous streaming
489/// (`until_current: false`) is an explicit opt-in; making it the default silently
490/// put a hand-written CDC config onto the never-terminating streaming path.
491fn default_true() -> bool {
492    true
493}
494
495#[derive(Debug, Deserialize, JsonSchema, Clone, Copy, PartialEq, Eq)]
496#[serde(rename_all = "snake_case")]
497pub enum ExportMode {
498    Full,
499    Incremental,
500    Chunked,
501    TimeWindow,
502    /// Log-based change data capture (see [`CdcExportConfig`]): stream
503    /// INSERT/UPDATE/DELETE from the source's transaction log instead of querying
504    /// the table. Reuses the export's `table` / `destination` / `format`.
505    Cdc,
506}
507
508/// Default PostgreSQL logical slot when `cdc.slot` is omitted — shared by the
509/// runner ([`crate::pipeline`]'s cdc job) and config validation, so the
510/// same-slot conflict check sees the value that will actually be used.
511pub const DEFAULT_PG_SLOT: &str = "rivet_slot";
512/// Default MySQL replica `server_id` when `cdc.server_id` is omitted (see
513/// [`DEFAULT_PG_SLOT`] for why this is a shared const).
514pub const DEFAULT_MYSQL_SERVER_ID: u32 = 4271;
515
516/// What the FIRST CDC run does before draining changes.
517#[derive(Debug, Deserialize, Serialize, JsonSchema, Clone, Copy, PartialEq, Eq)]
518#[serde(rename_all = "lowercase")]
519pub enum CdcInitialMode {
520    /// Anchor-then-snapshot: create the resume anchor (PostgreSQL slot /
521    /// MySQL binlog checkpoint / SQL Server LSN checkpoint) FIRST, then run a
522    /// full batch snapshot of each table into `<destination>[/<table>]/snapshot/`,
523    /// then drain CDC. Because the anchor predates the snapshot read, anything
524    /// changed during the snapshot also appears in the change stream — an
525    /// overlap (dedupe by PK + `__op`), never a gap. The safe switch ordering,
526    /// enforced by construction instead of operator discipline.
527    Snapshot,
528}
529
530/// Per-export CDC settings, required when `mode: cdc`. The output `table`,
531/// `destination`, and `format` come from the export itself; this carries only the
532/// CDC-specific knobs (resume + per-engine stream params).
533#[derive(Debug, Deserialize, Serialize, JsonSchema, Clone, Default)]
534pub struct CdcExportConfig {
535    /// First-run behaviour: `snapshot` = anchor → full snapshot → drain (see
536    /// [`CdcInitialMode`]). Omitted ⇒ capture changes only (the default; the
537    /// operator owns the initial load).
538    #[serde(default)]
539    pub initial: Option<CdcInitialMode>,
540    /// Persist/resume the source log position to this file. Omit to tail from the
541    /// current position without checkpointing.
542    pub checkpoint: Option<String>,
543    /// Catch up to the source's current end and exit (a bounded run), instead of
544    /// streaming indefinitely — ideal for a scheduler. For MySQL this is a
545    /// non-blocking binlog dump; PostgreSQL / SQL Server already drain-and-exit.
546    /// **Defaults to `true`** (bounded): the OSS model is scheduler-driven, and
547    /// omitting this must NOT silently start a never-terminating stream. Set it to
548    /// `false` to opt into continuous streaming explicitly.
549    #[serde(default = "default_true")]
550    pub until_current: bool,
551    /// Stop after N change events (default: until end of stream / interrupted).
552    pub max_events: Option<usize>,
553    /// Rows per output part file (default 100000). A part also rolls at a
554    /// transaction boundary, so it never splits a transaction. Larger ⇒ fewer,
555    /// bigger files but more drain memory — the PostgreSQL peek reads a part's
556    /// worth per batch, so drain RSS is O(rollover). Tune per workload: raise it
557    /// to cut file count, lower it to cap memory on a small extractor.
558    pub rollover: Option<usize>,
559    /// Roll a part once its buffered changes reach this many MB, whichever comes
560    /// first with `rollover`. Caps the in-memory buffer and the part file size by
561    /// bytes instead of a fixed row count — predictable for tables with wide
562    /// (large JSON / blob) rows, mirroring the batch path's `batch_size_memory_mb`.
563    pub rollover_memory_mb: Option<usize>,
564    /// MySQL replica server-id for the binlog connection (default 4271; must be
565    /// distinct from the source's and any other replica).
566    pub server_id: Option<u32>,
567    /// PostgreSQL logical replication slot name (default `rivet_slot`).
568    pub slot: Option<String>,
569    /// SQL Server CDC capture instance, e.g. `dbo_orders` — required for
570    /// `sqlserver://` sources.
571    pub capture_instance: Option<String>,
572}
573
574#[derive(Debug, Deserialize, Serialize, JsonSchema, Clone, Copy, PartialEq, Eq)]
575#[serde(rename_all = "lowercase")]
576pub enum TimeColumnType {
577    Timestamp,
578    Unix,
579}
580
581/// Calendar bucket width for date/timestamp output partitioning
582/// ([`ExportConfig::partition_by`]). The partition column must be a DATE or
583/// TIMESTAMP column; this picks how its range is split into contiguous Hive
584/// buckets. It is not a knob for partitioning by arbitrary column values.
585#[derive(Debug, Deserialize, Serialize, JsonSchema, Clone, Copy, PartialEq, Eq, Default)]
586#[serde(rename_all = "lowercase")]
587pub enum PartitionGranularity {
588    /// One bucket per calendar day (`col=2023-01-01/`). Default.
589    #[default]
590    Day,
591    /// One bucket per calendar month (`col=2023-01/`).
592    Month,
593    /// One bucket per calendar year (`col=2023/`).
594    Year,
595}
596
597/// Canonical fully-populated [`ExportConfig`] for tests across the crate.
598///
599/// One place lists every field, so adding a field is a single-site edit (the
600/// compiler still flags this literal if a field is missing). Test call sites
601/// take this baseline and override only the fields they exercise, rather than
602/// hand-writing the full struct — see `plan::build` and `preflight` tests.
603#[cfg(test)]
604pub(crate) fn sample_export(name: &str) -> ExportConfig {
605    ExportConfig {
606        name: name.into(),
607        target: None,
608        load: None,
609        verify: VerifyMode::Size,
610        query: Some("SELECT 1".into()),
611        query_file: None,
612        table: None,
613        tables: None,
614        mode: ExportMode::Full,
615        cdc: None,
616        cursor_column: None,
617        cursor_fallback_column: None,
618        incremental_cursor_mode: Default::default(),
619        chunk_column: None,
620        chunk_dense: false,
621        chunk_size: 100_000,
622        chunk_size_memory_mb: None,
623        chunk_count: None,
624        chunk_by_days: None,
625        chunk_by_key: None,
626        parallel: 1,
627        wave: None,
628        parallel_safe: None,
629        time_column: None,
630        time_column_type: TimeColumnType::Timestamp,
631        days_window: None,
632        partition_by: None,
633        partition_granularity: PartitionGranularity::Day,
634        format: FormatType::Parquet,
635        compression: CompressionType::None,
636        compression_level: None,
637        compression_profile: None,
638        skip_empty: false,
639        destination: crate::config::DestinationConfig {
640            destination_type: crate::config::DestinationType::Local,
641            path: Some("/tmp".into()),
642            ..Default::default()
643        },
644        meta_columns: MetaColumns::default(),
645        quality: None,
646        max_file_size: None,
647        chunk_checkpoint: false,
648        chunk_max_attempts: None,
649        tuning: None,
650        source_group: None,
651        reconcile_required: false,
652        columns: Default::default(),
653        on_schema_drift: Default::default(),
654        shape_drift_warn_factor: None,
655        parquet: None,
656    }
657}
658
659#[cfg(test)]
660mod tests {
661    use super::*;
662
663    // ── ExportConfig::max_file_size_bytes ───────────────────────────────────
664
665    fn make_export_yaml(name: &str, extra: &str) -> ExportConfig {
666        let yaml = format!(
667            "name: {name}\nquery: \"SELECT 1\"\nformat: parquet\ndestination:\n  type: local\n  path: /tmp\n{extra}"
668        );
669        serde_yaml_ng::from_str(&yaml).expect("parse ExportConfig")
670    }
671
672    #[test]
673    fn max_file_size_bytes_none_when_unset() {
674        let exp = make_export_yaml("no_limit", "");
675        assert!(exp.max_file_size_bytes().is_none());
676    }
677
678    #[test]
679    fn max_file_size_bytes_parses_mb() {
680        let exp = make_export_yaml("sized", "max_file_size: \"128MB\"\n");
681        assert_eq!(exp.max_file_size_bytes(), Some(128 * 1024 * 1024));
682    }
683
684    #[test]
685    fn max_file_size_bytes_parses_gb() {
686        let exp = make_export_yaml("sized_gb", "max_file_size: \"2GB\"\n");
687        assert_eq!(exp.max_file_size_bytes(), Some(2 * 1024 * 1024 * 1024));
688    }
689
690    #[test]
691    fn max_file_size_bytes_returns_none_on_invalid() {
692        let exp = make_export_yaml("bad_size", "max_file_size: \"notanumber\"\n");
693        assert!(exp.max_file_size_bytes().is_none());
694    }
695
696    // ── ExportConfig::resolve_query ─────────────────────────────────────────
697
698    // Build a minimal ExportConfig directly, bypassing Config::from_yaml validation.
699    // This lets us test the four branches inside resolve_query itself, including
700    // the (both-set / neither-set) error paths that are normally prevented by the
701    // top-level validator.
702    fn make_export_direct(query: Option<&str>, query_file: Option<&str>) -> ExportConfig {
703        ExportConfig {
704            query: query.map(|s| s.to_string()),
705            query_file: query_file.map(|s| s.to_string()),
706            ..sample_export("test")
707        }
708    }
709
710    fn params(pairs: &[(&str, &str)]) -> std::collections::HashMap<String, String> {
711        pairs
712            .iter()
713            .map(|(k, v)| (k.to_string(), v.to_string()))
714            .collect()
715    }
716
717    #[test]
718    fn resolve_query_inline_no_params_returns_query_as_is() {
719        let exp = make_export_direct(Some("SELECT id FROM orders"), None);
720        let q = exp.resolve_query(Path::new("/tmp"), None).unwrap();
721        assert_eq!(q, "SELECT id FROM orders");
722    }
723
724    #[test]
725    fn resolve_query_inline_with_params_substitutes_vars() {
726        let exp = make_export_direct(Some("SELECT ${col} FROM ${table}"), None);
727        let p = params(&[("col", "id"), ("table", "orders")]);
728        let q = exp.resolve_query(Path::new("/tmp"), Some(&p)).unwrap();
729        assert_eq!(q, "SELECT id FROM orders");
730    }
731
732    #[test]
733    fn resolve_query_inline_params_empty_map_is_noop() {
734        let exp = make_export_direct(Some("SELECT 1"), None);
735        let p = params(&[]);
736        let q = exp.resolve_query(Path::new("/tmp"), Some(&p)).unwrap();
737        assert_eq!(q, "SELECT 1");
738    }
739
740    #[test]
741    fn resolve_query_inline_missing_var_returns_error() {
742        // SAFETY: test-only; this binary is single-threaded in the test runner context.
743        unsafe { std::env::remove_var("UNSET_RIVET_TEST_VAR") };
744        let exp = make_export_direct(Some("SELECT ${UNSET_RIVET_TEST_VAR}"), None);
745        let p = params(&[]);
746        let result = exp.resolve_query(Path::new("/tmp"), Some(&p));
747        assert!(result.is_err());
748        let msg = format!("{:#}", result.unwrap_err());
749        assert!(
750            msg.contains("UNSET_RIVET_TEST_VAR") || msg.contains("not set"),
751            "got: {msg}"
752        );
753    }
754
755    #[test]
756    fn resolve_query_file_reads_content() {
757        let dir = tempfile::TempDir::new().unwrap();
758        let sql_path = dir.path().join("query.sql");
759        std::fs::write(&sql_path, "SELECT * FROM customers").unwrap();
760        let exp = make_export_direct(None, Some("query.sql"));
761        let q = exp.resolve_query(dir.path(), None).unwrap();
762        assert_eq!(q, "SELECT * FROM customers");
763    }
764
765    #[test]
766    fn resolve_query_file_with_params_substitutes() {
767        let dir = tempfile::TempDir::new().unwrap();
768        let sql_path = dir.path().join("q.sql");
769        std::fs::write(&sql_path, "SELECT ${col} FROM ${tbl}").unwrap();
770        let exp = make_export_direct(None, Some("q.sql"));
771        let p = params(&[("col", "name"), ("tbl", "users")]);
772        let q = exp.resolve_query(dir.path(), Some(&p)).unwrap();
773        assert_eq!(q, "SELECT name FROM users");
774    }
775
776    // ── `table:` shortcut ───────────────────────────────────────────────────
777
778    #[test]
779    fn resolve_query_table_shortcut_qualified() {
780        let mut exp = make_export_direct(None, None);
781        exp.table = Some("public.users".into());
782        let q = exp.resolve_query(Path::new("/tmp"), None).unwrap();
783        assert_eq!(q, "SELECT * FROM public.users");
784    }
785
786    #[test]
787    fn resolve_query_table_shortcut_unqualified() {
788        let mut exp = make_export_direct(None, None);
789        exp.table = Some("orders".into());
790        let q = exp.resolve_query(Path::new("/tmp"), None).unwrap();
791        assert_eq!(q, "SELECT * FROM orders");
792    }
793
794    #[test]
795    fn resolve_query_table_shortcut_rejects_three_part_name() {
796        let mut exp = make_export_direct(None, None);
797        exp.table = Some("db.public.users".into());
798        let err = exp.resolve_query(Path::new("/tmp"), None).unwrap_err();
799        let msg = format!("{err:#}");
800        assert!(msg.contains("<schema>.<name>"), "got: {msg}");
801    }
802
803    #[test]
804    fn resolve_query_table_shortcut_rejects_sql_injection() {
805        for bad in [
806            "users; DROP TABLE x",
807            "users--",
808            "users'",
809            "users\"",
810            "public.\"My Table\"",
811            "0starts_with_digit",
812            "",
813            ".trailing",
814            "leading.",
815            "two..dots",
816        ] {
817            let mut exp = make_export_direct(None, None);
818            exp.table = Some(bad.into());
819            assert!(
820                exp.resolve_query(Path::new("/tmp"), None).is_err(),
821                "should reject `table:` value '{bad}'",
822            );
823        }
824    }
825
826    #[test]
827    fn resolve_query_table_shortcut_takes_precedence_over_query() {
828        let mut exp = make_export_direct(Some("SELECT id FROM x"), None);
829        exp.table = Some("public.y".into());
830        let q = exp.resolve_query(Path::new("/tmp"), None).unwrap();
831        assert_eq!(q, "SELECT * FROM public.y");
832    }
833
834    #[test]
835    fn resolve_query_file_missing_returns_error() {
836        let dir = tempfile::TempDir::new().unwrap();
837        let exp = make_export_direct(None, Some("nonexistent.sql"));
838        let result = exp.resolve_query(dir.path(), None);
839        assert!(result.is_err());
840        let msg = format!("{:#}", result.unwrap_err());
841        assert!(
842            msg.contains("nonexistent.sql") || msg.contains("No such file"),
843            "got: {msg}"
844        );
845    }
846
847    #[test]
848    fn resolve_query_both_set_returns_error() {
849        let mut exp = make_export_direct(Some("SELECT 1"), None);
850        exp.query_file = Some("file.sql".into());
851        let result = exp.resolve_query(Path::new("/tmp"), None);
852        assert!(result.is_err());
853        let msg = format!("{:#}", result.unwrap_err());
854        assert!(
855            msg.contains("not both") || msg.contains("query_file"),
856            "got: {msg}"
857        );
858    }
859
860    #[test]
861    fn resolve_query_neither_set_returns_error() {
862        let exp = make_export_direct(None, None);
863        let result = exp.resolve_query(Path::new("/tmp"), None);
864        assert!(result.is_err());
865        let msg = format!("{:#}", result.unwrap_err());
866        assert!(
867            msg.contains("query") || msg.contains("query_file"),
868            "got: {msg}"
869        );
870    }
871
872    // ── SecOps: query_file path traversal prevention ──────────────────────────
873
874    #[test]
875    fn resolve_query_file_dotdot_is_rejected() {
876        let dir = tempfile::TempDir::new().unwrap();
877        let exp = make_export_direct(None, Some("../secret.sql"));
878        let result = exp.resolve_query(dir.path(), None);
879        assert!(result.is_err());
880        let msg = format!("{:#}", result.unwrap_err());
881        assert!(
882            msg.contains("..") || msg.contains("traversal"),
883            "got: {msg}"
884        );
885    }
886
887    #[test]
888    fn resolve_query_file_nested_dotdot_is_rejected() {
889        let dir = tempfile::TempDir::new().unwrap();
890        let exp = make_export_direct(None, Some("subdir/../../etc/passwd"));
891        let result = exp.resolve_query(dir.path(), None);
892        assert!(result.is_err());
893        let msg = format!("{:#}", result.unwrap_err());
894        assert!(
895            msg.contains("..") || msg.contains("traversal"),
896            "got: {msg}"
897        );
898    }
899
900    #[test]
901    fn resolve_query_file_absolute_path_is_rejected() {
902        let dir = tempfile::TempDir::new().unwrap();
903        let exp = make_export_direct(None, Some("/etc/passwd"));
904        let result = exp.resolve_query(dir.path(), None);
905        assert!(result.is_err());
906        let msg = format!("{:#}", result.unwrap_err());
907        assert!(
908            msg.contains("relative") || msg.contains("absolute"),
909            "got: {msg}"
910        );
911    }
912
913    #[test]
914    fn resolve_query_file_in_subdir_is_allowed() {
915        let dir = tempfile::TempDir::new().unwrap();
916        let subdir = dir.path().join("queries");
917        std::fs::create_dir(&subdir).unwrap();
918        std::fs::write(subdir.join("orders.sql"), "SELECT * FROM orders").unwrap();
919        let exp = make_export_direct(None, Some("queries/orders.sql"));
920        let q = exp.resolve_query(dir.path(), None).unwrap();
921        assert_eq!(q, "SELECT * FROM orders");
922    }
923}