Skip to main content

faucet_source_mssql/
config.rs

1//! Configuration for the MSSQL query source.
2
3use faucet_common_mssql::MssqlConnectionConfig;
4use faucet_core::{DEFAULT_BATCH_SIZE, FaucetError, validate_batch_size};
5use schemars::JsonSchema;
6use serde::{Deserialize, Serialize};
7use serde_json::Value;
8
9fn default_max_connections() -> u32 {
10    10
11}
12fn default_batch_size() -> usize {
13    DEFAULT_BATCH_SIZE
14}
15fn default_statement_timeout_secs() -> u64 {
16    300
17}
18
19/// How the source replicates rows across runs.
20///
21/// Serializes as `{ type: full }` or
22/// `{ type: incremental, column: "...", initial_value: ... }`.
23#[derive(Clone, Debug, Serialize, Deserialize, JsonSchema, Default, PartialEq)]
24#[serde(tag = "type", rename_all = "snake_case")]
25pub enum MssqlReplication {
26    /// Every run fetches the full result set (default).
27    #[default]
28    Full,
29    /// Only rows whose `column` is strictly greater than the stored bookmark
30    /// (or `initial_value` on the first run) are emitted.
31    ///
32    /// The bookmark is applied two ways: if the query contains the literal
33    /// token `@bookmark`, it is bound as a parameter so the server filters
34    /// (efficient); the source *also* filters client-side as a correctness
35    /// backstop. The new maximum of `column` is persisted on the final page.
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/// Configuration for [`MssqlSource`](crate::MssqlSource).
45#[derive(Clone, Serialize, Deserialize, JsonSchema)]
46pub struct MssqlSourceConfig {
47    /// Connection + TLS settings (`connection_url` or `connection_string`).
48    #[serde(flatten)]
49    pub connection: MssqlConnectionConfig,
50    /// SQL query to run. Use `@P1`, `@P2`, … for [`params`](Self::params), and
51    /// the literal `@bookmark` token to bind the incremental cursor server-side.
52    pub query: String,
53    /// Positional bind parameters for the query (`@P1`…`@Pn`). Defaults to empty.
54    #[serde(default)]
55    pub params: Vec<Value>,
56    /// Maximum pooled connections. Defaults to 10.
57    #[serde(default = "default_max_connections")]
58    pub max_connections: u32,
59    /// Records per emitted [`StreamPage`](faucet_core::StreamPage). `0` emits the
60    /// whole result set as a single page. Defaults to [`DEFAULT_BATCH_SIZE`].
61    #[serde(default = "default_batch_size")]
62    pub batch_size: usize,
63    /// Per-query timeout in seconds (`0` disables). Defaults to 300.
64    #[serde(default = "default_statement_timeout_secs")]
65    pub statement_timeout_secs: u64,
66    /// Replication mode. Defaults to [`MssqlReplication::Full`].
67    #[serde(default)]
68    pub replication: MssqlReplication,
69    /// Explicit state-store key for the bookmark. When unset, a key is derived
70    /// from the connection host and a query fingerprint.
71    #[serde(default, skip_serializing_if = "Option::is_none")]
72    pub state_key: Option<String>,
73    /// Optional primary-key range sharding for clustered (Mode B) execution.
74    ///
75    /// When set, the source advertises itself as shardable: the cluster
76    /// coordinator splits the query's `key` range into contiguous slices that
77    /// different workers process concurrently. Has **no effect** outside the
78    /// cluster coordinator (a plain `faucet run` streams the whole query), so
79    /// it is fully backward compatible. See [`ShardConfig`].
80    #[serde(default, skip_serializing_if = "Option::is_none")]
81    pub shard: Option<ShardConfig>,
82}
83
84/// Primary-key range sharding settings for the MSSQL source.
85///
86/// The source is split by contiguous ranges of an **integer-typed** column:
87/// each shard runs `SELECT * FROM (<query>) WHERE [key] >= lo AND [key] < hi`.
88/// The column must be present in the query's output and orderable as a 64-bit
89/// integer (e.g. a `BIGINT`/`INT` `IDENTITY` primary key).
90///
91/// **`ORDER BY`:** T-SQL forbids `ORDER BY` inside a derived table (without
92/// `TOP`/`OFFSET`), so a sharded query must not end in a top-level `ORDER BY`
93/// — ordering across concurrently-executing shards is meaningless anyway.
94///
95/// **Nullable keys:** if the key column contains NULLs, those rows are not
96/// visible to the `MIN`/`MAX` range enumeration. They are still read — exactly
97/// one shard (the last) additionally matches `[key] IS NULL`, so NULL-key rows
98/// are covered by precisely one shard with no loss and no duplication.
99#[derive(Clone, Debug, Serialize, Deserialize, JsonSchema, PartialEq, Eq)]
100pub struct ShardConfig {
101    /// Integer column to range-partition on. Quoted as an identifier (brackets)
102    /// before use, so it is safe against injection but must name a real output
103    /// column.
104    pub key: String,
105}
106
107impl std::fmt::Debug for MssqlSourceConfig {
108    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
109        f.debug_struct("MssqlSourceConfig")
110            .field("connection", &"***")
111            .field("query", &self.query)
112            .field("params", &self.params)
113            .field("max_connections", &self.max_connections)
114            .field("batch_size", &self.batch_size)
115            .field("statement_timeout_secs", &self.statement_timeout_secs)
116            .field("replication", &self.replication)
117            .field("state_key", &self.state_key)
118            .finish()
119    }
120}
121
122impl MssqlSourceConfig {
123    /// Build a config from a connection URL and query, with defaults elsewhere.
124    pub fn new(connection_url: impl Into<String>, query: impl Into<String>) -> Self {
125        Self {
126            connection: MssqlConnectionConfig {
127                connection_url: Some(connection_url.into()),
128                ..Default::default()
129            },
130            query: query.into(),
131            params: Vec::new(),
132            max_connections: default_max_connections(),
133            batch_size: default_batch_size(),
134            statement_timeout_secs: default_statement_timeout_secs(),
135            replication: MssqlReplication::Full,
136            state_key: None,
137            shard: None,
138        }
139    }
140
141    /// Validate connection source, batch size, and replication settings.
142    pub fn validate(&self) -> Result<(), FaucetError> {
143        self.connection.validate()?;
144        validate_batch_size(self.batch_size)?;
145        if let MssqlReplication::Incremental { column, .. } = &self.replication
146            && column.trim().is_empty()
147        {
148            return Err(FaucetError::Config(
149                "MSSQL incremental replication requires a non-empty `column`".into(),
150            ));
151        }
152        if self.incremental_without_bookmark_pushdown() {
153            tracing::warn!(
154                "MSSQL incremental replication query has no `@bookmark` token: the \
155                 cursor is applied client-side only, so the server returns the ENTIRE \
156                 table on every run (correctness is preserved, but it is a full re-scan). \
157                 Add `@bookmark` to the WHERE clause to push the cursor down, e.g. \
158                 `... WHERE {column} > @bookmark`",
159                column = match &self.replication {
160                    MssqlReplication::Incremental { column, .. } => column.as_str(),
161                    _ => "<column>",
162                }
163            );
164        }
165        Ok(())
166    }
167
168    /// `true` when replication is `Incremental` but the query omits the
169    /// `@bookmark` token, so the cursor cannot be pushed down to the server and
170    /// every run re-scans the whole table (F53). Pure predicate so the
171    /// load-time warning's condition is unit-testable.
172    pub(crate) fn incremental_without_bookmark_pushdown(&self) -> bool {
173        matches!(self.replication, MssqlReplication::Incremental { .. })
174            && !self.query.contains("@bookmark")
175    }
176}
177
178#[cfg(test)]
179mod tests {
180    use super::*;
181    use serde_json::json;
182
183    #[test]
184    fn replication_full_is_default_and_round_trips() {
185        let r: MssqlReplication = serde_json::from_value(json!({"type": "full"})).unwrap();
186        assert_eq!(r, MssqlReplication::Full);
187        assert_eq!(MssqlReplication::default(), MssqlReplication::Full);
188    }
189
190    #[test]
191    fn replication_incremental_parses_column_and_initial_value() {
192        let r: MssqlReplication = serde_json::from_value(json!({
193            "type": "incremental",
194            "column": "updated_at",
195            "initial_value": "1970-01-01T00:00:00Z"
196        }))
197        .unwrap();
198        assert_eq!(
199            r,
200            MssqlReplication::Incremental {
201                column: "updated_at".into(),
202                initial_value: json!("1970-01-01T00:00:00Z"),
203            }
204        );
205    }
206
207    #[test]
208    fn config_flattens_connection_fields() {
209        let cfg: MssqlSourceConfig = serde_json::from_value(json!({
210            "connection_url": "mssql://sa:pw@host:1433/db",
211            "query": "SELECT 1",
212        }))
213        .unwrap();
214        assert_eq!(
215            cfg.connection.connection_url.as_deref(),
216            Some("mssql://sa:pw@host:1433/db")
217        );
218        assert_eq!(cfg.batch_size, DEFAULT_BATCH_SIZE);
219        assert_eq!(cfg.max_connections, 10);
220        assert_eq!(cfg.statement_timeout_secs, 300);
221    }
222
223    #[test]
224    fn validate_rejects_incremental_without_column() {
225        let cfg = MssqlSourceConfig {
226            replication: MssqlReplication::Incremental {
227                column: "  ".into(),
228                initial_value: json!(0),
229            },
230            ..MssqlSourceConfig::new("mssql://sa:pw@h/db", "SELECT 1")
231        };
232        assert!(cfg.validate().is_err());
233    }
234
235    #[test]
236    fn incremental_without_bookmark_pushdown_flags_missing_token() {
237        // F53: Incremental + query lacking `@bookmark` → full re-scan warning.
238        let missing = MssqlSourceConfig {
239            replication: MssqlReplication::Incremental {
240                column: "updated_at".into(),
241                initial_value: json!("1970-01-01"),
242            },
243            ..MssqlSourceConfig::new("mssql://sa:pw@h/db", "SELECT * FROM t")
244        };
245        assert!(missing.incremental_without_bookmark_pushdown());
246        // validate still succeeds (warn, not hard-error).
247        assert!(missing.validate().is_ok());
248
249        // With the token present, no warning.
250        let with_token = MssqlSourceConfig {
251            replication: MssqlReplication::Incremental {
252                column: "updated_at".into(),
253                initial_value: json!("1970-01-01"),
254            },
255            ..MssqlSourceConfig::new(
256                "mssql://sa:pw@h/db",
257                "SELECT * FROM t WHERE updated_at > @bookmark",
258            )
259        };
260        assert!(!with_token.incremental_without_bookmark_pushdown());
261
262        // Full replication never warns.
263        let full = MssqlSourceConfig::new("mssql://sa:pw@h/db", "SELECT * FROM t");
264        assert!(!full.incremental_without_bookmark_pushdown());
265    }
266
267    #[test]
268    fn validate_rejects_bad_batch_size() {
269        let cfg = MssqlSourceConfig {
270            batch_size: faucet_core::MAX_BATCH_SIZE + 1,
271            ..MssqlSourceConfig::new("mssql://sa:pw@h/db", "SELECT 1")
272        };
273        assert!(cfg.validate().is_err());
274    }
275
276    #[test]
277    fn debug_masks_connection() {
278        let cfg = MssqlSourceConfig::new("mssql://sa:secret@h/db", "SELECT 1");
279        let dbg = format!("{cfg:?}");
280        assert!(dbg.contains("***"));
281        assert!(!dbg.contains("secret"));
282    }
283}