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

faucet_source_databricks/
config.rs
