Skip to main content

faucet_source_mysql/
config.rs

1//! MySQL source configuration.
2
3use faucet_core::DEFAULT_BATCH_SIZE;
4use schemars::JsonSchema;
5use serde::{Deserialize, Serialize};
6
7/// Configuration for the MySQL query source.
8#[derive(Clone, Serialize, Deserialize, JsonSchema)]
9pub struct MysqlSourceConfig {
10    /// MySQL connection URL (e.g. `mysql://user:pass@host/db`).
11    pub connection_url: String,
12    /// SQL query to execute.
13    pub query: String,
14    /// Maximum number of connections in the pool. Defaults to 10.
15    #[serde(default = "default_max_connections")]
16    pub max_connections: u32,
17    /// Records per emitted [`StreamPage`](faucet_core::StreamPage). Rows are
18    /// drained from the sqlx cursor and yielded whenever the buffer reaches
19    /// this size. Defaults to [`DEFAULT_BATCH_SIZE`].
20    ///
21    /// `batch_size = 0` is the "no batching" sentinel: the cursor is fully
22    /// drained and the entire result set is emitted in a single page. Useful
23    /// for small lookup tables or for sinks (e.g. SQL `COPY`, BigQuery load
24    /// jobs) that prefer one large request to many small ones.
25    #[serde(default = "default_batch_size")]
26    pub batch_size: usize,
27
28    /// Optional primary-key range sharding for clustered (Mode B) execution.
29    ///
30    /// When set, the source advertises itself as shardable: the cluster
31    /// coordinator splits the query's `key` range into contiguous slices that
32    /// different workers process concurrently. Has **no effect** outside the
33    /// cluster coordinator (a plain `faucet run` streams the whole query), so
34    /// it is fully backward compatible. See [`ShardConfig`].
35    #[serde(default, skip_serializing_if = "Option::is_none")]
36    pub shard: Option<ShardConfig>,
37}
38
39/// Primary-key range sharding settings for the MySQL source.
40///
41/// The source is split by contiguous ranges of an **integer-typed** column:
42/// each shard runs ``SELECT * FROM (<query>) WHERE `key` >= lo AND `key` < hi``.
43/// The column must be present in the query's output and orderable as a 64-bit
44/// integer (e.g. a `BIGINT`/`INT` auto-increment primary key).
45///
46/// **Nullable keys:** if the key column contains NULLs, those rows are not
47/// visible to the `MIN`/`MAX` range enumeration. They are still read — exactly
48/// one shard (the last) additionally matches ``` `key` IS NULL ```, so NULL-key
49/// rows are covered by precisely one shard with no loss and no duplication.
50#[derive(Clone, Debug, Serialize, Deserialize, JsonSchema, PartialEq, Eq)]
51pub struct ShardConfig {
52    /// Integer column to range-partition on. Quoted as an identifier (backticks)
53    /// before use, so it is safe against injection but must name a real output
54    /// column.
55    pub key: String,
56}
57
58fn default_max_connections() -> u32 {
59    10
60}
61
62fn default_batch_size() -> usize {
63    DEFAULT_BATCH_SIZE
64}
65
66impl std::fmt::Debug for MysqlSourceConfig {
67    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
68        f.debug_struct("MysqlSourceConfig")
69            .field("connection_url", &"***")
70            .field("query", &self.query)
71            .field("max_connections", &self.max_connections)
72            .field("batch_size", &self.batch_size)
73            .finish()
74    }
75}
76
77impl MysqlSourceConfig {
78    /// Create a new config with the required connection URL and query.
79    pub fn new(connection_url: impl Into<String>, query: impl Into<String>) -> Self {
80        Self {
81            connection_url: connection_url.into(),
82            query: query.into(),
83            max_connections: 10,
84            batch_size: DEFAULT_BATCH_SIZE,
85            shard: None,
86        }
87    }
88
89    /// Set the maximum number of connections in the pool.
90    pub fn with_max_connections(mut self, max_connections: u32) -> Self {
91        self.max_connections = max_connections;
92        self
93    }
94
95    /// Set the per-page row count for [`Source::stream_pages`](faucet_core::Source::stream_pages).
96    ///
97    /// Pass `0` to opt out of batching — the entire result set is emitted in
98    /// a single [`StreamPage`](faucet_core::StreamPage).
99    pub fn with_batch_size(mut self, batch_size: usize) -> Self {
100        self.batch_size = batch_size;
101        self
102    }
103}
104
105#[cfg(test)]
106mod tests {
107    use super::*;
108
109    #[test]
110    fn default_config() {
111        let config = MysqlSourceConfig::new("mysql://localhost/test", "SELECT * FROM events");
112        assert_eq!(config.query, "SELECT * FROM events");
113    }
114
115    #[test]
116    fn debug_masks_connection_url() {
117        let config = MysqlSourceConfig::new("mysql://secret:pass@host/db", "SELECT 1");
118        let debug = format!("{config:?}");
119        assert!(debug.contains("***"));
120        assert!(!debug.contains("secret"));
121        assert!(!debug.contains("pass"));
122    }
123
124    #[test]
125    fn batch_size_defaults_to_default_batch_size() {
126        let config = MysqlSourceConfig::new("mysql://localhost/test", "SELECT 1");
127        assert_eq!(config.batch_size, faucet_core::DEFAULT_BATCH_SIZE);
128    }
129
130    #[test]
131    fn with_batch_size_overrides_default() {
132        let config =
133            MysqlSourceConfig::new("mysql://localhost/test", "SELECT 1").with_batch_size(500);
134        assert_eq!(config.batch_size, 500);
135    }
136
137    #[test]
138    fn batch_size_zero_is_accepted_as_no_batching_sentinel() {
139        let config =
140            MysqlSourceConfig::new("mysql://localhost/test", "SELECT 1").with_batch_size(0);
141        assert_eq!(config.batch_size, 0);
142        assert!(faucet_core::validate_batch_size(config.batch_size).is_ok());
143    }
144
145    #[test]
146    fn batch_size_above_max_is_rejected_by_validate_batch_size() {
147        let config = MysqlSourceConfig::new("mysql://localhost/test", "SELECT 1")
148            .with_batch_size(faucet_core::MAX_BATCH_SIZE + 1);
149        assert!(faucet_core::validate_batch_size(config.batch_size).is_err());
150    }
151
152    #[test]
153    fn batch_size_deserializes_from_json() {
154        let json = r#"{
155            "connection_url": "mysql://localhost/test",
156            "query": "SELECT 1",
157            "batch_size": 250
158        }"#;
159        let config: MysqlSourceConfig = serde_json::from_str(json).unwrap();
160        assert_eq!(config.batch_size, 250);
161    }
162}