faucet_source_databricks/
config.rs

1//! Databricks SQL query source configuration.
2
3use faucet_core::{AuthSpec, DEFAULT_BATCH_SIZE, FaucetError};
4use schemars::JsonSchema;
5use serde::{Deserialize, Serialize};
6use serde_json::Value;
7
8/// How the row path receives a statement's result (ignored under
9/// `arrow_native`, which always uses external links).
10#[derive(Clone, Copy, Debug, Serialize, Deserialize, JsonSchema, Default, PartialEq, Eq)]
11#[serde(rename_all = "snake_case")]
12pub enum ResultDisposition {
13    /// Chunks are downloaded from presigned cloud-storage URLs (no size cap;
14    /// the client must be able to reach the workspace's storage). Default.
15    #[default]
16    ExternalLinks,
17    /// Chunks are returned inside the API response. Databricks fails a
18    /// statement whose result exceeds 25 MiB in this mode.
19    Inline,
20}
21
22/// How the source replicates rows across runs.
23///
24/// Serializes as `{ type: full }` or
25/// `{ type: incremental, column: "...", initial_value: ... }`.
26#[derive(Clone, Debug, Serialize, Deserialize, JsonSchema, Default, PartialEq)]
27#[serde(tag = "type", rename_all = "snake_case")]
28pub enum DatabricksReplication {
29    /// Every run fetches the full result set (default).
30    #[default]
31    Full,
32    /// Only rows not yet read, in `column` order, resuming from the stored
33    /// bookmark (or `initial_value` on the first run).
34    ///
35    /// The source wraps the SQL as
36    /// ``SELECT * FROM (<sql>) AS q WHERE q.`column` >= :_faucet_bookmark ORDER BY q.`column` ``
37    /// and drops the rows it already emitted at the bookmark value, so rows
38    /// that share a cursor value (ties, late commits) are never skipped.
39    /// Every page carries a bookmark, so a crash replays at most one page.
40    /// Write `${bookmark}` in the SQL (with `>=`) to apply the cursor further
41    /// inside it; a strict `>` is refused. Rows whose `column` is NULL are
42    /// never read.
43    Incremental {
44        /// Column whose value is the replication cursor (e.g. `updated_at`).
45        column: String,
46        /// Lower bound (inclusive) used on the first run, before any bookmark
47        /// is stored.
48        initial_value: Value,
49    },
50}
51
52fn default_wait_timeout() -> u64 {
53    50
54}
55
56fn default_statement_timeout() -> u64 {
57    3600
58}
59
60fn default_poll_interval() -> u64 {
61    1
62}
63
64fn default_batch_size() -> usize {
65    DEFAULT_BATCH_SIZE
66}
67
68/// Authentication for the Databricks SQL Statement Execution API.
69///
70/// Both variants send `Authorization: Bearer <token>` — Databricks accepts a
71/// Personal Access Token (PAT) or an OAuth machine-to-machine (M2M) access
72/// token in the same header. Uses the project-wide adjacently-tagged
73/// `{ type, config }` shape, e.g. `auth: { type: pat, config: { token: … } }`.
74/// A shared `auth: { ref: <name> }` provider that yields a `Bearer`/`Token`
75/// credential maps onto [`DatabricksAuth::Token`].
76///
77/// Same shape as [`faucet_common_databricks::DatabricksAuth`] (which the sink
78/// uses); kept as this crate's own type so its public API is unchanged, and
79/// converted with `From` where the shared client needs it.
80#[derive(Clone, Serialize, Deserialize, JsonSchema)]
81#[serde(tag = "type", content = "config", rename_all = "snake_case")]
82pub enum DatabricksAuth {
83    /// Databricks Personal Access Token.
84    Pat {
85        /// The token string (use `${env:…}` / `${vault:…}` to inject).
86        token: String,
87    },
88    /// A pre-obtained OAuth (M2M) bearer token.
89    Token {
90        /// The bearer token string.
91        token: String,
92    },
93}
94
95impl std::fmt::Debug for DatabricksAuth {
96    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
97        faucet_common_databricks::DatabricksAuth::from(self.clone()).fmt(f)
98    }
99}
100
101impl DatabricksAuth {
102    /// The `Authorization` header value (`Bearer <token>`).
103    pub fn authorization_value(&self) -> String {
104        faucet_common_databricks::DatabricksAuth::from(self.clone()).authorization_value()
105    }
106}
107
108impl From<DatabricksAuth> for faucet_common_databricks::DatabricksAuth {
109    fn from(a: DatabricksAuth) -> Self {
110        match a {
111            DatabricksAuth::Pat { token } => Self::Pat { token },
112            DatabricksAuth::Token { token } => Self::Token { token },
113        }
114    }
115}
116
117/// The shared-client form of an `auth` field.
118pub(crate) fn shared_auth(
119    auth: &AuthSpec<DatabricksAuth>,
120) -> AuthSpec<faucet_common_databricks::DatabricksAuth> {
121    match auth {
122        AuthSpec::Inline(a) => AuthSpec::Inline(a.clone().into()),
123        AuthSpec::Reference(r) => AuthSpec::Reference(r.clone()),
124    }
125}
126
127/// A named SQL parameter passed to the statement (`:name` markers in the SQL).
128#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
129#[serde(deny_unknown_fields)]
130pub struct DatabricksParam {
131    /// Parameter name (the `:name` marker in the SQL, without the colon).
132    pub name: String,
133    /// Parameter value. Sent to Databricks as its string form; `null` binds a
134    /// typed NULL.
135    #[serde(default)]
136    pub value: Value,
137    /// Optional Databricks SQL type (e.g. `INT`, `STRING`, `DATE`, `TIMESTAMP`).
138    /// When omitted, Databricks treats the value as `STRING`.
139    #[serde(rename = "type", default, skip_serializing_if = "Option::is_none")]
140    pub param_type: Option<String>,
141}
142
143/// Configuration for the Databricks SQL query source.
144///
145/// Runs `sql` against a Databricks SQL Warehouse via the Statement Execution
146/// REST API and streams the result rows as typed JSON objects.
147#[derive(Clone, Serialize, Deserialize, JsonSchema)]
148#[serde(deny_unknown_fields)]
149pub struct DatabricksSourceConfig {
150    /// Workspace base URL, e.g. `https://dbc-abc123.cloud.databricks.com`.
151    #[schemars(extend("x-faucet-resource" = "url"))]
152    pub workspace_url: String,
153    /// Target SQL Warehouse id.
154    pub warehouse_id: String,
155    /// The SQL statement to run. May contain `:name` parameter markers (see
156    /// [`parameters`](Self::parameters)) and a `${bookmark}` token for
157    /// incremental replication.
158    pub sql: String,
159    /// Authentication (PAT / OAuth bearer), inline or via a shared `auth: { ref }`.
160    pub auth: AuthSpec<DatabricksAuth>,
161    /// Default Unity Catalog catalog for the statement.
162    #[serde(default, skip_serializing_if = "Option::is_none")]
163    pub catalog: Option<String>,
164    /// Default schema for the statement.
165    #[serde(default, skip_serializing_if = "Option::is_none")]
166    pub schema: Option<String>,
167    /// Named SQL parameters (`:name` markers).
168    #[serde(default)]
169    pub parameters: Vec<DatabricksParam>,
170    /// Server-side wait before the statement goes async (seconds). Databricks
171    /// accepts `0` (fully async) or `5`–`50`. Defaults to `50`.
172    #[serde(default = "default_wait_timeout")]
173    pub wait_timeout_secs: u64,
174    /// Client poll cadence while the statement is `PENDING`/`RUNNING` (seconds).
175    #[serde(default = "default_poll_interval")]
176    pub poll_interval_secs: u64,
177    /// Client deadline for the statement to finish, in seconds; on expiry it is
178    /// cancelled and the read fails. `0` waits indefinitely. Defaults to 3600.
179    #[serde(default = "default_statement_timeout")]
180    pub statement_timeout_secs: u64,
181    /// Page size — rows accumulated before a `StreamPage` is emitted.
182    /// Defaults to [`DEFAULT_BATCH_SIZE`](faucet_core::DEFAULT_BATCH_SIZE).
183    ///
184    /// `batch_size = 0` is the "no batching" sentinel: result chunks are fully
185    /// drained and the entire result set is emitted in a single page. Useful
186    /// for small lookup tables or for sinks (e.g. SQL `COPY`, BigQuery load
187    /// jobs) that prefer one large request to many small ones.
188    #[serde(default = "default_batch_size")]
189    pub batch_size: usize,
190    /// Fetch results as Apache Arrow instead of JSON. When `true`, the
191    /// statement is submitted with `EXTERNAL_LINKS` disposition +
192    /// `ARROW_STREAM` format; each chunk's presigned link is fetched and
193    /// decoded as an Arrow IPC stream. This enables the **columnar** fast path
194    /// ([`Source::stream_batches`](faucet_core::Source::stream_batches)) so a
195    /// `databricks → parquet`/`delta` chain skips the per-cell JSON decode.
196    ///
197    /// Requires the crate-local `arrow` feature, and (in this release) only
198    /// [`DatabricksReplication::Full`] — the columnar path does not run the
199    /// per-row cursor bookmark. Defaults to `false` (the
200    /// `INLINE` + `JSON_ARRAY` row path, unchanged). RFC 0002 / #375.
201    #[serde(default)]
202    pub arrow_native: bool,
203    /// How the row path fetches results: `external_links` (default) downloads
204    /// each chunk from a presigned URL, so results of any size stream;
205    /// `inline` keeps them in the API response, which Databricks caps at
206    /// 25 MiB. Use `inline` only when the client cannot reach the workspace's
207    /// cloud storage.
208    #[serde(default)]
209    pub result_disposition: ResultDisposition,
210    /// Replication mode. Defaults to [`DatabricksReplication::Full`].
211    #[serde(default)]
212    pub replication: DatabricksReplication,
213    /// Explicit state-store key for the incremental bookmark. When unset, a key
214    /// is derived from the workspace/warehouse and a query fingerprint.
215    #[serde(default, skip_serializing_if = "Option::is_none")]
216    pub state_key: Option<String>,
217}
218
219impl DatabricksSourceConfig {
220    /// Validate the config; returns a human-readable error if invalid.
221    pub fn validate(&self) -> Result<(), FaucetError> {
222        if self.workspace_url.trim().is_empty() {
223            return Err(FaucetError::Config(
224                "databricks: `workspace_url` must not be empty".into(),
225            ));
226        }
227        if self.warehouse_id.trim().is_empty() {
228            return Err(FaucetError::Config(
229                "databricks: `warehouse_id` must not be empty".into(),
230            ));
231        }
232        if self.sql.trim().is_empty() {
233            return Err(FaucetError::Config(
234                "databricks: `sql` must not be empty".into(),
235            ));
236        }
237        // Databricks accepts wait_timeout of 0 (async) or 5..=50 seconds.
238        if self.wait_timeout_secs != 0 && !(5..=50).contains(&self.wait_timeout_secs) {
239            return Err(FaucetError::Config(format!(
240                "databricks: `wait_timeout_secs` must be 0 or between 5 and 50 (got {})",
241                self.wait_timeout_secs
242            )));
243        }
244        faucet_core::validate_batch_size(self.batch_size)?;
245        if let DatabricksReplication::Incremental {
246            column,
247            initial_value,
248        } = &self.replication
249        {
250            if column.trim().is_empty() {
251                return Err(FaucetError::Config(
252                    "databricks: incremental replication `column` must not be empty".into(),
253                ));
254            }
255            if initial_value.is_null() {
256                return Err(FaucetError::Config(
257                    "databricks: incremental replication requires a non-null `initial_value`"
258                        .into(),
259                ));
260            }
261            if faucet_core::replication::cursor::strict_comparison(&self.sql, BOOKMARK_TOKEN) {
262                return Err(FaucetError::Config(format!(
263                    "databricks: compare `{BOOKMARK_TOKEN}` with `>=`, not `>`: the source \
264                     re-reads the bookmark value and drops the rows it already emitted, and a \
265                     strict comparison skips rows that share it"
266                )));
267            }
268        }
269        if self.arrow_native {
270            if !cfg!(feature = "arrow") {
271                return Err(FaucetError::Config(
272                    "databricks: `arrow_native` requires the crate-local `arrow` feature to be \
273                     enabled"
274                        .into(),
275                ));
276            }
277            if !matches!(self.replication, DatabricksReplication::Full) {
278                return Err(FaucetError::Config(
279                    "databricks: `arrow_native` currently supports only `replication: full` — the \
280                     columnar path does not track the cursor bookmark"
281                        .into(),
282                ));
283            }
284        }
285        Ok(())
286    }
287}
288
289/// The token SQL uses to read the incremental bookmark
290/// (`WHERE updated_at >= ${bookmark}`).
291pub const BOOKMARK_TOKEN: &str = "${bookmark}";
292
293#[cfg(test)]
294mod tests {
295    use super::*;
296    use serde_json::json;
297
298    fn base() -> DatabricksSourceConfig {
299        DatabricksSourceConfig {
300            workspace_url: "https://x.cloud.databricks.com".into(),
301            warehouse_id: "wh1".into(),
302            sql: "SELECT 1".into(),
303            auth: AuthSpec::Inline(DatabricksAuth::Pat { token: "t".into() }),
304            catalog: None,
305            schema: None,
306            parameters: Vec::new(),
307            wait_timeout_secs: default_wait_timeout(),
308            poll_interval_secs: default_poll_interval(),
309            statement_timeout_secs: default_statement_timeout(),
310            batch_size: DEFAULT_BATCH_SIZE,
311            arrow_native: false,
312            result_disposition: ResultDisposition::default(),
313            replication: DatabricksReplication::Full,
314            state_key: None,
315        }
316    }
317
318    #[test]
319    fn valid_config_passes() {
320        base().validate().unwrap();
321    }
322
323    #[cfg(feature = "arrow")]
324    #[test]
325    fn arrow_native_full_passes_but_incremental_rejected() {
326        let mut c = base();
327        c.arrow_native = true;
328        c.validate().unwrap();
329        c.replication = DatabricksReplication::Incremental {
330            column: "ts".into(),
331            initial_value: json!("2026-01-01"),
332        };
333        let err = c.validate().unwrap_err();
334        assert!(err.to_string().contains("arrow_native"));
335    }
336
337    #[cfg(not(feature = "arrow"))]
338    #[test]
339    fn arrow_native_requires_feature() {
340        let mut c = base();
341        c.arrow_native = true;
342        let err = c.validate().unwrap_err();
343        assert!(err.to_string().contains("arrow"));
344    }
345
346    #[test]
347    fn auth_is_bearer_for_both_variants() {
348        assert_eq!(
349            DatabricksAuth::Pat {
350                token: "abc".into()
351            }
352            .authorization_value(),
353            "Bearer abc"
354        );
355        assert_eq!(
356            DatabricksAuth::Token {
357                token: "xyz".into()
358            }
359            .authorization_value(),
360            "Bearer xyz"
361        );
362        let dbg = format!(
363            "{:?}",
364            DatabricksAuth::Token {
365                token: "xyz".into()
366            }
367        );
368        assert!(!dbg.contains("xyz"));
369        let r = shared_auth(&AuthSpec::Reference(faucet_core::AuthReference {
370            name: "p".into(),
371        }));
372        assert!(matches!(r, AuthSpec::Reference(_)));
373    }
374
375    #[test]
376    fn rejects_empty_required_fields() {
377        let mut c = base();
378        c.workspace_url = "  ".into();
379        assert!(c.validate().is_err());
380        let mut c = base();
381        c.warehouse_id = "".into();
382        assert!(c.validate().is_err());
383        let mut c = base();
384        c.sql = "".into();
385        assert!(c.validate().is_err());
386    }
387
388    #[test]
389    fn rejects_bad_wait_timeout() {
390        let mut c = base();
391        c.wait_timeout_secs = 3; // not 0 and not in 5..=50
392        assert!(c.validate().is_err());
393        c.wait_timeout_secs = 51;
394        assert!(c.validate().is_err());
395        c.wait_timeout_secs = 0; // allowed (fully async)
396        assert!(c.validate().is_ok());
397    }
398
399    #[test]
400    fn rejects_bad_incremental_settings() {
401        for (sql, column, initial, needle) in [
402            ("SELECT * FROM t", " ", json!(0), "column"),
403            ("SELECT * FROM t", "ts", json!(null), "initial_value"),
404            (
405                "SELECT * FROM t WHERE ts > ${bookmark}",
406                "ts",
407                json!(0),
408                ">=",
409            ),
410        ] {
411            let mut c = base();
412            c.sql = sql.into();
413            c.replication = DatabricksReplication::Incremental {
414                column: column.into(),
415                initial_value: initial,
416            };
417            let err = c.validate().unwrap_err().to_string();
418            assert!(err.contains(needle), "{err}");
419        }
420    }
421
422    #[test]
423    fn rejects_oversized_batch() {
424        let mut c = base();
425        c.batch_size = faucet_core::MAX_BATCH_SIZE + 1;
426        assert!(c.validate().is_err());
427    }
428
429    #[test]
430    fn deserializes_full_shape() {
431        let v = json!({
432            "workspace_url": "https://x.cloud.databricks.com",
433            "warehouse_id": "wh1",
434            "sql": "SELECT * FROM t WHERE id > :min",
435            "auth": { "type": "pat", "config": { "token": "tok" } },
436            "catalog": "main",
437            "schema": "sales",
438            "parameters": [{ "name": "min", "value": 10, "type": "INT" }],
439            "wait_timeout_secs": 30,
440            "batch_size": 500
441        });
442        let c: DatabricksSourceConfig = serde_json::from_value(v).unwrap();
443        assert_eq!(c.warehouse_id, "wh1");
444        assert_eq!(c.catalog.as_deref(), Some("main"));
445        assert_eq!(c.parameters.len(), 1);
446        assert_eq!(c.parameters[0].name, "min");
447        assert_eq!(c.parameters[0].param_type.as_deref(), Some("INT"));
448        c.validate().unwrap();
449    }
450}
faucet_source_databricks/config.rs

faucet_source_databricks/
config.rs
