faucet_sink_snowflake/
config.rs

1//! Snowflake sink configuration.
2
3use faucet_core::{AuthSpec, DEFAULT_BATCH_SIZE};
4use schemars::JsonSchema;
5use serde::{Deserialize, Serialize};
6use std::collections::HashMap;
7use std::time::Duration;
8
9// Re-export the shared auth types so end-user imports remain stable
10// (`use faucet_sink_snowflake::SnowflakeAuth;` keeps working).
11pub use faucet_common_snowflake::SnowflakeAuth;
12
13/// Configuration for the Snowflake sink.
14#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
15pub struct SnowflakeSinkConfig {
16    /// Snowflake account identifier (e.g. `"xy12345.us-east-1"`).
17    pub account: String,
18    /// Warehouse to use for the session.
19    pub warehouse: String,
20    /// Database name.
21    pub database: String,
22    /// Schema name.
23    pub schema: String,
24    /// Target table name.
25    pub table: String,
26    /// Authentication: either inline (`{ type, config }`) or a `{ ref: <name> }`
27    /// pointer to a shared provider in the CLI's top-level `auth:` catalog.
28    /// A shared provider must yield a `Bearer` or `Token` credential, which
29    /// maps onto [`SnowflakeAuth::OAuth`]; key-pair JWT is always inline.
30    pub auth: AuthSpec<SnowflakeAuth>,
31    /// Maximum number of records sent per Snowflake SQL REST API request.
32    /// Defaults to [`DEFAULT_BATCH_SIZE`] (1000), which matches the
33    /// documented sweet spot for the SQL REST API.
34    ///
35    /// When `write_batch` is handed a slice larger than `batch_size`, the
36    /// sink re-chunks it into `batch_size` slices and issues one INSERT per
37    /// chunk. `batch_size = 0` is the **"no batching" sentinel** — the
38    /// records slice is forwarded as a single INSERT, no matter how large,
39    /// so upstream `StreamPage` framing flows through untouched.
40    #[serde(default = "default_batch_size")]
41    pub batch_size: usize,
42    /// Maximum wall-clock time to wait for an asynchronously-executed
43    /// INSERT to finish. Snowflake's SQL REST API answers an accepted but
44    /// not-yet-finished statement with HTTP 202 and a `statementHandle`;
45    /// the sink polls `GET /api/v2/statements/{handle}` until the statement
46    /// reports success before counting the rows as written. Without this
47    /// the sink would report success the moment Snowflake *accepted* the
48    /// statement, losing durability. Defaults to 300 seconds. Set to `0`
49    /// to poll indefinitely.
50    #[serde(
51        default = "default_poll_timeout",
52        with = "faucet_core::config::duration_secs"
53    )]
54    #[schemars(with = "u64")]
55    pub poll_timeout: Duration,
56    /// Arrow columnar **bulk-load** mode (#381): buffer Arrow `RecordBatch`es
57    /// to Parquet, upload to an external cloud stage's backing storage, then
58    /// `COPY INTO <table> FROM @stage FILE_FORMAT=(TYPE=PARQUET)` over the SQL
59    /// REST API. Requires a binary built with this crate's `arrow` feature; a
60    /// config that sets it on an `arrow`-off build is rejected at construction.
61    ///
62    /// This only drives the Arrow fast path the pipeline negotiates when the
63    /// source *and* sink are both columnar and no `Value`-shaped stage
64    /// (transforms, DLQ, exactly-once, masking, …) is configured. The regular
65    /// row path (`INSERT … PARSE_JSON`) and the exactly-once watermark MERGE
66    /// are unaffected — bulk-load is append-only.
67    #[serde(default, skip_serializing_if = "Option::is_none")]
68    pub bulk_load: Option<SnowflakeStageConfig>,
69}
70
71/// External-stage Parquet bulk-load configuration for the Arrow columnar path.
72///
73/// The named `stage` must already exist in Snowflake and point at the same
74/// cloud location as `url`; the sink uploads Parquet files to `url` (via
75/// `object_store`) and then references them as `@stage/<file>` in `COPY INTO`.
76#[derive(Clone, Serialize, Deserialize, JsonSchema)]
77pub struct SnowflakeStageConfig {
78    /// Named **external** stage in Snowflake — `MY_DB.MY_SCHEMA.MY_STAGE` or a
79    /// schema-relative `MY_STAGE`. Must already exist and reference `url`.
80    /// (Internal named stages use the `PUT` driver command, which the SQL REST
81    /// API does not support, so only external stages work here.)
82    pub stage: String,
83    /// Object-store URL of the stage's backing location, e.g.
84    /// `s3://bucket/prefix/`, `gs://bucket/prefix/`, or
85    /// `azure://container/prefix/`. Uploaded Parquet files land here and are
86    /// then loaded via `@stage/<file>`.
87    pub url: String,
88    /// Extra `object_store` config keys for the upload client (credentials,
89    /// region, endpoint, …), applied verbatim — e.g. `aws_access_key_id`,
90    /// `aws_secret_access_key`, `aws_region`, `google_service_account_key`.
91    /// Prefer `${secret:…}` / `${env:…}` interpolation for secret values.
92    #[serde(default, skip_serializing_if = "HashMap::is_empty")]
93    pub storage_options: HashMap<String, String>,
94    /// `MATCH_BY_COLUMN_NAME` COPY option — defaults to `CASE_INSENSITIVE` so
95    /// Parquet columns map to table columns by name. Set to `NONE` for
96    /// positional loading (rarely wanted for Parquet).
97    #[serde(default = "default_match_by_column_name")]
98    pub match_by_column_name: String,
99    /// Append `PURGE = TRUE` so Snowflake removes staged files after a
100    /// successful load. Default `false` (leave files for audit / debugging).
101    #[serde(default)]
102    pub purge: bool,
103}
104
105/// Manual `Debug` so `storage_options` values (which may carry cloud
106/// credentials) are never printed — only the key names are shown.
107impl std::fmt::Debug for SnowflakeStageConfig {
108    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
109        f.debug_struct("SnowflakeStageConfig")
110            .field("stage", &self.stage)
111            .field("url", &self.url)
112            .field(
113                "storage_options",
114                &self.storage_options.keys().collect::<Vec<_>>(),
115            )
116            .field("match_by_column_name", &self.match_by_column_name)
117            .field("purge", &self.purge)
118            .finish()
119    }
120}
121
122fn default_batch_size() -> usize {
123    DEFAULT_BATCH_SIZE
124}
125
126fn default_match_by_column_name() -> String {
127    "CASE_INSENSITIVE".to_string()
128}
129
130fn default_poll_timeout() -> Duration {
131    Duration::from_secs(300)
132}
133
134impl SnowflakeSinkConfig {
135    /// Create a new config with required fields and sensible defaults.
136    pub fn new(
137        account: impl Into<String>,
138        warehouse: impl Into<String>,
139        database: impl Into<String>,
140        schema: impl Into<String>,
141        table: impl Into<String>,
142        auth: SnowflakeAuth,
143    ) -> Self {
144        Self {
145            account: account.into(),
146            warehouse: warehouse.into(),
147            database: database.into(),
148            schema: schema.into(),
149            table: table.into(),
150            auth: AuthSpec::Inline(auth),
151            batch_size: DEFAULT_BATCH_SIZE,
152            poll_timeout: default_poll_timeout(),
153            bulk_load: None,
154        }
155    }
156
157    /// Enable Arrow columnar bulk-load via an external Parquet stage (#381).
158    pub fn with_bulk_load(mut self, stage: SnowflakeStageConfig) -> Self {
159        self.bulk_load = Some(stage);
160        self
161    }
162
163    /// Set the maximum wall-clock time spent polling an asynchronously
164    /// executed INSERT for completion. Pass `Duration::ZERO` to poll forever.
165    pub fn with_poll_timeout(mut self, timeout: Duration) -> Self {
166        self.poll_timeout = timeout;
167        self
168    }
169
170    /// Set the maximum number of records per Snowflake SQL REST API request.
171    ///
172    /// Pass `0` to opt out of re-chunking — the entire records slice handed
173    /// to `write_batch` is sent in a single INSERT request, preserving
174    /// upstream `StreamPage` framing.
175    pub fn with_batch_size(mut self, batch_size: usize) -> Self {
176        self.batch_size = batch_size;
177        self
178    }
179}
180
181#[cfg(test)]
182mod tests {
183    use super::*;
184
185    fn sample_auth() -> SnowflakeAuth {
186        SnowflakeAuth::OAuth {
187            token: "tok".into(),
188        }
189    }
190
191    fn sample_config() -> SnowflakeSinkConfig {
192        SnowflakeSinkConfig::new(
193            "xy12345",
194            "COMPUTE_WH",
195            "MY_DB",
196            "PUBLIC",
197            "events",
198            sample_auth(),
199        )
200    }
201
202    #[test]
203    fn default_config() {
204        let config = sample_config();
205        assert_eq!(config.account, "xy12345");
206        assert_eq!(config.warehouse, "COMPUTE_WH");
207        assert_eq!(config.database, "MY_DB");
208        assert_eq!(config.schema, "PUBLIC");
209        assert_eq!(config.table, "events");
210        assert_eq!(config.poll_timeout, Duration::from_secs(300));
211    }
212
213    #[test]
214    fn batch_size_defaults_to_default_batch_size() {
215        let config = sample_config();
216        assert_eq!(config.batch_size, faucet_core::DEFAULT_BATCH_SIZE);
217    }
218
219    #[test]
220    fn with_batch_size_overrides_default() {
221        let config = sample_config().with_batch_size(250);
222        assert_eq!(config.batch_size, 250);
223    }
224
225    #[test]
226    fn batch_size_zero_is_accepted_as_no_batching_sentinel() {
227        let config = sample_config().with_batch_size(0);
228        assert_eq!(config.batch_size, 0);
229        assert!(faucet_core::validate_batch_size(config.batch_size).is_ok());
230    }
231
232    #[test]
233    fn batch_size_above_max_is_rejected_by_validate_batch_size() {
234        let config = sample_config().with_batch_size(faucet_core::MAX_BATCH_SIZE + 1);
235        assert!(faucet_core::validate_batch_size(config.batch_size).is_err());
236    }
237
238    #[test]
239    fn with_bulk_load_sets_stage_and_defaults_deserialize() {
240        let cfg = sample_config().with_bulk_load(SnowflakeStageConfig {
241            stage: "STG".into(),
242            url: "s3://b/p/".into(),
243            storage_options: std::collections::HashMap::new(),
244            match_by_column_name: default_match_by_column_name(),
245            purge: false,
246        });
247        let stage = cfg.bulk_load.expect("bulk_load set");
248        assert_eq!(stage.stage, "STG");
249        assert_eq!(stage.match_by_column_name, "CASE_INSENSITIVE");
250        assert!(!stage.purge);
251
252        // JSON defaults: match_by_column_name + purge fall back sensibly.
253        let json = r#"{ "stage": "S", "url": "gs://b/" }"#;
254        let s: SnowflakeStageConfig = serde_json::from_str(json).unwrap();
255        assert_eq!(s.match_by_column_name, "CASE_INSENSITIVE");
256        assert!(!s.purge);
257        assert!(s.storage_options.is_empty());
258    }
259
260    #[test]
261    fn stage_debug_masks_storage_option_values() {
262        let mut opts = std::collections::HashMap::new();
263        opts.insert(
264            "aws_secret_access_key".to_string(),
265            "SUPER_SECRET".to_string(),
266        );
267        let stage = SnowflakeStageConfig {
268            stage: "STG".into(),
269            url: "s3://b/".into(),
270            storage_options: opts,
271            match_by_column_name: default_match_by_column_name(),
272            purge: true,
273        };
274        let dbg = format!("{stage:?}");
275        assert!(!dbg.contains("SUPER_SECRET"), "secret leaked: {dbg}");
276        assert!(
277            dbg.contains("aws_secret_access_key"),
278            "key name shown: {dbg}"
279        );
280        assert!(dbg.contains("purge: true"), "{dbg}");
281    }
282
283    #[test]
284    fn batch_size_deserializes_from_json() {
285        let json = r#"{
286            "account": "xy12345",
287            "warehouse": "COMPUTE_WH",
288            "database": "MY_DB",
289            "schema": "PUBLIC",
290            "table": "events",
291            "auth": {"type": "oauth", "config": {"token": "tok"}},
292            "batch_size": 250
293        }"#;
294        let config: SnowflakeSinkConfig = serde_json::from_str(json).unwrap();
295        assert_eq!(config.batch_size, 250);
296    }
297
298    #[test]
299    fn batch_size_defaults_when_absent_from_json() {
300        let json = r#"{
301            "account": "xy12345",
302            "warehouse": "COMPUTE_WH",
303            "database": "MY_DB",
304            "schema": "PUBLIC",
305            "table": "events",
306            "auth": {"type": "oauth", "config": {"token": "tok"}}
307        }"#;
308        let config: SnowflakeSinkConfig = serde_json::from_str(json).unwrap();
309        assert_eq!(config.batch_size, faucet_core::DEFAULT_BATCH_SIZE);
310    }
311}
faucet_sink_snowflake/config.rs

faucet_sink_snowflake/
config.rs
