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    /// Replication mode. Defaults to [`DatabricksReplication::Full`].
129    #[serde(default)]
130    pub replication: DatabricksReplication,
131    /// Explicit state-store key for the incremental bookmark. When unset, a key
132    /// is derived from the workspace/warehouse and a query fingerprint.
133    #[serde(default, skip_serializing_if = "Option::is_none")]
134    pub state_key: Option<String>,
135}
136
137impl DatabricksSourceConfig {
138    /// Validate the config; returns a human-readable error if invalid.
139    pub fn validate(&self) -> Result<(), FaucetError> {
140        if self.workspace_url.trim().is_empty() {
141            return Err(FaucetError::Config(
142                "databricks: `workspace_url` must not be empty".into(),
143            ));
144        }
145        if self.warehouse_id.trim().is_empty() {
146            return Err(FaucetError::Config(
147                "databricks: `warehouse_id` must not be empty".into(),
148            ));
149        }
150        if self.sql.trim().is_empty() {
151            return Err(FaucetError::Config(
152                "databricks: `sql` must not be empty".into(),
153            ));
154        }
155        // Databricks accepts wait_timeout of 0 (async) or 5..=50 seconds.
156        if self.wait_timeout_secs != 0 && !(5..=50).contains(&self.wait_timeout_secs) {
157            return Err(FaucetError::Config(format!(
158                "databricks: `wait_timeout_secs` must be 0 or between 5 and 50 (got {})",
159                self.wait_timeout_secs
160            )));
161        }
162        faucet_core::validate_batch_size(self.batch_size)?;
163        Ok(())
164    }
165}
166
167#[cfg(test)]
168mod tests {
169    use super::*;
170    use serde_json::json;
171
172    fn base() -> DatabricksSourceConfig {
173        DatabricksSourceConfig {
174            workspace_url: "https://x.cloud.databricks.com".into(),
175            warehouse_id: "wh1".into(),
176            sql: "SELECT 1".into(),
177            auth: AuthSpec::Inline(DatabricksAuth::Pat { token: "t".into() }),
178            catalog: None,
179            schema: None,
180            parameters: Vec::new(),
181            wait_timeout_secs: default_wait_timeout(),
182            poll_interval_secs: default_poll_interval(),
183            batch_size: DEFAULT_BATCH_SIZE,
184            replication: DatabricksReplication::Full,
185            state_key: None,
186        }
187    }
188
189    #[test]
190    fn valid_config_passes() {
191        base().validate().unwrap();
192    }
193
194    #[test]
195    fn auth_is_bearer_for_both_variants() {
196        assert_eq!(
197            DatabricksAuth::Pat {
198                token: "abc".into()
199            }
200            .authorization_value(),
201            "Bearer abc"
202        );
203        assert_eq!(
204            DatabricksAuth::Token {
205                token: "xyz".into()
206            }
207            .authorization_value(),
208            "Bearer xyz"
209        );
210    }
211
212    #[test]
213    fn rejects_empty_required_fields() {
214        let mut c = base();
215        c.workspace_url = "  ".into();
216        assert!(c.validate().is_err());
217        let mut c = base();
218        c.warehouse_id = "".into();
219        assert!(c.validate().is_err());
220        let mut c = base();
221        c.sql = "".into();
222        assert!(c.validate().is_err());
223    }
224
225    #[test]
226    fn rejects_bad_wait_timeout() {
227        let mut c = base();
228        c.wait_timeout_secs = 3; // not 0 and not in 5..=50
229        assert!(c.validate().is_err());
230        c.wait_timeout_secs = 51;
231        assert!(c.validate().is_err());
232        c.wait_timeout_secs = 0; // allowed (fully async)
233        assert!(c.validate().is_ok());
234    }
235
236    #[test]
237    fn rejects_oversized_batch() {
238        let mut c = base();
239        c.batch_size = faucet_core::MAX_BATCH_SIZE + 1;
240        assert!(c.validate().is_err());
241    }
242
243    #[test]
244    fn deserializes_full_shape() {
245        let v = json!({
246            "workspace_url": "https://x.cloud.databricks.com",
247            "warehouse_id": "wh1",
248            "sql": "SELECT * FROM t WHERE id > :min",
249            "auth": { "type": "pat", "config": { "token": "tok" } },
250            "catalog": "main",
251            "schema": "sales",
252            "parameters": [{ "name": "min", "value": 10, "type": "INT" }],
253            "wait_timeout_secs": 30,
254            "batch_size": 500
255        });
256        let c: DatabricksSourceConfig = serde_json::from_value(v).unwrap();
257        assert_eq!(c.warehouse_id, "wh1");
258        assert_eq!(c.catalog.as_deref(), Some("main"));
259        assert_eq!(c.parameters.len(), 1);
260        assert_eq!(c.parameters[0].name, "min");
261        assert_eq!(c.parameters[0].param_type.as_deref(), Some("INT"));
262        c.validate().unwrap();
263    }
264}
faucet_source_databricks/config.rs

faucet_source_databricks/
config.rs
