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#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
53#[serde(tag = "type", content = "config", rename_all = "snake_case")]
54pub enum DatabricksAuth {
55    /// Databricks Personal Access Token.
56    Pat {
57        /// The token string (use `${env:…}` / `${vault:…}` to inject).
58        token: String,
59    },
60    /// A pre-obtained OAuth (M2M) bearer token.
61    Token {
62        /// The bearer token string.
63        token: String,
64    },
65}
66
67impl DatabricksAuth {
68    /// The `Authorization` header value (`Bearer <token>`).
69    pub fn authorization_value(&self) -> String {
70        match self {
71            DatabricksAuth::Pat { token } | DatabricksAuth::Token { token } => {
72                format!("Bearer {token}")
73            }
74        }
75    }
76}
77
78/// A named SQL parameter passed to the statement (`:name` markers in the SQL).
79#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
80pub struct DatabricksParam {
81    /// Parameter name (the `:name` marker in the SQL, without the colon).
82    pub name: String,
83    /// Parameter value. Sent to Databricks as its string form; `null` binds a
84    /// typed NULL.
85    #[serde(default)]
86    pub value: Value,
87    /// Optional Databricks SQL type (e.g. `INT`, `STRING`, `DATE`, `TIMESTAMP`).
88    /// When omitted, Databricks treats the value as `STRING`.
89    #[serde(rename = "type", default, skip_serializing_if = "Option::is_none")]
90    pub param_type: Option<String>,
91}
92
93/// Configuration for the Databricks SQL query source.
94///
95/// Runs `sql` against a Databricks SQL Warehouse via the Statement Execution
96/// REST API and streams the result rows as typed JSON objects.
97#[derive(Clone, Serialize, Deserialize, JsonSchema)]
98pub struct DatabricksSourceConfig {
99    /// Workspace base URL, e.g. `https://dbc-abc123.cloud.databricks.com`.
100    pub workspace_url: String,
101    /// Target SQL Warehouse id.
102    pub warehouse_id: String,
103    /// The SQL statement to run. May contain `:name` parameter markers (see
104    /// [`parameters`](Self::parameters)) and a `${bookmark}` token for
105    /// incremental replication.
106    pub sql: String,
107    /// Authentication (PAT / OAuth bearer), inline or via a shared `auth: { ref }`.
108    pub auth: AuthSpec<DatabricksAuth>,
109    /// Default Unity Catalog catalog for the statement.
110    #[serde(default, skip_serializing_if = "Option::is_none")]
111    pub catalog: Option<String>,
112    /// Default schema for the statement.
113    #[serde(default, skip_serializing_if = "Option::is_none")]
114    pub schema: Option<String>,
115    /// Named SQL parameters (`:name` markers).
116    #[serde(default)]
117    pub parameters: Vec<DatabricksParam>,
118    /// Server-side wait before the statement goes async (seconds). Databricks
119    /// accepts `0` (fully async) or `5`–`50`. Defaults to `50`.
120    #[serde(default = "default_wait_timeout")]
121    pub wait_timeout_secs: u64,
122    /// Client poll cadence while the statement is `PENDING`/`RUNNING` (seconds).
123    #[serde(default = "default_poll_interval")]
124    pub poll_interval_secs: u64,
125    /// Page size — rows accumulated before a `StreamPage` is emitted.
126    #[serde(default = "default_batch_size")]
127    pub batch_size: usize,
128    /// Fetch results as Apache Arrow instead of JSON. When `true`, the
129    /// statement is submitted with `EXTERNAL_LINKS` disposition +
130    /// `ARROW_STREAM` format; each chunk's presigned link is fetched and
131    /// decoded as an Arrow IPC stream. This enables the **columnar** fast path
132    /// ([`Source::stream_batches`](faucet_core::Source::stream_batches)) so a
133    /// `databricks → parquet`/`delta` chain skips the per-cell JSON decode.
134    ///
135    /// Requires the crate-local `arrow` feature, and (in this release) only
136    /// [`DatabricksReplication::Full`] — the columnar path does not run the
137    /// per-row client-side incremental filter. Defaults to `false` (the
138    /// `INLINE` + `JSON_ARRAY` row path, unchanged). RFC 0002 / #375.
139    #[serde(default)]
140    pub arrow_native: bool,
141    /// Replication mode. Defaults to [`DatabricksReplication::Full`].
142    #[serde(default)]
143    pub replication: DatabricksReplication,
144    /// Explicit state-store key for the incremental bookmark. When unset, a key
145    /// is derived from the workspace/warehouse and a query fingerprint.
146    #[serde(default, skip_serializing_if = "Option::is_none")]
147    pub state_key: Option<String>,
148}
149
150impl DatabricksSourceConfig {
151    /// Validate the config; returns a human-readable error if invalid.
152    pub fn validate(&self) -> Result<(), FaucetError> {
153        if self.workspace_url.trim().is_empty() {
154            return Err(FaucetError::Config(
155                "databricks: `workspace_url` must not be empty".into(),
156            ));
157        }
158        if self.warehouse_id.trim().is_empty() {
159            return Err(FaucetError::Config(
160                "databricks: `warehouse_id` must not be empty".into(),
161            ));
162        }
163        if self.sql.trim().is_empty() {
164            return Err(FaucetError::Config(
165                "databricks: `sql` must not be empty".into(),
166            ));
167        }
168        // Databricks accepts wait_timeout of 0 (async) or 5..=50 seconds.
169        if self.wait_timeout_secs != 0 && !(5..=50).contains(&self.wait_timeout_secs) {
170            return Err(FaucetError::Config(format!(
171                "databricks: `wait_timeout_secs` must be 0 or between 5 and 50 (got {})",
172                self.wait_timeout_secs
173            )));
174        }
175        faucet_core::validate_batch_size(self.batch_size)?;
176        if self.arrow_native {
177            if !cfg!(feature = "arrow") {
178                return Err(FaucetError::Config(
179                    "databricks: `arrow_native` requires the crate-local `arrow` feature to be \
180                     enabled"
181                        .into(),
182                ));
183            }
184            if !matches!(self.replication, DatabricksReplication::Full) {
185                return Err(FaucetError::Config(
186                    "databricks: `arrow_native` currently supports only `replication: full` — the \
187                     columnar path does not run the client-side incremental filter"
188                        .into(),
189                ));
190            }
191        }
192        Ok(())
193    }
194}
195
196#[cfg(test)]
197mod tests {
198    use super::*;
199    use serde_json::json;
200
201    fn base() -> DatabricksSourceConfig {
202        DatabricksSourceConfig {
203            workspace_url: "https://x.cloud.databricks.com".into(),
204            warehouse_id: "wh1".into(),
205            sql: "SELECT 1".into(),
206            auth: AuthSpec::Inline(DatabricksAuth::Pat { token: "t".into() }),
207            catalog: None,
208            schema: None,
209            parameters: Vec::new(),
210            wait_timeout_secs: default_wait_timeout(),
211            poll_interval_secs: default_poll_interval(),
212            batch_size: DEFAULT_BATCH_SIZE,
213            arrow_native: false,
214            replication: DatabricksReplication::Full,
215            state_key: None,
216        }
217    }
218
219    #[test]
220    fn valid_config_passes() {
221        base().validate().unwrap();
222    }
223
224    #[cfg(feature = "arrow")]
225    #[test]
226    fn arrow_native_full_passes_but_incremental_rejected() {
227        let mut c = base();
228        c.arrow_native = true;
229        c.validate().unwrap();
230        c.replication = DatabricksReplication::Incremental {
231            column: "ts".into(),
232            initial_value: json!("2026-01-01"),
233        };
234        let err = c.validate().unwrap_err();
235        assert!(err.to_string().contains("arrow_native"));
236    }
237
238    #[cfg(not(feature = "arrow"))]
239    #[test]
240    fn arrow_native_requires_feature() {
241        let mut c = base();
242        c.arrow_native = true;
243        let err = c.validate().unwrap_err();
244        assert!(err.to_string().contains("arrow"));
245    }
246
247    #[test]
248    fn auth_is_bearer_for_both_variants() {
249        assert_eq!(
250            DatabricksAuth::Pat {
251                token: "abc".into()
252            }
253            .authorization_value(),
254            "Bearer abc"
255        );
256        assert_eq!(
257            DatabricksAuth::Token {
258                token: "xyz".into()
259            }
260            .authorization_value(),
261            "Bearer xyz"
262        );
263    }
264
265    #[test]
266    fn rejects_empty_required_fields() {
267        let mut c = base();
268        c.workspace_url = "  ".into();
269        assert!(c.validate().is_err());
270        let mut c = base();
271        c.warehouse_id = "".into();
272        assert!(c.validate().is_err());
273        let mut c = base();
274        c.sql = "".into();
275        assert!(c.validate().is_err());
276    }
277
278    #[test]
279    fn rejects_bad_wait_timeout() {
280        let mut c = base();
281        c.wait_timeout_secs = 3; // not 0 and not in 5..=50
282        assert!(c.validate().is_err());
283        c.wait_timeout_secs = 51;
284        assert!(c.validate().is_err());
285        c.wait_timeout_secs = 0; // allowed (fully async)
286        assert!(c.validate().is_ok());
287    }
288
289    #[test]
290    fn rejects_oversized_batch() {
291        let mut c = base();
292        c.batch_size = faucet_core::MAX_BATCH_SIZE + 1;
293        assert!(c.validate().is_err());
294    }
295
296    #[test]
297    fn deserializes_full_shape() {
298        let v = json!({
299            "workspace_url": "https://x.cloud.databricks.com",
300            "warehouse_id": "wh1",
301            "sql": "SELECT * FROM t WHERE id > :min",
302            "auth": { "type": "pat", "config": { "token": "tok" } },
303            "catalog": "main",
304            "schema": "sales",
305            "parameters": [{ "name": "min", "value": 10, "type": "INT" }],
306            "wait_timeout_secs": 30,
307            "batch_size": 500
308        });
309        let c: DatabricksSourceConfig = serde_json::from_value(v).unwrap();
310        assert_eq!(c.warehouse_id, "wh1");
311        assert_eq!(c.catalog.as_deref(), Some("main"));
312        assert_eq!(c.parameters.len(), 1);
313        assert_eq!(c.parameters[0].name, "min");
314        assert_eq!(c.parameters[0].param_type.as_deref(), Some("INT"));
315        c.validate().unwrap();
316    }
317}
faucet_source_databricks/config.rs

faucet_source_databricks/
config.rs
