faucet_source_snowflake/
config.rs

1//! Snowflake source configuration.
2
3use faucet_core::{AuthSpec, DEFAULT_BATCH_SIZE};
4use schemars::JsonSchema;
5use serde::{Deserialize, Serialize};
6use serde_json::Value;
7use std::time::Duration;
8
9// Re-export the shared auth types so end-user imports remain stable.
10pub use faucet_common_snowflake::SnowflakeAuth;
11
12fn default_statement_timeout() -> Duration {
13    Duration::from_secs(60)
14}
15
16fn default_poll_timeout() -> Duration {
17    Duration::from_secs(300)
18}
19
20fn default_batch_size() -> usize {
21    DEFAULT_BATCH_SIZE
22}
23
24/// Configuration for the Snowflake query source.
25#[derive(Clone, Serialize, Deserialize, JsonSchema)]
26pub struct SnowflakeSourceConfig {
27    /// Snowflake account identifier (e.g. `"xy12345.us-east-1"`).
28    pub account: String,
29    /// Warehouse to use for the session.
30    pub warehouse: String,
31    /// Database name.
32    pub database: String,
33    /// Schema name.
34    pub schema: String,
35    /// Optional role to assume for the session.
36    #[serde(default, skip_serializing_if = "Option::is_none")]
37    pub role: Option<String>,
38    /// Authentication: either inline (`{ type, config }`) or a `{ ref: <name> }`
39    /// pointer to a shared provider in the CLI's top-level `auth:` catalog.
40    /// A shared provider must yield a `Bearer` or `Token` credential, which
41    /// maps onto [`SnowflakeAuth::OAuth`]; key-pair JWT is always inline.
42    pub auth: AuthSpec<SnowflakeAuth>,
43    /// SQL statement to execute. May contain `${field.path}` placeholders that
44    /// are resolved against the parent-record context at runtime; resolved
45    /// values are sent as positional bind parameters appended after
46    /// [`params`](Self::params).
47    pub query: String,
48    /// Positional bind parameters for the query, applied in order before any
49    /// context-derived values. Snowflake's SQL REST API uses 1-based positional
50    /// binds in the JSON request body (see the `bindings` field in the
51    /// [SQL API docs](https://docs.snowflake.com/en/developer-guide/sql-api/submitting-requests#using-bind-variables-in-a-statement)).
52    #[serde(default)]
53    pub params: Vec<Value>,
54    /// Per-statement server-side timeout. Defaults to 60 seconds. Passed
55    /// through as the `timeout` field on the `POST /api/v2/statements` request
56    /// body. The HTTP-level timeout for each individual request is configured
57    /// separately by the source via the underlying `reqwest` client defaults.
58    #[serde(
59        default = "default_statement_timeout",
60        with = "faucet_core::config::duration_secs"
61    )]
62    #[schemars(with = "u64")]
63    pub statement_timeout: Duration,
64    /// Maximum wall-clock time the source will spend polling an asynchronous
65    /// statement (one for which the initial `POST /api/v2/statements`
66    /// returns HTTP 202) before giving up with
67    /// [`FaucetError::Source`](faucet_core::FaucetError::Source).
68    /// Without this cap a statement that never finishes would loop forever.
69    /// Defaults to 300 seconds. Set to `0` to disable the cap and poll
70    /// indefinitely.
71    #[serde(
72        default = "default_poll_timeout",
73        with = "faucet_core::config::duration_secs"
74    )]
75    #[schemars(with = "u64")]
76    pub poll_timeout: Duration,
77    /// Records per emitted [`StreamPage`](faucet_core::StreamPage).
78    ///
79    /// Snowflake's SQL REST API splits large result sets into *partitions* (one
80    /// chunk of rows per HTTP response). The source re-frames partitions into
81    /// `batch_size`-sized pages: rows accumulate in a buffer and are yielded as
82    /// soon as the buffer reaches `batch_size`. Defaults to
83    /// [`DEFAULT_BATCH_SIZE`].
84    ///
85    /// `batch_size = 0` is the **"no batching" sentinel**: the entire result
86    /// set is buffered and emitted in a single page. Useful for small lookup
87    /// tables, or for sinks (e.g. SQL `COPY`, BigQuery load jobs) that prefer
88    /// one large request to many small ones.
89    #[serde(default = "default_batch_size")]
90    pub batch_size: usize,
91}
92
93impl std::fmt::Debug for SnowflakeSourceConfig {
94    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
95        f.debug_struct("SnowflakeSourceConfig")
96            .field("account", &self.account)
97            .field("warehouse", &self.warehouse)
98            .field("database", &self.database)
99            .field("schema", &self.schema)
100            .field("role", &self.role)
101            .field("auth", &self.auth)
102            .field("query", &self.query)
103            .field("params", &self.params)
104            .field("statement_timeout", &self.statement_timeout)
105            .field("poll_timeout", &self.poll_timeout)
106            .field("batch_size", &self.batch_size)
107            .finish()
108    }
109}
110
111impl SnowflakeSourceConfig {
112    /// Create a new config with required fields and sensible defaults.
113    pub fn new(
114        account: impl Into<String>,
115        warehouse: impl Into<String>,
116        database: impl Into<String>,
117        schema: impl Into<String>,
118        auth: SnowflakeAuth,
119        query: impl Into<String>,
120    ) -> Self {
121        Self {
122            account: account.into(),
123            warehouse: warehouse.into(),
124            database: database.into(),
125            schema: schema.into(),
126            role: None,
127            auth: AuthSpec::Inline(auth),
128            query: query.into(),
129            params: Vec::new(),
130            statement_timeout: default_statement_timeout(),
131            poll_timeout: default_poll_timeout(),
132            batch_size: DEFAULT_BATCH_SIZE,
133        }
134    }
135
136    /// Set the session role.
137    pub fn with_role(mut self, role: impl Into<String>) -> Self {
138        self.role = Some(role.into());
139        self
140    }
141
142    /// Set positional bind parameters for the query.
143    pub fn with_params(mut self, params: Vec<Value>) -> Self {
144        self.params = params;
145        self
146    }
147
148    /// Set the per-statement server-side timeout.
149    pub fn with_statement_timeout(mut self, timeout: Duration) -> Self {
150        self.statement_timeout = timeout;
151        self
152    }
153
154    /// Set the maximum wall-clock time spent polling an asynchronous
155    /// statement before giving up. Pass `Duration::ZERO` to poll forever.
156    pub fn with_poll_timeout(mut self, timeout: Duration) -> Self {
157        self.poll_timeout = timeout;
158        self
159    }
160
161    /// Set the records-per-page hint for [`Source::stream_pages`](faucet_core::Source::stream_pages).
162    ///
163    /// Pass `0` to opt out of batching — the entire result set is emitted in
164    /// a single page.
165    pub fn with_batch_size(mut self, batch_size: usize) -> Self {
166        self.batch_size = batch_size;
167        self
168    }
169}
170
171#[cfg(test)]
172mod tests {
173    use super::*;
174    use serde_json::json;
175
176    fn sample_config() -> SnowflakeSourceConfig {
177        SnowflakeSourceConfig::new(
178            "xy12345",
179            "COMPUTE_WH",
180            "MY_DB",
181            "PUBLIC",
182            SnowflakeAuth::OAuth { token: "t".into() },
183            "SELECT * FROM events",
184        )
185    }
186
187    #[test]
188    fn default_config() {
189        let cfg = sample_config();
190        assert_eq!(cfg.account, "xy12345");
191        assert_eq!(cfg.warehouse, "COMPUTE_WH");
192        assert_eq!(cfg.database, "MY_DB");
193        assert_eq!(cfg.schema, "PUBLIC");
194        assert!(cfg.role.is_none());
195        assert!(cfg.params.is_empty());
196        assert_eq!(cfg.statement_timeout, Duration::from_secs(60));
197        assert_eq!(cfg.poll_timeout, Duration::from_secs(300));
198        assert_eq!(cfg.batch_size, faucet_core::DEFAULT_BATCH_SIZE);
199    }
200
201    #[test]
202    fn builders_compose() {
203        let cfg = sample_config()
204            .with_role("ANALYST")
205            .with_params(vec![json!(42)])
206            .with_statement_timeout(Duration::from_secs(15))
207            .with_batch_size(500);
208        assert_eq!(cfg.role.as_deref(), Some("ANALYST"));
209        assert_eq!(cfg.params, vec![json!(42)]);
210        assert_eq!(cfg.statement_timeout, Duration::from_secs(15));
211        assert_eq!(cfg.batch_size, 500);
212    }
213
214    #[test]
215    fn batch_size_zero_is_accepted_as_no_batching_sentinel() {
216        let cfg = sample_config().with_batch_size(0);
217        assert_eq!(cfg.batch_size, 0);
218        assert!(faucet_core::validate_batch_size(cfg.batch_size).is_ok());
219    }
220
221    #[test]
222    fn batch_size_above_max_is_rejected_by_validate_batch_size() {
223        let cfg = sample_config().with_batch_size(faucet_core::MAX_BATCH_SIZE + 1);
224        assert!(faucet_core::validate_batch_size(cfg.batch_size).is_err());
225    }
226
227    #[test]
228    fn deserializes_minimal_json() {
229        let json = r#"{
230            "account": "xy12345",
231            "warehouse": "WH",
232            "database": "DB",
233            "schema": "PUBLIC",
234            "auth": {"type": "oauth", "config": {"token": "t"}},
235            "query": "SELECT 1"
236        }"#;
237        let cfg: SnowflakeSourceConfig = serde_json::from_str(json).unwrap();
238        assert_eq!(cfg.batch_size, faucet_core::DEFAULT_BATCH_SIZE);
239        assert_eq!(cfg.statement_timeout, Duration::from_secs(60));
240        assert!(cfg.role.is_none());
241        assert!(cfg.params.is_empty());
242    }
243
244    #[test]
245    fn deserializes_all_optional_fields() {
246        let json = r#"{
247            "account": "xy12345",
248            "warehouse": "WH",
249            "database": "DB",
250            "schema": "PUBLIC",
251            "role": "ANALYST",
252            "auth": {"type": "oauth", "config": {"token": "t"}},
253            "query": "SELECT $1",
254            "params": [42],
255            "statement_timeout": 15,
256            "batch_size": 250
257        }"#;
258        let cfg: SnowflakeSourceConfig = serde_json::from_str(json).unwrap();
259        assert_eq!(cfg.role.as_deref(), Some("ANALYST"));
260        assert_eq!(cfg.params, vec![json!(42)]);
261        assert_eq!(cfg.statement_timeout, Duration::from_secs(15));
262        assert_eq!(cfg.batch_size, 250);
263    }
264
265    #[test]
266    fn debug_masks_auth_secrets() {
267        let cfg = SnowflakeSourceConfig::new(
268            "acct",
269            "wh",
270            "db",
271            "schema",
272            SnowflakeAuth::KeyPair {
273                user: "alice".into(),
274                private_key_pem: "PRIVATE-KEY-DATA".into(),
275            },
276            "SELECT 1",
277        );
278        let dbg = format!("{cfg:?}");
279        assert!(dbg.contains("***"));
280        assert!(!dbg.contains("PRIVATE-KEY-DATA"));
281    }
282}
faucet_source_snowflake/config.rs

faucet_source_snowflake/
config.rs
