Skip to main content

laminar_connectors/postgres/
mod.rs

1//! PostgreSQL connector-specific configuration and implementations.
2
3#[cfg(feature = "postgres-cdc")]
4pub mod cdc;
5#[cfg(feature = "postgres-cdc")]
6pub mod lookup;
7#[cfg(feature = "postgres-cdc")]
8pub mod reference;
9#[cfg(feature = "postgres-sink")]
10pub mod sink;
11#[cfg(feature = "postgres-sink")]
12pub mod sink_config;
13#[cfg(feature = "postgres-sink")]
14pub mod sink_metrics;
15mod tls;
16#[cfg(feature = "postgres-sink")]
17pub mod types;
18
19pub(crate) use tls::make_rustls_connector;
20/// `PostgreSQL` connection security policy.
21pub use tls::SslMode;
22
23#[cfg(feature = "postgres-cdc")]
24pub use cdc::{register_postgres_cdc_source, Lsn, PostgresCdcConfig, PostgresCdcSource};
25#[cfg(feature = "postgres-cdc")]
26pub use lookup::{PostgresLookupSource, PostgresLookupSourceConfig};
27#[cfg(feature = "postgres-cdc")]
28pub use reference::PostgresReferenceTableSource;
29
30// Re-export primary sink types at module level.
31#[cfg(feature = "postgres-sink")]
32pub use sink::PostgresSink;
33#[cfg(feature = "postgres-sink")]
34pub use sink_config::{PostgresSinkConfig, WriteMode};
35#[cfg(feature = "postgres-sink")]
36pub use sink_metrics::PostgresSinkMetrics;
37
38#[cfg(feature = "postgres-cdc")]
39use std::future::Future;
40#[cfg(feature = "postgres-sink")]
41use std::sync::Arc;
42
43#[cfg(feature = "postgres-sink")]
44use crate::config::{ConfigKeySpec, ConnectorInfo};
45#[cfg(feature = "postgres-sink")]
46use crate::registry::ConnectorRegistry;
47
48/// Poll a driver future from a task whose lifetime is independent of its caller.
49/// Dropping the waiter detaches the task; it does not cancel the driver operation.
50#[cfg(feature = "postgres-cdc")]
51async fn await_owned_driver<T, E>(
52    future: impl Future<Output = Result<T, E>> + Send + 'static,
53    join_error: impl FnOnce(tokio::task::JoinError) -> E + Send + 'static,
54) -> Result<T, E>
55where
56    T: Send + 'static,
57    E: Send + 'static,
58{
59    match tokio::spawn(future).await {
60        Ok(result) => result,
61        Err(error) => Err(join_error(error)),
62    }
63}
64
65/// Registers the `PostgreSQL` sink connector with the given registry.
66///
67/// # Errors
68///
69/// Returns an error if the connector name is already registered or the registry is frozen.
70#[cfg(feature = "postgres-sink")]
71pub fn register_postgres_sink(
72    registry: &ConnectorRegistry,
73) -> Result<(), crate::error::ConnectorError> {
74    let info = ConnectorInfo {
75        name: "postgres-sink".to_string(),
76        display_name: "PostgreSQL Sink".to_string(),
77        version: env!("CARGO_PKG_VERSION").to_string(),
78        is_source: false,
79        is_sink: true,
80        config_keys: postgres_sink_config_keys(),
81    };
82
83    registry.register_sink(
84        "postgres-sink",
85        info,
86        Arc::new(|config, registry: Option<&Arc<prometheus::Registry>>| {
87            Ok(Box::new(PostgresSink::from_connector_config(
88                config,
89                registry.map(Arc::as_ref),
90            )?))
91        }),
92    )
93}
94
95#[cfg(feature = "postgres-sink")]
96fn postgres_sink_config_keys() -> Vec<ConfigKeySpec> {
97    vec![
98        ConfigKeySpec::required("hostname", "PostgreSQL server hostname"),
99        ConfigKeySpec::required("database", "Target database name"),
100        ConfigKeySpec::required("username", "Authentication username"),
101        ConfigKeySpec::required("table.name", "Target table name"),
102        ConfigKeySpec::optional("password", "Authentication password", ""),
103        ConfigKeySpec::optional("port", "PostgreSQL port", "5432"),
104        ConfigKeySpec::optional("schema.name", "Target schema name", "public"),
105        ConfigKeySpec::optional(
106            "write.mode",
107            "Write mode: 'append' (COPY BINARY) or 'upsert' (ON CONFLICT)",
108            "append",
109        ),
110        ConfigKeySpec::optional(
111            "primary.key",
112            "Comma-separated primary key columns (required for upsert mode)",
113            "",
114        ),
115        ConfigKeySpec::optional("flush.interval.ms", "Max time before flush (ms)", "250"),
116        ConfigKeySpec::optional("connect.timeout.ms", "Connection timeout (ms)", "10000"),
117        ConfigKeySpec::optional("statement.timeout.ms", "Statement timeout (ms)", "30000"),
118        ConfigKeySpec::optional(
119            "ssl.mode",
120            "Connection security: verify-full or explicit disable",
121            "verify-full",
122        ),
123        ConfigKeySpec::optional(
124            "ssl.ca.cert.path",
125            "PEM file with trusted CA certificates; defaults to webpki roots",
126            "",
127        ),
128        ConfigKeySpec::optional(
129            "auto.create.table",
130            "Create target table from Arrow schema if missing",
131            "false",
132        ),
133        ConfigKeySpec::optional(
134            "changelog.mode",
135            "Handle Z-set records (split INSERT/DELETE by _op)",
136            "false",
137        ),
138    ]
139}
140
141#[cfg(all(test, feature = "postgres-cdc"))]
142mod driver_tests {
143    use super::await_owned_driver;
144
145    #[tokio::test]
146    async fn owned_driver_outlives_a_cancelled_waiter() {
147        let (started_tx, started_rx) = tokio::sync::oneshot::channel();
148        let (release_tx, release_rx) = tokio::sync::oneshot::channel();
149        let (completed_tx, completed_rx) = tokio::sync::oneshot::channel();
150
151        let waiter = tokio::spawn(async move {
152            await_owned_driver(
153                async move {
154                    let _ = started_tx.send(());
155                    let _ = release_rx.await;
156                    let _ = completed_tx.send(());
157                    Ok::<(), ()>(())
158                },
159                |_| (),
160            )
161            .await
162        });
163
164        started_rx.await.expect("owned task started");
165        waiter.abort();
166        assert!(waiter
167            .await
168            .expect_err("waiter must be cancelled")
169            .is_cancelled());
170        release_tx
171            .send(())
172            .expect("owned task still receives release");
173        tokio::time::timeout(std::time::Duration::from_secs(1), completed_rx)
174            .await
175            .expect("owned task must finish after waiter cancellation")
176            .expect("completion signal");
177    }
178}
179
180#[cfg(all(test, feature = "postgres-sink"))]
181mod tests {
182    use super::*;
183    use arrow_schema::{DataType, Field, Schema, SchemaRef};
184
185    fn base_factory_config() -> crate::config::ConnectorConfig {
186        let mut config = crate::config::ConnectorConfig::new("postgres-sink");
187        config.set("hostname", "localhost");
188        config.set("database", "analytics");
189        config.set("username", "writer");
190        config.set("table.name", "events");
191        config
192    }
193
194    fn factory_config(schema: &SchemaRef) -> crate::config::ConnectorConfig {
195        let mut config = base_factory_config();
196        config.set(
197            "_arrow_schema",
198            crate::config::encode_arrow_schema_ipc(schema.as_ref()),
199        );
200        config
201    }
202
203    #[test]
204    fn test_register_postgres_sink() {
205        let registry = ConnectorRegistry::new();
206        register_postgres_sink(&registry).unwrap();
207
208        let info = registry.sink_info("postgres-sink");
209        assert!(info.is_some());
210        let info = info.unwrap();
211        assert_eq!(info.name, "postgres-sink");
212        assert!(info.is_sink);
213        assert!(!info.is_source);
214        assert!(!info.config_keys.is_empty());
215    }
216
217    #[test]
218    fn test_config_keys_required() {
219        let keys = postgres_sink_config_keys();
220        let required: Vec<&str> = keys
221            .iter()
222            .filter(|k| k.required)
223            .map(|k| k.key.as_str())
224            .collect();
225        assert!(required.contains(&"hostname"));
226        assert!(required.contains(&"database"));
227        assert!(required.contains(&"username"));
228        assert!(required.contains(&"table.name"));
229    }
230
231    #[test]
232    fn test_config_keys_optional_present() {
233        let keys = postgres_sink_config_keys();
234        let optional: Vec<&str> = keys
235            .iter()
236            .filter(|k| !k.required)
237            .map(|k| k.key.as_str())
238            .collect();
239        assert!(optional.contains(&"port"));
240        assert!(optional.contains(&"write.mode"));
241        assert!(optional.contains(&"primary.key"));
242        assert!(!optional.contains(&"batch.size"));
243        assert!(!optional.contains(&"pool.size"));
244        assert!(!optional.contains(&"delivery.guarantee"));
245        assert!(optional.contains(&"changelog.mode"));
246        assert!(optional.contains(&"ssl.mode"));
247        assert!(optional.contains(&"ssl.ca.cert.path"));
248        assert!(optional.contains(&"statement.timeout.ms"));
249    }
250
251    #[test]
252    fn test_factory_creates_sink() {
253        let registry = ConnectorRegistry::new();
254        register_postgres_sink(&registry).unwrap();
255
256        let schema = Arc::new(Schema::new(vec![
257            Field::new("tenant", DataType::Utf8, false),
258            Field::new("sequence", DataType::Int64, false),
259            Field::new("enabled", DataType::Boolean, true),
260        ]));
261        let sink = registry
262            .create_sink(&factory_config(&schema), None)
263            .unwrap();
264        assert_eq!(sink.schema(), schema);
265    }
266
267    #[test]
268    fn test_factory_rejects_missing_or_malformed_arrow_schema() {
269        let registry = ConnectorRegistry::new();
270        register_postgres_sink(&registry).unwrap();
271
272        let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int64, false)]));
273        let missing = base_factory_config();
274        let missing_error = registry
275            .create_sink(&missing, None)
276            .err()
277            .expect("missing schema must fail")
278            .to_string();
279        assert!(missing_error.contains("_arrow_schema"), "{missing_error}");
280
281        let mut malformed = factory_config(&schema);
282        malformed.set("_arrow_schema", "not-arrow-ipc");
283        let malformed_error = registry
284            .create_sink(&malformed, None)
285            .err()
286            .expect("malformed schema must fail")
287            .to_string();
288        assert!(
289            malformed_error.contains("invalid") && malformed_error.contains("_arrow_schema"),
290            "{malformed_error}"
291        );
292    }
293}