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 not yet read, in `column` order, resuming from the stored
26    /// bookmark (or `initial_value` on the first run).
27    ///
28    /// The source wraps the query as
29    /// ``SELECT * FROM (<query>) AS q WHERE q.`column` >= CAST(@bookmark AS <type>) ORDER BY q.`column` ``
30    /// (`<type>` is the column's type, read from the query's result metadata)
31    /// and drops the rows it already emitted at the bookmark value, so rows
32    /// that share a cursor value (ties, late commits) are never skipped.
33    /// Every page carries a bookmark, so a crash replays at most one page.
34    /// Write `@bookmark` in the query (with `>=`) to apply the cursor further
35    /// inside it; a strict `>` is refused. Rows whose `column` is NULL are
36    /// never read.
37    ///
38    /// `@bookmark` inside the query is bound by its JSON type (a timestamp
39    /// bookmark is a STRING), and Spanner does not coerce parameter types: cast
40    /// it where you use it, e.g. `WHERE updated_at >= TIMESTAMP(@bookmark)`.
41    Incremental {
42        /// Column whose value is the replication cursor (e.g. `updated_at`).
43        column: String,
44        /// Lower bound (inclusive) used on the first run, before any bookmark
45        /// is stored.
46        initial_value: Value,
47    },
48}
49
50/// Primary-key range sharding settings for the Spanner source (clustered
51/// Mode B execution).
52///
53/// The source is split by contiguous ranges of an **INT64-typed** column:
54/// each shard runs ``SELECT * FROM (<query>) AS _faucet_shard WHERE `key` >=
55/// lo AND `key` < hi``. The column must be present in the query's output.
56///
57/// **Nullable keys:** rows whose key is NULL are invisible to the `MIN`/`MAX`
58/// range enumeration but are still read — exactly one shard (the last)
59/// additionally matches ``key IS NULL``, so NULL-key rows are covered by
60/// precisely one shard with no loss and no duplication.
61#[derive(Clone, Debug, Serialize, Deserialize, JsonSchema, PartialEq, Eq)]
62pub struct ShardConfig {
63    /// INT64 column to range-partition on. Backtick-quoted before use, so it
64    /// is safe against injection but must name a real output column.
65    pub key: String,
66}
67
68/// Configuration for [`SpannerSource`](crate::SpannerSource).
69#[derive(Clone, Serialize, Deserialize, JsonSchema)]
70pub struct SpannerSourceConfig {
71    /// Connection settings (`project_id` / `instance` / `database` / `auth` /
72    /// `max_sessions` / `emulator_host`), flattened.
73    #[serde(flatten)]
74    pub connection: SpannerConnection,
75    /// GoogleSQL query to run. Use named parameters (`@name`) for
76    /// [`params`](Self::params), and the named parameter `@bookmark` to bind
77    /// the incremental cursor server-side.
78    pub query: String,
79    /// Named bind parameters for the query. Values must be scalars (string,
80    /// integer, float, or boolean) — Spanner parameters are typed, and
81    /// null/array/object values have no unambiguous parameter type.
82    #[serde(default)]
83    pub params: HashMap<String, Value>,
84    /// Records per emitted [`StreamPage`](faucet_core::StreamPage). `0` emits
85    /// the whole result set as a single page. Defaults to
86    /// [`DEFAULT_BATCH_SIZE`].
87    #[serde(default = "default_batch_size")]
88    pub batch_size: usize,
89    /// Read at a timestamp this many seconds in the past instead of a strong
90    /// read. Stale reads can be served by any replica, offloading the leader
91    /// — at the cost of missing rows committed within the staleness window.
92    #[serde(default, skip_serializing_if = "Option::is_none")]
93    pub exact_staleness_secs: Option<u64>,
94    /// Replication mode. Defaults to [`SpannerReplication::Full`].
95    #[serde(default)]
96    pub replication: SpannerReplication,
97    /// Explicit state-store key for the bookmark. When unset, a key is
98    /// derived from the database path and a query fingerprint.
99    #[serde(default, skip_serializing_if = "Option::is_none")]
100    pub state_key: Option<String>,
101    /// Optional primary-key range sharding for clustered (Mode B) execution.
102    ///
103    /// When set, the source advertises itself as shardable: the cluster
104    /// coordinator splits the query's `key` range into contiguous slices that
105    /// different workers process concurrently. Has **no effect** outside the
106    /// cluster coordinator (a plain `faucet run` streams the whole query).
107    #[serde(default, skip_serializing_if = "Option::is_none")]
108    pub shard: Option<ShardConfig>,
109}
110
111impl std::fmt::Debug for SpannerSourceConfig {
112    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
113        // `SpannerConnection`'s Debug masks inline key material already; keep
114        // the connector-level Debug conservative anyway.
115        f.debug_struct("SpannerSourceConfig")
116            .field("connection", &self.connection)
117            .field("query", &self.query)
118            .field("params", &self.params)
119            .field("batch_size", &self.batch_size)
120            .field("exact_staleness_secs", &self.exact_staleness_secs)
121            .field("replication", &self.replication)
122            .field("state_key", &self.state_key)
123            .field("shard", &self.shard)
124            .finish()
125    }
126}
127
128/// `true` when `name` is a valid Spanner parameter name
129/// (`[A-Za-z_][A-Za-z0-9_]*`).
130fn valid_param_name(name: &str) -> bool {
131    let mut chars = name.chars();
132    match chars.next() {
133        Some(c) if c.is_ascii_alphabetic() || c == '_' => {}
134        _ => return false,
135    }
136    chars.all(|c| c.is_ascii_alphanumeric() || c == '_')
137}
138
139impl SpannerSourceConfig {
140    /// Build a config from connection identifiers and a query, with defaults
141    /// elsewhere.
142    pub fn new(
143        project_id: impl Into<String>,
144        instance: impl Into<String>,
145        database: impl Into<String>,
146        query: impl Into<String>,
147    ) -> Self {
148        Self {
149            connection: SpannerConnection {
150                project_id: project_id.into(),
151                instance: instance.into(),
152                database: database.into(),
153                auth: Default::default(),
154                max_sessions: 100,
155                emulator_host: None,
156            },
157            query: query.into(),
158            params: HashMap::new(),
159            batch_size: default_batch_size(),
160            exact_staleness_secs: None,
161            replication: SpannerReplication::Full,
162            state_key: None,
163            shard: None,
164        }
165    }
166
167    /// Validate connection, batch size, params, and replication settings.
168    pub fn validate(&self) -> Result<(), FaucetError> {
169        self.connection.validate()?;
170        validate_batch_size(self.batch_size)?;
171        if self.query.trim().is_empty() {
172            return Err(FaucetError::Config(
173                "spanner: `query` must not be empty".into(),
174            ));
175        }
176        for (name, value) in &self.params {
177            if !valid_param_name(name) {
178                return Err(FaucetError::Config(format!(
179                    "spanner: invalid parameter name `{name}` \
180                     (must match [A-Za-z_][A-Za-z0-9_]*)"
181                )));
182            }
183            match value {
184                Value::String(_) | Value::Bool(_) => {}
185                Value::Number(n) => {
186                    if n.as_u64().is_some_and(|u| i64::try_from(u).is_err()) {
187                        return Err(FaucetError::Config(format!(
188                            "spanner: parameter `{name}` ({n}) overflows INT64"
189                        )));
190                    }
191                }
192                _ => {
193                    return Err(FaucetError::Config(format!(
194                        "spanner: parameter `{name}` must be a scalar \
195                         (string / integer / float / boolean) — Spanner \
196                         parameters are typed and null/array/object values \
197                         have no unambiguous type"
198                    )));
199                }
200            }
201        }
202        if let SpannerReplication::Incremental {
203            column,
204            initial_value,
205        } = &self.replication
206        {
207            if column.trim().is_empty() {
208                return Err(FaucetError::Config(
209                    "spanner: incremental replication requires a non-empty `column`".into(),
210                ));
211            }
212            if !matches!(
213                initial_value,
214                Value::String(_) | Value::Number(_) | Value::Bool(_)
215            ) {
216                return Err(FaucetError::Config(
217                    "spanner: incremental replication requires a scalar `initial_value` \
218                     (string / integer / float / boolean)"
219                        .into(),
220                ));
221            }
222            if self.params.contains_key("bookmark") {
223                return Err(FaucetError::Config(
224                    "spanner: the parameter name `bookmark` is reserved for the \
225                     incremental replication cursor"
226                        .into(),
227                ));
228            }
229            if strict_bookmark_comparison(&self.query) {
230                return Err(FaucetError::Config(format!(
231                    "spanner: compare `{BOOKMARK_TOKEN}` with `>=`, not `>`: the source \
232                     re-reads the bookmark value and drops the rows it already emitted, and a \
233                     strict comparison skips rows that share it"
234                )));
235            }
236        }
237        Ok(())
238    }
239}
240
241/// The named parameter a query uses to read the incremental bookmark
242/// (`WHERE updated_at >= TIMESTAMP(@bookmark)`).
243pub const BOOKMARK_TOKEN: &str = "@bookmark";
244
245/// `true` when the query compares `@bookmark` strictly (`> @bookmark`,
246/// `> TIMESTAMP(@bookmark)`, `> CAST(@bookmark AS NUMERIC)`, `@bookmark < c`).
247fn strict_bookmark_comparison(query: &str) -> bool {
248    let ident = |c: char| c.is_ascii_alphanumeric() || c == '_';
249    let mut unwrapped = String::with_capacity(query.len());
250    let mut rest = query;
251    while let Some(i) = rest.find(BOOKMARK_TOKEN) {
252        let end = i + BOOKMARK_TOKEN.len();
253        if rest[end..].starts_with(ident) {
254            // Another parameter (`@bookmark_x`): keep it out of the check.
255            unwrapped.push_str(&rest[..i]);
256            unwrapped.push_str("@_param");
257            rest = &rest[end..];
258            continue;
259        }
260        let (mut start, mut stop) = (i, end);
261        let before = rest[..i].trim_end();
262        if let Some(open) = before.strip_suffix('(') {
263            let name = open.trim_end();
264            let fname = name.trim_end_matches(ident);
265            let after = rest[end..].trim_start();
266            let after = match after.get(..3) {
267                Some(kw) if kw.eq_ignore_ascii_case("AS ") => after[3..]
268                    .trim_start()
269                    .trim_start_matches(ident)
270                    .trim_start(),
271                _ => after,
272            };
273            if let Some(tail) = after.strip_prefix(')') {
274                start = fname.len();
275                stop = rest.len() - tail.len();
276            }
277        }
278        unwrapped.push_str(&rest[..start]);
279        unwrapped.push_str(BOOKMARK_TOKEN);
280        rest = &rest[stop..];
281    }
282    unwrapped.push_str(rest);
283    faucet_core::replication::cursor::strict_comparison(&unwrapped, BOOKMARK_TOKEN)
284}
285
286#[cfg(test)]
287mod tests {
288    use super::*;
289    use serde_json::json;
290
291    fn base() -> SpannerSourceConfig {
292        SpannerSourceConfig::new("p", "i", "d", "SELECT * FROM t")
293    }
294
295    #[test]
296    fn replication_full_is_default_and_round_trips() {
297        let r: SpannerReplication = serde_json::from_value(json!({"type": "full"})).unwrap();
298        assert_eq!(r, SpannerReplication::Full);
299        assert_eq!(SpannerReplication::default(), SpannerReplication::Full);
300    }
301
302    #[test]
303    fn replication_incremental_parses_column_and_initial_value() {
304        let r: SpannerReplication = serde_json::from_value(json!({
305            "type": "incremental",
306            "column": "updated_at",
307            "initial_value": "1970-01-01T00:00:00Z"
308        }))
309        .unwrap();
310        assert_eq!(
311            r,
312            SpannerReplication::Incremental {
313                column: "updated_at".into(),
314                initial_value: json!("1970-01-01T00:00:00Z"),
315            }
316        );
317    }
318
319    #[test]
320    fn config_flattens_connection_fields() {
321        let cfg: SpannerSourceConfig = serde_json::from_value(json!({
322            "project_id": "p",
323            "instance": "i",
324            "database": "d",
325            "query": "SELECT 1",
326        }))
327        .unwrap();
328        assert_eq!(cfg.connection.project_id, "p");
329        assert_eq!(cfg.connection.max_sessions, 100);
330        assert_eq!(cfg.batch_size, DEFAULT_BATCH_SIZE);
331        assert!(cfg.exact_staleness_secs.is_none());
332        assert!(cfg.validate().is_ok());
333    }
334
335    #[test]
336    fn validate_rejects_empty_query_and_bad_batch() {
337        let mut cfg = base();
338        cfg.query = "  ".into();
339        assert!(cfg.validate().is_err());
340        let mut cfg = base();
341        cfg.batch_size = faucet_core::MAX_BATCH_SIZE + 1;
342        assert!(cfg.validate().is_err());
343    }
344
345    #[test]
346    fn validate_rejects_non_scalar_and_ill_named_params() {
347        let mut cfg = base();
348        cfg.params.insert("ok".into(), json!("x"));
349        assert!(cfg.validate().is_ok());
350
351        let mut cfg = base();
352        cfg.params.insert("bad".into(), json!(null));
353        assert!(cfg.validate().is_err());
354        let mut cfg = base();
355        cfg.params.insert("bad".into(), json!([1]));
356        assert!(cfg.validate().is_err());
357        let mut cfg = base();
358        cfg.params.insert("bad".into(), json!({"a": 1}));
359        assert!(cfg.validate().is_err());
360        let mut cfg = base();
361        cfg.params.insert("9lives".into(), json!(1));
362        assert!(cfg.validate().is_err(), "names must not start with a digit");
363        let mut cfg = base();
364        cfg.params.insert("we-ird".into(), json!(1));
365        assert!(cfg.validate().is_err(), "names must be [A-Za-z0-9_]");
366        let mut cfg = base();
367        cfg.params.insert("big".into(), json!(u64::MAX));
368        assert!(
369            cfg.validate().is_err(),
370            "u64 above i64::MAX overflows INT64"
371        );
372    }
373
374    #[test]
375    fn validate_rejects_incremental_without_column_or_with_reserved_param() {
376        let mut cfg = base();
377        cfg.replication = SpannerReplication::Incremental {
378            column: " ".into(),
379            initial_value: json!(0),
380        };
381        assert!(cfg.validate().is_err());
382
383        let mut cfg = base();
384        cfg.query = "SELECT * FROM t WHERE c > @bookmark".into();
385        cfg.replication = SpannerReplication::Incremental {
386            column: "c".into(),
387            initial_value: json!(0),
388        };
389        cfg.params.insert("bookmark".into(), json!(1));
390        assert!(cfg.validate().is_err(), "`bookmark` is reserved");
391    }
392
393    #[test]
394    fn validate_rejects_bad_incremental_settings() {
395        for (query, column, initial, needle) in [
396            (
397                "SELECT * FROM t",
398                "updated_at",
399                json!(null),
400                "initial_value",
401            ),
402            ("SELECT * FROM t", "updated_at", json!([1]), "initial_value"),
403            ("SELECT * FROM t WHERE c > @bookmark", "c", json!(0), ">="),
404            ("SELECT * FROM t WHERE @bookmark < c", "c", json!(0), ">="),
405            (
406                "SELECT * FROM t WHERE ts > TIMESTAMP(@bookmark)",
407                "ts",
408                json!("1970-01-01T00:00:00Z"),
409                ">=",
410            ),
411            (
412                "SELECT * FROM t WHERE n > CAST( @bookmark AS NUMERIC ) ORDER BY n",
413                "n",
414                json!("0"),
415                ">=",
416            ),
417            ("SELECT * FROM t WHERE c > (@bookmark)", "c", json!(0), ">="),
418        ] {
419            let mut cfg = base();
420            cfg.query = query.into();
421            cfg.replication = SpannerReplication::Incremental {
422                column: column.into(),
423                initial_value: initial,
424            };
425            let err = cfg.validate().unwrap_err().to_string();
426            assert!(err.contains(needle), "{err}");
427        }
428        let mut ok = base();
429        ok.query = "SELECT * FROM t WHERE c >= @bookmark".into();
430        ok.replication = SpannerReplication::Incremental {
431            column: "c".into(),
432            initial_value: json!(0),
433        };
434        assert!(ok.validate().is_ok());
435        for query in [
436            "SELECT * FROM t WHERE ts >= TIMESTAMP(@bookmark)",
437            "SELECT * FROM t WHERE n >= CAST(@bookmark AS NUMERIC)",
438            "SELECT * FROM t WHERE c > @bookmark_other AND c >= @bookmark",
439            "SELECT * FROM t WHERE c > f (x) AND c >= @bookmark",
440        ] {
441            assert!(!strict_bookmark_comparison(query), "{query}");
442        }
443    }
444
445    #[test]
446    fn valid_param_names() {
447        assert!(valid_param_name("cursor"));
448        assert!(valid_param_name("_x9"));
449        assert!(valid_param_name("A_b_1"));
450        assert!(!valid_param_name(""));
451        assert!(!valid_param_name("1abc"));
452        assert!(!valid_param_name("a-b"));
453        assert!(!valid_param_name("a b"));
454    }
455
456    #[test]
457    fn debug_does_not_leak_inline_key() {
458        let mut cfg = base();
459        cfg.connection.auth = faucet_common_spanner::SpannerCredentials::ServiceAccountKey {
460            json: "{\"private_key\":\"SECRET\"}".into(),
461        };
462        let dbg = format!("{cfg:?}");
463        assert!(!dbg.contains("SECRET"));
464    }
465}
faucet_source_spanner/config.rs

faucet_source_spanner/
config.rs
