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

faucet_source_databricks/
config.rs
