Skip to main content

laminar_connectors/postgres/
sink_config.rs

1//! `PostgreSQL` sink connector configuration.
2//!
3//! [`PostgresSinkConfig`] encapsulates all settings for writing Arrow
4//! `RecordBatch` data to `PostgreSQL`, parsed from SQL `WITH (...)` clauses.
5
6use std::path::PathBuf;
7use std::time::Duration;
8
9use crate::config::ConnectorConfig;
10use crate::error::ConnectorError;
11use crate::postgres::SslMode;
12
13const MAX_POSTGRES_IDENTIFIER_BYTES: usize = 63;
14const MAX_POSTGRES_STATEMENT_TIMEOUT_MS: u128 = 2_147_483_647;
15
16const REMOVED_CONFIG_KEYS: &[&str] = &["batch.size", "pool.size"];
17
18const ALLOWED_CONFIG_KEYS: &[&str] = &[
19    "_arrow_schema",
20    "auto.create.table",
21    "changelog.mode",
22    "connect.timeout.ms",
23    "database",
24    "flush.interval.ms",
25    "hostname",
26    "password",
27    "port",
28    "primary.key",
29    "schema.name",
30    "ssl.ca.cert.path",
31    "ssl.mode",
32    "statement.timeout.ms",
33    "table.name",
34    "username",
35    "write.mode",
36];
37
38/// Configuration for the `PostgreSQL` sink connector.
39///
40/// Parsed from SQL `WITH (...)` clause options via [`from_config`](Self::from_config).
41#[derive(Debug, Clone)]
42pub struct PostgresSinkConfig {
43    /// `PostgreSQL` hostname.
44    pub hostname: String,
45
46    /// `PostgreSQL` port (default: 5432).
47    pub port: u16,
48
49    /// Database name.
50    pub database: String,
51
52    /// Username for authentication.
53    pub username: String,
54
55    /// Password for authentication.
56    pub password: String,
57
58    /// Target schema name (default: `"public"`).
59    pub schema_name: String,
60
61    /// Target table name.
62    pub table_name: String,
63
64    /// Write mode: append (COPY BINARY) or upsert (ON CONFLICT).
65    pub write_mode: WriteMode,
66
67    /// Primary key columns (required for upsert mode).
68    pub primary_key_columns: Vec<String>,
69
70    /// Maximum time to buffer before flushing.
71    pub flush_interval: Duration,
72
73    /// Connection timeout.
74    pub connect_timeout: Duration,
75
76    /// SSL mode for connections.
77    pub ssl_mode: SslMode,
78
79    /// Optional PEM file containing trusted CA certificates.
80    pub ssl_ca_cert_path: Option<PathBuf>,
81
82    /// Whether to create the target table if it doesn't exist.
83    pub auto_create_table: bool,
84
85    /// Whether to handle changelog/retraction records.
86    pub changelog_mode: bool,
87
88    /// Per-query statement timeout (default: 30s).
89    pub statement_timeout: Duration,
90}
91
92impl Default for PostgresSinkConfig {
93    fn default() -> Self {
94        Self {
95            hostname: "localhost".to_string(),
96            port: 5432,
97            database: String::new(),
98            username: "postgres".to_string(),
99            password: String::new(),
100            schema_name: "public".to_string(),
101            table_name: String::new(),
102            write_mode: WriteMode::Append,
103            primary_key_columns: Vec::new(),
104            flush_interval: Duration::from_millis(250),
105            connect_timeout: Duration::from_secs(10),
106            ssl_mode: SslMode::VerifyFull,
107            ssl_ca_cert_path: None,
108            auto_create_table: false,
109            changelog_mode: false,
110            statement_timeout: Duration::from_secs(30),
111        }
112    }
113}
114
115impl PostgresSinkConfig {
116    /// Creates a minimal config for testing.
117    #[must_use]
118    pub fn new(hostname: &str, database: &str, table_name: &str) -> Self {
119        Self {
120            hostname: hostname.to_string(),
121            database: database.to_string(),
122            table_name: table_name.to_string(),
123            ..Default::default()
124        }
125    }
126
127    /// Parses a sink config from a [`ConnectorConfig`] (SQL WITH clause).
128    ///
129    /// # Required keys
130    ///
131    /// - `hostname` - `PostgreSQL` server hostname
132    /// - `database` - Target database name
133    /// - `username` - Authentication username
134    /// - `table.name` - Target table name
135    ///
136    /// # Errors
137    ///
138    /// Returns `ConnectorError::MissingConfig` if required keys are absent,
139    /// or `ConnectorError::ConfigurationError` on invalid values.
140    #[allow(clippy::field_reassign_with_default)]
141    pub fn from_config(config: &ConnectorConfig) -> Result<Self, ConnectorError> {
142        Self::reject_removed_keys(config)?;
143        config.reject_unknown_properties(ALLOWED_CONFIG_KEYS, "PostgreSQL sink")?;
144
145        let mut cfg = Self::default();
146
147        cfg.hostname = config.require("hostname")?.to_string();
148        cfg.database = config.require("database")?.to_string();
149        cfg.username = config.require("username")?.to_string();
150        cfg.table_name = config.require("table.name")?.to_string();
151
152        if let Some(v) = config.get("password") {
153            cfg.password = v.to_string();
154        }
155        if let Some(v) = config.get("port") {
156            cfg.port = crate::config::parse_port(v)?;
157        }
158        if let Some(v) = config.get("schema.name") {
159            cfg.schema_name = v.to_string();
160        }
161        if let Some(v) = config.get("write.mode") {
162            cfg.write_mode = v.parse().map_err(|_| {
163                ConnectorError::ConfigurationError(format!(
164                    "invalid write.mode: '{v}' (expected 'append' or 'upsert')"
165                ))
166            })?;
167        }
168        if let Some(v) = config.get("primary.key") {
169            if v.trim().is_empty() {
170                cfg.primary_key_columns.clear();
171            } else {
172                let columns: Vec<_> = v.split(',').map(str::trim).collect();
173                if columns.iter().any(|column| column.is_empty()) {
174                    return Err(ConnectorError::ConfigurationError(
175                        "primary.key contains an empty column name".into(),
176                    ));
177                }
178                cfg.primary_key_columns = columns.into_iter().map(str::to_string).collect();
179            }
180        }
181        if let Some(v) = config.get("flush.interval.ms") {
182            let ms: u64 = v.parse().map_err(|_| {
183                ConnectorError::ConfigurationError(format!("invalid flush.interval.ms: '{v}'"))
184            })?;
185            cfg.flush_interval = Duration::from_millis(ms);
186        }
187        if let Some(v) = config.get("connect.timeout.ms") {
188            let ms: u64 = v.parse().map_err(|_| {
189                ConnectorError::ConfigurationError(format!("invalid connect.timeout.ms: '{v}'"))
190            })?;
191            cfg.connect_timeout = Duration::from_millis(ms);
192        }
193        if let Some(v) = config.get("ssl.mode") {
194            cfg.ssl_mode = v.parse().map_err(|_| {
195                ConnectorError::ConfigurationError(format!(
196                    "invalid ssl.mode: '{v}' (expected 'disable' or 'verify-full')"
197                ))
198            })?;
199        }
200        cfg.ssl_ca_cert_path = config.get("ssl.ca.cert.path").map(PathBuf::from);
201        if let Some(v) = config.get("auto.create.table") {
202            cfg.auto_create_table = v.eq_ignore_ascii_case("true");
203        }
204        if let Some(v) = config.get("changelog.mode") {
205            cfg.changelog_mode = v.eq_ignore_ascii_case("true");
206        }
207        if let Some(v) = config.get("statement.timeout.ms") {
208            let ms: u64 = v.parse().map_err(|_| {
209                ConnectorError::ConfigurationError(format!("invalid statement.timeout.ms: '{v}'"))
210            })?;
211            cfg.statement_timeout = Duration::from_millis(ms);
212        }
213        cfg.validate()?;
214        Ok(cfg)
215    }
216
217    fn reject_removed_keys(config: &ConnectorConfig) -> Result<(), ConnectorError> {
218        if let Some(key) = REMOVED_CONFIG_KEYS
219            .iter()
220            .find(|key| config.get(key).is_some())
221        {
222            return Err(ConnectorError::ConfigurationError(format!(
223                "PostgreSQL sink property '{key}' is not supported: batching is owned by the runtime and retained-byte limit"
224            )));
225        }
226        Ok(())
227    }
228
229    /// Validates the configuration for consistency.
230    ///
231    /// # Errors
232    ///
233    /// Returns `ConnectorError::ConfigurationError` on invalid combinations.
234    pub fn validate(&self) -> Result<(), ConnectorError> {
235        crate::config::require_non_empty(&self.hostname, "hostname")?;
236        crate::config::require_non_empty(&self.database, "database")?;
237        crate::config::require_non_empty(&self.username, "username")?;
238        crate::config::require_non_empty(&self.schema_name, "schema.name")?;
239        crate::config::require_non_empty(&self.table_name, "table.name")?;
240        validate_sql_identifier(&self.schema_name, "schema.name")?;
241        validate_sql_identifier(&self.table_name, "table.name")?;
242        let mut primary_keys = std::collections::HashSet::new();
243        for column in &self.primary_key_columns {
244            validate_sql_identifier(column, "primary.key column")?;
245            if !primary_keys.insert(column) {
246                return Err(ConnectorError::ConfigurationError(format!(
247                    "primary.key contains duplicate column '{column}'"
248                )));
249            }
250        }
251        if self.port == 0 {
252            return Err(ConnectorError::ConfigurationError(
253                "port must be > 0".into(),
254            ));
255        }
256        if self
257            .ssl_ca_cert_path
258            .as_ref()
259            .is_some_and(|path| path.as_os_str().is_empty())
260        {
261            return Err(ConnectorError::ConfigurationError(
262                "ssl.ca.cert.path must not be empty".into(),
263            ));
264        }
265        if self.ssl_mode == SslMode::Disable && self.ssl_ca_cert_path.is_some() {
266            return Err(ConnectorError::ConfigurationError(
267                "ssl.ca.cert.path requires ssl.mode=verify-full".into(),
268            ));
269        }
270        if self.write_mode == WriteMode::Upsert && self.primary_key_columns.is_empty() {
271            return Err(ConnectorError::ConfigurationError(
272                "upsert mode requires 'primary.key' to be set".into(),
273            ));
274        }
275        if self.changelog_mode && self.write_mode != WriteMode::Upsert {
276            return Err(ConnectorError::ConfigurationError(
277                "changelog mode requires write.mode = 'upsert'".into(),
278            ));
279        }
280        if self.flush_interval.is_zero() {
281            return Err(ConnectorError::ConfigurationError(
282                "flush.interval.ms must be > 0".into(),
283            ));
284        }
285        if self.connect_timeout.is_zero() {
286            return Err(ConnectorError::ConfigurationError(
287                "connect.timeout.ms must be > 0".into(),
288            ));
289        }
290        if self.statement_timeout < Duration::from_secs(1) {
291            return Err(ConnectorError::ConfigurationError(
292                "statement.timeout.ms must be >= 1000 (1 second)".into(),
293            ));
294        }
295        if self.statement_timeout.as_millis() > MAX_POSTGRES_STATEMENT_TIMEOUT_MS {
296            return Err(ConnectorError::ConfigurationError(format!(
297                "statement.timeout.ms must be <= {MAX_POSTGRES_STATEMENT_TIMEOUT_MS}"
298            )));
299        }
300        if !self
301            .statement_timeout
302            .subsec_nanos()
303            .is_multiple_of(1_000_000)
304        {
305            return Err(ConnectorError::ConfigurationError(
306                "statement timeout must be an integer number of milliseconds".into(),
307            ));
308        }
309        Ok(())
310    }
311
312    /// `PostgreSQL` startup option applied to every connection created by the pool.
313    #[must_use]
314    pub(super) fn statement_timeout_startup_option(&self) -> String {
315        format!(
316            "-c statement_timeout={}",
317            self.statement_timeout.as_millis()
318        )
319    }
320
321    /// Returns the safely quoted fully qualified table name (`"schema"."table"`).
322    #[must_use]
323    pub fn qualified_table_name(&self) -> String {
324        format!(
325            "{}.{}",
326            quote_sql_identifier(&self.schema_name),
327            quote_sql_identifier(&self.table_name)
328        )
329    }
330}
331
332/// Validates one `PostgreSQL` identifier segment before it can reach generated SQL.
333pub(super) fn validate_sql_identifier(identifier: &str, label: &str) -> Result<(), ConnectorError> {
334    if identifier.is_empty() {
335        return Err(ConnectorError::ConfigurationError(format!(
336            "{label} must not be empty"
337        )));
338    }
339    if identifier.contains('\0') {
340        return Err(ConnectorError::ConfigurationError(format!(
341            "{label} must not contain NUL"
342        )));
343    }
344    if identifier.len() > MAX_POSTGRES_IDENTIFIER_BYTES {
345        return Err(ConnectorError::ConfigurationError(format!(
346            "{label} exceeds PostgreSQL's {MAX_POSTGRES_IDENTIFIER_BYTES}-byte identifier limit"
347        )));
348    }
349    Ok(())
350}
351
352/// Quotes one already-validated identifier segment. `PostgreSQL` escapes embedded quotes by
353/// doubling them; dots remain literal characters within the segment.
354pub(super) fn quote_sql_identifier(identifier: &str) -> String {
355    format!("\"{}\"", identifier.replace('"', "\"\""))
356}
357
358/// Write mode for the `PostgreSQL` sink.
359#[derive(Debug, Clone, Copy, PartialEq, Eq)]
360pub enum WriteMode {
361    /// Append-only: uses COPY BINARY for maximum throughput.
362    /// No deduplication — every record is inserted.
363    Append,
364    /// Upsert: `INSERT ... ON CONFLICT DO UPDATE`.
365    /// Requires primary key columns. Deduplicates on key.
366    Upsert,
367}
368
369str_enum!(WriteMode, lowercase_nodash, String, "unknown write mode",
370    Append => "append", "copy";
371    Upsert => "upsert", "insert"
372);
373
374#[cfg(test)]
375mod tests {
376    use super::*;
377
378    fn make_config(pairs: &[(&str, &str)]) -> ConnectorConfig {
379        let mut config = ConnectorConfig::new("postgres-sink");
380        for (k, v) in pairs {
381            config.set(*k, *v);
382        }
383        config
384    }
385
386    fn required_pairs() -> Vec<(&'static str, &'static str)> {
387        vec![
388            ("hostname", "localhost"),
389            ("database", "mydb"),
390            ("username", "writer"),
391            ("table.name", "events"),
392        ]
393    }
394
395    #[test]
396    fn test_parse_required_fields() {
397        let config = make_config(&required_pairs());
398        let cfg = PostgresSinkConfig::from_config(&config).unwrap();
399        assert_eq!(cfg.hostname, "localhost");
400        assert_eq!(cfg.database, "mydb");
401        assert_eq!(cfg.username, "writer");
402        assert_eq!(cfg.table_name, "events");
403        assert_eq!(cfg.port, 5432);
404        assert_eq!(cfg.schema_name, "public");
405        assert_eq!(cfg.write_mode, WriteMode::Append);
406    }
407
408    #[test]
409    fn test_missing_hostname() {
410        let config = make_config(&[("database", "db"), ("username", "u"), ("table.name", "t")]);
411        assert!(PostgresSinkConfig::from_config(&config).is_err());
412    }
413
414    #[test]
415    fn test_missing_database() {
416        let config = make_config(&[("hostname", "h"), ("username", "u"), ("table.name", "t")]);
417        assert!(PostgresSinkConfig::from_config(&config).is_err());
418    }
419
420    #[test]
421    fn test_missing_username() {
422        let config = make_config(&[("hostname", "h"), ("database", "db"), ("table.name", "t")]);
423        assert!(PostgresSinkConfig::from_config(&config).is_err());
424    }
425
426    #[test]
427    fn test_missing_table_name() {
428        let config = make_config(&[("hostname", "h"), ("database", "db"), ("username", "u")]);
429        assert!(PostgresSinkConfig::from_config(&config).is_err());
430    }
431
432    #[test]
433    fn test_parse_all_optional_fields() {
434        let mut pairs = required_pairs();
435        pairs.extend_from_slice(&[
436            ("password", "secret"),
437            ("port", "5433"),
438            ("schema.name", "analytics"),
439            ("write.mode", "upsert"),
440            ("primary.key", "id, region"),
441            ("flush.interval.ms", "500"),
442            ("connect.timeout.ms", "5000"),
443            ("ssl.mode", "verify-full"),
444            ("ssl.ca.cert.path", "/certs/ca.pem"),
445            ("auto.create.table", "true"),
446            ("changelog.mode", "true"),
447        ]);
448        let config = make_config(&pairs);
449        let cfg = PostgresSinkConfig::from_config(&config).unwrap();
450
451        assert_eq!(cfg.password, "secret");
452        assert_eq!(cfg.port, 5433);
453        assert_eq!(cfg.schema_name, "analytics");
454        assert_eq!(cfg.write_mode, WriteMode::Upsert);
455        assert_eq!(cfg.primary_key_columns, vec!["id", "region"]);
456        assert_eq!(cfg.flush_interval, Duration::from_millis(500));
457        assert_eq!(cfg.connect_timeout, Duration::from_secs(5));
458        assert_eq!(cfg.ssl_mode, SslMode::VerifyFull);
459        assert_eq!(cfg.ssl_ca_cert_path, Some(PathBuf::from("/certs/ca.pem")));
460        assert!(cfg.auto_create_table);
461        assert!(cfg.changelog_mode);
462    }
463
464    #[test]
465    fn test_upsert_requires_primary_key() {
466        let mut pairs = required_pairs();
467        pairs.push(("write.mode", "upsert"));
468        let config = make_config(&pairs);
469        let result = PostgresSinkConfig::from_config(&config);
470        assert!(result.is_err());
471        let err = result.unwrap_err().to_string();
472        assert!(err.contains("primary.key"), "error: {err}");
473    }
474
475    #[test]
476    fn test_changelog_requires_upsert() {
477        let mut pairs = required_pairs();
478        pairs.push(("changelog.mode", "true"));
479        let error = PostgresSinkConfig::from_config(&make_config(&pairs)).unwrap_err();
480        assert!(error.to_string().contains("write.mode = 'upsert'"));
481    }
482
483    #[test]
484    fn removed_buffer_and_pool_options_are_rejected() {
485        for key in REMOVED_CONFIG_KEYS {
486            let mut pairs = required_pairs();
487            pairs.push((*key, "1"));
488            let error = PostgresSinkConfig::from_config(&make_config(&pairs)).unwrap_err();
489            assert!(error.to_string().contains(key));
490        }
491    }
492
493    #[test]
494    fn test_qualified_table_name() {
495        let cfg = PostgresSinkConfig::new("localhost", "db", "events");
496        assert_eq!(cfg.qualified_table_name(), "\"public\".\"events\"");
497
498        let mut cfg2 = cfg;
499        cfg2.schema_name = "analytics".to_string();
500        cfg2.table_name = "order\"items".to_string();
501        assert_eq!(
502            cfg2.qualified_table_name(),
503            "\"analytics\".\"order\"\"items\""
504        );
505    }
506
507    #[test]
508    fn identifiers_fail_closed_before_sql_generation() {
509        for invalid in ["", "bad\0name"] {
510            let mut cfg = PostgresSinkConfig::new("localhost", "db", invalid);
511            assert!(cfg.validate().is_err(), "table identifier {invalid:?}");
512            cfg.table_name = "events".into();
513            cfg.schema_name = invalid.into();
514            assert!(cfg.validate().is_err(), "schema identifier {invalid:?}");
515        }
516
517        let mut cfg = PostgresSinkConfig::new("localhost", "db", "events");
518        cfg.write_mode = WriteMode::Upsert;
519        cfg.primary_key_columns = vec!["id".into(), "id".into()];
520        assert!(cfg
521            .validate()
522            .unwrap_err()
523            .to_string()
524            .contains("duplicate"));
525    }
526
527    #[test]
528    fn test_defaults() {
529        let cfg = PostgresSinkConfig::default();
530        assert_eq!(cfg.hostname, "localhost");
531        assert_eq!(cfg.port, 5432);
532        assert_eq!(cfg.schema_name, "public");
533        assert_eq!(cfg.write_mode, WriteMode::Append);
534        assert_eq!(cfg.flush_interval, Duration::from_millis(250));
535        assert_eq!(cfg.ssl_mode, SslMode::VerifyFull);
536        assert!(!cfg.auto_create_table);
537        assert!(!cfg.changelog_mode);
538    }
539
540    #[test]
541    fn test_write_mode_parse() {
542        assert_eq!("append".parse::<WriteMode>().unwrap(), WriteMode::Append);
543        assert_eq!("copy".parse::<WriteMode>().unwrap(), WriteMode::Append);
544        assert_eq!("upsert".parse::<WriteMode>().unwrap(), WriteMode::Upsert);
545        assert_eq!("insert".parse::<WriteMode>().unwrap(), WriteMode::Upsert);
546        assert!("unknown".parse::<WriteMode>().is_err());
547    }
548
549    #[test]
550    fn test_write_mode_display() {
551        assert_eq!(WriteMode::Append.to_string(), "append");
552        assert_eq!(WriteMode::Upsert.to_string(), "upsert");
553    }
554
555    #[test]
556    fn legacy_ssl_modes_are_rejected() {
557        for mode in ["off", "prefer", "require", "verify-ca"] {
558            let mut pairs = required_pairs();
559            pairs.push(("ssl.mode", mode));
560            let error = PostgresSinkConfig::from_config(&make_config(&pairs)).unwrap_err();
561            let message = error.to_string();
562            assert!(message.contains("invalid ssl.mode"), "{message}");
563        }
564    }
565
566    #[test]
567    fn plaintext_rejects_unused_ca_path() {
568        let mut pairs = required_pairs();
569        pairs.extend_from_slice(&[
570            ("ssl.mode", "disable"),
571            ("ssl.ca.cert.path", "/certs/ca.pem"),
572        ]);
573        let error = PostgresSinkConfig::from_config(&make_config(&pairs)).unwrap_err();
574        assert!(error.to_string().contains("ssl.mode=verify-full"));
575    }
576
577    #[test]
578    fn test_invalid_port() {
579        let mut pairs = required_pairs();
580        pairs.push(("port", "not_a_number"));
581        let config = make_config(&pairs);
582        assert!(PostgresSinkConfig::from_config(&config).is_err());
583    }
584
585    #[test]
586    fn test_zero_flush_interval_rejected() {
587        let mut pairs = required_pairs();
588        pairs.push(("flush.interval.ms", "0"));
589        let error = PostgresSinkConfig::from_config(&make_config(&pairs)).unwrap_err();
590        assert!(error.to_string().contains("flush.interval.ms must be > 0"));
591    }
592
593    #[test]
594    fn zero_connect_timeout_is_rejected() {
595        let mut pairs = required_pairs();
596        pairs.push(("connect.timeout.ms", "0"));
597        let error = PostgresSinkConfig::from_config(&make_config(&pairs)).unwrap_err();
598        assert!(error.to_string().contains("connect.timeout.ms must be > 0"));
599    }
600
601    #[test]
602    fn statement_timeout_is_a_bounded_integer_startup_setting() {
603        let mut cfg = PostgresSinkConfig::new("localhost", "db", "events");
604        cfg.statement_timeout = Duration::from_millis(30_000);
605        cfg.validate().unwrap();
606        assert_eq!(
607            cfg.statement_timeout_startup_option(),
608            "-c statement_timeout=30000"
609        );
610
611        cfg.statement_timeout =
612            Duration::from_millis(u64::try_from(MAX_POSTGRES_STATEMENT_TIMEOUT_MS).unwrap());
613        cfg.validate().unwrap();
614        cfg.statement_timeout =
615            Duration::from_millis(u64::try_from(MAX_POSTGRES_STATEMENT_TIMEOUT_MS + 1).unwrap());
616        assert!(cfg.validate().is_err());
617
618        cfg.statement_timeout = Duration::from_secs(1) + Duration::from_micros(1);
619        assert!(cfg.validate().is_err());
620    }
621}