Skip to main content

faucet_source_spanner/
config.rs

1//! Configuration for the Cloud Spanner query source.
2
3use std::collections::HashMap;
4
5use faucet_common_spanner::SpannerConnection;
6use faucet_core::{DEFAULT_BATCH_SIZE, FaucetError, validate_batch_size};
7use schemars::JsonSchema;
8use serde::{Deserialize, Serialize};
9use serde_json::Value;
10
11fn default_batch_size() -> usize {
12    DEFAULT_BATCH_SIZE
13}
14
15/// How the source replicates rows across runs.
16///
17/// Serializes as `{ type: full }` or
18/// `{ type: incremental, column: "...", initial_value: ... }`.
19#[derive(Clone, Debug, Serialize, Deserialize, JsonSchema, Default, PartialEq)]
20#[serde(tag = "type", rename_all = "snake_case")]
21pub enum SpannerReplication {
22    /// Every run fetches the full result set (default).
23    #[default]
24    Full,
25    /// Only rows whose `column` is strictly greater than the stored bookmark
26    /// (or `initial_value` on the first run) are emitted.
27    ///
28    /// The bookmark is applied two ways: if the query contains the named
29    /// parameter `@bookmark`, it is bound so the server filters (efficient);
30    /// the source *also* filters client-side as a correctness backstop. The
31    /// new maximum of `column` is persisted on the final page.
32    ///
33    /// Spanner does not implicitly coerce parameter types: a string bookmark
34    /// compared against a `TIMESTAMP` column needs an explicit cast in the
35    /// query, e.g. `WHERE updated_at > TIMESTAMP(@bookmark)`.
36    Incremental {
37        /// Column whose value is the replication cursor (e.g. `updated_at`).
38        column: String,
39        /// Lower bound used on the first run, before any bookmark is stored.
40        initial_value: Value,
41    },
42}
43
44/// Primary-key range sharding settings for the Spanner source (clustered
45/// Mode B execution).
46///
47/// The source is split by contiguous ranges of an **INT64-typed** column:
48/// each shard runs ``SELECT * FROM (<query>) AS _faucet_shard WHERE `key` >=
49/// lo AND `key` < hi``. The column must be present in the query's output.
50///
51/// **Nullable keys:** rows whose key is NULL are invisible to the `MIN`/`MAX`
52/// range enumeration but are still read — exactly one shard (the last)
53/// additionally matches ``key IS NULL``, so NULL-key rows are covered by
54/// precisely one shard with no loss and no duplication.
55#[derive(Clone, Debug, Serialize, Deserialize, JsonSchema, PartialEq, Eq)]
56pub struct ShardConfig {
57    /// INT64 column to range-partition on. Backtick-quoted before use, so it
58    /// is safe against injection but must name a real output column.
59    pub key: String,
60}
61
62/// Configuration for [`SpannerSource`](crate::SpannerSource).
63#[derive(Clone, Serialize, Deserialize, JsonSchema)]
64pub struct SpannerSourceConfig {
65    /// Connection settings (`project_id` / `instance` / `database` / `auth` /
66    /// `max_sessions` / `emulator_host`), flattened.
67    #[serde(flatten)]
68    pub connection: SpannerConnection,
69    /// GoogleSQL query to run. Use named parameters (`@name`) for
70    /// [`params`](Self::params), and the named parameter `@bookmark` to bind
71    /// the incremental cursor server-side.
72    pub query: String,
73    /// Named bind parameters for the query. Values must be scalars (string,
74    /// integer, float, or boolean) — Spanner parameters are typed, and
75    /// null/array/object values have no unambiguous parameter type.
76    #[serde(default)]
77    pub params: HashMap<String, Value>,
78    /// Records per emitted [`StreamPage`](faucet_core::StreamPage). `0` emits
79    /// the whole result set as a single page. Defaults to
80    /// [`DEFAULT_BATCH_SIZE`].
81    #[serde(default = "default_batch_size")]
82    pub batch_size: usize,
83    /// Read at a timestamp this many seconds in the past instead of a strong
84    /// read. Stale reads can be served by any replica, offloading the leader
85    /// — at the cost of missing rows committed within the staleness window.
86    #[serde(default, skip_serializing_if = "Option::is_none")]
87    pub exact_staleness_secs: Option<u64>,
88    /// Replication mode. Defaults to [`SpannerReplication::Full`].
89    #[serde(default)]
90    pub replication: SpannerReplication,
91    /// Explicit state-store key for the bookmark. When unset, a key is
92    /// derived from the database path and a query fingerprint.
93    #[serde(default, skip_serializing_if = "Option::is_none")]
94    pub state_key: Option<String>,
95    /// Optional primary-key range sharding for clustered (Mode B) execution.
96    ///
97    /// When set, the source advertises itself as shardable: the cluster
98    /// coordinator splits the query's `key` range into contiguous slices that
99    /// different workers process concurrently. Has **no effect** outside the
100    /// cluster coordinator (a plain `faucet run` streams the whole query).
101    #[serde(default, skip_serializing_if = "Option::is_none")]
102    pub shard: Option<ShardConfig>,
103}
104
105impl std::fmt::Debug for SpannerSourceConfig {
106    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
107        // `SpannerConnection`'s Debug masks inline key material already; keep
108        // the connector-level Debug conservative anyway.
109        f.debug_struct("SpannerSourceConfig")
110            .field("connection", &self.connection)
111            .field("query", &self.query)
112            .field("params", &self.params)
113            .field("batch_size", &self.batch_size)
114            .field("exact_staleness_secs", &self.exact_staleness_secs)
115            .field("replication", &self.replication)
116            .field("state_key", &self.state_key)
117            .field("shard", &self.shard)
118            .finish()
119    }
120}
121
122/// `true` when `name` is a valid Spanner parameter name
123/// (`[A-Za-z_][A-Za-z0-9_]*`).
124fn valid_param_name(name: &str) -> bool {
125    let mut chars = name.chars();
126    match chars.next() {
127        Some(c) if c.is_ascii_alphabetic() || c == '_' => {}
128        _ => return false,
129    }
130    chars.all(|c| c.is_ascii_alphanumeric() || c == '_')
131}
132
133impl SpannerSourceConfig {
134    /// Build a config from connection identifiers and a query, with defaults
135    /// elsewhere.
136    pub fn new(
137        project_id: impl Into<String>,
138        instance: impl Into<String>,
139        database: impl Into<String>,
140        query: impl Into<String>,
141    ) -> Self {
142        Self {
143            connection: SpannerConnection {
144                project_id: project_id.into(),
145                instance: instance.into(),
146                database: database.into(),
147                auth: Default::default(),
148                max_sessions: 100,
149                emulator_host: None,
150            },
151            query: query.into(),
152            params: HashMap::new(),
153            batch_size: default_batch_size(),
154            exact_staleness_secs: None,
155            replication: SpannerReplication::Full,
156            state_key: None,
157            shard: None,
158        }
159    }
160
161    /// Validate connection, batch size, params, and replication settings.
162    pub fn validate(&self) -> Result<(), FaucetError> {
163        self.connection.validate()?;
164        validate_batch_size(self.batch_size)?;
165        if self.query.trim().is_empty() {
166            return Err(FaucetError::Config(
167                "spanner: `query` must not be empty".into(),
168            ));
169        }
170        for (name, value) in &self.params {
171            if !valid_param_name(name) {
172                return Err(FaucetError::Config(format!(
173                    "spanner: invalid parameter name `{name}` \
174                     (must match [A-Za-z_][A-Za-z0-9_]*)"
175                )));
176            }
177            match value {
178                Value::String(_) | Value::Bool(_) => {}
179                Value::Number(n) => {
180                    if n.as_u64().is_some_and(|u| i64::try_from(u).is_err()) {
181                        return Err(FaucetError::Config(format!(
182                            "spanner: parameter `{name}` ({n}) overflows INT64"
183                        )));
184                    }
185                }
186                _ => {
187                    return Err(FaucetError::Config(format!(
188                        "spanner: parameter `{name}` must be a scalar \
189                         (string / integer / float / boolean) — Spanner \
190                         parameters are typed and null/array/object values \
191                         have no unambiguous type"
192                    )));
193                }
194            }
195        }
196        if let SpannerReplication::Incremental { column, .. } = &self.replication {
197            if column.trim().is_empty() {
198                return Err(FaucetError::Config(
199                    "spanner: incremental replication requires a non-empty `column`".into(),
200                ));
201            }
202            if self.params.contains_key("bookmark") {
203                return Err(FaucetError::Config(
204                    "spanner: the parameter name `bookmark` is reserved for the \
205                     incremental replication cursor"
206                        .into(),
207                ));
208            }
209        }
210        if self.incremental_without_bookmark_pushdown() {
211            tracing::warn!(
212                "spanner incremental replication query has no `@bookmark` parameter: the \
213                 cursor is applied client-side only, so the server returns the ENTIRE \
214                 table on every run (correctness is preserved, but it is a full re-scan). \
215                 Add `@bookmark` to the WHERE clause to push the cursor down, e.g. \
216                 `... WHERE {column} > @bookmark`",
217                column = match &self.replication {
218                    SpannerReplication::Incremental { column, .. } => column.as_str(),
219                    _ => "<column>",
220                }
221            );
222        }
223        Ok(())
224    }
225
226    /// `true` when replication is `Incremental` but the query omits the
227    /// `@bookmark` parameter, so the cursor cannot be pushed down to the
228    /// server and every run re-scans the whole table. Pure predicate so the
229    /// load-time warning's condition is unit-testable.
230    pub(crate) fn incremental_without_bookmark_pushdown(&self) -> bool {
231        matches!(self.replication, SpannerReplication::Incremental { .. })
232            && !self.query.contains("@bookmark")
233    }
234}
235
236#[cfg(test)]
237mod tests {
238    use super::*;
239    use serde_json::json;
240
241    fn base() -> SpannerSourceConfig {
242        SpannerSourceConfig::new("p", "i", "d", "SELECT * FROM t")
243    }
244
245    #[test]
246    fn replication_full_is_default_and_round_trips() {
247        let r: SpannerReplication = serde_json::from_value(json!({"type": "full"})).unwrap();
248        assert_eq!(r, SpannerReplication::Full);
249        assert_eq!(SpannerReplication::default(), SpannerReplication::Full);
250    }
251
252    #[test]
253    fn replication_incremental_parses_column_and_initial_value() {
254        let r: SpannerReplication = serde_json::from_value(json!({
255            "type": "incremental",
256            "column": "updated_at",
257            "initial_value": "1970-01-01T00:00:00Z"
258        }))
259        .unwrap();
260        assert_eq!(
261            r,
262            SpannerReplication::Incremental {
263                column: "updated_at".into(),
264                initial_value: json!("1970-01-01T00:00:00Z"),
265            }
266        );
267    }
268
269    #[test]
270    fn config_flattens_connection_fields() {
271        let cfg: SpannerSourceConfig = serde_json::from_value(json!({
272            "project_id": "p",
273            "instance": "i",
274            "database": "d",
275            "query": "SELECT 1",
276        }))
277        .unwrap();
278        assert_eq!(cfg.connection.project_id, "p");
279        assert_eq!(cfg.connection.max_sessions, 100);
280        assert_eq!(cfg.batch_size, DEFAULT_BATCH_SIZE);
281        assert!(cfg.exact_staleness_secs.is_none());
282        assert!(cfg.validate().is_ok());
283    }
284
285    #[test]
286    fn validate_rejects_empty_query_and_bad_batch() {
287        let mut cfg = base();
288        cfg.query = "  ".into();
289        assert!(cfg.validate().is_err());
290        let mut cfg = base();
291        cfg.batch_size = faucet_core::MAX_BATCH_SIZE + 1;
292        assert!(cfg.validate().is_err());
293    }
294
295    #[test]
296    fn validate_rejects_non_scalar_and_ill_named_params() {
297        let mut cfg = base();
298        cfg.params.insert("ok".into(), json!("x"));
299        assert!(cfg.validate().is_ok());
300
301        let mut cfg = base();
302        cfg.params.insert("bad".into(), json!(null));
303        assert!(cfg.validate().is_err());
304        let mut cfg = base();
305        cfg.params.insert("bad".into(), json!([1]));
306        assert!(cfg.validate().is_err());
307        let mut cfg = base();
308        cfg.params.insert("bad".into(), json!({"a": 1}));
309        assert!(cfg.validate().is_err());
310        let mut cfg = base();
311        cfg.params.insert("9lives".into(), json!(1));
312        assert!(cfg.validate().is_err(), "names must not start with a digit");
313        let mut cfg = base();
314        cfg.params.insert("we-ird".into(), json!(1));
315        assert!(cfg.validate().is_err(), "names must be [A-Za-z0-9_]");
316        let mut cfg = base();
317        cfg.params.insert("big".into(), json!(u64::MAX));
318        assert!(
319            cfg.validate().is_err(),
320            "u64 above i64::MAX overflows INT64"
321        );
322    }
323
324    #[test]
325    fn validate_rejects_incremental_without_column_or_with_reserved_param() {
326        let mut cfg = base();
327        cfg.replication = SpannerReplication::Incremental {
328            column: " ".into(),
329            initial_value: json!(0),
330        };
331        assert!(cfg.validate().is_err());
332
333        let mut cfg = base();
334        cfg.query = "SELECT * FROM t WHERE c > @bookmark".into();
335        cfg.replication = SpannerReplication::Incremental {
336            column: "c".into(),
337            initial_value: json!(0),
338        };
339        cfg.params.insert("bookmark".into(), json!(1));
340        assert!(cfg.validate().is_err(), "`bookmark` is reserved");
341    }
342
343    #[test]
344    fn incremental_without_bookmark_pushdown_flags_missing_token() {
345        let mut missing = base();
346        missing.replication = SpannerReplication::Incremental {
347            column: "updated_at".into(),
348            initial_value: json!("1970-01-01"),
349        };
350        assert!(missing.incremental_without_bookmark_pushdown());
351        // validate still succeeds (warn, not hard-error).
352        assert!(missing.validate().is_ok());
353
354        let mut with_token = missing.clone();
355        with_token.query = "SELECT * FROM t WHERE updated_at > @bookmark".into();
356        assert!(!with_token.incremental_without_bookmark_pushdown());
357
358        assert!(!base().incremental_without_bookmark_pushdown());
359    }
360
361    #[test]
362    fn valid_param_names() {
363        assert!(valid_param_name("cursor"));
364        assert!(valid_param_name("_x9"));
365        assert!(valid_param_name("A_b_1"));
366        assert!(!valid_param_name(""));
367        assert!(!valid_param_name("1abc"));
368        assert!(!valid_param_name("a-b"));
369        assert!(!valid_param_name("a b"));
370    }
371
372    #[test]
373    fn debug_does_not_leak_inline_key() {
374        let mut cfg = base();
375        cfg.connection.auth = faucet_common_spanner::SpannerCredentials::ServiceAccountKey {
376            json: "{\"private_key\":\"SECRET\"}".into(),
377        };
378        let dbg = format!("{cfg:?}");
379        assert!(!dbg.contains("SECRET"));
380    }
381}