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    #[schemars(extend("x-faucet-sql" = true))]
159    pub sql: String,
160    /// Authentication (PAT / OAuth bearer), inline or via a shared `auth: { ref }`.
161    pub auth: AuthSpec<DatabricksAuth>,
162    /// Default Unity Catalog catalog for the statement.
163    #[serde(default, skip_serializing_if = "Option::is_none")]
164    pub catalog: Option<String>,
165    /// Default schema for the statement.
166    #[serde(default, skip_serializing_if = "Option::is_none")]
167    pub schema: Option<String>,
168    /// Named SQL parameters (`:name` markers).
169    #[serde(default)]
170    pub parameters: Vec<DatabricksParam>,
171    /// Server-side wait before the statement goes async (seconds). Databricks
172    /// accepts `0` (fully async) or `5`–`50`. Defaults to `50`.
173    #[serde(default = "default_wait_timeout")]
174    pub wait_timeout_secs: u64,
175    /// Client poll cadence while the statement is `PENDING`/`RUNNING` (seconds).
176    #[serde(default = "default_poll_interval")]
177    pub poll_interval_secs: u64,
178    /// Client deadline for the statement to finish, in seconds; on expiry it is
179    /// cancelled and the read fails. `0` waits indefinitely. Defaults to 3600.
180    #[serde(default = "default_statement_timeout")]
181    pub statement_timeout_secs: u64,
182    /// Page size — rows accumulated before a `StreamPage` is emitted.
183    /// Defaults to [`DEFAULT_BATCH_SIZE`](faucet_core::DEFAULT_BATCH_SIZE).
184    ///
185    /// `batch_size = 0` is the "no batching" sentinel: result chunks are fully
186    /// drained and the entire result set is emitted in a single page. Useful
187    /// for small lookup tables or for sinks (e.g. SQL `COPY`, BigQuery load
188    /// jobs) that prefer one large request to many small ones.
189    #[serde(default = "default_batch_size")]
190    pub batch_size: usize,
191    /// Fetch results as Apache Arrow instead of JSON. When `true`, the
192    /// statement is submitted with `EXTERNAL_LINKS` disposition +
193    /// `ARROW_STREAM` format; each chunk's presigned link is fetched and
194    /// decoded as an Arrow IPC stream. This enables the **columnar** fast path
195    /// ([`Source::stream_batches`](faucet_core::Source::stream_batches)) so a
196    /// `databricks → parquet`/`delta` chain skips the per-cell JSON decode.
197    ///
198    /// Requires the crate-local `arrow` feature, and (in this release) only
199    /// [`DatabricksReplication::Full`] — the columnar path does not run the
200    /// per-row cursor bookmark. Defaults to `false` (the
201    /// `INLINE` + `JSON_ARRAY` row path, unchanged). RFC 0002 / #375.
202    #[serde(default)]
203    pub arrow_native: bool,
204    /// How the row path fetches results: `external_links` (default) downloads
205    /// each chunk from a presigned URL, so results of any size stream;
206    /// `inline` keeps them in the API response, which Databricks caps at
207    /// 25 MiB. Use `inline` only when the client cannot reach the workspace's
208    /// cloud storage.
209    #[serde(default)]
210    pub result_disposition: ResultDisposition,
211    /// Replication mode. Defaults to [`DatabricksReplication::Full`].
212    #[serde(default)]
213    pub replication: DatabricksReplication,
214    /// Explicit state-store key for the incremental bookmark. When unset, a key
215    /// is derived from the workspace/warehouse and a query fingerprint.
216    #[serde(default, skip_serializing_if = "Option::is_none")]
217    pub state_key: Option<String>,
218}
219
220impl DatabricksSourceConfig {
221    /// Validate the config; returns a human-readable error if invalid.
222    pub fn validate(&self) -> Result<(), FaucetError> {
223        if self.workspace_url.trim().is_empty() {
224            return Err(FaucetError::Config(
225                "databricks: `workspace_url` must not be empty".into(),
226            ));
227        }
228        if self.warehouse_id.trim().is_empty() {
229            return Err(FaucetError::Config(
230                "databricks: `warehouse_id` must not be empty".into(),
231            ));
232        }
233        if self.sql.trim().is_empty() {
234            return Err(FaucetError::Config(
235                "databricks: `sql` must not be empty".into(),
236            ));
237        }
238        // Databricks accepts wait_timeout of 0 (async) or 5..=50 seconds.
239        if self.wait_timeout_secs != 0 && !(5..=50).contains(&self.wait_timeout_secs) {
240            return Err(FaucetError::Config(format!(
241                "databricks: `wait_timeout_secs` must be 0 or between 5 and 50 (got {})",
242                self.wait_timeout_secs
243            )));
244        }
245        faucet_core::validate_batch_size(self.batch_size)?;
246        if let DatabricksReplication::Incremental {
247            column,
248            initial_value,
249        } = &self.replication
250        {
251            if column.trim().is_empty() {
252                return Err(FaucetError::Config(
253                    "databricks: incremental replication `column` must not be empty".into(),
254                ));
255            }
256            if initial_value.is_null() {
257                return Err(FaucetError::Config(
258                    "databricks: incremental replication requires a non-null `initial_value`"
259                        .into(),
260                ));
261            }
262            if faucet_core::replication::cursor::strict_comparison(&self.sql, BOOKMARK_TOKEN) {
263                return Err(FaucetError::Config(format!(
264                    "databricks: compare `{BOOKMARK_TOKEN}` with `>=`, not `>`: the source \
265                     re-reads the bookmark value and drops the rows it already emitted, and a \
266                     strict comparison skips rows that share it"
267                )));
268            }
269        }
270        if self.arrow_native {
271            if !cfg!(feature = "arrow") {
272                return Err(FaucetError::Config(
273                    "databricks: `arrow_native` requires the crate-local `arrow` feature to be \
274                     enabled"
275                        .into(),
276                ));
277            }
278            if !matches!(self.replication, DatabricksReplication::Full) {
279                return Err(FaucetError::Config(
280                    "databricks: `arrow_native` currently supports only `replication: full` — the \
281                     columnar path does not track the cursor bookmark"
282                        .into(),
283                ));
284            }
285        }
286        Ok(())
287    }
288}
289
290/// The token SQL uses to read the incremental bookmark
291/// (`WHERE updated_at >= ${bookmark}`).
292pub const BOOKMARK_TOKEN: &str = "${bookmark}";
293
294#[cfg(test)]
295mod tests {
296    use super::*;
297    use serde_json::json;
298
299    fn base() -> DatabricksSourceConfig {
300        DatabricksSourceConfig {
301            workspace_url: "https://x.cloud.databricks.com".into(),
302            warehouse_id: "wh1".into(),
303            sql: "SELECT 1".into(),
304            auth: AuthSpec::Inline(DatabricksAuth::Pat { token: "t".into() }),
305            catalog: None,
306            schema: None,
307            parameters: Vec::new(),
308            wait_timeout_secs: default_wait_timeout(),
309            poll_interval_secs: default_poll_interval(),
310            statement_timeout_secs: default_statement_timeout(),
311            batch_size: DEFAULT_BATCH_SIZE,
312            arrow_native: false,
313            result_disposition: ResultDisposition::default(),
314            replication: DatabricksReplication::Full,
315            state_key: None,
316        }
317    }
318
319    #[test]
320    fn valid_config_passes() {
321        base().validate().unwrap();
322    }
323
324    #[cfg(feature = "arrow")]
325    #[test]
326    fn arrow_native_full_passes_but_incremental_rejected() {
327        let mut c = base();
328        c.arrow_native = true;
329        c.validate().unwrap();
330        c.replication = DatabricksReplication::Incremental {
331            column: "ts".into(),
332            initial_value: json!("2026-01-01"),
333        };
334        let err = c.validate().unwrap_err();
335        assert!(err.to_string().contains("arrow_native"));
336    }
337
338    #[cfg(not(feature = "arrow"))]
339    #[test]
340    fn arrow_native_requires_feature() {
341        let mut c = base();
342        c.arrow_native = true;
343        let err = c.validate().unwrap_err();
344        assert!(err.to_string().contains("arrow"));
345    }
346
347    #[test]
348    fn auth_is_bearer_for_both_variants() {
349        assert_eq!(
350            DatabricksAuth::Pat {
351                token: "abc".into()
352            }
353            .authorization_value(),
354            "Bearer abc"
355        );
356        assert_eq!(
357            DatabricksAuth::Token {
358                token: "xyz".into()
359            }
360            .authorization_value(),
361            "Bearer xyz"
362        );
363        let dbg = format!(
364            "{:?}",
365            DatabricksAuth::Token {
366                token: "xyz".into()
367            }
368        );
369        assert!(!dbg.contains("xyz"));
370        let r = shared_auth(&AuthSpec::Reference(faucet_core::AuthReference {
371            name: "p".into(),
372        }));
373        assert!(matches!(r, AuthSpec::Reference(_)));
374    }
375
376    #[test]
377    fn rejects_empty_required_fields() {
378        let mut c = base();
379        c.workspace_url = "  ".into();
380        assert!(c.validate().is_err());
381        let mut c = base();
382        c.warehouse_id = "".into();
383        assert!(c.validate().is_err());
384        let mut c = base();
385        c.sql = "".into();
386        assert!(c.validate().is_err());
387    }
388
389    #[test]
390    fn rejects_bad_wait_timeout() {
391        let mut c = base();
392        c.wait_timeout_secs = 3; // not 0 and not in 5..=50
393        assert!(c.validate().is_err());
394        c.wait_timeout_secs = 51;
395        assert!(c.validate().is_err());
396        c.wait_timeout_secs = 0; // allowed (fully async)
397        assert!(c.validate().is_ok());
398    }
399
400    #[test]
401    fn rejects_bad_incremental_settings() {
402        for (sql, column, initial, needle) in [
403            ("SELECT * FROM t", " ", json!(0), "column"),
404            ("SELECT * FROM t", "ts", json!(null), "initial_value"),
405            (
406                "SELECT * FROM t WHERE ts > ${bookmark}",
407                "ts",
408                json!(0),
409                ">=",
410            ),
411        ] {
412            let mut c = base();
413            c.sql = sql.into();
414            c.replication = DatabricksReplication::Incremental {
415                column: column.into(),
416                initial_value: initial,
417            };
418            let err = c.validate().unwrap_err().to_string();
419            assert!(err.contains(needle), "{err}");
420        }
421    }
422
423    #[test]
424    fn rejects_oversized_batch() {
425        let mut c = base();
426        c.batch_size = faucet_core::MAX_BATCH_SIZE + 1;
427        assert!(c.validate().is_err());
428    }
429
430    #[test]
431    fn deserializes_full_shape() {
432        let v = json!({
433            "workspace_url": "https://x.cloud.databricks.com",
434            "warehouse_id": "wh1",
435            "sql": "SELECT * FROM t WHERE id > :min",
436            "auth": { "type": "pat", "config": { "token": "tok" } },
437            "catalog": "main",
438            "schema": "sales",
439            "parameters": [{ "name": "min", "value": 10, "type": "INT" }],
440            "wait_timeout_secs": 30,
441            "batch_size": 500
442        });
443        let c: DatabricksSourceConfig = serde_json::from_value(v).unwrap();
444        assert_eq!(c.warehouse_id, "wh1");
445        assert_eq!(c.catalog.as_deref(), Some("main"));
446        assert_eq!(c.parameters.len(), 1);
447        assert_eq!(c.parameters[0].name, "min");
448        assert_eq!(c.parameters[0].param_type.as_deref(), Some("INT"));
449        c.validate().unwrap();
450    }
451}
faucet_source_databricks/config.rs

faucet_source_databricks/
config.rs
