use crate::destination::gcs::GcsStore;
use crate::types::target::{TargetColumnSpec, TargetStatus};
use anyhow::{Context, Result, bail};
mod bigquery;
pub mod cdc;
pub mod plan;
pub mod reconcile;
mod snowflake;
pub use bigquery::BigQueryLoader;
pub use snowflake::SnowflakeLoader;
#[derive(Debug, Clone)]
pub struct LoadReport {
pub rows_loaded: u64,
pub target_table: String,
pub source_cleaned: bool,
}
#[derive(Debug, Clone)]
pub struct CdcLoadReport {
pub rows_appended: u64,
pub changes_table: String,
pub view: String,
pub source_cleaned: bool,
}
pub trait TargetLoader {
fn fqtn(&self, table: &str) -> String;
fn materialize(&self, table: &str, specs: &[TargetColumnSpec], uris: &[String]) -> Result<u64>;
fn append_changelog(
&self,
table: &str,
specs: &[TargetColumnSpec],
uris: &[String],
pk: &[String],
) -> Result<u64>;
fn warehouse(&self) -> cdc::Warehouse;
fn create_view(&self, table: &str, view_sql: &str) -> Result<()>;
}
fn validate_specs(table: &str, specs: &[TargetColumnSpec]) -> Result<()> {
if specs.is_empty() {
bail!("no column specs for `{table}` — nothing to build a schema from");
}
let failed: Vec<&str> = specs
.iter()
.filter(|s| s.status == TargetStatus::Fail)
.map(|s| s.column_name.as_str())
.collect();
if !failed.is_empty() {
bail!(
"cannot load `{table}`: {} column(s) do not map to the warehouse: {}",
failed.len(),
failed.join(", ")
);
}
Ok(())
}
fn maybe_cleanup(cleanup: Option<(&GcsStore, &str)>) -> bool {
match cleanup {
Some((store, prefix)) => match delete_under(store, prefix) {
Ok(()) => true,
Err(e) => {
eprintln!("warning: source cleanup failed (data is safely loaded): {e:#}");
false
}
},
None => false,
}
}
#[allow(private_interfaces)]
pub fn run_load(
loader: &dyn TargetLoader,
table: &str,
specs: &[TargetColumnSpec],
uris: &[String],
expected_rows: Option<u64>,
cleanup: Option<(&GcsStore, &str)>,
) -> Result<LoadReport> {
if uris.is_empty() {
bail!("no Parquet URIs to load into `{table}`");
}
validate_specs(table, specs)?;
let rows_loaded = loader.materialize(table, specs, uris)?;
if let Some(expected) = expected_rows
&& rows_loaded != expected
{
bail!(
"count validation failed for `{}`: loaded {rows_loaded} rows, expected {expected} — \
NOT cleaning up source; investigate before re-running",
loader.fqtn(table)
);
}
let source_cleaned = maybe_cleanup(cleanup);
Ok(LoadReport {
rows_loaded,
target_table: loader.fqtn(table),
source_cleaned,
})
}
#[allow(clippy::too_many_arguments, private_interfaces)]
fn append_and_view(
loader: &dyn TargetLoader,
table: &str,
specs: &[TargetColumnSpec],
uris: &[String],
pk: &[String],
expected_delta: Option<u64>,
cleanup: Option<(&GcsStore, &str)>,
label: &str,
build_view: impl FnOnce(&dyn TargetLoader) -> Result<()>,
) -> Result<CdcLoadReport> {
if uris.is_empty() {
bail!("no Parquet URIs to append into `{table}__changes`");
}
if pk.is_empty() {
bail!("{label} load of `{table}` needs a primary key for the dedup view (pass --pk)");
}
validate_specs(&format!("{table}__changes"), specs)?;
let rows_appended = loader.append_changelog(table, specs, uris, pk)?;
if let Some(expected) = expected_delta
&& rows_appended != expected
{
bail!(
"{label} count validation failed for `{}__changes`: appended {rows_appended} rows, \
expected {expected} from the run manifests — investigate before trusting the view",
table
);
}
build_view(loader)?;
let source_cleaned = maybe_cleanup(cleanup);
Ok(CdcLoadReport {
rows_appended,
changes_table: loader.fqtn(&format!("{table}__changes")),
view: loader.fqtn(table),
source_cleaned,
})
}
#[allow(clippy::too_many_arguments, private_interfaces)]
pub fn run_load_cdc(
loader: &dyn TargetLoader,
table: &str,
specs: &[TargetColumnSpec],
uris: &[String],
pk: &[String],
engine: cdc::SourceEngine,
expected_delta: Option<u64>,
cleanup: Option<(&GcsStore, &str)>,
) -> Result<CdcLoadReport> {
append_and_view(
loader,
table,
specs,
uris,
pk,
expected_delta,
cleanup,
"CDC",
|l| {
let pk_refs: Vec<&str> = pk.iter().map(String::as_str).collect();
let sql = cdc::dedup_view_sql(
l.warehouse(),
&l.fqtn(table),
&l.fqtn(&format!("{table}__changes")),
&pk_refs,
engine,
);
l.create_view(table, &sql)
},
)
}
#[allow(clippy::too_many_arguments, private_interfaces)]
pub fn run_load_incremental(
loader: &dyn TargetLoader,
table: &str,
specs: &[TargetColumnSpec],
uris: &[String],
pk: &[String],
cursor_column: &str,
expected_delta: Option<u64>,
cleanup: Option<(&GcsStore, &str)>,
) -> Result<CdcLoadReport> {
if cursor_column.is_empty() {
bail!(
"incremental load of `{table}` needs a cursor column (the export's `cursor_column:`) \
for the dedup view's latest-per-PK ordering"
);
}
if !specs.iter().any(|s| s.column_name == cursor_column) {
let cols: Vec<&str> = specs.iter().map(|s| s.column_name.as_str()).collect();
bail!(
"incremental load of `{table}`: cursor_column `{cursor_column}` is not one of the \
exported columns [{}] — add it to the export's SELECT so the dedup view can order \
the change log by it",
cols.join(", ")
);
}
append_and_view(
loader,
table,
specs,
uris,
pk,
expected_delta,
cleanup,
"incremental",
|l| {
let pk_refs: Vec<&str> = pk.iter().map(String::as_str).collect();
let sql = cdc::inc_dedup_view_sql(
l.warehouse(),
&l.fqtn(table),
&l.fqtn(&format!("{table}__changes")),
&pk_refs,
cursor_column,
);
l.create_view(table, &sql)
},
)
}
pub(crate) fn split_gs_uri(uri: &str) -> Result<(&str, &str)> {
uri.strip_prefix("gs://")
.and_then(|rest| rest.split_once('/'))
.with_context(|| format!("not a `gs://bucket/path` URI: {uri}"))
}
pub(crate) fn delete_under(store: &GcsStore, gs_prefix: &str) -> Result<()> {
let (_, rel) = split_gs_uri(gs_prefix)?;
store
.remove_all(rel)
.with_context(|| format!("source cleanup (recursive delete of {gs_prefix}) failed"))
}
#[allow(private_interfaces)]
pub fn open_store(dest: &crate::config::DestinationConfig) -> Result<GcsStore> {
GcsStore::new(dest)
}
pub fn build_loader(plan: &plan::LoadPlan, run_id: &str) -> Box<dyn TargetLoader> {
use plan::LoadTarget;
let load = &plan.load;
match &load.target {
LoadTarget::Bigquery { project, dataset } => Box::new(build_bigquery_loader(
project,
dataset,
plan.partition_by.as_deref(),
&load.cluster_by,
run_id,
)),
LoadTarget::Snowflake {
connection,
warehouse,
database,
schema,
storage_integration,
} => {
let mut l = SnowflakeLoader::new(connection.clone());
l.warehouse = warehouse.clone();
l.database = database.clone();
l.schema = schema.clone();
l.storage_integration = storage_integration.clone();
l.cluster_by = load.cluster_by.clone();
l.run_id = Some(run_id.to_string());
l.gcs_url = plan.gcs_prefix.replacen("gs://", "gcs://", 1);
l.private_key_path = std::env::var("RIVET_SNOWFLAKE_KEY").ok();
Box::new(l)
}
}
}
fn build_bigquery_loader(
project: &str,
dataset: &str,
partition_by: Option<&str>,
cluster_by: &[String],
run_id: &str,
) -> BigQueryLoader {
let mut l = BigQueryLoader::new(project, dataset).run_id(run_id);
if let Some(part) = partition_by {
l = l.partition_by(part);
}
if !cluster_by.is_empty() {
l = l.cluster_by(cluster_by.to_vec());
}
l
}
#[cfg(test)]
mod tests {
use super::*;
use std::cell::RefCell;
#[derive(Default)]
struct FakeLoader {
rows: u64,
materialized: RefCell<Vec<String>>,
appended: RefCell<Vec<String>>,
views: RefCell<Vec<String>>,
}
impl TargetLoader for FakeLoader {
fn fqtn(&self, table: &str) -> String {
format!("db.{table}")
}
fn materialize(&self, table: &str, _: &[TargetColumnSpec], _: &[String]) -> Result<u64> {
self.materialized.borrow_mut().push(table.into());
Ok(self.rows)
}
fn append_changelog(
&self,
table: &str,
_: &[TargetColumnSpec],
_: &[String],
_: &[String],
) -> Result<u64> {
self.appended.borrow_mut().push(table.into());
Ok(self.rows)
}
fn warehouse(&self) -> cdc::Warehouse {
cdc::Warehouse::BigQuery
}
fn create_view(&self, table: &str, _view_sql: &str) -> Result<()> {
self.views.borrow_mut().push(table.into());
Ok(())
}
}
fn fs_store_with_prefix(dir: &tempfile::TempDir, rel: &str) -> GcsStore {
let obj = dir.path().join(rel).join("x.parquet");
std::fs::create_dir_all(obj.parent().unwrap()).unwrap();
std::fs::write(obj, b"x").unwrap();
GcsStore::open_fs(dir.path().to_str().unwrap()).unwrap()
}
fn prefix_populated(store: &GcsStore, rel: &str) -> bool {
!store.list_files(rel).unwrap().is_empty()
}
fn spec(status: TargetStatus) -> Vec<TargetColumnSpec> {
vec![TargetColumnSpec {
column_name: "id".into(),
target_type: "INT64".into(),
autoload_type: String::new(),
status,
note: None,
cast_sql: None,
}]
}
fn uris() -> Vec<String> {
vec!["gs://b/p/x.parquet".into()]
}
const PREFIX: &str = "gs://b/p";
const REL: &str = "p";
#[test]
fn empty_uris_bail_before_materialize() {
let f = FakeLoader {
rows: 10,
..Default::default()
};
assert!(run_load(&f, "t", &spec(TargetStatus::Ok), &[], Some(10), None).is_err());
assert!(f.materialized.borrow().is_empty());
}
#[test]
fn fail_spec_bails_before_materialize() {
let f = FakeLoader::default();
assert!(run_load(&f, "t", &spec(TargetStatus::Fail), &uris(), Some(10), None).is_err());
assert!(f.materialized.borrow().is_empty());
}
#[test]
fn count_mismatch_bails_without_cleanup() {
let f = FakeLoader {
rows: 7,
..Default::default()
};
let dir = tempfile::tempdir().unwrap();
let store = fs_store_with_prefix(&dir, REL);
let err = run_load(
&f,
"t",
&spec(TargetStatus::Ok),
&uris(),
Some(10),
Some((&store, PREFIX)),
)
.unwrap_err()
.to_string();
assert!(err.contains("count validation failed"), "{err}");
assert!(
prefix_populated(&store, REL),
"cleanup must not run on a failed gate — the source prefix stays intact"
);
}
#[test]
fn match_with_prefix_cleans_once() {
let f = FakeLoader {
rows: 10,
..Default::default()
};
let dir = tempfile::tempdir().unwrap();
let store = fs_store_with_prefix(&dir, REL);
let r = run_load(
&f,
"t",
&spec(TargetStatus::Ok),
&uris(),
Some(10),
Some((&store, PREFIX)),
)
.unwrap();
assert!(r.source_cleaned);
assert!(
!prefix_populated(&store, REL),
"a passed gate drains the source prefix through the injected store"
);
assert_eq!(r.target_table, "db.t");
}
#[test]
fn match_without_prefix_does_not_clean() {
let f = FakeLoader {
rows: 10,
..Default::default()
};
let r = run_load(&f, "t", &spec(TargetStatus::Ok), &uris(), Some(10), None).unwrap();
assert!(!r.source_cleaned);
}
#[test]
fn none_expected_skips_the_gate() {
let f = FakeLoader {
rows: 999,
..Default::default()
};
assert!(run_load(&f, "t", &spec(TargetStatus::Ok), &uris(), None, None).is_ok());
}
#[test]
fn cdc_delta_mismatch_bails_without_view() {
let f = FakeLoader {
rows: 3,
..Default::default()
};
let dir = tempfile::tempdir().unwrap();
let store = fs_store_with_prefix(&dir, REL);
let err = run_load_cdc(
&f,
"t",
&spec(TargetStatus::Ok),
&uris(),
&["id".into()],
cdc::SourceEngine::MySql,
Some(5),
Some((&store, PREFIX)),
)
.unwrap_err()
.to_string();
assert!(err.contains("CDC count validation failed"), "{err}");
assert!(
f.views.borrow().is_empty(),
"view must not be built on a failed gate"
);
assert!(
prefix_populated(&store, REL),
"cleanup must not run on a failed gate"
);
}
#[test]
fn cdc_match_builds_view_then_cleans() {
let f = FakeLoader {
rows: 5,
..Default::default()
};
let dir = tempfile::tempdir().unwrap();
let store = fs_store_with_prefix(&dir, REL);
let r = run_load_cdc(
&f,
"t",
&spec(TargetStatus::Ok),
&uris(),
&["id".into()],
cdc::SourceEngine::MySql,
Some(5),
Some((&store, PREFIX)),
)
.unwrap();
assert_eq!(r.rows_appended, 5);
assert_eq!(*f.views.borrow(), vec!["t".to_string()]);
assert!(
!prefix_populated(&store, REL),
"a passed CDC gate drains the source prefix after the view is built"
);
assert_eq!(r.changes_table, "db.t__changes");
}
#[test]
fn delete_under_drains_the_prefix_through_the_store() {
let dir = tempfile::tempdir().unwrap();
let store = fs_store_with_prefix(&dir, REL);
assert!(prefix_populated(&store, REL), "seeded object is present");
delete_under(&store, PREFIX).unwrap();
assert!(
!prefix_populated(&store, REL),
"delete_under recursively removes the bucket-relative prefix behind the gs:// URI"
);
}
#[test]
fn build_bigquery_loader_wires_partition_and_cluster_keys() {
let l = build_bigquery_loader(
"proj",
"ds",
Some("day"),
&["customer_id".to_string(), "region".to_string()],
"run-1",
);
assert_eq!(l.cluster_by, ["customer_id", "region"]);
assert_eq!(l.partition_by.as_deref(), Some("day"));
let bare = build_bigquery_loader("proj", "ds", None, &[], "run-1");
assert!(bare.cluster_by.is_empty());
assert!(bare.partition_by.is_none());
}
#[test]
fn split_gs_uri_parses_bucket_and_bucket_relative_key() {
assert_eq!(split_gs_uri("gs://b/p").unwrap(), ("b", "p"));
assert_eq!(
split_gs_uri("gs://bucket/a/b/c.parquet").unwrap(),
("bucket", "a/b/c.parquet"),
"only the FIRST '/' splits bucket from key; the rest is the key"
);
assert!(
split_gs_uri("s3://b/p").is_err(),
"a non-gs scheme is rejected"
);
assert!(
split_gs_uri("gs://bucket-only").is_err(),
"a bucket with no '/' has no (bucket, key) split"
);
}
#[test]
fn cdc_empty_pk_bails() {
let f = FakeLoader::default();
assert!(
run_load_cdc(
&f,
"t",
&spec(TargetStatus::Ok),
&uris(),
&[],
cdc::SourceEngine::MySql,
None,
None
)
.is_err()
);
assert!(f.appended.borrow().is_empty());
}
#[test]
fn incremental_cursor_not_in_specs_bails_before_append() {
let f = FakeLoader::default();
let err = run_load_incremental(
&f,
"t",
&spec(TargetStatus::Ok),
&uris(),
&["id".to_string()],
"updated_at",
None,
None,
)
.unwrap_err()
.to_string();
assert!(
err.contains("updated_at") && err.contains("not one of the exported columns"),
"{err}"
);
assert!(
f.appended.borrow().is_empty(),
"nothing appended before the bail"
);
}
}