Skip to main content

faucet_source_sqlite/
config.rs

1//! SQLite source configuration.
2
3use faucet_core::DEFAULT_BATCH_SIZE;
4use schemars::JsonSchema;
5use serde::{Deserialize, Serialize};
6
7/// Configuration for the SQLite query source.
8#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
9pub struct SqliteSourceConfig {
10    /// SQLite database URL (file path or `sqlite::memory:`).
11    pub database_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    ///
36    /// Note: sharding a SQLite source across cluster workers requires every
37    /// worker to reach the same database file (e.g. a shared volume).
38    #[serde(default, skip_serializing_if = "Option::is_none")]
39    pub shard: Option<ShardConfig>,
40}
41
42/// Primary-key range sharding settings for the SQLite source.
43///
44/// The source is split by contiguous ranges of an **integer-typed** column:
45/// each shard runs `SELECT * FROM (<query>) WHERE "key" >= lo AND "key" < hi`.
46/// The column must be present in the query's output and orderable as a 64-bit
47/// integer (e.g. an `INTEGER PRIMARY KEY` rowid alias).
48///
49/// **Nullable keys:** if the key column contains NULLs, those rows are not
50/// visible to the `MIN`/`MAX` range enumeration. They are still read — exactly
51/// one shard (the last) additionally matches `"key" IS NULL`, so NULL-key rows
52/// are covered by precisely one shard with no loss and no duplication.
53#[derive(Clone, Debug, Serialize, Deserialize, JsonSchema, PartialEq, Eq)]
54pub struct ShardConfig {
55    /// Integer column to range-partition on. Quoted as an identifier before use,
56    /// so it is safe against injection but must name a real output column.
57    pub key: String,
58}
59
60fn default_max_connections() -> u32 {
61    10
62}
63
64fn default_batch_size() -> usize {
65    DEFAULT_BATCH_SIZE
66}
67
68impl SqliteSourceConfig {
69    /// Create a new config with the required database URL and query.
70    pub fn new(database_url: impl Into<String>, query: impl Into<String>) -> Self {
71        Self {
72            database_url: database_url.into(),
73            query: query.into(),
74            max_connections: 10,
75            batch_size: DEFAULT_BATCH_SIZE,
76            shard: None,
77        }
78    }
79
80    /// Set the maximum number of connections in the pool.
81    pub fn with_max_connections(mut self, max_connections: u32) -> Self {
82        self.max_connections = max_connections;
83        self
84    }
85
86    /// Set the per-page row count for [`Source::stream_pages`](faucet_core::Source::stream_pages).
87    ///
88    /// Pass `0` to opt out of batching — the entire result set is emitted in
89    /// a single [`StreamPage`](faucet_core::StreamPage).
90    pub fn with_batch_size(mut self, batch_size: usize) -> Self {
91        self.batch_size = batch_size;
92        self
93    }
94}
95
96#[cfg(test)]
97mod tests {
98    use super::*;
99
100    #[test]
101    fn default_config() {
102        let config = SqliteSourceConfig::new("sqlite:test.db", "SELECT * FROM events");
103        assert_eq!(config.database_url, "sqlite:test.db");
104        assert_eq!(config.query, "SELECT * FROM events");
105    }
106
107    #[test]
108    fn memory_database() {
109        let config = SqliteSourceConfig::new("sqlite::memory:", "SELECT 1");
110        assert_eq!(config.database_url, "sqlite::memory:");
111    }
112
113    #[test]
114    fn batch_size_defaults_to_default_batch_size() {
115        let config = SqliteSourceConfig::new("sqlite::memory:", "SELECT 1");
116        assert_eq!(config.batch_size, faucet_core::DEFAULT_BATCH_SIZE);
117    }
118
119    #[test]
120    fn with_batch_size_overrides_default() {
121        let config = SqliteSourceConfig::new("sqlite::memory:", "SELECT 1").with_batch_size(500);
122        assert_eq!(config.batch_size, 500);
123    }
124
125    #[test]
126    fn batch_size_zero_is_accepted_as_no_batching_sentinel() {
127        let config = SqliteSourceConfig::new("sqlite::memory:", "SELECT 1").with_batch_size(0);
128        assert_eq!(config.batch_size, 0);
129        assert!(faucet_core::validate_batch_size(config.batch_size).is_ok());
130    }
131
132    #[test]
133    fn batch_size_above_max_is_rejected_by_validate_batch_size() {
134        let config = SqliteSourceConfig::new("sqlite::memory:", "SELECT 1")
135            .with_batch_size(faucet_core::MAX_BATCH_SIZE + 1);
136        assert!(faucet_core::validate_batch_size(config.batch_size).is_err());
137    }
138
139    #[test]
140    fn batch_size_deserializes_from_json() {
141        let json = r#"{
142            "database_url": "sqlite::memory:",
143            "query": "SELECT 1",
144            "batch_size": 250
145        }"#;
146        let config: SqliteSourceConfig = serde_json::from_str(json).unwrap();
147        assert_eq!(config.batch_size, 250);
148    }
149}