faucet_source_postgres/
config.rs

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

faucet_source_postgres/
config.rs
