rivet/load/
snowflake.rs

1//! Snowflake loader.
2//!
3//! Loads Rivet Parquet from GCS into a native-typed Snowflake table via a
4//! `COPY INTO` off an external stage (a GCS `STORAGE INTEGRATION`). Unlike
5//! BigQuery — where `LOAD DATA` with a declared schema coerces the Parquet for
6//! free — Snowflake:
7//!   * bills warehouse compute for the `COPY` (there is no free load), and
8//!   * does NOT parse a Parquet JSON string into a navigable `VARIANT`: a plain
9//!     `COPY` lands it as a `VARIANT`-wrapped string (`meta:key` → NULL). So a
10//!     `VARIANT` column is loaded through a `PARSE_JSON($1:col)` transform in
11//!     the same `COPY` (one billed pass), which yields a navigable `OBJECT`.
12//!
13//! Both facts were verified live before this was written.
14//!
15//! Cost attribution rides a `QUERY_TAG` (Snowflake's analogue of BigQuery job
16//! labels): the tag shows up in `ACCOUNT_USAGE.QUERY_HISTORY`, so per-`rivet_op`
17//! warehouse credits can be summed after the fact.
18
19use super::TargetLoader;
20use crate::types::target::TargetColumnSpec;
21use anyhow::{Context, Result, bail};
22use std::process::Command;
23
24/// Loads Rivet Parquet from GCS into Snowflake.
25#[derive(Debug, Default, Clone)]
26pub struct SnowflakeLoader {
27    /// The `snow` CLI connection name (e.g. `rivet`).
28    pub connection: String,
29    pub warehouse: String,
30    pub database: String,
31    pub schema: String,
32    /// A pre-created GCS `STORAGE INTEGRATION` (grants Snowflake read on the
33    /// bucket). The external stage is created per-load using it.
34    pub storage_integration: String,
35    /// `gcs://bucket/prefix/` — the external stage's URL (Snowflake wants the
36    /// `gcs://` scheme, not `gs://`).
37    pub gcs_url: String,
38    /// Clustering key — column(s)/expression(s) for `CLUSTER BY`, enabling
39    /// background auto-clustering. Empty = no clustering key. Applies only at
40    /// table creation.
41    pub cluster_by: Vec<String>,
42    /// Absolute path to the connection's private key. The `snow` CLI does not
43    /// expand `~`, so a `~`-relative `private_key_path` in the connection file
44    /// must be overridden with an absolute path via env.
45    pub private_key_path: Option<String>,
46    /// Load-run correlation id, emitted in the `QUERY_TAG` JSON as `rivet_run`
47    /// so every statement of one `rivet load` invocation shares a run key —
48    /// cost slices per run (across tables) as well as per table. `None` omits it.
49    pub run_id: Option<String>,
50}
51
52impl SnowflakeLoader {
53    pub fn new(connection: impl Into<String>) -> Self {
54        Self {
55            connection: connection.into(),
56            ..Default::default()
57        }
58    }
59
60    /// Fully-qualified `db.schema.table`. Identifiers are passed **unquoted**:
61    /// the warehouse/database/schema are pre-existing objects, and quoting a
62    /// lowercase name would miss an unquoted-created (upper-cased) object. The
63    /// tradeoff is that a reserved-word / special-char column is not handled —
64    /// a hardening TODO once a real source needs it.
65    fn fqtn(&self, table: &str) -> String {
66        format!("{}.{}.{}", self.database, self.schema, table)
67    }
68
69    /// `  id NUMBER(38,0),\n  meta VARIANT` — the native column DDL.
70    fn build_schema_ddl(specs: &[TargetColumnSpec]) -> String {
71        specs
72            .iter()
73            .map(|s| format!("  {} {}", s.column_name, s.target_type))
74            .collect::<Vec<_>>()
75            .join(",\n")
76    }
77
78    /// The `COPY` transform projection: a `VARIANT` column is parsed with
79    /// `PARSE_JSON` (a plain `COPY` would leave it a string), everything else is
80    /// passed through `$1:col` (the path key preserves the Parquet field case)
81    /// and coerced by the target column's type.
82    fn build_copy_select(specs: &[TargetColumnSpec]) -> String {
83        specs
84            .iter()
85            .map(|s| {
86                let path = format!("$1:{}", s.column_name);
87                if needs_parse_json(s) {
88                    format!("PARSE_JSON({path})")
89                } else {
90                    path
91                }
92            })
93            .collect::<Vec<_>>()
94            .join(", ")
95    }
96
97    /// `id, meta` — the explicit column list the transform loads into.
98    fn build_column_list(specs: &[TargetColumnSpec]) -> String {
99        specs
100            .iter()
101            .map(|s| s.column_name.clone())
102            .collect::<Vec<_>>()
103            .join(", ")
104    }
105
106    /// ` CLUSTER BY (created, region)` — the clustering-key clause (empty when
107    /// no key). Snowflake wraps the key in parentheses (unlike BigQuery).
108    fn cluster_clause(cluster_by: &[String]) -> String {
109        if cluster_by.is_empty() {
110            String::new()
111        } else {
112            format!(" CLUSTER BY ({})", cluster_by.join(", "))
113        }
114    }
115
116    /// A JSON query tag for post-hoc cost attribution in `QUERY_HISTORY`.
117    /// Carries `rivet_run` too when a load-run id is set, so credits summed from
118    /// `QUERY_ATTRIBUTION_HISTORY` slice per run as well as per table.
119    fn query_tag(&self, table: &str) -> String {
120        let run = self
121            .run_id
122            .as_deref()
123            .map(|r| format!(r#","rivet_run":"{}""#, sanitize_tag(r)))
124            .unwrap_or_default();
125        format!(
126            r#"{{"managed_by":"rivet","rivet_op":"load","rivet_table":"{}"{run}}}"#,
127            sanitize_tag(table)
128        )
129    }
130
131    /// Run a SQL script through `snow sql`, returning parsed JSON blocks.
132    fn run_snow(&self, sql: &str) -> Result<serde_json::Value> {
133        let mut cmd = Command::new("snow");
134        cmd.args(["sql", "-c", &self.connection, "--format", "json", "-q", sql]);
135        if let Some(key) = &self.private_key_path {
136            // snow reads SNOWFLAKE_CONNECTIONS_<CONN>_PRIVATE_KEY_PATH.
137            let env_key = format!(
138                "SNOWFLAKE_CONNECTIONS_{}_PRIVATE_KEY_PATH",
139                self.connection.to_uppercase()
140            );
141            cmd.env(env_key, key);
142        }
143        let out = cmd
144            .output()
145            .context("running `snow sql` — is the Snowflake CLI installed?")?;
146        if !out.status.success() {
147            bail!(
148                "snow sql failed: {}",
149                String::from_utf8_lossy(&out.stderr).trim()
150            );
151        }
152        serde_json::from_slice(&out.stdout).with_context(|| {
153            format!(
154                "parsing snow sql JSON output: {}",
155                String::from_utf8_lossy(&out.stdout)
156            )
157        })
158    }
159}
160
161impl TargetLoader for SnowflakeLoader {
162    fn fqtn(&self, table: &str) -> String {
163        format!("{}.{}.{}", self.database, self.schema, table)
164    }
165
166    fn materialize(&self, table: &str, specs: &[TargetColumnSpec], uris: &[String]) -> Result<u64> {
167        let fqtn = self.fqtn(table);
168        let ddl = Self::build_schema_ddl(specs);
169        let select = Self::build_copy_select(specs);
170        let columns = Self::build_column_list(specs);
171        // A per-load external stage over the export's GCS prefix; the COPY loads
172        // exactly the driver-selected files (`FILES=(…)`), NOT every Parquet under
173        // the prefix — so the mode-aware/ledger per-run selection is honored (a
174        // `PATTERN` over the prefix would load stale runs and fail the count gate).
175        let stage = format!("rivet_stage_{}", sanitize_tag(table));
176        let files = copy_files_clause(&self.gcs_url, uris)?;
177        let cluster = Self::cluster_clause(&self.cluster_by);
178
179        // `CREATE OR REPLACE` (overwrite): storage is the source of truth. Pin
180        // the session to UTC before the COPY — Snowflake otherwise stamps a
181        // Parquet timestamp with the session offset, shifting a `timestamptz`.
182        let sql = format!(
183            "ALTER SESSION SET QUERY_TAG = '{tag}';\n\
184             ALTER SESSION SET TIMEZONE = 'UTC';\n\
185             USE WAREHOUSE {wh};\n\
186             USE SCHEMA {db}.{sc};\n\
187             CREATE FILE FORMAT IF NOT EXISTS rivet_pq TYPE=PARQUET BINARY_AS_TEXT=FALSE;\n\
188             CREATE OR REPLACE STAGE {stage} URL='{url}' STORAGE_INTEGRATION={si} FILE_FORMAT=rivet_pq;\n\
189             CREATE OR REPLACE TABLE {fqtn} (\n{ddl}\n){cluster};\n\
190             COPY INTO {fqtn} ({columns})\n\
191             \x20 FROM (SELECT {select} FROM @{stage})\n\
192             \x20 FILE_FORMAT=(FORMAT_NAME=rivet_pq) {files};\n\
193             SELECT COUNT(*) AS ROWS_ FROM {fqtn};",
194            tag = self.query_tag(table),
195            wh = self.warehouse,
196            db = self.database,
197            sc = self.schema,
198            si = self.storage_integration,
199            url = self.gcs_url,
200        );
201
202        let result = self.run_snow(&sql)?;
203        // ponytail: rows via COUNT(*); can become the COPY's `rows_loaded`
204        // (metadata) behind this seam, no driver change.
205        extract_count(&result)
206            .context("COPY ran but the row count could not be read from snow output")
207    }
208
209    fn append_changelog(
210        &self,
211        table: &str,
212        specs: &[TargetColumnSpec],
213        uris: &[String],
214        pk: &[String],
215    ) -> Result<u64> {
216        use crate::load::cdc::Warehouse;
217        // Full change-log schema: rivet's `__op`/`__pos`/`__seq` meta columns
218        // (not reported by `rivet check`) ahead of the resolved data columns.
219        let mut full = crate::load::cdc::meta_column_specs(Warehouse::Snowflake);
220        full.extend(
221            specs
222                .iter()
223                .filter(|s| !is_meta_column(&s.column_name))
224                .cloned(),
225        );
226
227        let changes = format!("{table}__changes");
228        let changes_fqtn = self.fqtn(&changes);
229        let ddl = Self::build_schema_ddl(&full);
230        let select = Self::build_copy_select(&full);
231        let columns = Self::build_column_list(&full);
232        let cluster = Self::cluster_clause(pk);
233        let stage = format!("rivet_stage_{}", sanitize_tag(&changes));
234        let files = copy_files_clause(&self.gcs_url, uris)?;
235
236        // Ensure the log exists (clustered on PK), COUNT before, append via COPY,
237        // COUNT after — the delta is what THIS load added; the driver gates it.
238        let sql = format!(
239            "ALTER SESSION SET QUERY_TAG = '{tag}';\n\
240             ALTER SESSION SET TIMEZONE = 'UTC';\n\
241             USE WAREHOUSE {wh};\n\
242             USE SCHEMA {db}.{sc};\n\
243             CREATE FILE FORMAT IF NOT EXISTS rivet_pq TYPE=PARQUET BINARY_AS_TEXT=FALSE;\n\
244             CREATE OR REPLACE STAGE {stage} URL='{url}' STORAGE_INTEGRATION={si} FILE_FORMAT=rivet_pq;\n\
245             CREATE TABLE IF NOT EXISTS {changes_fqtn} (\n{ddl}\n){cluster};\n\
246             SELECT COUNT(*) AS BEFORE_ FROM {changes_fqtn};\n\
247             COPY INTO {changes_fqtn} ({columns})\n\
248             \x20 FROM (SELECT {select} FROM @{stage})\n\
249             \x20 FILE_FORMAT=(FORMAT_NAME=rivet_pq) {files};\n\
250             SELECT COUNT(*) AS AFTER_ FROM {changes_fqtn};",
251            tag = self.query_tag(&changes),
252            wh = self.warehouse,
253            db = self.database,
254            sc = self.schema,
255            si = self.storage_integration,
256            url = self.gcs_url,
257        );
258
259        let result = self.run_snow(&sql)?;
260        let before = extract_named(&result, "BEFORE_")
261            .context("CDC load ran but the pre-append count (BEFORE_) could not be read")?;
262        let after = extract_named(&result, "AFTER_")
263            .context("CDC load ran but the post-append count (AFTER_) could not be read")?;
264        Ok(after.saturating_sub(before))
265    }
266
267    fn warehouse(&self) -> crate::load::cdc::Warehouse {
268        crate::load::cdc::Warehouse::Snowflake
269    }
270
271    fn create_view(&self, table: &str, view_sql: &str) -> Result<()> {
272        // Fully-qualified DDL; a QUERY_TAG keeps it cost-attributable. CREATE VIEW
273        // is metadata — no warehouse compute needed.
274        let sql = format!(
275            "ALTER SESSION SET QUERY_TAG = '{tag}';\n{view_sql}",
276            tag = self.query_tag(table),
277        );
278        self.run_snow(&sql)?;
279        Ok(())
280    }
281}
282
283/// Whether a column name is one of rivet's CDC meta columns.
284fn is_meta_column(name: &str) -> bool {
285    matches!(name, "__op" | "__pos" | "__seq")
286}
287
288/// A column is loaded through `PARSE_JSON` iff its native type is `VARIANT`
289/// (Rivet's Snowflake resolver maps JSON → `VARIANT`); a plain `COPY` would
290/// leave it a string.
291fn needs_parse_json(spec: &TargetColumnSpec) -> bool {
292    spec.target_type.eq_ignore_ascii_case("VARIANT")
293}
294
295/// Keep a table name safe for a stage name / query tag (alnum + underscore).
296fn sanitize_tag(s: &str) -> String {
297    s.chars()
298        .map(|c| if c.is_ascii_alphanumeric() { c } else { '_' })
299        .collect()
300}
301
302/// Strip either GCS URI scheme — `gcs://` (Snowflake's stage scheme) or `gs://`
303/// (what the load driver hands us) — so a stage URL and a driver-selected URI
304/// compare on bucket/key alone, not scheme. Without this the two never matched.
305fn strip_gcs_scheme(s: &str) -> &str {
306    s.strip_prefix("gcs://")
307        .or_else(|| s.strip_prefix("gs://"))
308        .unwrap_or(s)
309}
310
311/// Build the COPY `FILES=('a.parquet', 'b/c.parquet', …)` clause from the
312/// driver-selected `uris`, each made relative to the stage URL (`gcs_url`). This
313/// is what makes Snowflake honor the mode-aware/ledger per-run selection instead
314/// of loading every Parquet under the prefix — a `PATTERN` over the prefix would
315/// re-load stale/already-loaded runs and fail the count gate.
316fn copy_files_clause(gcs_url: &str, uris: &[String]) -> Result<String> {
317    if uris.is_empty() {
318        bail!("Snowflake COPY: no Parquet URIs selected to load");
319    }
320    // Snowflake caps an explicit FILES=() list at 1000 entries; a normal run is a
321    // handful. Batching past that is a follow-up — fail loud rather than silently
322    // fall back to a whole-prefix PATTERN (the bug this fix closes).
323    if uris.len() > 1000 {
324        bail!(
325            "Snowflake COPY FILES=() caps at 1000 files, got {} — batch the load or reduce parallelism",
326            uris.len()
327        );
328    }
329    // Scheme-blind: the stage URL is `gcs://` (Snowflake's scheme) but the driver
330    // hands `gs://` uris. Strip both before matching so FILES entries come out
331    // stage-RELATIVE — else strip_prefix never matches, the FULL uri leaks into
332    // FILES=(), and Snowflake resolves it relative to the stage → a doubled
333    // `gcs://…/gs://…` path (the live-caught "file not found" bug).
334    let base = strip_gcs_scheme(gcs_url).trim_end_matches('/');
335    let files = uris
336        .iter()
337        .map(|u| {
338            let stripped = strip_gcs_scheme(u);
339            let rel = stripped
340                .strip_prefix(base)
341                .unwrap_or(stripped)
342                .trim_start_matches('/');
343            format!("'{rel}'")
344        })
345        .collect::<Vec<_>>()
346        .join(", ");
347    Ok(format!("FILES=({files})"))
348}
349
350/// Pull the `ROWS_` count out of snow's JSON (array of statement result blocks).
351fn extract_count(value: &serde_json::Value) -> Option<u64> {
352    extract_named(value, "ROWS_")
353}
354
355/// Pull a named integer column (e.g. `BEFORE_` / `AFTER_` / `ROWS_`) out of
356/// snow's JSON — an array of statement result blocks, each an array of row
357/// objects. Returns the first block carrying the key.
358fn extract_named(value: &serde_json::Value, key: &str) -> Option<u64> {
359    let blocks = value.as_array()?;
360    for block in blocks {
361        if let Some(rows) = block.as_array() {
362            for row in rows {
363                if let Some(n) = row.get(key).and_then(|v| v.as_u64()) {
364                    return Some(n);
365                }
366            }
367        }
368    }
369    None
370}
371
372#[cfg(test)]
373mod tests {
374    use super::*;
375    use crate::types::target::TargetStatus;
376
377    fn spec(name: &str, ty: &str) -> TargetColumnSpec {
378        TargetColumnSpec {
379            column_name: name.to_string(),
380            target_type: ty.to_string(),
381            autoload_type: String::new(),
382            status: TargetStatus::Ok,
383            note: None,
384            cast_sql: None,
385        }
386    }
387
388    #[test]
389    fn copy_files_clause_lists_the_selected_files_relative_to_the_stage() {
390        // PRODUCTION REALITY: the stage URL is `gcs://` (Snowflake's scheme) while
391        // the driver hands `gs://` uris. FILES must still come out stage-RELATIVE
392        // across the scheme gap — a scheme-blind strip_prefix leaks the FULL uri,
393        // which Snowflake resolves relative to the stage → a doubled
394        // `gcs://…/gs://…` path ("file not found", caught live). This test now
395        // mixes the schemes so it reproduces that bug.
396        let clause = copy_files_clause(
397            "gcs://bucket/exports/orders/",
398            &[
399                "gs://bucket/exports/orders/part-0.parquet".to_string(),
400                "gs://bucket/exports/orders/snapshot/part-1.parquet".to_string(),
401            ],
402        )
403        .unwrap();
404        assert_eq!(
405            clause, "FILES=('part-0.parquet', 'snapshot/part-1.parquet')",
406            "each uri stripped to a stage-relative path across schemes; no PATTERN"
407        );
408        assert!(!clause.contains("PATTERN"));
409        assert!(
410            !clause.contains("gs://") && !clause.contains("gcs://"),
411            "no absolute URI may leak into FILES=() — that doubles the stage prefix"
412        );
413        // A `gcs://` stage URL without a trailing slash still strips a `gs://` uri.
414        assert_eq!(
415            copy_files_clause("gcs://b/p", &["gs://b/p/f.parquet".to_string()]).unwrap(),
416            "FILES=('f.parquet')"
417        );
418        // Empty selection and the 1000-file cap both bail (never a silent
419        // whole-prefix fallback).
420        assert!(copy_files_clause("gcs://b/p/", &[]).is_err());
421        let many: Vec<String> = (0..1001)
422            .map(|i| format!("gs://b/p/f{i}.parquet"))
423            .collect();
424        assert!(copy_files_clause("gcs://b/p/", &many).is_err());
425    }
426
427    #[test]
428    fn variant_columns_are_parsed_scalars_pass_through() {
429        let specs = [
430            spec("id", "NUMBER(38,0)"),
431            spec("meta", "VARIANT"),
432            spec("created", "DATE"),
433        ];
434        let sel = SnowflakeLoader::build_copy_select(&specs);
435        assert_eq!(sel, "$1:id, PARSE_JSON($1:meta), $1:created");
436    }
437
438    #[test]
439    fn schema_ddl_and_column_list_are_unquoted() {
440        let specs = [spec("id", "NUMBER(38,0)"), spec("meta", "VARIANT")];
441        assert_eq!(
442            SnowflakeLoader::build_schema_ddl(&specs),
443            "  id NUMBER(38,0),\n  meta VARIANT"
444        );
445        assert_eq!(SnowflakeLoader::build_column_list(&specs), "id, meta");
446    }
447
448    #[test]
449    fn cluster_clause_wraps_key_in_parens_and_is_empty_when_unset() {
450        assert_eq!(SnowflakeLoader::cluster_clause(&[]), "");
451        assert_eq!(
452            SnowflakeLoader::cluster_clause(&["created".to_string(), "customer".to_string()]),
453            " CLUSTER BY (created, customer)"
454        );
455    }
456
457    #[test]
458    fn is_meta_column_matches_only_cdc_meta() {
459        assert!(is_meta_column("__op"));
460        assert!(is_meta_column("__pos"));
461        assert!(is_meta_column("__seq"));
462        assert!(!is_meta_column("id"));
463        assert!(!is_meta_column("__other"));
464    }
465
466    #[test]
467    fn fqtn_qualifies_database_schema_table() {
468        let mut l = SnowflakeLoader::new("c");
469        l.database = "DB".into();
470        l.schema = "SC".into();
471        assert_eq!(l.fqtn("orders"), "DB.SC.orders");
472    }
473
474    #[test]
475    fn query_tag_carries_run_id_when_set_and_omits_it_otherwise() {
476        let mut l = SnowflakeLoader::new("c");
477        assert_eq!(
478            l.query_tag("Orders"),
479            r#"{"managed_by":"rivet","rivet_op":"load","rivet_table":"Orders"}"#
480        );
481        // Non-alphanumerics in the id are coerced to `_` so QUERY_TAG stays
482        // valid JSON. The generated run id is pure hex, so this only bites a
483        // user-supplied `--run-id` with punctuation.
484        l.run_id = Some("r-7".to_string());
485        assert_eq!(
486            l.query_tag("Orders"),
487            r#"{"managed_by":"rivet","rivet_op":"load","rivet_table":"Orders","rivet_run":"r_7"}"#
488        );
489    }
490
491    #[test]
492    fn count_is_extracted_from_snow_json() {
493        let v = serde_json::json!([
494            [{"status": "ok"}],
495            [{"ROWS_": 50}]
496        ]);
497        assert_eq!(extract_count(&v), Some(50));
498    }
499
500    #[test]
501    fn cdc_before_and_after_counts_are_extracted_by_name() {
502        // The CDC script emits BEFORE_ and AFTER_ counts in separate blocks.
503        let v = serde_json::json!([
504            [{"status": "Statement executed successfully."}],
505            [{"BEFORE_": 10}],
506            [{"status": "rows loaded"}],
507            [{"AFTER_": 35}]
508        ]);
509        assert_eq!(extract_named(&v, "BEFORE_"), Some(10));
510        assert_eq!(extract_named(&v, "AFTER_"), Some(35));
511        assert_eq!(extract_named(&v, "MISSING_"), None);
512    }
513}
rivet/load/snowflake.rs

rivet/load/
snowflake.rs
