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

faucet_source_spanner/
config.rs
