1use faucet_core::DEFAULT_BATCH_SIZE;
4use schemars::JsonSchema;
5use serde::{Deserialize, Serialize};
6
7#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
9#[serde(rename_all = "snake_case")]
10pub enum PostgresColumnMapping {
11 Jsonb { column: String },
14 AutoMap,
17}
18
19impl Default for PostgresColumnMapping {
20 fn default() -> Self {
21 Self::Jsonb {
22 column: "data".into(),
23 }
24 }
25}
26
27#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize, JsonSchema)]
29#[serde(rename_all = "snake_case")]
30pub enum PostgresWriteMethod {
31 #[default]
34 Insert,
35 Copy,
48}
49
50#[derive(Clone, Serialize, Deserialize, JsonSchema)]
52pub struct PostgresSinkConfig {
53 pub connection_url: String,
55 pub table_name: String,
57 #[serde(default)]
66 pub schema: Option<String>,
67 #[serde(default)]
70 pub column_mapping: PostgresColumnMapping,
71 #[serde(default = "default_batch_size")]
92 pub batch_size: usize,
93 #[serde(default = "default_max_connections")]
99 pub max_connections: u32,
100 #[serde(flatten)]
104 pub write: faucet_core::WriteSpec,
105 #[serde(default)]
111 pub write_method: PostgresWriteMethod,
112 #[serde(default, skip_serializing_if = "Option::is_none")]
116 pub scope: Option<faucet_core::OverwriteScope>,
117}
118
119fn default_batch_size() -> usize {
120 DEFAULT_BATCH_SIZE
121}
122
123fn default_max_connections() -> u32 {
124 5
125}
126
127impl std::fmt::Debug for PostgresSinkConfig {
128 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
129 f.debug_struct("PostgresSinkConfig")
130 .field("connection_url", &"***")
131 .field("table_name", &self.table_name)
132 .field("schema", &self.schema)
133 .field("column_mapping", &self.column_mapping)
134 .field("batch_size", &self.batch_size)
135 .field("max_connections", &self.max_connections)
136 .field("write_method", &self.write_method)
137 .finish()
138 }
139}
140
141impl PostgresSinkConfig {
142 pub fn new(connection_url: impl Into<String>, table_name: impl Into<String>) -> Self {
144 Self {
145 connection_url: connection_url.into(),
146 table_name: table_name.into(),
147 schema: None,
148 column_mapping: PostgresColumnMapping::default(),
149 batch_size: DEFAULT_BATCH_SIZE,
150 max_connections: 5,
151 write: faucet_core::WriteSpec::default(),
152 write_method: PostgresWriteMethod::default(),
153 scope: None,
154 }
155 }
156
157 pub fn with_schema(mut self, schema: impl Into<String>) -> Self {
160 self.schema = Some(schema.into());
161 self
162 }
163
164 pub fn column_mapping(mut self, mapping: PostgresColumnMapping) -> Self {
166 self.column_mapping = mapping;
167 self
168 }
169
170 pub fn with_batch_size(mut self, batch_size: usize) -> Self {
176 self.batch_size = batch_size;
177 self
178 }
179
180 pub fn max_connections(mut self, n: u32) -> Self {
182 self.max_connections = n;
183 self
184 }
185
186 pub fn with_write_method(mut self, method: PostgresWriteMethod) -> Self {
189 self.write_method = method;
190 self
191 }
192}
193
194#[cfg(test)]
195mod tests {
196 use super::*;
197
198 #[test]
199 fn default_config() {
200 let config = PostgresSinkConfig::new("postgres://localhost/test", "events");
201 assert_eq!(config.table_name, "events");
202 assert_eq!(config.batch_size, DEFAULT_BATCH_SIZE);
203 assert!(matches!(
204 config.column_mapping,
205 PostgresColumnMapping::Jsonb { ref column } if column == "data"
206 ));
207 }
208
209 #[test]
210 fn builder_methods() {
211 let config = PostgresSinkConfig::new("postgres://localhost/test", "events")
212 .column_mapping(PostgresColumnMapping::AutoMap)
213 .with_batch_size(100);
214 assert_eq!(config.batch_size, 100);
215 assert!(matches!(
216 config.column_mapping,
217 PostgresColumnMapping::AutoMap
218 ));
219 }
220
221 #[test]
222 fn jsonb_custom_column() {
223 let config = PostgresSinkConfig::new("postgres://localhost/test", "events").column_mapping(
224 PostgresColumnMapping::Jsonb {
225 column: "payload".into(),
226 },
227 );
228 assert!(matches!(
229 config.column_mapping,
230 PostgresColumnMapping::Jsonb { ref column } if column == "payload"
231 ));
232 }
233
234 #[test]
235 fn with_batch_size_overrides_default() {
236 let config =
237 PostgresSinkConfig::new("postgres://localhost/test", "events").with_batch_size(250);
238 assert_eq!(config.batch_size, 250);
239 }
240
241 #[test]
242 fn batch_size_zero_is_accepted_as_no_batching_sentinel() {
243 let config =
244 PostgresSinkConfig::new("postgres://localhost/test", "events").with_batch_size(0);
245 assert_eq!(config.batch_size, 0);
246 assert!(faucet_core::validate_batch_size(config.batch_size).is_ok());
247 }
248
249 #[test]
250 fn batch_size_above_max_is_rejected_by_validate_batch_size() {
251 let config = PostgresSinkConfig::new("postgres://localhost/test", "events")
252 .with_batch_size(faucet_core::MAX_BATCH_SIZE + 1);
253 assert!(faucet_core::validate_batch_size(config.batch_size).is_err());
254 }
255
256 #[test]
257 fn batch_size_deserializes_from_json() {
258 let json = r#"{
259 "connection_url": "postgres://localhost/test",
260 "table_name": "events",
261 "column_mapping": {"jsonb": {"column": "data"}},
262 "batch_size": 250,
263 "max_connections": 5
264 }"#;
265 let config: PostgresSinkConfig = serde_json::from_str(json).unwrap();
266 assert_eq!(config.batch_size, 250);
267 }
268
269 #[test]
270 fn batch_size_defaults_when_absent_in_json() {
271 let json = r#"{
272 "connection_url": "postgres://localhost/test",
273 "table_name": "events",
274 "column_mapping": {"jsonb": {"column": "data"}},
275 "max_connections": 5
276 }"#;
277 let config: PostgresSinkConfig = serde_json::from_str(json).unwrap();
278 assert_eq!(config.batch_size, DEFAULT_BATCH_SIZE);
279 }
280
281 #[test]
282 fn config_builder_chaining() {
283 let config = PostgresSinkConfig::new("postgres://localhost/test", "events")
284 .with_batch_size(100)
285 .with_batch_size(250);
286 assert_eq!(config.batch_size, 250);
287 }
288
289 #[test]
290 fn max_connections_and_column_mapping_default_when_absent_in_json() {
291 let json = r#"{
295 "connection_url": "postgres://localhost/test",
296 "table_name": "events"
297 }"#;
298 let config: PostgresSinkConfig = serde_json::from_str(json).unwrap();
299 assert_eq!(config.max_connections, 5);
300 assert!(matches!(
301 config.column_mapping,
302 PostgresColumnMapping::Jsonb { .. }
303 ));
304 assert_eq!(config.batch_size, DEFAULT_BATCH_SIZE);
305 }
306
307 #[test]
308 fn max_connections_deserializes_when_present() {
309 let json = r#"{
310 "connection_url": "postgres://localhost/test",
311 "table_name": "events",
312 "max_connections": 20
313 }"#;
314 let config: PostgresSinkConfig = serde_json::from_str(json).unwrap();
315 assert_eq!(config.max_connections, 20);
316 }
317
318 #[test]
319 fn write_method_defaults_to_insert() {
320 let config = PostgresSinkConfig::new("postgres://localhost/test", "events");
321 assert_eq!(config.write_method, PostgresWriteMethod::Insert);
322
323 let json = r#"{
325 "connection_url": "postgres://localhost/test",
326 "table_name": "events"
327 }"#;
328 let config: PostgresSinkConfig = serde_json::from_str(json).unwrap();
329 assert_eq!(config.write_method, PostgresWriteMethod::Insert);
330 }
331
332 #[test]
333 fn write_method_serde_round_trips() {
334 let json = r#"{
335 "connection_url": "postgres://localhost/test",
336 "table_name": "events",
337 "write_method": "copy"
338 }"#;
339 let config: PostgresSinkConfig = serde_json::from_str(json).unwrap();
340 assert_eq!(config.write_method, PostgresWriteMethod::Copy);
341 let text = serde_json::to_string(&config).unwrap();
342 assert!(text.contains("\"write_method\":\"copy\""));
343 }
344
345 #[test]
346 fn with_write_method_builder() {
347 let config = PostgresSinkConfig::new("postgres://localhost/test", "events")
348 .with_write_method(PostgresWriteMethod::Copy);
349 assert_eq!(config.write_method, PostgresWriteMethod::Copy);
350 }
351}