use crate::types::target::{TargetColumnSpec, TargetStatus};
use anyhow::{Context, Result, bail};
use serde::Deserialize;
use std::process::Command;
#[derive(Debug, Clone, Deserialize)]
pub struct LoadSection {
#[serde(flatten)]
pub target: LoadTarget,
#[serde(default)]
pub cleanup_source: bool,
#[serde(default)]
pub pk: Vec<String>,
#[serde(default)]
pub allow_source_drift: bool,
#[serde(default)]
pub gc_orphans: bool,
#[serde(default)]
pub cluster_by: Vec<String>,
}
#[derive(Debug, Deserialize)]
#[serde(deny_unknown_fields)]
struct LoadOverride {
#[serde(default)]
pk: Option<Vec<String>>,
#[serde(default)]
cleanup_source: Option<bool>,
#[serde(default)]
gc_orphans: Option<bool>,
#[serde(default)]
cluster_by: Option<Vec<String>>,
#[serde(default)]
allow_source_drift: Option<bool>,
#[serde(default)]
target: Option<serde_json::Value>,
}
impl LoadSection {
fn with_override(&self, o: &LoadOverride) -> LoadSection {
let mut eff = self.clone();
if let Some(pk) = &o.pk {
eff.pk = pk.clone();
}
if let Some(c) = o.cleanup_source {
eff.cleanup_source = c;
}
if let Some(g) = o.gc_orphans {
eff.gc_orphans = g;
}
if let Some(cb) = &o.cluster_by {
eff.cluster_by = cb.clone();
}
if let Some(d) = o.allow_source_drift {
eff.allow_source_drift = d;
}
eff
}
}
#[derive(Debug, Clone, Deserialize)]
#[serde(tag = "target", rename_all = "lowercase")]
pub enum LoadTarget {
Bigquery {
project: String,
dataset: String,
},
Snowflake {
connection: String,
warehouse: String,
database: String,
schema: String,
storage_integration: String,
},
}
impl LoadTarget {
pub fn name(&self) -> &'static str {
match self {
LoadTarget::Bigquery { .. } => "bigquery",
LoadTarget::Snowflake { .. } => "snowflake",
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum LoadMode {
Full,
Incremental,
Cdc,
}
impl LoadMode {
pub fn ledger_str(self) -> &'static str {
match self {
LoadMode::Full => "full",
LoadMode::Incremental => "incremental",
LoadMode::Cdc => "cdc",
}
}
}
#[derive(Debug, Clone)]
pub struct LoadPlan {
pub table: String,
pub partition_by: Option<String>,
pub specs: Vec<TargetColumnSpec>,
pub gcs_prefix: String,
pub destination: crate::config::DestinationConfig,
pub load: LoadSection,
pub mode: LoadMode,
pub cursor_column: Option<String>,
}
#[derive(Deserialize)]
struct ExportReport {
export: String,
columns: Vec<ColReport>,
}
#[derive(Deserialize)]
struct ColReport {
column: String,
target_type: String,
target_status: String,
}
const LOAD_KEYS: &[&str] = &[
"target",
"cleanup_source",
"pk",
"allow_source_drift",
"gc_orphans",
"cluster_by",
"project",
"dataset",
"connection",
"warehouse",
"database",
"schema",
"storage_integration",
];
fn check_load_keys(value: &serde_json::Value, whose: &str) -> Result<()> {
if let Some(obj) = value.as_object() {
for k in obj.keys() {
if !LOAD_KEYS.contains(&k.as_str()) {
bail!(
"unknown key `{k}` in the {whose} `load:` block — valid keys are: {}",
LOAD_KEYS.join(", ")
);
}
}
}
Ok(())
}
pub fn plan_loads(config_path: &str, rivet_bin: &str) -> Result<Vec<LoadPlan>> {
let yaml = std::fs::read_to_string(config_path)
.with_context(|| format!("reading config {config_path}"))?;
let cfg = crate::config::Config::from_yaml(&yaml).context("parsing rivet config")?;
if cfg.exports.is_empty() {
bail!("config has no exports");
}
let load_value = cfg.load.clone().context(
"config has no top-level `load:` block — add `load: { target, ... }` to load into a warehouse",
)?;
check_load_keys(&load_value, "top-level")?;
let load: LoadSection =
serde_json::from_value(load_value).context("parsing the top-level `load:` block")?;
let out = Command::new(rivet_bin)
.args([
"check",
"-c",
config_path,
"--target",
load.target.name(),
"--json",
])
.output()
.with_context(|| {
format!("running `{rivet_bin} check` — is rivet on PATH? pass --rivet-bin")
})?;
if !out.status.success() {
bail!(
"rivet check failed: {}",
String::from_utf8_lossy(&out.stderr).trim()
);
}
let reports: Vec<ExportReport> = serde_json::Deserializer::from_slice(&out.stdout)
.into_iter::<ExportReport>()
.collect::<Result<_, _>>()
.context("parsing `rivet check --json` (one document per export)")?;
build_plans(&cfg, &load, reports)
}
fn build_plans(
cfg: &crate::config::Config,
load: &LoadSection,
reports: Vec<ExportReport>,
) -> Result<Vec<LoadPlan>> {
let mut plans = Vec::with_capacity(reports.len());
for report in reports {
let export = cfg
.exports
.iter()
.find(|e| e.name == report.export)
.with_context(|| {
format!(
"rivet check reported export `{}` not found in config",
report.export
)
})?;
let table = export.table.clone().unwrap_or_else(|| export.name.clone());
let dest = &export.destination;
let bucket = dest.bucket.as_deref().with_context(|| {
format!(
"export `{}` has no destination `bucket` — a GCS destination is required",
export.name
)
})?;
let prefix = dest.prefix.as_deref().unwrap_or("");
let base = prefix.split("{partition}").next().unwrap_or(prefix);
let gcs_prefix = format!("gs://{bucket}/{base}");
let specs = report
.columns
.into_iter()
.map(|c| TargetColumnSpec {
column_name: c.column,
target_type: c.target_type,
autoload_type: String::new(),
status: match c.target_status.as_str() {
"fail" => TargetStatus::Fail,
"warn" => TargetStatus::Warn,
_ => TargetStatus::Ok,
},
note: None,
cast_sql: None,
})
.collect();
let mode = match export.mode {
crate::config::ExportMode::Cdc => LoadMode::Cdc,
crate::config::ExportMode::Incremental => LoadMode::Incremental,
crate::config::ExportMode::Full => LoadMode::Full, crate::config::ExportMode::Chunked => LoadMode::Full, crate::config::ExportMode::TimeWindow => LoadMode::Full, };
let eff_load = match &export.load {
Some(v) => {
let o: LoadOverride = serde_json::from_value(v.clone()).with_context(|| {
format!("parsing export `{}` `load:` override", export.name)
})?;
if o.target.is_some() {
bail!(
"export `{}`: a per-export `load:` cannot override `target:` — the \
warehouse is shared; set `target:` in the top-level `load:` block only",
export.name
);
}
load.with_override(&o)
}
None => load.clone(),
};
plans.push(LoadPlan {
table,
partition_by: export.partition_by.clone(),
specs,
gcs_prefix,
destination: export.destination.clone(),
load: eff_load,
mode,
cursor_column: export.cursor_column.clone(),
});
}
reject_duplicate_target_tables(&plans.iter().map(|p| p.table.as_str()).collect::<Vec<_>>())?;
Ok(plans)
}
fn reject_duplicate_target_tables(tables: &[&str]) -> Result<()> {
let mut seen = std::collections::HashSet::new();
for t in tables {
if !seen.insert(*t) {
bail!(
"two exports resolve to the same load target table `{t}` — each would clobber \
the other (a full OVERWRITE vs a cdc/incremental append share the table and \
its ledger). Give each export its own `table:` or destination."
);
}
}
Ok(())
}
pub fn source_engine(config_path: &str) -> Result<crate::load::cdc::SourceEngine> {
use crate::config::SourceType;
use crate::load::cdc::SourceEngine;
let yaml = std::fs::read_to_string(config_path)
.with_context(|| format!("reading config {config_path}"))?;
let cfg = crate::config::Config::from_yaml(&yaml).context("parsing rivet config")?;
match cfg.source.source_type {
SourceType::Postgres => Ok(SourceEngine::Postgres),
SourceType::Mysql => Ok(SourceEngine::MySql),
SourceType::Mssql => Ok(SourceEngine::SqlServer),
SourceType::Mongo => Ok(SourceEngine::Mongo),
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn ledger_str_names_each_mode_stably() {
assert_eq!(LoadMode::Full.ledger_str(), "full");
assert_eq!(LoadMode::Incremental.ledger_str(), "incremental");
assert_eq!(LoadMode::Cdc.ledger_str(), "cdc");
}
fn col(name: &str, status: &str) -> ColReport {
ColReport {
column: name.into(),
target_type: "STRING".into(),
target_status: status.into(),
}
}
#[test]
fn build_plans_matches_by_name_maps_statuses_and_mode() {
let cfg = crate::config::Config::from_yaml(
r#"
source:
type: postgres
url: "postgresql://localhost/test"
exports:
- name: alpha
table: alpha_tbl
mode: full
format: parquet
destination:
type: gcs
bucket: b1
prefix: exports/alpha/
- name: beta
table: beta_tbl
mode: incremental
cursor_column: updated_at
format: parquet
destination:
type: gcs
bucket: b2
prefix: exports/beta/
load:
target: bigquery
project: p
dataset: d
"#,
)
.unwrap();
let load: LoadSection = serde_json::from_value(cfg.load.clone().unwrap()).unwrap();
let reports = vec![
ExportReport {
export: "beta".into(),
columns: vec![col("id", "ok"), col("f", "fail"), col("w", "warn")],
},
ExportReport {
export: "alpha".into(),
columns: vec![col("id", "ok")],
},
];
let plans = build_plans(&cfg, &load, reports).unwrap();
assert_eq!(plans.len(), 2);
assert_eq!(
plans[0].table, "beta_tbl",
"found beta by name, not position"
);
assert_eq!(plans[0].mode, LoadMode::Incremental);
assert_eq!(plans[0].cursor_column.as_deref(), Some("updated_at"));
assert_eq!(plans[0].gcs_prefix, "gs://b2/exports/beta/");
let statuses: Vec<_> = plans[0].specs.iter().map(|s| s.status).collect();
assert_eq!(
statuses,
vec![TargetStatus::Ok, TargetStatus::Fail, TargetStatus::Warn]
);
assert_eq!(plans[1].table, "alpha_tbl");
assert_eq!(plans[1].mode, LoadMode::Full);
assert_eq!(plans[1].gcs_prefix, "gs://b1/exports/alpha/");
}
#[test]
fn build_plans_bails_on_a_report_for_an_unknown_export() {
let cfg = crate::config::Config::from_yaml(
"source:\n type: postgres\n url: \"postgresql://localhost/test\"\n\
exports:\n - name: a\n query: \"SELECT 1\"\n format: parquet\n \
destination:\n type: gcs\n bucket: b\n prefix: p/\nload:\n \
target: bigquery\n project: p\n dataset: d\n",
)
.unwrap();
let load: LoadSection = serde_json::from_value(cfg.load.clone().unwrap()).unwrap();
let reports = vec![ExportReport {
export: "ghost".into(),
columns: vec![],
}];
let err = build_plans(&cfg, &load, reports).unwrap_err().to_string();
assert!(err.contains("ghost") && err.contains("not found"), "{err}");
}
#[test]
fn reject_duplicate_target_tables_catches_a_collision() {
assert!(reject_duplicate_target_tables(&["orders", "events", "orders"]).is_err());
assert!(reject_duplicate_target_tables(&["orders", "events"]).is_ok());
assert!(reject_duplicate_target_tables(&[]).is_ok());
}
#[test]
fn multi_export_check_json_parses_as_a_stream() {
let raw = concat!(
"{\"export\":\"orders\",\"columns\":[{\"column\":\"id\",\"target_type\":\"NUMBER\",\"target_status\":\"ok\"}]}\n",
"{\"export\":\"customers\",\"columns\":[{\"column\":\"cid\",\"target_type\":\"NUMBER\",\"target_status\":\"ok\"}]}\n"
);
let reports: Vec<ExportReport> = serde_json::Deserializer::from_slice(raw.as_bytes())
.into_iter::<ExportReport>()
.collect::<Result<_, _>>()
.unwrap();
assert_eq!(reports.len(), 2);
assert_eq!(reports[0].export, "orders");
assert_eq!(reports[1].export, "customers");
assert_eq!(reports[1].columns[0].column, "cid");
}
#[test]
fn bigquery_load_section_deserializes_into_its_variant() {
let value = serde_json::json!({
"target": "bigquery", "project": "p", "dataset": "d",
"cleanup_source": true, "cluster_by": ["customer"]
});
let load: LoadSection = serde_json::from_value(value).unwrap();
assert_eq!(load.target.name(), "bigquery");
assert!(load.cleanup_source);
assert_eq!(load.cluster_by, vec!["customer"]);
match load.target {
LoadTarget::Bigquery { project, dataset } => {
assert_eq!((project.as_str(), dataset.as_str()), ("p", "d"));
}
_ => panic!("expected Bigquery variant"),
}
}
#[test]
fn snowflake_missing_field_is_unrepresentable_no_runtime_validate() {
let full = serde_json::json!({
"target": "snowflake", "connection": "rivet", "warehouse": "wh",
"database": "db", "schema": "sc", "storage_integration": "si"
});
let load: LoadSection = serde_json::from_value(full).unwrap();
assert_eq!(load.target.name(), "snowflake");
let partial = serde_json::json!({
"target": "snowflake", "connection": "rivet", "warehouse": "wh",
"database": "db", "schema": "sc"
});
let err = serde_json::from_value::<LoadSection>(partial).unwrap_err();
assert!(
err.to_string().contains("storage_integration"),
"error should name the missing field: {err}"
);
}
#[test]
fn unknown_target_is_rejected_at_deserialize() {
let value = serde_json::json!({ "target": "redshift", "project": "p" });
assert!(serde_json::from_value::<LoadSection>(value).is_err());
}
fn top_level_load() -> LoadSection {
serde_json::from_value(serde_json::json!({
"target": "bigquery", "project": "p", "dataset": "d",
"pk": ["top"], "cleanup_source": true, "gc_orphans": false,
"cluster_by": ["c0"], "allow_source_drift": false,
}))
.unwrap()
}
#[test]
fn with_override_replaces_some_fields_and_inherits_the_rest() {
let top = top_level_load();
let o: LoadOverride =
serde_json::from_value(serde_json::json!({ "pk": ["id"], "gc_orphans": true }))
.unwrap();
let eff = top.with_override(&o);
assert_eq!(eff.pk, vec!["id"], "pk replaced");
assert!(eff.gc_orphans, "gc_orphans replaced");
assert!(
eff.cleanup_source,
"cleanup_source inherited (top-level true)"
);
assert_eq!(eff.cluster_by, vec!["c0"], "cluster_by inherited");
assert!(!eff.allow_source_drift, "allow_source_drift inherited");
}
#[test]
fn override_parsing_leaves_omitted_fields_none() {
let o: LoadOverride = serde_json::from_value(serde_json::json!({ "pk": ["id"] })).unwrap();
assert_eq!(o.pk.as_deref(), Some(&["id".to_string()][..]));
assert!(o.cleanup_source.is_none());
assert!(o.gc_orphans.is_none());
assert!(o.cluster_by.is_none());
assert!(o.allow_source_drift.is_none());
assert!(o.target.is_none());
}
#[test]
fn empty_override_is_distinct_from_inherit() {
let top = top_level_load();
let cleared: LoadOverride =
serde_json::from_value(serde_json::json!({ "pk": [] })).unwrap();
assert!(
top.with_override(&cleared).pk.is_empty(),
"explicit [] clears"
);
let inherit: LoadOverride = serde_json::from_value(serde_json::json!({})).unwrap();
assert_eq!(
top.with_override(&inherit).pk,
vec!["top"],
"omitted inherits"
);
}
#[test]
fn override_carrying_target_is_detected() {
let o: LoadOverride =
serde_json::from_value(serde_json::json!({ "target": "snowflake" })).unwrap();
assert!(
o.target.is_some(),
"target captured for the plan_loads guard"
);
}
#[test]
fn unknown_top_level_load_key_is_rejected() {
let typo = serde_json::json!({
"target": "bigquery", "project": "p", "dataset": "d", "gc_orphan": true
});
let err = check_load_keys(&typo, "top-level").unwrap_err().to_string();
assert!(err.contains("gc_orphan"), "{err}");
let ok = serde_json::json!({
"target": "bigquery", "project": "p", "dataset": "d",
"gc_orphans": true, "cleanup_source": false, "pk": ["id"],
"allow_source_drift": true, "cluster_by": ["a"]
});
assert!(check_load_keys(&ok, "top-level").is_ok());
}
#[test]
fn unknown_per_export_override_key_is_rejected() {
let typo = serde_json::json!({ "pk": ["id"], "cluster_bye": ["x"] });
assert!(
serde_json::from_value::<LoadOverride>(typo).is_err(),
"a typo'd override key must fail to parse"
);
let ok = serde_json::json!({ "pk": ["id"], "cleanup_source": true });
assert!(serde_json::from_value::<LoadOverride>(ok).is_ok());
}
}