faucet_source_snowflake/
stream.rs

1//! Snowflake SQL REST API source.
2//!
3//! Executes the configured SQL statement against
4//! [`POST /api/v2/statements`](https://docs.snowflake.com/en/developer-guide/sql-api/submitting-requests)
5//! and streams the rows back as JSON. Large result sets are split server-side
6//! into partitions; the source fetches each subsequent partition via
7//! `GET /api/v2/statements/{handle}?partition={n}` and re-frames the rows
8//! into pages of [`SnowflakeSourceConfig::batch_size`].
9
10use crate::config::SnowflakeSourceConfig;
11use crate::convert::{ColumnMeta, row_to_json};
12use async_trait::async_trait;
13use faucet_common_snowflake::{
14    SnowflakeAuth, authorization_header, credential_to_auth, snowflake_token_type,
15};
16use faucet_core::util::substitute_context_bind_params;
17use faucet_core::{AuthSpec, FaucetError, SharedAuthProvider, Stream, StreamPage};
18use reqwest::Client;
19use serde::Deserialize;
20use serde_json::{Map, Value, json};
21use std::collections::HashMap;
22use std::pin::Pin;
23use std::time::Duration;
24
25/// Snowflake API response envelope. Only the fields we need are pulled out;
26/// everything else (`createdOn`, `numRows`, `databaseProvider`, ...) is
27/// ignored.
28#[derive(Debug, Deserialize)]
29struct StatementResponse {
30    /// Snowflake SQL state code. `"090001"` means "statement executed
31    /// successfully"; anything else is an error.
32    #[serde(default)]
33    code: Option<String>,
34    /// Human-readable message paired with `code`.
35    #[serde(default)]
36    message: Option<String>,
37    /// Opaque handle for fetching additional partitions and re-polling
38    /// asynchronous results.
39    #[serde(rename = "statementHandle", default)]
40    statement_handle: Option<String>,
41    /// Result schema and partition manifest. Only present on the first
42    /// (synchronous) response or on the resync GET after async submission;
43    /// follow-up partition fetches omit it.
44    #[serde(rename = "resultSetMetaData", default)]
45    result_set_metadata: Option<ResultSetMetadata>,
46    /// One element per row, each a JSON array of stringified cells.
47    #[serde(default)]
48    data: Option<Vec<Vec<Value>>>,
49}
50
51#[derive(Debug, Deserialize)]
52struct ResultSetMetadata {
53    /// Per-column metadata. Required to map cell arrays back to JSON keys.
54    #[serde(rename = "rowType", default)]
55    row_type: Vec<ColumnMeta>,
56    /// One entry per partition. The first entry covers the rows in the
57    /// initial response; subsequent entries must be fetched individually.
58    #[serde(rename = "partitionInfo", default)]
59    partition_info: Vec<PartitionInfo>,
60}
61
62#[derive(Debug, Deserialize)]
63#[allow(dead_code)] // `rowCount` is informational; we drain partitions until empty.
64struct PartitionInfo {
65    #[serde(rename = "rowCount", default)]
66    row_count: u64,
67}
68
69/// A source that streams the rows of a Snowflake SQL statement via the SQL
70/// REST API.
71pub struct SnowflakeSource {
72    config: SnowflakeSourceConfig,
73    client: Client,
74    /// Optional explicit endpoint override (full base URL up to but not
75    /// including `/api/v2/statements`). When `None`, derived from
76    /// `config.account`. Intended for wiremock tests and private-link
77    /// deployments.
78    endpoint_base: Option<String>,
79    /// Optional shared auth provider. When set, takes precedence over inline
80    /// auth; the provider yields a `Bearer` or `Token` credential mapped onto
81    /// [`SnowflakeAuth::OAuth`]. Set via [`Self::with_auth_provider`].
82    auth_provider: Option<SharedAuthProvider>,
83}
84
85impl SnowflakeSource {
86    /// Create a new Snowflake source. Initialises the underlying HTTP client and
87    /// does no I/O; it fails only on an invalid config (an out-of-range
88    /// `batch_size`).
89    pub fn new(config: SnowflakeSourceConfig) -> Result<Self, FaucetError> {
90        faucet_core::validate_batch_size(config.batch_size)?;
91        Ok(Self {
92            config,
93            client: Client::new(),
94            endpoint_base: None,
95            auth_provider: None,
96        })
97    }
98
99    /// Attach a shared [`AuthProvider`](faucet_core::AuthProvider). When set,
100    /// the provider supplies the credential for every request (taking
101    /// precedence over inline auth), so several sources can share one OAuth
102    /// token with single-flight refresh. Used by the CLI to resolve
103    /// `auth: { ref }`, and by library callers who inject a provider directly.
104    ///
105    /// The provider must yield a `Bearer` or `Token` credential, which maps
106    /// onto [`SnowflakeAuth::OAuth`]. Key-pair JWT cannot be supplied via a
107    /// provider (JWT is minted locally from the RSA key).
108    pub fn with_auth_provider(mut self, provider: SharedAuthProvider) -> Self {
109        self.auth_provider = Some(provider);
110        self
111    }
112
113    /// Override the base URL used to reach the SQL REST API.
114    ///
115    /// Pass `http://127.0.0.1:1234` to point the source at a mock server —
116    /// the source will issue requests to `{base}/api/v2/statements` and
117    /// `{base}/api/v2/statements/{handle}?partition=N`. Useful for tests
118    /// and proxy/private-link setups.
119    pub fn with_endpoint_base(mut self, base: impl Into<String>) -> Self {
120        self.endpoint_base = Some(base.into());
121        self
122    }
123
124    fn base_url(&self) -> String {
125        match &self.endpoint_base {
126            Some(b) => b.trim_end_matches('/').to_owned(),
127            None => format!("https://{}.snowflakecomputing.com", self.config.account),
128        }
129    }
130
131    fn statements_url(&self) -> String {
132        format!("{}/api/v2/statements", self.base_url())
133    }
134
135    fn partition_url(&self, handle: &str, partition: usize) -> String {
136        format!(
137            "{}/api/v2/statements/{}?partition={}",
138            self.base_url(),
139            handle,
140            partition
141        )
142    }
143
144    /// Resolve the effective [`SnowflakeAuth`] for this request.
145    ///
146    /// Resolution order:
147    /// 1. If a shared provider is attached, call it and map the credential.
148    /// 2. Otherwise, use the inline auth from the config.
149    /// 3. If the config holds an unresolved `Reference` with no provider,
150    ///    return [`FaucetError::Auth`].
151    async fn resolve_auth(&self) -> Result<SnowflakeAuth, FaucetError> {
152        if let Some(p) = &self.auth_provider {
153            return credential_to_auth(p.credential().await?);
154        }
155        match &self.config.auth {
156            AuthSpec::Inline(a) => Ok(a.clone()),
157            AuthSpec::Reference(r) => Err(FaucetError::Auth(format!(
158                "auth references provider '{}' but no provider was supplied",
159                r.name
160            ))),
161        }
162    }
163
164    /// Build the JSON body for `POST /api/v2/statements`.
165    ///
166    /// Bindings are sent as 1-based positional values under the documented
167    /// `bindings: {"1": {"type": "<TYPE>", "value": "..."}}` shape. The value
168    /// is always a string; the `type` is inferred from the JSON value
169    /// (`FIXED` for integers, `REAL` for floats, `BOOLEAN` for bools, `TEXT`
170    /// for strings/arrays/objects) so a numeric/boolean bind compares against
171    /// a typed column rather than being forced to TEXT (#78/#34). A JSON
172    /// `null` is bound as an explicit `{"type":"TEXT","value":null}` so it
173    /// keeps its positional slot rather than shifting later parameters.
174    fn build_request_body(&self, bindings: &[Value]) -> Value {
175        let mut body = json!({
176            "statement": self.config.query,
177            "timeout": self.config.statement_timeout.as_secs(),
178            "database": self.config.database,
179            "schema": self.config.schema,
180            "warehouse": self.config.warehouse,
181        });
182
183        if let Some(role) = &self.config.role {
184            body["role"] = json!(role);
185        }
186
187        if !bindings.is_empty() {
188            let mut map = Map::with_capacity(bindings.len());
189            for (i, v) in bindings.iter().enumerate() {
190                // Bindings are 1-based positional (`?` markers). A NULL must be
191                // bound as an explicit `{"type":"TEXT","value":null}` rather
192                // than skipped — skipping leaves a gap that shifts every later
193                // parameter onto the wrong marker and corrupts the query
194                // (#78/#18).
195                //
196                // The Snowflake type is inferred from the JSON value so a
197                // numeric or boolean bind compares correctly against a
198                // NUMBER/BOOLEAN column instead of being forced to TEXT
199                // (#78/#34). The value is always sent as a string (or null).
200                let (ty, value) = match v {
201                    Value::Null => ("TEXT", Value::Null),
202                    Value::Bool(b) => ("BOOLEAN", Value::String(b.to_string())),
203                    Value::Number(n) => {
204                        let ty = if n.is_i64() || n.is_u64() {
205                            "FIXED"
206                        } else {
207                            "REAL"
208                        };
209                        (ty, Value::String(n.to_string()))
210                    }
211                    Value::String(s) => ("TEXT", Value::String(s.clone())),
212                    other => ("TEXT", Value::String(other.to_string())),
213                };
214                map.insert((i + 1).to_string(), json!({"type": ty, "value": value}));
215            }
216            body["bindings"] = Value::Object(map);
217        }
218
219        body
220    }
221
222    /// Resolve the final SQL statement and ordered bind values for a given
223    /// parent-record context.
224    ///
225    /// `{key}` tokens whose key matches `context` are replaced with `?`
226    /// markers (Snowflake's positional binding form) and their values are
227    /// appended to the returned `Vec` after any [`config.params`](SnowflakeSourceConfig::params).
228    fn resolve_query(&self, context: &HashMap<String, Value>) -> (String, Vec<Value>) {
229        let mut bindings = self.config.params.clone();
230        let (rewritten, context_values) = if context.is_empty() {
231            (self.config.query.clone(), Vec::new())
232        } else {
233            substitute_context_bind_params(&self.config.query, context, bindings.len() + 1, |_| {
234                "?".to_string()
235            })
236        };
237        bindings.extend(context_values);
238        (rewritten, bindings)
239    }
240
241    /// Issue the initial `POST /api/v2/statements` request.
242    async fn submit_statement(
243        &self,
244        context: &HashMap<String, Value>,
245    ) -> Result<StatementResponse, FaucetError> {
246        let (query, bindings) = self.resolve_query(context);
247        let mut body = self.build_request_body(&bindings);
248        body["statement"] = Value::String(query);
249
250        let url = self.statements_url();
251        let effective = self.resolve_auth().await?;
252        let auth = authorization_header(&effective, &self.config.account)?;
253        let token_type = snowflake_token_type(&effective);
254
255        let resp = self
256            .client
257            .post(&url)
258            .header("Authorization", &auth)
259            .header("Content-Type", "application/json")
260            .header("Accept", "application/json")
261            .header("X-Snowflake-Authorization-Token-Type", token_type)
262            .json(&body)
263            .send()
264            .await
265            .map_err(|e| FaucetError::Source(format!("Snowflake request failed: {e}")))?;
266
267        let status = resp.status();
268        let async_pending = status.as_u16() == 202;
269        if !status.is_success() && !async_pending {
270            let text = resp.text().await.unwrap_or_default();
271            return Err(FaucetError::Source(format!(
272                "Snowflake SQL API returned HTTP {status}: {text}"
273            )));
274        }
275
276        let parsed: StatementResponse = resp
277            .json()
278            .await
279            .map_err(|e| FaucetError::Source(format!("failed to parse Snowflake response: {e}")))?;
280
281        if async_pending {
282            // Statement still running on the server. Poll until it completes.
283            let handle = parsed.statement_handle.clone().ok_or_else(|| {
284                FaucetError::Source(
285                    "Snowflake returned 202 without a statementHandle to poll".into(),
286                )
287            })?;
288            self.poll_until_ready(&handle, &auth, token_type).await
289        } else {
290            check_code(&parsed)?;
291            Ok(parsed)
292        }
293    }
294
295    /// Poll the same handle as a partition-less GET until the response is
296    /// 200 + `code: "090001"`. Used after a 202 from the initial POST.
297    async fn poll_until_ready(
298        &self,
299        handle: &str,
300        auth: &str,
301        token_type: &'static str,
302    ) -> Result<StatementResponse, FaucetError> {
303        let url = format!("{}/api/v2/statements/{}", self.base_url(), handle);
304        let poll_timeout = self.config.poll_timeout;
305        let started = std::time::Instant::now();
306        loop {
307            let resp = self
308                .client
309                .get(&url)
310                .header("Authorization", auth)
311                .header("Accept", "application/json")
312                .header("X-Snowflake-Authorization-Token-Type", token_type)
313                .send()
314                .await
315                .map_err(|e| FaucetError::Source(format!("Snowflake poll request failed: {e}")))?;
316
317            let status = resp.status();
318            if status.as_u16() == 202 {
319                // `poll_timeout == 0` disables the cap (poll forever).
320                if !poll_timeout.is_zero() && started.elapsed() >= poll_timeout {
321                    return Err(FaucetError::Source(format!(
322                        "Snowflake statement '{handle}' did not finish within poll_timeout ({}s); still HTTP 202",
323                        poll_timeout.as_secs()
324                    )));
325                }
326                tokio::time::sleep(Duration::from_millis(500)).await;
327                continue;
328            }
329            if !status.is_success() {
330                let text = resp.text().await.unwrap_or_default();
331                return Err(FaucetError::Source(format!(
332                    "Snowflake poll returned HTTP {status}: {text}"
333                )));
334            }
335            let parsed: StatementResponse = resp.json().await.map_err(|e| {
336                FaucetError::Source(format!("failed to parse Snowflake poll response: {e}"))
337            })?;
338            check_code(&parsed)?;
339            return Ok(parsed);
340        }
341    }
342
343    /// Fetch one additional partition's worth of rows.
344    async fn fetch_partition(
345        &self,
346        handle: &str,
347        partition: usize,
348        auth: &str,
349        token_type: &'static str,
350    ) -> Result<Vec<Vec<Value>>, FaucetError> {
351        let url = self.partition_url(handle, partition);
352        let resp = self
353            .client
354            .get(&url)
355            .header("Authorization", auth)
356            .header("Accept", "application/json")
357            .header("X-Snowflake-Authorization-Token-Type", token_type)
358            .send()
359            .await
360            .map_err(|e| {
361                FaucetError::Source(format!(
362                    "Snowflake partition fetch failed (partition {partition}): {e}"
363                ))
364            })?;
365
366        let status = resp.status();
367        if !status.is_success() {
368            let text = resp.text().await.unwrap_or_default();
369            return Err(FaucetError::Source(format!(
370                "Snowflake partition fetch returned HTTP {status} (partition {partition}): {text}"
371            )));
372        }
373
374        let parsed: StatementResponse = resp.json().await.map_err(|e| {
375            FaucetError::Source(format!(
376                "failed to parse Snowflake partition response (partition {partition}): {e}"
377            ))
378        })?;
379        check_code(&parsed)?;
380        Ok(parsed.data.unwrap_or_default())
381    }
382}
383
384/// Validate `code: "090001"` (statement executed successfully). Any other
385/// code surfaces as `FaucetError::Source` carrying Snowflake's message.
386fn check_code(resp: &StatementResponse) -> Result<(), FaucetError> {
387    if let Some(code) = &resp.code
388        && code != "090001"
389    {
390        return Err(FaucetError::Source(format!(
391            "Snowflake error {}: {}",
392            code,
393            resp.message.clone().unwrap_or_default()
394        )));
395    }
396    Ok(())
397}
398
399#[async_trait]
400impl faucet_core::Source for SnowflakeSource {
401    fn connector_name(&self) -> &'static str {
402        "snowflake"
403    }
404
405    fn config_schema(&self) -> Value {
406        serde_json::to_value(faucet_core::schema_for!(SnowflakeSourceConfig))
407            .expect("schema serialization")
408    }
409
410    fn dataset_uri(&self) -> String {
411        format!(
412            "snowflake://{}/{}/{}?query={}",
413            self.config.account, self.config.database, self.config.schema, self.config.query
414        )
415    }
416
417    /// Preflight probe for `faucet doctor`. Overrides the default (which pulls a
418    /// page via `stream_pages` and would **execute the configured query**).
419    /// Instead submits a cheap `SELECT 1` so doctor validates auth, warehouse,
420    /// and connectivity without running the user's real statement.
421    async fn check(
422        &self,
423        ctx: &faucet_core::check::CheckContext,
424    ) -> Result<faucet_core::check::CheckReport, FaucetError> {
425        use faucet_core::check::{CheckReport, Probe};
426        let start = std::time::Instant::now();
427        let probe = async {
428            let mut body = self.build_request_body(&[]);
429            body["statement"] = Value::String("SELECT 1".to_string());
430            let url = self.statements_url();
431            let effective = self.resolve_auth().await.map_err(|e| {
432                Probe::fail_hint(
433                    "auth",
434                    start.elapsed(),
435                    e.to_string(),
436                    "verify the Snowflake credentials / shared auth provider",
437                )
438            })?;
439            let auth = authorization_header(&effective, &self.config.account)
440                .map_err(|e| Probe::fail("auth", start.elapsed(), e.to_string()))?;
441            let token_type = snowflake_token_type(&effective);
442            let resp = self
443                .client
444                .post(&url)
445                .header("Authorization", &auth)
446                .header("Content-Type", "application/json")
447                .header("Accept", "application/json")
448                .header("X-Snowflake-Authorization-Token-Type", token_type)
449                .json(&body)
450                .send()
451                .await
452                .map_err(|e| {
453                    Probe::fail_hint(
454                        "query",
455                        start.elapsed(),
456                        format!("Snowflake request failed: {e}"),
457                        "verify the account endpoint is reachable",
458                    )
459                })?;
460            let status = resp.status();
461            if status.is_success() || status.as_u16() == 202 {
462                Ok::<Probe, Probe>(Probe::pass("query", start.elapsed()))
463            } else {
464                let text = resp.text().await.unwrap_or_default();
465                Err(Probe::fail_hint(
466                    "query",
467                    start.elapsed(),
468                    format!("Snowflake SQL API returned HTTP {status}: {text}"),
469                    "verify credentials, warehouse, database/schema, and role",
470                ))
471            }
472        };
473        let probe = match tokio::time::timeout(ctx.timeout, probe).await {
474            Ok(Ok(p)) | Ok(Err(p)) => p,
475            Err(_elapsed) => Probe::fail_hint(
476                "query",
477                start.elapsed(),
478                "Snowflake probe timed out",
479                "Snowflake did not respond within the check timeout",
480            ),
481        };
482        Ok(CheckReport::single(probe))
483    }
484
485    async fn fetch_with_context(
486        &self,
487        context: &HashMap<String, Value>,
488    ) -> Result<Vec<Value>, FaucetError> {
489        let initial = self.submit_statement(context).await?;
490        let columns = initial
491            .result_set_metadata
492            .as_ref()
493            .map(|m| m.row_type.clone())
494            .unwrap_or_default();
495
496        if columns.is_empty() {
497            tracing::info!(
498                rows = 0,
499                query = %self.config.query,
500                "Snowflake source fetch returned no schema (likely no rows)",
501            );
502            return Ok(Vec::new());
503        }
504
505        let mut rows: Vec<Value> = initial
506            .data
507            .unwrap_or_default()
508            .iter()
509            .map(|r| row_to_json(r, &columns))
510            .collect();
511
512        let partition_count = initial
513            .result_set_metadata
514            .as_ref()
515            .map(|m| m.partition_info.len())
516            .unwrap_or(0);
517
518        if partition_count > 1 {
519            let handle = initial.statement_handle.ok_or_else(|| {
520                FaucetError::Source(
521                    "Snowflake reported >1 partition without a statementHandle to fetch them"
522                        .into(),
523                )
524            })?;
525            let effective = self.resolve_auth().await?;
526            let auth = authorization_header(&effective, &self.config.account)?;
527            let token_type = snowflake_token_type(&effective);
528
529            for i in 1..partition_count {
530                let raw = self.fetch_partition(&handle, i, &auth, token_type).await?;
531                for r in raw {
532                    rows.push(row_to_json(&r, &columns));
533                }
534            }
535        }
536
537        tracing::info!(
538            rows = rows.len(),
539            query = %self.config.query,
540            "Snowflake source fetch complete",
541        );
542        Ok(rows)
543    }
544
545    /// Stream rows partition-by-partition without buffering the full result
546    /// set. Rows are accumulated into a per-page buffer of
547    /// [`SnowflakeSourceConfig::batch_size`] entries and yielded as soon as
548    /// the buffer is full.
549    ///
550    /// The trait-level `batch_size` argument is ignored in favour of the
551    /// config field — the config is the user-facing knob the README
552    /// documents, and routing the pipeline-supplied hint through it would
553    /// silently override an explicit config value.
554    ///
555    /// `batch_size = 0` drains every partition and emits the full result
556    /// set as a single page. The source has no incremental-replication mode
557    /// today, so every emitted page carries `bookmark: None`.
558    fn stream_pages<'a>(
559        &'a self,
560        context: &'a HashMap<String, Value>,
561        _batch_size: usize,
562    ) -> Pin<Box<dyn Stream<Item = Result<StreamPage, FaucetError>> + Send + 'a>> {
563        let batch_size = self.config.batch_size;
564
565        Box::pin(async_stream::try_stream! {
566            let initial = self.submit_statement(context).await?;
567            let columns = initial
568                .result_set_metadata
569                .as_ref()
570                .map(|m| m.row_type.clone())
571                .unwrap_or_default();
572
573            if columns.is_empty() {
574                // Either the query returned no schema or an empty result set
575                // with no metadata. Either way nothing to emit.
576                return;
577            }
578
579            let partition_count = initial
580                .result_set_metadata
581                .as_ref()
582                .map(|m| m.partition_info.len())
583                .unwrap_or(1);
584
585            let chunk = if batch_size == 0 { usize::MAX } else { batch_size };
586            let initial_capacity = if batch_size == 0 {
587                // Sum the row counts across partitions for a tight pre-allocation.
588                initial
589                    .result_set_metadata
590                    .as_ref()
591                    .map(|m| m.partition_info.iter().map(|p| p.row_count as usize).sum())
592                    .unwrap_or(1024)
593            } else {
594                batch_size
595            };
596
597            let mut buffer: Vec<Value> = Vec::with_capacity(initial_capacity);
598            let mut total = 0usize;
599
600            // First partition: rows come back in the POST response itself.
601            for raw in initial.data.unwrap_or_default() {
602                buffer.push(row_to_json(&raw, &columns));
603                if buffer.len() >= chunk {
604                    let page = std::mem::replace(&mut buffer, Vec::with_capacity(initial_capacity));
605                    total += page.len();
606                    yield StreamPage { records: page, bookmark: None };
607                }
608            }
609
610            // Remaining partitions: one GET each.
611            if partition_count > 1 {
612                let handle = initial.statement_handle.ok_or_else(|| {
613                    FaucetError::Source(
614                        "Snowflake reported >1 partition without a statementHandle".into(),
615                    )
616                })?;
617                let effective = self.resolve_auth().await?;
618                let auth = authorization_header(&effective, &self.config.account)?;
619                let token_type = snowflake_token_type(&effective);
620
621                for i in 1..partition_count {
622                    let raw = self.fetch_partition(&handle, i, &auth, token_type).await?;
623                    for r in raw {
624                        buffer.push(row_to_json(&r, &columns));
625                        if buffer.len() >= chunk {
626                            let page = std::mem::replace(
627                                &mut buffer,
628                                Vec::with_capacity(initial_capacity),
629                            );
630                            total += page.len();
631                            yield StreamPage { records: page, bookmark: None };
632                        }
633                    }
634                }
635            }
636
637            if !buffer.is_empty() {
638                total += buffer.len();
639                yield StreamPage { records: buffer, bookmark: None };
640            }
641
642            tracing::info!(
643                rows = total,
644                batch_size,
645                query = %self.config.query,
646                "Snowflake source stream complete",
647            );
648        })
649    }
650}
651
652#[cfg(test)]
653mod tests {
654    use super::*;
655    use crate::config::SnowflakeAuth;
656
657    fn cfg() -> SnowflakeSourceConfig {
658        SnowflakeSourceConfig::new(
659            "xy12345.us-east-1",
660            "WH",
661            "DB",
662            "PUBLIC",
663            SnowflakeAuth::OAuth { token: "t".into() },
664            "SELECT 1",
665        )
666    }
667
668    #[test]
669    fn new_rejects_out_of_range_batch_size() {
670        let mut config = cfg();
671        config.batch_size = faucet_core::MAX_BATCH_SIZE + 1;
672        match SnowflakeSource::new(config) {
673            Err(FaucetError::Config(m)) => assert!(m.contains("batch_size"), "got: {m}"),
674            _ => panic!("expected a batch_size Config error"),
675        }
676    }
677
678    #[test]
679    fn statements_url_uses_account_when_no_override() {
680        let src = SnowflakeSource::new(cfg()).unwrap();
681        assert_eq!(
682            src.statements_url(),
683            "https://xy12345.us-east-1.snowflakecomputing.com/api/v2/statements"
684        );
685    }
686
687    #[test]
688    fn statements_url_uses_endpoint_override() {
689        let src = SnowflakeSource::new(cfg())
690            .unwrap()
691            .with_endpoint_base("http://127.0.0.1:9999");
692        assert_eq!(
693            src.statements_url(),
694            "http://127.0.0.1:9999/api/v2/statements"
695        );
696    }
697
698    #[test]
699    fn partition_url_includes_handle_and_index() {
700        let src = SnowflakeSource::new(cfg())
701            .unwrap()
702            .with_endpoint_base("http://srv");
703        assert_eq!(
704            src.partition_url("abc-123", 2),
705            "http://srv/api/v2/statements/abc-123?partition=2"
706        );
707    }
708
709    #[test]
710    fn build_request_body_minimal() {
711        let src = SnowflakeSource::new(cfg()).unwrap();
712        let body = src.build_request_body(&[]);
713        assert_eq!(body["statement"], "SELECT 1");
714        assert_eq!(body["timeout"], 60);
715        assert_eq!(body["database"], "DB");
716        assert_eq!(body["schema"], "PUBLIC");
717        assert_eq!(body["warehouse"], "WH");
718        assert!(body.get("bindings").is_none());
719        assert!(body.get("role").is_none());
720    }
721
722    #[test]
723    fn build_request_body_includes_role_when_set() {
724        let mut c = cfg();
725        c.role = Some("ANALYST".into());
726        let src = SnowflakeSource::new(c).unwrap();
727        let body = src.build_request_body(&[]);
728        assert_eq!(body["role"], "ANALYST");
729    }
730
731    #[test]
732    fn build_request_body_infers_binding_types() {
733        // #78/#34: types are inferred from the JSON value (value still a
734        // string), so numeric/boolean binds compare against typed columns.
735        let src = SnowflakeSource::new(cfg()).unwrap();
736        let body = src.build_request_body(&[
737            Value::String("alice".into()),
738            json!(42),
739            json!(true),
740            json!(3.5),
741        ]);
742        let b = &body["bindings"];
743        assert_eq!(b["1"]["type"], "TEXT");
744        assert_eq!(b["1"]["value"], "alice");
745        assert_eq!(b["2"]["type"], "FIXED");
746        assert_eq!(b["2"]["value"], "42");
747        assert_eq!(b["3"]["type"], "BOOLEAN");
748        assert_eq!(b["3"]["value"], "true");
749        assert_eq!(b["4"]["type"], "REAL");
750        assert_eq!(b["4"]["value"], "3.5");
751    }
752
753    #[test]
754    fn build_request_body_array_and_object_bindings_fall_back_to_text() {
755        // The catch-all arm (`other => ("TEXT", ...)`) stringifies array/object
756        // JSON values into a TEXT bind so they keep their positional slot.
757        let src = SnowflakeSource::new(cfg()).unwrap();
758        let body = src.build_request_body(&[json!([1, 2, 3]), json!({"k": "v"})]);
759        let b = &body["bindings"];
760        assert_eq!(b["1"]["type"], "TEXT");
761        assert_eq!(b["1"]["value"], "[1,2,3]");
762        assert_eq!(b["2"]["type"], "TEXT");
763        assert_eq!(b["2"]["value"], r#"{"k":"v"}"#);
764    }
765
766    #[test]
767    fn connector_name_is_snowflake() {
768        use faucet_core::Source;
769        let src = SnowflakeSource::new(cfg()).unwrap();
770        assert_eq!(src.connector_name(), "snowflake");
771    }
772
773    #[test]
774    fn config_schema_reports_required_fields() {
775        use faucet_core::Source;
776        let src = SnowflakeSource::new(cfg()).unwrap();
777        let schema = src.config_schema();
778        // The generated JSON Schema should describe the config's properties.
779        assert!(schema["properties"]["account"].is_object());
780        assert!(schema["properties"]["query"].is_object());
781        let required = schema["required"].as_array().expect("required array");
782        assert!(required.iter().any(|v| v == "account"));
783        assert!(required.iter().any(|v| v == "query"));
784    }
785
786    #[test]
787    fn build_request_body_null_binding_preserves_positional_alignment() {
788        // Regression for #78/#18: a NULL must occupy its positional slot as an
789        // explicit null-valued binding, not be skipped (which would shift "42"
790        // onto marker 1 and leave marker 2 unbound).
791        let src = SnowflakeSource::new(cfg()).unwrap();
792        let body = src.build_request_body(&[Value::Null, json!(42)]);
793        let b = &body["bindings"];
794        assert_eq!(b["1"]["type"], "TEXT");
795        assert_eq!(
796            b["1"]["value"],
797            Value::Null,
798            "position 1 must be a NULL binding"
799        );
800        assert_eq!(b["2"]["value"], "42", "position 2 must still be 42");
801    }
802
803    #[test]
804    fn resolve_query_with_no_context_returns_input() {
805        let src = SnowflakeSource::new(cfg().with_params(vec![json!(7)])).unwrap();
806        let (q, binds) = src.resolve_query(&HashMap::new());
807        assert_eq!(q, "SELECT 1");
808        assert_eq!(binds, vec![json!(7)]);
809    }
810
811    #[test]
812    fn resolve_query_substitutes_context_with_question_mark_markers() {
813        let mut c = cfg();
814        c.query = "SELECT * FROM t WHERE id = {parent.id}".into();
815        let src = SnowflakeSource::new(c).unwrap();
816        let mut ctx = HashMap::new();
817        ctx.insert("parent.id".to_string(), json!(7));
818        let (q, binds) = src.resolve_query(&ctx);
819        assert_eq!(q, "SELECT * FROM t WHERE id = ?");
820        assert_eq!(binds, vec![json!(7)]);
821    }
822
823    #[test]
824    fn dataset_uri_includes_account_db_schema_and_query() {
825        use faucet_core::Source;
826        let src = SnowflakeSource::new(cfg()).unwrap();
827        assert_eq!(
828            src.dataset_uri(),
829            "snowflake://xy12345.us-east-1/DB/PUBLIC?query=SELECT 1"
830        );
831    }
832}
faucet_source_snowflake/stream.rs

faucet_source_snowflake/
stream.rs
