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`](Self::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`](Self::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 /// Policy applied when structural schema drift is detected (column added, removed, or retyped).
264 /// Defaults to `warn`: log a warning and continue.
265 #[serde(default)]
266 pub on_schema_drift: SchemaDriftPolicy,
267
268 /// Growth-factor threshold for data shape drift warnings (Epic 8).
269 /// When a string/binary column's max observed byte length in the current run
270 /// exceeds `stored_max * shape_drift_warn_factor`, Rivet logs a warning.
271 /// `None` uses the default of 2.0. Set to `0.0` to disable shape tracking.
272 #[serde(default)]
273 pub shape_drift_warn_factor: Option<f64>,
274
275 /// Parquet row group tuning. Only meaningful when `format: parquet`.
276 /// When absent, the parquet library default (1,048,576 rows/group) is used.
277 #[serde(default)]
278 pub parquet: Option<ParquetConfig>,
279}
280
281impl ExportConfig {
282 /// Resolve the effective `(CompressionType, level)` for this export.
283 /// `compression_profile` takes precedence over `compression` + `compression_level`.
284 ///
285 /// L24: when a profile is set *and* a conflicting explicit codec/level was
286 /// written, warn once that the profile wins rather than silently dropping the
287 /// explicit choice. An explicit codec is only detectable when it differs from
288 /// the `#[serde(default)]` (Zstd) — a literal `compression: zstd` alongside a
289 /// profile is indistinguishable from an omitted field and stays silent.
290 pub fn effective_compression(&self) -> (CompressionType, Option<u32>) {
291 if let Some(profile) = self.compression_profile {
292 let explicit_codec =
293 (self.compression != CompressionType::default()).then_some(self.compression);
294 if let Some(msg) = super::format::compression_profile_override_warning(
295 profile,
296 explicit_codec,
297 self.compression_level,
298 ) {
299 log::warn!("export '{}': {}", self.name, msg);
300 }
301 profile.to_codec()
302 } else {
303 (self.compression, self.compression_level)
304 }
305 }
306
307 pub fn max_file_size_bytes(&self) -> Option<u64> {
308 self.max_file_size
309 .as_ref()
310 .and_then(|s| parse_file_size(s).ok())
311 }
312
313 pub fn resolve_query(
314 &self,
315 config_dir: &Path,
316 params: Option<&std::collections::HashMap<String, String>>,
317 ) -> crate::error::Result<String> {
318 // table: shortcut takes precedence — already validated by
319 // `validate_business_rules` to be mutually exclusive with query/query_file.
320 if let Some(tbl) = &self.table {
321 validate_table_shortcut_ident(&self.name, tbl)?;
322 return Ok(format!("SELECT * FROM {tbl}"));
323 }
324 match (&self.query, &self.query_file) {
325 (Some(q), None) => {
326 if params.is_some() {
327 resolve_vars(q, params)
328 } else {
329 Ok(q.clone())
330 }
331 }
332 (None, Some(file)) => {
333 let file_path = std::path::Path::new(file);
334 // SecOps: block absolute paths and `..` traversal components.
335 if file_path.is_absolute() {
336 anyhow::bail!(
337 "export '{}': query_file must be a relative path: '{}'",
338 self.name,
339 file
340 );
341 }
342 if file_path
343 .components()
344 .any(|c| c == std::path::Component::ParentDir)
345 {
346 anyhow::bail!(
347 "export '{}': query_file path must not contain '..': '{}'",
348 self.name,
349 file
350 );
351 }
352 let joined = config_dir.join(file);
353 // Canonicalize-based check catches symlink-based evasion for files
354 // that already exist on disk.
355 if let Ok(canonical) = joined.canonicalize() {
356 let base = config_dir
357 .canonicalize()
358 .unwrap_or_else(|_| config_dir.to_path_buf());
359 if !canonical.starts_with(&base) {
360 anyhow::bail!(
361 "export '{}': query_file '{}' resolves outside the config directory",
362 self.name,
363 file
364 );
365 }
366 }
367 let raw = std::fs::read_to_string(&joined)?;
368 resolve_vars(&raw, params)
369 }
370 (Some(_), Some(_)) => {
371 anyhow::bail!(
372 "export '{}': specify either 'query' or 'query_file', not both",
373 self.name
374 )
375 }
376 (None, None) => {
377 anyhow::bail!(
378 "export '{}': must specify exactly one of 'query', 'query_file', or 'table'",
379 self.name
380 )
381 }
382 }
383 }
384}
385
386/// Validate the value of the `table:` YAML shortcut.
387///
388/// Accepts ASCII identifiers in the form `<table>` or `<schema>.<table>`. Each
389/// segment must match `[A-Za-z_][A-Za-z0-9_]*`. Anything else (quoted
390/// identifiers, exotic chars, three-part names, SQL injection attempts) is
391/// rejected — the user should fall back to `query:` for those cases.
392///
393/// The bound on identifier shape keeps generated SQL safe to interpolate
394/// without quoting and ensures the generated `SELECT * FROM <ident>` form is
395/// recognised by the PG catalog-hint parser ([src/source/postgres.rs]).
396fn validate_table_shortcut_ident(export_name: &str, raw: &str) -> crate::error::Result<()> {
397 let trimmed = raw.trim();
398 if trimmed.is_empty() {
399 anyhow::bail!("export '{export_name}': 'table' is empty");
400 }
401 let parts: Vec<&str> = trimmed.split('.').collect();
402 if parts.len() > 2 {
403 anyhow::bail!(
404 "export '{export_name}': 'table' must be `<name>` or `<schema>.<name>` (got '{raw}')"
405 );
406 }
407 for part in &parts {
408 if part.is_empty() {
409 anyhow::bail!("export '{export_name}': 'table' has an empty segment in '{raw}'");
410 }
411 let mut chars = part.chars();
412 let first = chars.next().unwrap();
413 if !(first.is_ascii_alphabetic() || first == '_') {
414 anyhow::bail!(
415 "export '{export_name}': 'table' segment '{part}' must start with a letter or underscore (use 'query:' for quoted identifiers)"
416 );
417 }
418 if !chars.all(|c| c.is_ascii_alphanumeric() || c == '_') {
419 anyhow::bail!(
420 "export '{export_name}': 'table' segment '{part}' contains non-identifier characters (use 'query:' for quoted identifiers)"
421 );
422 }
423 }
424 Ok(())
425}
426
427#[derive(Debug, Deserialize, Serialize, JsonSchema, Clone)]
428#[serde(deny_unknown_fields)]
429pub struct QualityConfig {
430 pub row_count_min: Option<usize>,
431 pub row_count_max: Option<usize>,
432 #[serde(default)]
433 pub null_ratio_max: std::collections::HashMap<String, f64>,
434 #[serde(default)]
435 pub unique_columns: Vec<String>,
436 /// Cap on the number of distinct values tracked per column during uniqueness checks.
437 /// When the limit is hit, a Warn issue is emitted and tracking stops for that column.
438 /// Prevents unbounded HashSet growth on high-cardinality columns.
439 pub unique_max_entries: Option<usize>,
440}
441
442#[derive(Debug, Deserialize, Serialize, JsonSchema, Clone, Default)]
443#[serde(deny_unknown_fields)]
444pub struct MetaColumns {
445 #[serde(default)]
446 pub exported_at: bool,
447 #[serde(default)]
448 pub row_hash: bool,
449}
450
451fn default_mode() -> ExportMode {
452 ExportMode::Full
453}
454
455fn default_chunk_size() -> usize {
456 100_000
457}
458
459fn default_parallel() -> usize {
460 1
461}
462
463fn default_time_column_type() -> TimeColumnType {
464 TimeColumnType::Timestamp
465}
466
467#[derive(Debug, Deserialize, JsonSchema, Clone, Copy, PartialEq, Eq)]
468#[serde(rename_all = "snake_case")]
469pub enum ExportMode {
470 Full,
471 Incremental,
472 Chunked,
473 TimeWindow,
474 /// Log-based change data capture (see [`CdcExportConfig`]): stream
475 /// INSERT/UPDATE/DELETE from the source's transaction log instead of querying
476 /// the table. Reuses the export's `table` / `destination` / `format`.
477 Cdc,
478}
479
480/// Default PostgreSQL logical slot when `cdc.slot` is omitted — shared by the
481/// runner ([`crate::pipeline`]'s cdc job) and config validation, so the
482/// same-slot conflict check sees the value that will actually be used.
483pub const DEFAULT_PG_SLOT: &str = "rivet_slot";
484/// Default MySQL replica `server_id` when `cdc.server_id` is omitted (see
485/// [`DEFAULT_PG_SLOT`] for why this is a shared const).
486pub const DEFAULT_MYSQL_SERVER_ID: u32 = 4271;
487
488/// What the FIRST CDC run does before draining changes.
489#[derive(Debug, Deserialize, Serialize, JsonSchema, Clone, Copy, PartialEq, Eq)]
490#[serde(rename_all = "lowercase")]
491pub enum CdcInitialMode {
492 /// Anchor-then-snapshot: create the resume anchor (PostgreSQL slot /
493 /// MySQL binlog checkpoint / SQL Server LSN checkpoint) FIRST, then run a
494 /// full batch snapshot of each table into `<destination>[/<table>]/snapshot/`,
495 /// then drain CDC. Because the anchor predates the snapshot read, anything
496 /// changed during the snapshot also appears in the change stream — an
497 /// overlap (dedupe by PK + `__op`), never a gap. The safe switch ordering,
498 /// enforced by construction instead of operator discipline.
499 Snapshot,
500}
501
502/// Per-export CDC settings, required when `mode: cdc`. The output `table`,
503/// `destination`, and `format` come from the export itself; this carries only the
504/// CDC-specific knobs (resume + per-engine stream params).
505#[derive(Debug, Deserialize, Serialize, JsonSchema, Clone, Default)]
506pub struct CdcExportConfig {
507 /// First-run behaviour: `snapshot` = anchor → full snapshot → drain (see
508 /// [`CdcInitialMode`]). Omitted ⇒ capture changes only (the default; the
509 /// operator owns the initial load).
510 #[serde(default)]
511 pub initial: Option<CdcInitialMode>,
512 /// Persist/resume the source log position to this file. Omit to tail from the
513 /// current position without checkpointing.
514 pub checkpoint: Option<String>,
515 /// Catch up to the source's current end and exit (a bounded run), instead of
516 /// streaming indefinitely — ideal for a scheduler. For MySQL this is a
517 /// non-blocking binlog dump; PostgreSQL / SQL Server already drain-and-exit.
518 #[serde(default)]
519 pub until_current: bool,
520 /// Stop after N change events (default: until end of stream / interrupted).
521 pub max_events: Option<usize>,
522 /// Rows per output part file (default 10000). A part also rolls at a
523 /// transaction boundary, so it never splits a transaction.
524 pub rollover: Option<usize>,
525 /// Roll a part once its buffered changes reach this many MB, whichever comes
526 /// first with `rollover`. Caps the in-memory buffer and the part file size by
527 /// bytes instead of a fixed row count — predictable for tables with wide
528 /// (large JSON / blob) rows, mirroring the batch path's `batch_size_memory_mb`.
529 pub rollover_memory_mb: Option<usize>,
530 /// MySQL replica server-id for the binlog connection (default 4271; must be
531 /// distinct from the source's and any other replica).
532 pub server_id: Option<u32>,
533 /// PostgreSQL logical replication slot name (default `rivet_slot`).
534 pub slot: Option<String>,
535 /// SQL Server CDC capture instance, e.g. `dbo_orders` — required for
536 /// `sqlserver://` sources.
537 pub capture_instance: Option<String>,
538}
539
540#[derive(Debug, Deserialize, Serialize, JsonSchema, Clone, Copy, PartialEq, Eq)]
541#[serde(rename_all = "lowercase")]
542pub enum TimeColumnType {
543 Timestamp,
544 Unix,
545}
546
547/// Calendar bucket width for date/timestamp output partitioning
548/// ([`ExportConfig::partition_by`]). The partition column must be a DATE or
549/// TIMESTAMP column; this picks how its range is split into contiguous Hive
550/// buckets. It is not a knob for partitioning by arbitrary column values.
551#[derive(Debug, Deserialize, Serialize, JsonSchema, Clone, Copy, PartialEq, Eq, Default)]
552#[serde(rename_all = "lowercase")]
553pub enum PartitionGranularity {
554 /// One bucket per calendar day (`col=2023-01-01/`). Default.
555 #[default]
556 Day,
557 /// One bucket per calendar month (`col=2023-01/`).
558 Month,
559 /// One bucket per calendar year (`col=2023/`).
560 Year,
561}
562
563/// Canonical fully-populated [`ExportConfig`] for tests across the crate.
564///
565/// One place lists every field, so adding a field is a single-site edit (the
566/// compiler still flags this literal if a field is missing). Test call sites
567/// take this baseline and override only the fields they exercise, rather than
568/// hand-writing the full struct — see `plan::build` and `preflight` tests.
569#[cfg(test)]
570pub(crate) fn sample_export(name: &str) -> ExportConfig {
571 ExportConfig {
572 name: name.into(),
573 target: None,
574 verify: VerifyMode::Size,
575 query: Some("SELECT 1".into()),
576 query_file: None,
577 table: None,
578 tables: None,
579 mode: ExportMode::Full,
580 cdc: None,
581 cursor_column: None,
582 cursor_fallback_column: None,
583 incremental_cursor_mode: Default::default(),
584 chunk_column: None,
585 chunk_dense: false,
586 chunk_size: 100_000,
587 chunk_size_memory_mb: None,
588 chunk_count: None,
589 chunk_by_days: None,
590 chunk_by_key: None,
591 parallel: 1,
592 wave: None,
593 parallel_safe: None,
594 time_column: None,
595 time_column_type: TimeColumnType::Timestamp,
596 days_window: None,
597 partition_by: None,
598 partition_granularity: PartitionGranularity::Day,
599 format: FormatType::Parquet,
600 compression: CompressionType::None,
601 compression_level: None,
602 compression_profile: None,
603 skip_empty: false,
604 destination: crate::config::DestinationConfig {
605 destination_type: crate::config::DestinationType::Local,
606 path: Some("/tmp".into()),
607 ..Default::default()
608 },
609 meta_columns: MetaColumns::default(),
610 quality: None,
611 max_file_size: None,
612 chunk_checkpoint: false,
613 chunk_max_attempts: None,
614 tuning: None,
615 source_group: None,
616 reconcile_required: false,
617 columns: Default::default(),
618 on_schema_drift: Default::default(),
619 shape_drift_warn_factor: None,
620 parquet: None,
621 }
622}
623
624#[cfg(test)]
625mod tests {
626 use super::*;
627
628 // ── ExportConfig::max_file_size_bytes ───────────────────────────────────
629
630 fn make_export_yaml(name: &str, extra: &str) -> ExportConfig {
631 let yaml = format!(
632 "name: {name}\nquery: \"SELECT 1\"\nformat: parquet\ndestination:\n type: local\n path: /tmp\n{extra}"
633 );
634 serde_yaml_ng::from_str(&yaml).expect("parse ExportConfig")
635 }
636
637 #[test]
638 fn max_file_size_bytes_none_when_unset() {
639 let exp = make_export_yaml("no_limit", "");
640 assert!(exp.max_file_size_bytes().is_none());
641 }
642
643 #[test]
644 fn max_file_size_bytes_parses_mb() {
645 let exp = make_export_yaml("sized", "max_file_size: \"128MB\"\n");
646 assert_eq!(exp.max_file_size_bytes(), Some(128 * 1024 * 1024));
647 }
648
649 #[test]
650 fn max_file_size_bytes_parses_gb() {
651 let exp = make_export_yaml("sized_gb", "max_file_size: \"2GB\"\n");
652 assert_eq!(exp.max_file_size_bytes(), Some(2 * 1024 * 1024 * 1024));
653 }
654
655 #[test]
656 fn max_file_size_bytes_returns_none_on_invalid() {
657 let exp = make_export_yaml("bad_size", "max_file_size: \"notanumber\"\n");
658 assert!(exp.max_file_size_bytes().is_none());
659 }
660
661 // ── ExportConfig::resolve_query ─────────────────────────────────────────
662
663 // Build a minimal ExportConfig directly, bypassing Config::from_yaml validation.
664 // This lets us test the four branches inside resolve_query itself, including
665 // the (both-set / neither-set) error paths that are normally prevented by the
666 // top-level validator.
667 fn make_export_direct(query: Option<&str>, query_file: Option<&str>) -> ExportConfig {
668 ExportConfig {
669 query: query.map(|s| s.to_string()),
670 query_file: query_file.map(|s| s.to_string()),
671 ..sample_export("test")
672 }
673 }
674
675 fn params(pairs: &[(&str, &str)]) -> std::collections::HashMap<String, String> {
676 pairs
677 .iter()
678 .map(|(k, v)| (k.to_string(), v.to_string()))
679 .collect()
680 }
681
682 #[test]
683 fn resolve_query_inline_no_params_returns_query_as_is() {
684 let exp = make_export_direct(Some("SELECT id FROM orders"), None);
685 let q = exp.resolve_query(Path::new("/tmp"), None).unwrap();
686 assert_eq!(q, "SELECT id FROM orders");
687 }
688
689 #[test]
690 fn resolve_query_inline_with_params_substitutes_vars() {
691 let exp = make_export_direct(Some("SELECT ${col} FROM ${table}"), None);
692 let p = params(&[("col", "id"), ("table", "orders")]);
693 let q = exp.resolve_query(Path::new("/tmp"), Some(&p)).unwrap();
694 assert_eq!(q, "SELECT id FROM orders");
695 }
696
697 #[test]
698 fn resolve_query_inline_params_empty_map_is_noop() {
699 let exp = make_export_direct(Some("SELECT 1"), None);
700 let p = params(&[]);
701 let q = exp.resolve_query(Path::new("/tmp"), Some(&p)).unwrap();
702 assert_eq!(q, "SELECT 1");
703 }
704
705 #[test]
706 fn resolve_query_inline_missing_var_returns_error() {
707 // SAFETY: test-only; this binary is single-threaded in the test runner context.
708 unsafe { std::env::remove_var("UNSET_RIVET_TEST_VAR") };
709 let exp = make_export_direct(Some("SELECT ${UNSET_RIVET_TEST_VAR}"), None);
710 let p = params(&[]);
711 let result = exp.resolve_query(Path::new("/tmp"), Some(&p));
712 assert!(result.is_err());
713 let msg = format!("{:#}", result.unwrap_err());
714 assert!(
715 msg.contains("UNSET_RIVET_TEST_VAR") || msg.contains("not set"),
716 "got: {msg}"
717 );
718 }
719
720 #[test]
721 fn resolve_query_file_reads_content() {
722 let dir = tempfile::TempDir::new().unwrap();
723 let sql_path = dir.path().join("query.sql");
724 std::fs::write(&sql_path, "SELECT * FROM customers").unwrap();
725 let exp = make_export_direct(None, Some("query.sql"));
726 let q = exp.resolve_query(dir.path(), None).unwrap();
727 assert_eq!(q, "SELECT * FROM customers");
728 }
729
730 #[test]
731 fn resolve_query_file_with_params_substitutes() {
732 let dir = tempfile::TempDir::new().unwrap();
733 let sql_path = dir.path().join("q.sql");
734 std::fs::write(&sql_path, "SELECT ${col} FROM ${tbl}").unwrap();
735 let exp = make_export_direct(None, Some("q.sql"));
736 let p = params(&[("col", "name"), ("tbl", "users")]);
737 let q = exp.resolve_query(dir.path(), Some(&p)).unwrap();
738 assert_eq!(q, "SELECT name FROM users");
739 }
740
741 // ── `table:` shortcut ───────────────────────────────────────────────────
742
743 #[test]
744 fn resolve_query_table_shortcut_qualified() {
745 let mut exp = make_export_direct(None, None);
746 exp.table = Some("public.users".into());
747 let q = exp.resolve_query(Path::new("/tmp"), None).unwrap();
748 assert_eq!(q, "SELECT * FROM public.users");
749 }
750
751 #[test]
752 fn resolve_query_table_shortcut_unqualified() {
753 let mut exp = make_export_direct(None, None);
754 exp.table = Some("orders".into());
755 let q = exp.resolve_query(Path::new("/tmp"), None).unwrap();
756 assert_eq!(q, "SELECT * FROM orders");
757 }
758
759 #[test]
760 fn resolve_query_table_shortcut_rejects_three_part_name() {
761 let mut exp = make_export_direct(None, None);
762 exp.table = Some("db.public.users".into());
763 let err = exp.resolve_query(Path::new("/tmp"), None).unwrap_err();
764 let msg = format!("{err:#}");
765 assert!(msg.contains("<schema>.<name>"), "got: {msg}");
766 }
767
768 #[test]
769 fn resolve_query_table_shortcut_rejects_sql_injection() {
770 for bad in [
771 "users; DROP TABLE x",
772 "users--",
773 "users'",
774 "users\"",
775 "public.\"My Table\"",
776 "0starts_with_digit",
777 "",
778 ".trailing",
779 "leading.",
780 "two..dots",
781 ] {
782 let mut exp = make_export_direct(None, None);
783 exp.table = Some(bad.into());
784 assert!(
785 exp.resolve_query(Path::new("/tmp"), None).is_err(),
786 "should reject `table:` value '{bad}'",
787 );
788 }
789 }
790
791 #[test]
792 fn resolve_query_table_shortcut_takes_precedence_over_query() {
793 let mut exp = make_export_direct(Some("SELECT id FROM x"), None);
794 exp.table = Some("public.y".into());
795 let q = exp.resolve_query(Path::new("/tmp"), None).unwrap();
796 assert_eq!(q, "SELECT * FROM public.y");
797 }
798
799 #[test]
800 fn resolve_query_file_missing_returns_error() {
801 let dir = tempfile::TempDir::new().unwrap();
802 let exp = make_export_direct(None, Some("nonexistent.sql"));
803 let result = exp.resolve_query(dir.path(), None);
804 assert!(result.is_err());
805 let msg = format!("{:#}", result.unwrap_err());
806 assert!(
807 msg.contains("nonexistent.sql") || msg.contains("No such file"),
808 "got: {msg}"
809 );
810 }
811
812 #[test]
813 fn resolve_query_both_set_returns_error() {
814 let mut exp = make_export_direct(Some("SELECT 1"), None);
815 exp.query_file = Some("file.sql".into());
816 let result = exp.resolve_query(Path::new("/tmp"), None);
817 assert!(result.is_err());
818 let msg = format!("{:#}", result.unwrap_err());
819 assert!(
820 msg.contains("not both") || msg.contains("query_file"),
821 "got: {msg}"
822 );
823 }
824
825 #[test]
826 fn resolve_query_neither_set_returns_error() {
827 let exp = make_export_direct(None, None);
828 let result = exp.resolve_query(Path::new("/tmp"), None);
829 assert!(result.is_err());
830 let msg = format!("{:#}", result.unwrap_err());
831 assert!(
832 msg.contains("query") || msg.contains("query_file"),
833 "got: {msg}"
834 );
835 }
836
837 // ── SecOps: query_file path traversal prevention ──────────────────────────
838
839 #[test]
840 fn resolve_query_file_dotdot_is_rejected() {
841 let dir = tempfile::TempDir::new().unwrap();
842 let exp = make_export_direct(None, Some("../secret.sql"));
843 let result = exp.resolve_query(dir.path(), None);
844 assert!(result.is_err());
845 let msg = format!("{:#}", result.unwrap_err());
846 assert!(
847 msg.contains("..") || msg.contains("traversal"),
848 "got: {msg}"
849 );
850 }
851
852 #[test]
853 fn resolve_query_file_nested_dotdot_is_rejected() {
854 let dir = tempfile::TempDir::new().unwrap();
855 let exp = make_export_direct(None, Some("subdir/../../etc/passwd"));
856 let result = exp.resolve_query(dir.path(), None);
857 assert!(result.is_err());
858 let msg = format!("{:#}", result.unwrap_err());
859 assert!(
860 msg.contains("..") || msg.contains("traversal"),
861 "got: {msg}"
862 );
863 }
864
865 #[test]
866 fn resolve_query_file_absolute_path_is_rejected() {
867 let dir = tempfile::TempDir::new().unwrap();
868 let exp = make_export_direct(None, Some("/etc/passwd"));
869 let result = exp.resolve_query(dir.path(), None);
870 assert!(result.is_err());
871 let msg = format!("{:#}", result.unwrap_err());
872 assert!(
873 msg.contains("relative") || msg.contains("absolute"),
874 "got: {msg}"
875 );
876 }
877
878 #[test]
879 fn resolve_query_file_in_subdir_is_allowed() {
880 let dir = tempfile::TempDir::new().unwrap();
881 let subdir = dir.path().join("queries");
882 std::fs::create_dir(&subdir).unwrap();
883 std::fs::write(subdir.join("orders.sql"), "SELECT * FROM orders").unwrap();
884 let exp = make_export_direct(None, Some("queries/orders.sql"));
885 let q = exp.resolve_query(dir.path(), None).unwrap();
886 assert_eq!(q, "SELECT * FROM orders");
887 }
888}