faucet_sink_databricks/
config.rs

1//! Databricks SQL warehouse sink configuration.
2
3use faucet_core::staging::{StagingCleanup, StagingCredentials};
4use faucet_core::{AuthSpec, DEFAULT_BATCH_SIZE, FaucetError, WriteMode, WriteSpec};
5use schemars::JsonSchema;
6use serde::{Deserialize, Serialize};
7
8pub use faucet_common_databricks::DatabricksAuth;
9
10/// How a page is loaded into the table.
11#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize, JsonSchema)]
12#[serde(rename_all = "snake_case")]
13pub enum DatabricksLoadMethod {
14    /// Stage the page when `staging` is configured and the page is at least
15    /// `copy_threshold_bytes`; otherwise a multi-row `INSERT` (the default).
16    #[default]
17    Auto,
18    /// Always a multi-row `INSERT … SELECT … FROM VALUES` through the
19    /// Statement Execution API. No staging location needed.
20    Insert,
21    /// Always stage the page as Parquet, then `COPY INTO` (append) or read it
22    /// with `read_files` (upsert / delete / exactly-once). Requires `staging`.
23    CopyInto,
24}
25
26/// Where staged Parquet files are written before the warehouse loads them.
27#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, JsonSchema)]
28#[serde(deny_unknown_fields)]
29pub struct DatabricksStagingConfig {
30    /// Staging root. Either a Unity Catalog volume path
31    /// (`/Volumes/<catalog>/<schema>/<volume>[/prefix]`, uploaded through the
32    /// Files API with the sink's own credentials) or a cloud location the
33    /// warehouse can read through an external location: `s3://bucket/prefix`,
34    /// `gs://bucket/prefix`, or
35    /// `abfss://container@account.dfs.core.windows.net/prefix`. Cloud locations
36    /// need the crate's `staging` feature and ambient cloud credentials for the
37    /// upload.
38    #[schemars(extend("x-faucet-resource" = "url"))]
39    pub location: String,
40    /// When staged files are deleted after a load. Default `always`.
41    #[serde(default)]
42    pub cleanup: StagingCleanup,
43    /// Extra `COPY_OPTIONS` entries appended verbatim to the append path's
44    /// `COPY INTO`, e.g. `'mergeSchema' = 'false'`. The operator owns their
45    /// correctness.
46    #[serde(default, skip_serializing_if = "Option::is_none")]
47    pub copy_options: Option<String>,
48    /// How faucet authenticates to a cloud location when it uploads the
49    /// files (default: the store's default chain, i.e. the identity of the
50    /// host faucet runs on). Must stay unset for a volume, which uploads with
51    /// the sink's own credentials.
52    #[serde(default, skip_serializing_if = "StagingCredentials::is_default")]
53    pub credentials: StagingCredentials,
54}
55
56fn default_true() -> bool {
57    true
58}
59fn default_batch_size() -> usize {
60    DEFAULT_BATCH_SIZE
61}
62fn default_copy_threshold() -> usize {
63    1024 * 1024
64}
65fn default_max_statement_bytes() -> usize {
66    8 * 1024 * 1024
67}
68fn default_wait_timeout() -> u64 {
69    50
70}
71fn default_poll_interval_ms() -> u64 {
72    1000
73}
74fn default_statement_timeout() -> u64 {
75    3600
76}
77fn default_max_retries() -> u32 {
78    5
79}
80fn default_retry_backoff_ms() -> u64 {
81    1000
82}
83
84/// Hard ceiling on one Statement Execution API request's statement text.
85pub const STATEMENT_TEXT_LIMIT: usize = 16 * 1024 * 1024;
86
87/// Configuration for the Databricks SQL warehouse sink.
88///
89/// Writes into a Delta table through a SQL warehouse's
90/// [Statement Execution API](https://docs.databricks.com/api/workspace/statementexecution).
91/// `write_mode` / `key` / `delete_marker` are flattened in from
92/// [`WriteSpec`](faucet_core::WriteSpec).
93#[derive(Clone, Serialize, Deserialize, JsonSchema)]
94pub struct DatabricksSinkConfig {
95    /// Workspace base URL, e.g. `https://dbc-abc123.cloud.databricks.com`.
96    #[schemars(extend("x-faucet-resource" = "url"))]
97    pub workspace_url: String,
98    /// Target SQL warehouse id.
99    pub warehouse_id: String,
100    /// Authentication (PAT / OAuth bearer), inline or via a shared `auth: { ref }`
101    /// (use an `oauth2` client-credentials provider for an M2M service principal).
102    pub auth: AuthSpec<DatabricksAuth>,
103    /// Unity Catalog catalog. When unset, the warehouse's default catalog.
104    #[serde(default, skip_serializing_if = "Option::is_none")]
105    pub catalog: Option<String>,
106    /// Target schema.
107    pub schema: String,
108    /// Target table.
109    pub table: String,
110    /// Create the table from the first page's inferred schema when it does not
111    /// exist. Default `true`.
112    #[serde(default = "default_true")]
113    pub create_table: bool,
114    /// Load path. Default `auto`.
115    #[serde(default)]
116    pub load_method: DatabricksLoadMethod,
117    /// Staging location for `copy_into` / `auto`.
118    #[serde(default, skip_serializing_if = "Option::is_none")]
119    pub staging: Option<DatabricksStagingConfig>,
120    /// `auto` stages a page whose estimated size is at least this many bytes.
121    /// Default 1 MiB.
122    #[serde(default = "default_copy_threshold")]
123    pub copy_threshold_bytes: usize,
124    /// Rows per `INSERT` / `MERGE` statement on the insert path (`0` = no row
125    /// limit — statements are then bounded only by `max_statement_bytes`).
126    #[serde(default = "default_batch_size")]
127    pub batch_size: usize,
128    /// Upper bound on one statement's SQL text on the insert path; a page is
129    /// split into several statements above it. Default 8 MiB (the API accepts
130    /// up to 16 MiB).
131    #[serde(default = "default_max_statement_bytes")]
132    pub max_statement_bytes: usize,
133    /// Server-side wait before a statement goes async (`0` or `5`–`50` s).
134    #[serde(default = "default_wait_timeout")]
135    pub wait_timeout_secs: u64,
136    /// Poll cadence while a statement is queued or running (a serverless
137    /// warehouse may take a minute to start). Default 1000 ms.
138    #[serde(default = "default_poll_interval_ms")]
139    pub poll_interval_ms: u64,
140    /// Client deadline for one statement; the statement is cancelled when it
141    /// is exceeded. `0` waits indefinitely. Default 3600 s.
142    #[serde(default = "default_statement_timeout")]
143    pub statement_timeout_secs: u64,
144    /// Retries on `429` / `503` and on Delta concurrent-write conflicts of
145    /// idempotent statements. Default 5.
146    #[serde(default = "default_max_retries")]
147    pub max_retries: u32,
148    /// Base of the exponential retry backoff. Default 1000 ms.
149    #[serde(default = "default_retry_backoff_ms")]
150    pub retry_backoff_ms: u64,
151    /// `write_mode` (`append` / `upsert` / `delete` / `overwrite`), `key`,
152    /// `delete_marker`.
153    #[serde(flatten)]
154    pub write: WriteSpec,
155}
156
157impl std::fmt::Debug for DatabricksSinkConfig {
158    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
159        f.debug_struct("DatabricksSinkConfig")
160            .field("workspace_url", &self.workspace_url)
161            .field("warehouse_id", &self.warehouse_id)
162            .field("catalog", &self.catalog)
163            .field("schema", &self.schema)
164            .field("table", &self.table)
165            .field("load_method", &self.load_method)
166            .field("staging", &self.staging)
167            .field("write", &self.write)
168            .finish_non_exhaustive()
169    }
170}
171
172impl DatabricksSinkConfig {
173    /// What a failed write leaves behind: a page may span several statements,
174    /// and the statements before the failing one have committed.
175    pub fn batch_atomicity(&self) -> faucet_core::BatchAtomicity {
176        faucet_core::BatchAtomicity::BestEffort
177    }
178
179    /// A config with every optional field at its default.
180    pub fn new(
181        workspace_url: impl Into<String>,
182        warehouse_id: impl Into<String>,
183        schema: impl Into<String>,
184        table: impl Into<String>,
185        auth: DatabricksAuth,
186    ) -> Self {
187        Self {
188            workspace_url: workspace_url.into(),
189            warehouse_id: warehouse_id.into(),
190            auth: AuthSpec::Inline(auth),
191            catalog: None,
192            schema: schema.into(),
193            table: table.into(),
194            create_table: true,
195            load_method: DatabricksLoadMethod::Auto,
196            staging: None,
197            copy_threshold_bytes: default_copy_threshold(),
198            batch_size: default_batch_size(),
199            max_statement_bytes: default_max_statement_bytes(),
200            wait_timeout_secs: default_wait_timeout(),
201            poll_interval_ms: default_poll_interval_ms(),
202            statement_timeout_secs: default_statement_timeout(),
203            max_retries: default_max_retries(),
204            retry_backoff_ms: default_retry_backoff_ms(),
205            write: WriteSpec::default(),
206        }
207    }
208
209    /// Validate the config at load time.
210    pub fn validate(&self) -> Result<(), FaucetError> {
211        for (name, v) in [
212            ("workspace_url", &self.workspace_url),
213            ("warehouse_id", &self.warehouse_id),
214            ("schema", &self.schema),
215            ("table", &self.table),
216        ] {
217            if v.trim().is_empty() {
218                return Err(FaucetError::Config(format!(
219                    "databricks sink: `{name}` must not be empty"
220                )));
221            }
222        }
223        if self.catalog.as_deref().is_some_and(|c| c.trim().is_empty()) {
224            return Err(FaucetError::Config(
225                "databricks sink: `catalog` must not be empty when set".into(),
226            ));
227        }
228        if self.wait_timeout_secs != 0 && !(5..=50).contains(&self.wait_timeout_secs) {
229            return Err(FaucetError::Config(format!(
230                "databricks sink: `wait_timeout_secs` must be 0 or between 5 and 50 (got {})",
231                self.wait_timeout_secs
232            )));
233        }
234        faucet_core::validate_batch_size(self.batch_size)?;
235        if self.poll_interval_ms == 0 {
236            return Err(FaucetError::Config(
237                "databricks sink: `poll_interval_ms` must be > 0 (0 would poll the workspace \
238                 in a tight loop)"
239                    .into(),
240            ));
241        }
242        if self.max_statement_bytes == 0 || self.max_statement_bytes > STATEMENT_TEXT_LIMIT {
243            return Err(FaucetError::Config(format!(
244                "databricks sink: `max_statement_bytes` must be between 1 and {STATEMENT_TEXT_LIMIT}"
245            )));
246        }
247        if self.load_method == DatabricksLoadMethod::CopyInto && self.staging.is_none() {
248            return Err(FaucetError::Config(
249                "databricks sink: `load_method: copy_into` requires a `staging` block".into(),
250            ));
251        }
252        if let Some(s) = &self.staging {
253            let loc = crate::stage::StageLocation::parse(&s.location)?;
254            if let Some(scheme) = loc.scheme() {
255                s.credentials.check_scheme(scheme)?;
256            } else if !s.credentials.is_default() {
257                return Err(FaucetError::Config(
258                    "databricks sink: a volume `staging.location` uploads with the sink's own \
259                     credentials; remove `staging.credentials`"
260                        .into(),
261                ));
262            }
263        }
264        self.write.validate()?;
265        if self.write.write_mode == WriteMode::Overwrite && self.table.ends_with(OVW_SUFFIX) {
266            return Err(FaucetError::Config(format!(
267                "databricks sink: table names ending in `{OVW_SUFFIX}` are reserved for overwrite staging"
268            )));
269        }
270        Ok(())
271    }
272}
273
274/// Overwrite staging-table suffix.
275pub(crate) const OVW_SUFFIX: &str = faucet_core::idempotency::OVERWRITE_STAGING_SUFFIX;
276
277#[cfg(test)]
278mod tests {
279    use super::*;
280    use serde_json::json;
281
282    fn base() -> DatabricksSinkConfig {
283        DatabricksSinkConfig::new(
284            "https://x.cloud.databricks.com",
285            "wh",
286            "sales",
287            "orders",
288            DatabricksAuth::Pat { token: "t".into() },
289        )
290    }
291
292    #[test]
293    fn defaults_validate() {
294        assert_eq!(
295            base().batch_atomicity(),
296            faucet_core::BatchAtomicity::BestEffort
297        );
298        let c = base();
299        c.validate().unwrap();
300        assert!(c.create_table);
301        assert_eq!(c.load_method, DatabricksLoadMethod::Auto);
302        assert_eq!(c.batch_size, DEFAULT_BATCH_SIZE);
303        let d = format!("{c:?}");
304        assert!(d.contains("orders") && !d.contains("\"t\""));
305    }
306
307    #[test]
308    fn deserializes_full_shape_with_flattened_write_spec() {
309        let c: DatabricksSinkConfig = serde_json::from_value(json!({
310            "workspace_url": "https://x",
311            "warehouse_id": "wh",
312            "auth": {"type": "pat", "config": {"token": "t"}},
313            "catalog": "main",
314            "schema": "s",
315            "table": "t",
316            "load_method": "copy_into",
317            "staging": {"location": "/Volumes/main/s/stage/faucet", "cleanup": "on_success"},
318            "write_mode": "upsert",
319            "key": ["id"]
320        }))
321        .unwrap();
322        c.validate().unwrap();
323        assert_eq!(c.write.write_mode, WriteMode::Upsert);
324        assert_eq!(c.staging.unwrap().cleanup, StagingCleanup::OnSuccess);
325    }
326
327    #[test]
328    fn rejects_invalid_configs() {
329        type Mutate = Box<dyn Fn(&mut DatabricksSinkConfig)>;
330        let cases: Vec<Mutate> = vec![
331            Box::new(|c| c.workspace_url = " ".into()),
332            Box::new(|c| c.warehouse_id = "".into()),
333            Box::new(|c| c.schema = "".into()),
334            Box::new(|c| c.table = "".into()),
335            Box::new(|c| c.catalog = Some(" ".into())),
336            Box::new(|c| c.wait_timeout_secs = 3),
337            Box::new(|c| c.batch_size = faucet_core::MAX_BATCH_SIZE + 1),
338            Box::new(|c| c.max_statement_bytes = 0),
339            Box::new(|c| c.max_statement_bytes = STATEMENT_TEXT_LIMIT + 1),
340            Box::new(|c| c.load_method = DatabricksLoadMethod::CopyInto),
341            Box::new(|c| {
342                c.staging = Some(DatabricksStagingConfig {
343                    location: "ftp://nope".into(),
344                    cleanup: StagingCleanup::Always,
345                    copy_options: None,
346                    credentials: Default::default(),
347                })
348            }),
349            Box::new(|c| c.write.write_mode = WriteMode::Upsert),
350            Box::new(|c| {
351                c.staging = Some(DatabricksStagingConfig {
352                    location: "/Volumes/a/b/c".into(),
353                    cleanup: StagingCleanup::Always,
354                    copy_options: None,
355                    credentials: StagingCredentials::AccountKey {
356                        account_key: "k".into(),
357                    },
358                })
359            }),
360            Box::new(|c| {
361                c.staging = Some(DatabricksStagingConfig {
362                    location: "gs://b/p".into(),
363                    cleanup: StagingCleanup::Always,
364                    copy_options: None,
365                    credentials: StagingCredentials::AccountKey {
366                        account_key: "k".into(),
367                    },
368                })
369            }),
370            Box::new(|c| {
371                c.write.write_mode = WriteMode::Overwrite;
372                c.table = "t__faucet_ovw".into();
373            }),
374        ];
375        for (i, mutate) in cases.iter().enumerate() {
376            let mut c = base();
377            mutate(&mut c);
378            assert!(c.validate().is_err(), "case {i} should be rejected");
379        }
380        let mut ok = base();
381        ok.wait_timeout_secs = 0;
382        ok.validate().unwrap();
383    }
384}
faucet_sink_databricks/config.rs

faucet_sink_databricks/
config.rs
