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

faucet_source_databricks/
config.rs
