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_core::util::substitute_context_bind_params;
14use faucet_core::{AuthSpec, FaucetError, SharedAuthProvider, Stream, StreamPage};
15use faucet_common_snowflake::{
16    SnowflakeAuth, authorization_header, credential_to_auth, snowflake_token_type,
17};
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    /// Preflight probe for `faucet doctor`. Overrides the default (which pulls a
411    /// page via `stream_pages` and would **execute the configured query**).
412    /// Instead submits a cheap `SELECT 1` so doctor validates auth, warehouse,
413    /// and connectivity without running the user's real statement.
414    async fn check(
415        &self,
416        ctx: &faucet_core::check::CheckContext,
417    ) -> Result<faucet_core::check::CheckReport, FaucetError> {
418        use faucet_core::check::{CheckReport, Probe};
419        let start = std::time::Instant::now();
420        let probe = async {
421            let mut body = self.build_request_body(&[]);
422            body["statement"] = Value::String("SELECT 1".to_string());
423            let url = self.statements_url();
424            let effective = self.resolve_auth().await.map_err(|e| {
425                Probe::fail_hint(
426                    "auth",
427                    start.elapsed(),
428                    e.to_string(),
429                    "verify the Snowflake credentials / shared auth provider",
430                )
431            })?;
432            let auth = authorization_header(&effective, &self.config.account)
433                .map_err(|e| Probe::fail("auth", start.elapsed(), e.to_string()))?;
434            let token_type = snowflake_token_type(&effective);
435            let resp = self
436                .client
437                .post(&url)
438                .header("Authorization", &auth)
439                .header("Content-Type", "application/json")
440                .header("Accept", "application/json")
441                .header("X-Snowflake-Authorization-Token-Type", token_type)
442                .json(&body)
443                .send()
444                .await
445                .map_err(|e| {
446                    Probe::fail_hint(
447                        "query",
448                        start.elapsed(),
449                        format!("Snowflake request failed: {e}"),
450                        "verify the account endpoint is reachable",
451                    )
452                })?;
453            let status = resp.status();
454            if status.is_success() || status.as_u16() == 202 {
455                Ok::<Probe, Probe>(Probe::pass("query", start.elapsed()))
456            } else {
457                let text = resp.text().await.unwrap_or_default();
458                Err(Probe::fail_hint(
459                    "query",
460                    start.elapsed(),
461                    format!("Snowflake SQL API returned HTTP {status}: {text}"),
462                    "verify credentials, warehouse, database/schema, and role",
463                ))
464            }
465        };
466        let probe = match tokio::time::timeout(ctx.timeout, probe).await {
467            Ok(Ok(p)) | Ok(Err(p)) => p,
468            Err(_elapsed) => Probe::fail_hint(
469                "query",
470                start.elapsed(),
471                "Snowflake probe timed out",
472                "Snowflake did not respond within the check timeout",
473            ),
474        };
475        Ok(CheckReport::single(probe))
476    }
477
478    async fn fetch_with_context(
479        &self,
480        context: &HashMap<String, Value>,
481    ) -> Result<Vec<Value>, FaucetError> {
482        let initial = self.submit_statement(context).await?;
483        let columns = initial
484            .result_set_metadata
485            .as_ref()
486            .map(|m| m.row_type.clone())
487            .unwrap_or_default();
488
489        if columns.is_empty() {
490            tracing::info!(
491                rows = 0,
492                query = %self.config.query,
493                "Snowflake source fetch returned no schema (likely no rows)",
494            );
495            return Ok(Vec::new());
496        }
497
498        let mut rows: Vec<Value> = initial
499            .data
500            .unwrap_or_default()
501            .iter()
502            .map(|r| row_to_json(r, &columns))
503            .collect();
504
505        let partition_count = initial
506            .result_set_metadata
507            .as_ref()
508            .map(|m| m.partition_info.len())
509            .unwrap_or(0);
510
511        if partition_count > 1 {
512            let handle = initial.statement_handle.ok_or_else(|| {
513                FaucetError::Source(
514                    "Snowflake reported >1 partition without a statementHandle to fetch them"
515                        .into(),
516                )
517            })?;
518            let effective = self.resolve_auth().await?;
519            let auth = authorization_header(&effective, &self.config.account)?;
520            let token_type = snowflake_token_type(&effective);
521
522            for i in 1..partition_count {
523                let raw = self.fetch_partition(&handle, i, &auth, token_type).await?;
524                for r in raw {
525                    rows.push(row_to_json(&r, &columns));
526                }
527            }
528        }
529
530        tracing::info!(
531            rows = rows.len(),
532            query = %self.config.query,
533            "Snowflake source fetch complete",
534        );
535        Ok(rows)
536    }
537
538    /// Stream rows partition-by-partition without buffering the full result
539    /// set. Rows are accumulated into a per-page buffer of
540    /// [`SnowflakeSourceConfig::batch_size`] entries and yielded as soon as
541    /// the buffer is full.
542    ///
543    /// The trait-level `batch_size` argument is ignored in favour of the
544    /// config field — the config is the user-facing knob the README
545    /// documents, and routing the pipeline-supplied hint through it would
546    /// silently override an explicit config value.
547    ///
548    /// `batch_size = 0` drains every partition and emits the full result
549    /// set as a single page. The source has no incremental-replication mode
550    /// today, so every emitted page carries `bookmark: None`.
551    fn stream_pages<'a>(
552        &'a self,
553        context: &'a HashMap<String, Value>,
554        _batch_size: usize,
555    ) -> Pin<Box<dyn Stream<Item = Result<StreamPage, FaucetError>> + Send + 'a>> {
556        let batch_size = self.config.batch_size;
557
558        Box::pin(async_stream::try_stream! {
559            let initial = self.submit_statement(context).await?;
560            let columns = initial
561                .result_set_metadata
562                .as_ref()
563                .map(|m| m.row_type.clone())
564                .unwrap_or_default();
565
566            if columns.is_empty() {
567                // Either the query returned no schema or an empty result set
568                // with no metadata. Either way nothing to emit.
569                return;
570            }
571
572            let partition_count = initial
573                .result_set_metadata
574                .as_ref()
575                .map(|m| m.partition_info.len())
576                .unwrap_or(1);
577
578            let chunk = if batch_size == 0 { usize::MAX } else { batch_size };
579            let initial_capacity = if batch_size == 0 {
580                // Sum the row counts across partitions for a tight pre-allocation.
581                initial
582                    .result_set_metadata
583                    .as_ref()
584                    .map(|m| m.partition_info.iter().map(|p| p.row_count as usize).sum())
585                    .unwrap_or(1024)
586            } else {
587                batch_size
588            };
589
590            let mut buffer: Vec<Value> = Vec::with_capacity(initial_capacity);
591            let mut total = 0usize;
592
593            // First partition: rows come back in the POST response itself.
594            for raw in initial.data.unwrap_or_default() {
595                buffer.push(row_to_json(&raw, &columns));
596                if buffer.len() >= chunk {
597                    let page = std::mem::replace(&mut buffer, Vec::with_capacity(initial_capacity));
598                    total += page.len();
599                    yield StreamPage { records: page, bookmark: None };
600                }
601            }
602
603            // Remaining partitions: one GET each.
604            if partition_count > 1 {
605                let handle = initial.statement_handle.ok_or_else(|| {
606                    FaucetError::Source(
607                        "Snowflake reported >1 partition without a statementHandle".into(),
608                    )
609                })?;
610                let effective = self.resolve_auth().await?;
611                let auth = authorization_header(&effective, &self.config.account)?;
612                let token_type = snowflake_token_type(&effective);
613
614                for i in 1..partition_count {
615                    let raw = self.fetch_partition(&handle, i, &auth, token_type).await?;
616                    for r in raw {
617                        buffer.push(row_to_json(&r, &columns));
618                        if buffer.len() >= chunk {
619                            let page = std::mem::replace(
620                                &mut buffer,
621                                Vec::with_capacity(initial_capacity),
622                            );
623                            total += page.len();
624                            yield StreamPage { records: page, bookmark: None };
625                        }
626                    }
627                }
628            }
629
630            if !buffer.is_empty() {
631                total += buffer.len();
632                yield StreamPage { records: buffer, bookmark: None };
633            }
634
635            tracing::info!(
636                rows = total,
637                batch_size,
638                query = %self.config.query,
639                "Snowflake source stream complete",
640            );
641        })
642    }
643}
644
645#[cfg(test)]
646mod tests {
647    use super::*;
648    use crate::config::SnowflakeAuth;
649
650    fn cfg() -> SnowflakeSourceConfig {
651        SnowflakeSourceConfig::new(
652            "xy12345.us-east-1",
653            "WH",
654            "DB",
655            "PUBLIC",
656            SnowflakeAuth::OAuth { token: "t".into() },
657            "SELECT 1",
658        )
659    }
660
661    #[test]
662    fn new_rejects_out_of_range_batch_size() {
663        let mut config = cfg();
664        config.batch_size = faucet_core::MAX_BATCH_SIZE + 1;
665        match SnowflakeSource::new(config) {
666            Err(FaucetError::Config(m)) => assert!(m.contains("batch_size"), "got: {m}"),
667            _ => panic!("expected a batch_size Config error"),
668        }
669    }
670
671    #[test]
672    fn statements_url_uses_account_when_no_override() {
673        let src = SnowflakeSource::new(cfg()).unwrap();
674        assert_eq!(
675            src.statements_url(),
676            "https://xy12345.us-east-1.snowflakecomputing.com/api/v2/statements"
677        );
678    }
679
680    #[test]
681    fn statements_url_uses_endpoint_override() {
682        let src = SnowflakeSource::new(cfg())
683            .unwrap()
684            .with_endpoint_base("http://127.0.0.1:9999");
685        assert_eq!(
686            src.statements_url(),
687            "http://127.0.0.1:9999/api/v2/statements"
688        );
689    }
690
691    #[test]
692    fn partition_url_includes_handle_and_index() {
693        let src = SnowflakeSource::new(cfg())
694            .unwrap()
695            .with_endpoint_base("http://srv");
696        assert_eq!(
697            src.partition_url("abc-123", 2),
698            "http://srv/api/v2/statements/abc-123?partition=2"
699        );
700    }
701
702    #[test]
703    fn build_request_body_minimal() {
704        let src = SnowflakeSource::new(cfg()).unwrap();
705        let body = src.build_request_body(&[]);
706        assert_eq!(body["statement"], "SELECT 1");
707        assert_eq!(body["timeout"], 60);
708        assert_eq!(body["database"], "DB");
709        assert_eq!(body["schema"], "PUBLIC");
710        assert_eq!(body["warehouse"], "WH");
711        assert!(body.get("bindings").is_none());
712        assert!(body.get("role").is_none());
713    }
714
715    #[test]
716    fn build_request_body_includes_role_when_set() {
717        let mut c = cfg();
718        c.role = Some("ANALYST".into());
719        let src = SnowflakeSource::new(c).unwrap();
720        let body = src.build_request_body(&[]);
721        assert_eq!(body["role"], "ANALYST");
722    }
723
724    #[test]
725    fn build_request_body_infers_binding_types() {
726        // #78/#34: types are inferred from the JSON value (value still a
727        // string), so numeric/boolean binds compare against typed columns.
728        let src = SnowflakeSource::new(cfg()).unwrap();
729        let body = src.build_request_body(&[
730            Value::String("alice".into()),
731            json!(42),
732            json!(true),
733            json!(3.5),
734        ]);
735        let b = &body["bindings"];
736        assert_eq!(b["1"]["type"], "TEXT");
737        assert_eq!(b["1"]["value"], "alice");
738        assert_eq!(b["2"]["type"], "FIXED");
739        assert_eq!(b["2"]["value"], "42");
740        assert_eq!(b["3"]["type"], "BOOLEAN");
741        assert_eq!(b["3"]["value"], "true");
742        assert_eq!(b["4"]["type"], "REAL");
743        assert_eq!(b["4"]["value"], "3.5");
744    }
745
746    #[test]
747    fn build_request_body_null_binding_preserves_positional_alignment() {
748        // Regression for #78/#18: a NULL must occupy its positional slot as an
749        // explicit null-valued binding, not be skipped (which would shift "42"
750        // onto marker 1 and leave marker 2 unbound).
751        let src = SnowflakeSource::new(cfg()).unwrap();
752        let body = src.build_request_body(&[Value::Null, json!(42)]);
753        let b = &body["bindings"];
754        assert_eq!(b["1"]["type"], "TEXT");
755        assert_eq!(
756            b["1"]["value"],
757            Value::Null,
758            "position 1 must be a NULL binding"
759        );
760        assert_eq!(b["2"]["value"], "42", "position 2 must still be 42");
761    }
762
763    #[test]
764    fn resolve_query_with_no_context_returns_input() {
765        let src = SnowflakeSource::new(cfg().with_params(vec![json!(7)])).unwrap();
766        let (q, binds) = src.resolve_query(&HashMap::new());
767        assert_eq!(q, "SELECT 1");
768        assert_eq!(binds, vec![json!(7)]);
769    }
770
771    #[test]
772    fn resolve_query_substitutes_context_with_question_mark_markers() {
773        let mut c = cfg();
774        c.query = "SELECT * FROM t WHERE id = {parent.id}".into();
775        let src = SnowflakeSource::new(c).unwrap();
776        let mut ctx = HashMap::new();
777        ctx.insert("parent.id".to_string(), json!(7));
778        let (q, binds) = src.resolve_query(&ctx);
779        assert_eq!(q, "SELECT * FROM t WHERE id = ?");
780        assert_eq!(binds, vec![json!(7)]);
781    }
782}
faucet_source_snowflake/stream.rs

faucet_source_snowflake/
stream.rs
