Skip to main content

laminar_connectors/postgres/sink/
mod.rs

1//! `PostgreSQL` sink connector implementation.
2//!
3//! [`PostgresSink`] implements [`SinkConnector`], writing Arrow `RecordBatch`
4//! data to `PostgreSQL` tables via two strategies:
5//!
6//! - **Append mode**: COPY BINARY for maximum throughput (>500K rows/sec)
7//! - **Upsert mode**: `INSERT ... ON CONFLICT DO UPDATE` with UNNEST arrays
8//!
9//! The connector provides durable at-least-once delivery. Typed runtime
10//! admission excludes it from exactly-once pipelines because it has no
11//! coordinated external checkpoint committer.
12//!
13//! # Ring Architecture
14//!
15//! - **Ring 0**: No sink code. Data arrives via SPSC channel (~5ns push).
16//! - **Ring 1**: Batch buffering, COPY/INSERT writes, transaction management.
17//! - **Ring 2**: Connection pool and table creation.
18
19use std::time::Duration;
20
21use arrow_array::RecordBatch;
22use arrow_schema::SchemaRef;
23use async_trait::async_trait;
24use tracing::{debug, info};
25
26#[cfg(feature = "postgres-sink")]
27use crate::changelog::collapse_changelog;
28use crate::config::{ConnectorConfig, ConnectorState};
29use crate::connector::{
30    SinkConnector, SinkConsistency, SinkContract, SinkInputMode, SinkTopology, WriteResult,
31};
32use crate::error::ConnectorError;
33
34use super::sink_config::{PostgresSinkConfig, WriteMode};
35use super::sink_metrics::PostgresSinkMetrics;
36
37mod input;
38mod statements;
39
40use input::{build_user_schema, validate_sink_schema};
41#[cfg(any(feature = "postgres-sink", test))]
42use input::{collapse_upsert_batch, requires_preflush, retained_batch_bytes};
43#[cfg(feature = "postgres-sink")]
44use input::{
45    execute_unnest, retained_batch_bytes_u64, strip_metadata_columns, validate_changelog_input,
46    validate_input_batch,
47};
48
49#[cfg(feature = "postgres-sink")]
50fn postgres_dispatched_write_error(
51    operation: &str,
52    error: &tokio_postgres::Error,
53) -> ConnectorError {
54    classify_postgres_write_failure(operation, error, error.as_db_error().is_some())
55}
56
57#[cfg(any(feature = "postgres-sink", test))]
58fn classify_postgres_write_failure(
59    operation: &str,
60    error: &dyn std::fmt::Display,
61    server_rejected: bool,
62) -> ConnectorError {
63    if server_rejected {
64        ConnectorError::WriteError(format!("{operation}: {error}"))
65    } else {
66        ConnectorError::outcome_unknown(
67            format!(
68                "PostgreSQL {operation} failed without a server response and may have committed: {error}"
69            ),
70            true,
71        )
72    }
73}
74
75#[cfg(any(feature = "postgres-sink", test))]
76fn resolve_uncommitted_transaction_error(error: ConnectorError) -> ConnectorError {
77    match error {
78        ConnectorError::OutcomeUnknown { message, retryable } if retryable => {
79            ConnectorError::ConnectionFailed(format!(
80                "PostgreSQL transaction mutation failed before COMMIT; no commit was dispatched: {message}"
81            ))
82        }
83        ConnectorError::OutcomeUnknown { message, .. } => ConnectorError::TransactionError(
84            format!(
85                "PostgreSQL transaction mutation failed before COMMIT; no commit was dispatched: {message}"
86            ),
87        ),
88        error => error,
89    }
90}
91#[cfg(feature = "postgres-sink")]
92use super::types::{arrow_column_to_pg_array, postgres_copy_batch};
93#[cfg(feature = "postgres-sink")]
94use bytes::BytesMut;
95#[cfg(feature = "postgres-sink")]
96use deadpool_postgres::Pool;
97
98// A flush temporarily owns the buffered Arrow arrays plus a concatenated batch and either COPY
99// wire bytes or owned UNNEST parameters. Keeping the retained input at 16 MiB leaves conservative
100// headroom for those transient copies without adding another deployment tuning dimension.
101#[cfg(feature = "postgres-sink")]
102const MAX_BUFFERED_RETAINED_BYTES: usize = 16 * 1024 * 1024;
103#[cfg(feature = "postgres-sink")]
104const COPY_ENCODE_INITIAL_CAPACITY: usize = 64 * 1024;
105#[cfg(feature = "postgres-sink")]
106const SINK_POOL_SIZE: usize = 1;
107
108#[cfg(feature = "postgres-sink")]
109fn build_pool(config: &PostgresSinkConfig) -> Result<Pool, ConnectorError> {
110    config.validate()?;
111
112    let mut pool_config = deadpool_postgres::Config::new();
113    pool_config.host = Some(config.hostname.clone());
114    pool_config.port = Some(config.port);
115    pool_config.dbname = Some(config.database.clone());
116    pool_config.user = Some(config.username.clone());
117    pool_config.password = Some(config.password.clone());
118    pool_config.options = Some(config.statement_timeout_startup_option());
119    pool_config.connect_timeout = Some(config.connect_timeout);
120    pool_config.ssl_mode = Some(match config.ssl_mode {
121        crate::postgres::SslMode::Disable => deadpool_postgres::SslMode::Disable,
122        crate::postgres::SslMode::VerifyFull => deadpool_postgres::SslMode::Require,
123    });
124    let mut deadpool_config = deadpool_postgres::PoolConfig::new(SINK_POOL_SIZE);
125    deadpool_config.timeouts.wait = Some(config.connect_timeout);
126    deadpool_config.timeouts.create = Some(config.connect_timeout);
127    pool_config.pool = Some(deadpool_config);
128
129    let runtime = Some(deadpool_postgres::Runtime::Tokio1);
130    match config.ssl_mode {
131        crate::postgres::SslMode::Disable => pool_config
132            .create_pool(runtime, tokio_postgres::NoTls)
133            .map_err(|error| {
134                ConnectorError::ConnectionFailed(format!("pool creation failed: {error}"))
135            }),
136        crate::postgres::SslMode::VerifyFull => {
137            let tls = crate::postgres::make_rustls_connector(config.ssl_ca_cert_path.as_deref())?;
138            pool_config.create_pool(runtime, tls).map_err(|error| {
139                ConnectorError::ConnectionFailed(format!("TLS pool creation failed: {error}"))
140            })
141        }
142    }
143}
144
145/// `PostgreSQL` sink connector.
146///
147/// Writes Arrow `RecordBatch` to `PostgreSQL` tables using COPY BINARY
148/// (append) or UNNEST-based upsert with durable at-least-once semantics.
149pub struct PostgresSink {
150    /// Sink configuration.
151    config: PostgresSinkConfig,
152    /// Arrow schema for input batches.
153    schema: SchemaRef,
154    /// User-visible schema (metadata columns stripped).
155    user_schema: SchemaRef,
156    /// Connector lifecycle state.
157    state: ConnectorState,
158    /// Buffered records awaiting flush.
159    buffer: Vec<RecordBatch>,
160    /// Total rows in buffer.
161    buffered_rows: usize,
162    /// Arrow-reported retained bytes held by `buffer`.
163    buffered_retained_bytes: usize,
164    /// Sink metrics.
165    metrics: PostgresSinkMetrics,
166    /// Cached upsert SQL statement (for upsert mode).
167    upsert_sql: Option<String>,
168    /// Cached COPY SQL statement (for append mode).
169    copy_sql: Option<String>,
170    /// Cached CREATE TABLE SQL (for auto-create).
171    create_table_sql: Option<String>,
172    /// Cached DELETE SQL (for changelog mode).
173    delete_sql: Option<String>,
174    /// Connection pool (`None` until `open()` is called).
175    #[cfg(feature = "postgres-sink")]
176    pool: Option<Pool>,
177}
178
179impl PostgresSink {
180    /// Creates a new `PostgreSQL` sink connector.
181    #[must_use]
182    pub fn new(
183        schema: SchemaRef,
184        config: PostgresSinkConfig,
185        registry: Option<&prometheus::Registry>,
186    ) -> Self {
187        let user_schema = build_user_schema(&schema);
188        Self {
189            config,
190            schema,
191            user_schema,
192            state: ConnectorState::Created,
193            buffer: Vec::with_capacity(4),
194            buffered_rows: 0,
195            buffered_retained_bytes: 0,
196            metrics: PostgresSinkMetrics::new(registry),
197            upsert_sql: None,
198            copy_sql: None,
199            create_table_sql: None,
200            delete_sql: None,
201            #[cfg(feature = "postgres-sink")]
202            pool: None,
203        }
204    }
205
206    /// Builds the engine-managed sink from its complete runtime configuration.
207    /// Engine-managed sinks must carry the upstream Arrow schema; using a
208    /// placeholder here would generate the wrong table and write statements.
209    pub(crate) fn from_connector_config(
210        config: &ConnectorConfig,
211        registry: Option<&prometheus::Registry>,
212    ) -> Result<Self, ConnectorError> {
213        let (sink_config, schema) = Self::decode_connector_config(config)?;
214        Ok(Self::new(schema, sink_config, registry))
215    }
216
217    fn decode_connector_config(
218        config: &ConnectorConfig,
219    ) -> Result<(PostgresSinkConfig, SchemaRef), ConnectorError> {
220        let sink_config = PostgresSinkConfig::from_config(config)?;
221        if config.get("_arrow_schema").is_none() {
222            return Err(ConnectorError::ConfigurationError(
223                "PostgreSQL sink requires the engine-injected '_arrow_schema'".into(),
224            ));
225        }
226        let schema = config.arrow_schema().ok_or_else(|| {
227            ConnectorError::ConfigurationError(
228                "invalid PostgreSQL sink '_arrow_schema' encoding".into(),
229            )
230        })?;
231        validate_sink_schema(&schema, &sink_config)?;
232        Ok((sink_config, schema))
233    }
234
235    fn apply_connector_config(&mut self, config: &ConnectorConfig) -> Result<(), ConnectorError> {
236        let (sink_config, schema) = Self::decode_connector_config(config)?;
237        let user_schema = build_user_schema(&schema);
238        self.config = sink_config;
239        self.schema = schema;
240        self.user_schema = user_schema;
241        Ok(())
242    }
243
244    /// Returns the current connector state.
245    #[must_use]
246    pub fn state(&self) -> ConnectorState {
247        self.state
248    }
249
250    /// Returns the number of buffered rows pending flush.
251    #[must_use]
252    pub fn buffered_rows(&self) -> usize {
253        self.buffered_rows
254    }
255
256    #[cfg(feature = "postgres-sink")]
257    fn retain_batch(&mut self, batch: &RecordBatch, retained_bytes: usize) {
258        self.buffer.push(batch.clone());
259        self.buffered_rows = self.buffered_rows.saturating_add(batch.num_rows());
260        self.buffered_retained_bytes = self.buffered_retained_bytes.saturating_add(retained_bytes);
261    }
262
263    /// Moves pending batches out before any await so timeout cancellation cannot leave stale
264    /// records or accounting in the connector.
265    fn take_buffer(&mut self) -> Vec<RecordBatch> {
266        self.buffered_rows = 0;
267        self.buffered_retained_bytes = 0;
268        std::mem::take(&mut self.buffer)
269    }
270
271    /// Returns a reference to the sink metrics.
272    #[must_use]
273    pub fn sink_metrics(&self) -> &PostgresSinkMetrics {
274        &self.metrics
275    }
276
277    // ── Internal helpers ────────────────────────────────────────────
278
279    /// Prepares cached SQL statements based on schema and config.
280    fn prepare_statements(&mut self) -> Result<(), ConnectorError> {
281        self.copy_sql = (self.config.write_mode == WriteMode::Append)
282            .then(|| Self::build_copy_sql(&self.schema, &self.config))
283            .transpose()?;
284        self.upsert_sql = (self.config.write_mode == WriteMode::Upsert)
285            .then(|| Self::build_upsert_sql(&self.schema, &self.config))
286            .transpose()?;
287        self.create_table_sql = self
288            .config
289            .auto_create_table
290            .then(|| Self::build_create_table_sql(&self.schema, &self.config))
291            .transpose()?;
292        self.delete_sql = self
293            .config
294            .changelog_mode
295            .then(|| Self::build_delete_sql(&self.schema, &self.config))
296            .transpose()?;
297        Ok(())
298    }
299
300    /// Returns a reference to the connection pool, or an error if not initialized.
301    #[cfg(feature = "postgres-sink")]
302    fn pool(&self) -> Result<&Pool, ConnectorError> {
303        self.pool.as_ref().ok_or(ConnectorError::InvalidState {
304            expected: "pool initialized (call open() first)".into(),
305            actual: "pool not initialized".into(),
306        })
307    }
308
309    /// Concatenates all buffered batches and strips metadata columns.
310    #[cfg(feature = "postgres-sink")]
311    fn concat_buffer(&self, buffer: &[RecordBatch]) -> Result<RecordBatch, ConnectorError> {
312        if buffer.is_empty() {
313            return Ok(RecordBatch::new_empty(self.user_schema.clone()));
314        }
315
316        let stripped: Result<Vec<RecordBatch>, ConnectorError> =
317            buffer.iter().map(strip_metadata_columns).collect();
318        let stripped = stripped?;
319
320        arrow_select::concat::concat_batches(&self.user_schema, &stripped)
321            .map_err(|e| ConnectorError::Internal(format!("batch concat failed: {e}")))
322    }
323
324    /// Flushes buffered data to `PostgreSQL` using the COPY BINARY protocol.
325    #[cfg(feature = "postgres-sink")]
326    async fn flush_append(
327        &mut self,
328        client: &tokio_postgres::Client,
329        buffer: &[RecordBatch],
330    ) -> Result<WriteResult, ConnectorError> {
331        let user_batch = self.concat_buffer(buffer)?;
332        if user_batch.num_rows() == 0 {
333            return Ok(WriteResult::new(0, 0));
334        }
335        let copy_batch = postgres_copy_batch(&user_batch)?;
336        let mut encoder = pgpq::ArrowToPostgresBinaryEncoder::try_new(copy_batch.schema().as_ref())
337            .map_err(|e| ConnectorError::Internal(format!("pgpq encoder init: {e}")))?;
338
339        // `freeze()` transfers ownership to the COPY sink, so keeping this local also guarantees
340        // cancellation releases partially encoded data.
341        let mut encode_buf = BytesMut::with_capacity(COPY_ENCODE_INITIAL_CAPACITY);
342        encoder.write_header(&mut encode_buf);
343        encoder
344            .write_batch(&copy_batch, &mut encode_buf)
345            .map_err(|e| ConnectorError::Internal(format!("pgpq encode: {e}")))?;
346        encoder
347            .write_footer(&mut encode_buf)
348            .map_err(|e| ConnectorError::Internal(format!("pgpq footer: {e}")))?;
349
350        let encoded_bytes = encode_buf.len();
351        let bytes_to_send = encode_buf.freeze();
352
353        let copy_sql = self
354            .copy_sql
355            .as_deref()
356            .ok_or_else(|| ConnectorError::Internal("COPY SQL not prepared".into()))?;
357
358        let sink = client
359            .copy_in(copy_sql)
360            .await
361            .map_err(|error| postgres_dispatched_write_error("COPY start", &error))?;
362
363        {
364            use futures_util::SinkExt;
365            futures_util::pin_mut!(sink);
366            sink.send(bytes_to_send)
367                .await
368                .map_err(|error| postgres_dispatched_write_error("COPY send", &error))?;
369            sink.close()
370                .await
371                .map_err(|error| postgres_dispatched_write_error("COPY finish", &error))?;
372        }
373
374        let rows = user_batch.num_rows();
375        self.metrics.record_write(rows as u64, encoded_bytes as u64);
376        self.metrics.record_flush();
377        self.metrics.record_copy();
378
379        Ok(WriteResult::new(rows, encoded_bytes as u64))
380    }
381
382    /// Flushes buffered data to `PostgreSQL` using UNNEST-based upsert.
383    #[cfg(feature = "postgres-sink")]
384    #[allow(clippy::cast_possible_truncation)]
385    async fn flush_upsert(
386        &mut self,
387        client: &mut tokio_postgres::Client,
388        buffer: &[RecordBatch],
389    ) -> Result<WriteResult, ConnectorError> {
390        if self.config.changelog_mode {
391            return self.flush_changelog(client, buffer).await;
392        }
393
394        let user_batch = self.concat_buffer(buffer)?;
395        if user_batch.num_rows() == 0 {
396            return Ok(WriteResult::new(0, 0));
397        }
398        let user_batch = collapse_upsert_batch(&user_batch, &self.config.primary_key_columns)?;
399
400        let upsert_sql = self
401            .upsert_sql
402            .as_deref()
403            .ok_or_else(|| ConnectorError::Internal("upsert SQL not prepared".into()))?;
404
405        let rows = execute_unnest(client, upsert_sql, &user_batch).await?;
406
407        let byte_estimate = retained_batch_bytes_u64(&user_batch);
408        self.metrics.record_write(rows, byte_estimate);
409        self.metrics.record_flush();
410        self.metrics.record_upsert();
411
412        Ok(WriteResult::new(rows as usize, byte_estimate))
413    }
414
415    /// Flushes changelog batches per primary key: one collapsed terminal op each, then upserts
416    /// before deletes.
417    #[cfg(feature = "postgres-sink")]
418    #[allow(clippy::cast_possible_truncation)]
419    async fn flush_changelog(
420        &mut self,
421        client: &mut tokio_postgres::Client,
422        buffer: &[RecordBatch],
423    ) -> Result<WriteResult, ConnectorError> {
424        if buffer.is_empty() {
425            return Ok(WriteResult::new(0, 0));
426        }
427
428        // Validate the entire flush window before issuing either the UPSERT or DELETE. A malformed
429        // later batch must not be discovered after an earlier key has already mutated PostgreSQL.
430        for batch in buffer {
431            validate_changelog_input(batch)?;
432        }
433
434        // Collapse the whole buffered flush window per primary key into a cardinality-safe
435        // `{U,D}` batch that `split_changelog_batch` understands. A Z-set changelog (`__weight`,
436        // no `_op`) MUST be collapsed or its many retract+insert events per key fail the split /
437        // violate ON CONFLICT cardinality. A CDC changelog (`_op`) MUST be collapsed too:
438        // otherwise a delete-then-reinsert of a key in one flush window splits into both bins
439        // and, applied upserts-then-deletes, wrongly ends deleted (CN-3). `collapse_changelog`
440        // keeps the last arrival per key and normalizes `_op`, so each key contributes exactly
441        // one terminal U or D.
442        let split_input: Vec<RecordBatch> = {
443            let schema = buffer[0].schema();
444            let combined = arrow_select::concat::concat_batches(&schema, buffer)
445                .map_err(|e| ConnectorError::Internal(format!("concat changelog: {e}")))?;
446            vec![collapse_changelog(
447                &combined,
448                &self.config.primary_key_columns,
449            )?]
450        };
451
452        // Split each batch into inserts/deletes.
453        let mut all_inserts = Vec::new();
454        let mut all_deletes = Vec::new();
455        for batch in &split_input {
456            let (ins, del) = Self::split_changelog_batch(batch)?;
457            if ins.num_rows() > 0 {
458                all_inserts.push(ins);
459            }
460            if del.num_rows() > 0 {
461                all_deletes.push(del);
462            }
463        }
464
465        let transaction = client.transaction().await.map_err(|error| {
466            ConnectorError::ConnectionFailed(format!(
467                "begin PostgreSQL changelog transaction: {error}"
468            ))
469        })?;
470        let mutation = async {
471            let mut upserted = 0_u64;
472            let mut deleted = 0_usize;
473            let mut bytes = 0_u64;
474
475            if !all_inserts.is_empty() {
476                let insert_batch =
477                    arrow_select::concat::concat_batches(&self.user_schema, &all_inserts)
478                        .map_err(|e| ConnectorError::Internal(format!("concat inserts: {e}")))?;
479                let upsert_sql = self
480                    .upsert_sql
481                    .as_deref()
482                    .ok_or_else(|| ConnectorError::Internal("upsert SQL not prepared".into()))?;
483                upserted = execute_unnest(&transaction, upsert_sql, &insert_batch).await?;
484                bytes = retained_batch_bytes_u64(&insert_batch);
485            }
486
487            if !all_deletes.is_empty() {
488                let delete_batch =
489                    arrow_select::concat::concat_batches(&self.user_schema, &all_deletes)
490                        .map_err(|e| ConnectorError::Internal(format!("concat deletes: {e}")))?;
491                bytes = bytes.saturating_add(retained_batch_bytes_u64(&delete_batch));
492                deleted = self.execute_deletes(&transaction, &delete_batch).await?;
493            }
494
495            Ok::<_, ConnectorError>((upserted, deleted, bytes))
496        }
497        .await;
498
499        let (upserted, deleted, total_bytes) = match mutation {
500            Ok(result) => {
501                transaction.commit().await.map_err(|error| {
502                    postgres_dispatched_write_error("transaction COMMIT", &error)
503                })?;
504                result
505            }
506            Err(error) => {
507                if let Err(rollback_error) = transaction.rollback().await {
508                    tracing::warn!(
509                        %rollback_error,
510                        "PostgreSQL changelog rollback failed after a mutation error; no COMMIT was dispatched"
511                    );
512                }
513                return Err(resolve_uncommitted_transaction_error(error));
514            }
515        };
516
517        let total_rows = upserted.saturating_add(deleted as u64);
518        if total_rows != 0 {
519            self.metrics.record_write(total_rows, total_bytes);
520        }
521        if upserted != 0 {
522            self.metrics.record_upsert();
523        }
524        if deleted != 0 {
525            self.metrics.record_deletes(deleted as u64);
526        }
527        self.metrics.record_flush();
528        Ok(WriteResult::new(total_rows as usize, total_bytes))
529    }
530
531    /// Executes batched DELETE for changelog delete records.
532    #[cfg(feature = "postgres-sink")]
533    #[allow(clippy::cast_possible_truncation)]
534    async fn execute_deletes<C>(
535        &self,
536        client: &C,
537        delete_batch: &RecordBatch,
538    ) -> Result<usize, ConnectorError>
539    where
540        C: tokio_postgres::GenericClient + Sync,
541    {
542        if delete_batch.num_rows() == 0 {
543            return Ok(0);
544        }
545
546        let delete_sql = self
547            .delete_sql
548            .as_deref()
549            .ok_or_else(|| ConnectorError::Internal("DELETE SQL not prepared".into()))?;
550
551        let pk_params: Vec<Box<dyn postgres_types::ToSql + Sync + Send>> = self
552            .config
553            .primary_key_columns
554            .iter()
555            .map(|col| {
556                let idx = delete_batch.schema().index_of(col).map_err(|_| {
557                    ConnectorError::ConfigurationError(format!(
558                        "primary key column '{col}' not in delete batch"
559                    ))
560                })?;
561                arrow_column_to_pg_array(delete_batch.column(idx))
562            })
563            .collect::<Result<_, _>>()?;
564
565        let pk_refs: Vec<&(dyn postgres_types::ToSql + Sync)> = pk_params
566            .iter()
567            .map(|p| p.as_ref() as &(dyn postgres_types::ToSql + Sync))
568            .collect();
569
570        let rows = client
571            .execute(delete_sql, &pk_refs)
572            .await
573            .map_err(|error| postgres_dispatched_write_error("DELETE", &error))?;
574
575        Ok(rows as usize)
576    }
577
578    /// Dispatches flush to the appropriate write mode.
579    #[cfg(feature = "postgres-sink")]
580    async fn flush_to_client(
581        &mut self,
582        client: &mut tokio_postgres::Client,
583        buffer: &[RecordBatch],
584    ) -> Result<WriteResult, ConnectorError> {
585        match self.config.write_mode {
586            WriteMode::Append => self.flush_append(client, buffer).await,
587            WriteMode::Upsert => self.flush_upsert(client, buffer).await,
588        }
589    }
590
591    /// Flushes the current buffer, dropping the drained batches on every success, error, or
592    /// cancellation path. The empty vector allocation is recovered after non-cancelled I/O.
593    #[cfg(feature = "postgres-sink")]
594    async fn flush_buffer(&mut self) -> Result<WriteResult, ConnectorError> {
595        let mut pending = self.take_buffer();
596        if pending.is_empty() {
597            self.buffer = pending;
598            return Ok(WriteResult::new(0, 0));
599        }
600
601        let result = async {
602            let mut client = self
603                .pool()?
604                .get()
605                .await
606                .map_err(|e| ConnectorError::ConnectionFailed(format!("pool checkout: {e}")))?;
607            self.flush_to_client(&mut client, &pending).await
608        }
609        .await;
610
611        pending.clear();
612        self.buffer = pending;
613        result
614    }
615
616    #[cfg(feature = "postgres-sink")]
617    async fn write_batch_with_retained_limit(
618        &mut self,
619        batch: &RecordBatch,
620        retained_limit: usize,
621    ) -> Result<WriteResult, ConnectorError> {
622        if self.state != ConnectorState::Running {
623            return Err(ConnectorError::InvalidState {
624                expected: "Running".into(),
625                actual: self.state.to_string(),
626            });
627        }
628
629        validate_input_batch(batch, &self.schema, &self.config)?;
630
631        if batch.num_rows() == 0 {
632            return Ok(WriteResult::new(0, 0));
633        }
634
635        let batch_retained_bytes = retained_batch_bytes(batch);
636        let preflush = requires_preflush(
637            self.buffered_retained_bytes,
638            batch_retained_bytes,
639            retained_limit,
640        )?;
641        let result = if preflush {
642            self.flush_buffer().await?
643        } else {
644            WriteResult::new(0, 0)
645        };
646
647        self.retain_batch(batch, batch_retained_bytes);
648        Ok(result)
649    }
650}
651
652// ── SinkConnector implementation ────────────────────────────────────
653
654#[cfg(feature = "postgres-sink")]
655#[async_trait]
656impl SinkConnector for PostgresSink {
657    fn contract(&self, config: &ConnectorConfig) -> Result<SinkContract, ConnectorError> {
658        let cfg = if config.properties().is_empty() {
659            self.config.clone()
660        } else {
661            Self::decode_connector_config(config)?.0
662        };
663        cfg.validate()?;
664        if config.properties().is_empty() {
665            validate_sink_schema(&self.schema, &cfg)?;
666        }
667        let input_mode = if cfg.changelog_mode {
668            SinkInputMode::FullChangelog
669        } else if cfg.write_mode == WriteMode::Upsert {
670            SinkInputMode::KeyedUpsert
671        } else {
672            SinkInputMode::AppendOnly
673        };
674        // Append-only writes commute across independent runtime writers. Mutable writes do not:
675        // without key-affine placement and a fenced handoff, an older upsert/delete from one node
676        // can land after a newer value from another node. Keep those modes singleton until the
677        // runtime can prove that stronger topology protocol.
678        let topology = if cfg.write_mode == WriteMode::Append {
679            SinkTopology::MultiWriter
680        } else {
681            SinkTopology::Singleton
682        };
683        Ok(SinkContract::new(
684            SinkConsistency::DurableAtLeastOnce,
685            topology,
686            input_mode,
687        ))
688    }
689
690    async fn open(&mut self, config: &ConnectorConfig) -> Result<(), ConnectorError> {
691        if config.properties().is_empty() {
692            // Direct programmatic construction supplies the schema to `new`.
693            self.config.validate()?;
694            validate_sink_schema(&self.schema, &self.config)?;
695        } else {
696            self.apply_connector_config(config)?;
697        }
698
699        // Complete all deterministic admission before touching the network.
700        self.prepare_statements()?;
701        let pool = build_pool(&self.config)?;
702        self.state = ConnectorState::Initializing;
703
704        info!(
705            table = %self.config.qualified_table_name(),
706            mode = %self.config.write_mode,
707            "opening PostgreSQL sink connector"
708        );
709
710        // Validate connectivity.
711        let client = pool.get().await.map_err(|e| {
712            ConnectorError::ConnectionFailed(format!("initial connection failed: {e}"))
713        })?;
714
715        // Auto-create target table.
716        if self.config.auto_create_table {
717            if let Some(ddl) = &self.create_table_sql {
718                client.batch_execute(ddl.as_str()).await.map_err(|e| {
719                    ConnectorError::Internal(format!("auto-create table failed: {e}"))
720                })?;
721                debug!(table = %self.config.qualified_table_name(), "target table ensured");
722            }
723        }
724
725        self.pool = Some(pool);
726        self.state = ConnectorState::Running;
727
728        info!(table = %self.config.qualified_table_name(), "PostgreSQL sink connector opened");
729
730        Ok(())
731    }
732
733    async fn write_batch(&mut self, batch: &RecordBatch) -> Result<WriteResult, ConnectorError> {
734        self.write_batch_with_retained_limit(batch, MAX_BUFFERED_RETAINED_BYTES)
735            .await
736    }
737
738    fn schema(&self) -> SchemaRef {
739        self.schema.clone()
740    }
741
742    fn suggested_write_timeout(&self) -> Duration {
743        let statement_budget = if self.config.changelog_mode {
744            self.config.statement_timeout.saturating_mul(2)
745        } else {
746            self.config.statement_timeout
747        };
748        statement_budget + Duration::from_secs(5)
749    }
750
751    fn flush_interval(&self) -> Duration {
752        self.config.flush_interval
753    }
754
755    async fn flush(&mut self) -> Result<(), ConnectorError> {
756        self.flush_buffer().await.map(|_| ())
757    }
758
759    async fn close(&mut self) -> Result<(), ConnectorError> {
760        info!("closing PostgreSQL sink connector");
761
762        let flush_result = self.flush_buffer().await.map(|_| ());
763        self.pool = None;
764        self.state = ConnectorState::Closed;
765
766        info!(
767            table = %self.config.qualified_table_name(),
768            records = self.metrics.records_written.get(),
769            "PostgreSQL sink connector closed"
770        );
771
772        flush_result
773    }
774}
775
776#[cfg(not(feature = "postgres-sink"))]
777#[async_trait]
778impl SinkConnector for PostgresSink {
779    fn contract(&self, _config: &ConnectorConfig) -> Result<SinkContract, ConnectorError> {
780        Err(ConnectorError::ConfigurationError(
781            "PostgreSQL sink requires the 'postgres-sink' feature".into(),
782        ))
783    }
784
785    async fn open(&mut self, _config: &ConnectorConfig) -> Result<(), ConnectorError> {
786        Err(ConnectorError::ConfigurationError(
787            "PostgreSQL sink requires the 'postgres-sink' feature".into(),
788        ))
789    }
790
791    async fn write_batch(&mut self, _batch: &RecordBatch) -> Result<WriteResult, ConnectorError> {
792        Err(ConnectorError::ConfigurationError(
793            "PostgreSQL sink requires the 'postgres-sink' feature".into(),
794        ))
795    }
796
797    fn schema(&self) -> SchemaRef {
798        self.schema.clone()
799    }
800
801    fn suggested_write_timeout(&self) -> Duration {
802        self.config.statement_timeout + Duration::from_secs(5)
803    }
804
805    fn flush_interval(&self) -> Duration {
806        self.config.flush_interval
807    }
808
809    async fn close(&mut self) -> Result<(), ConnectorError> {
810        drop(self.take_buffer());
811        self.state = ConnectorState::Closed;
812        Ok(())
813    }
814}
815
816impl std::fmt::Debug for PostgresSink {
817    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
818        f.debug_struct("PostgresSink")
819            .field("state", &self.state)
820            .field("table", &self.config.qualified_table_name())
821            .field("mode", &self.config.write_mode)
822            .field("buffered_rows", &self.buffered_rows)
823            .field("buffered_retained_bytes", &self.buffered_retained_bytes)
824            .finish_non_exhaustive()
825    }
826}
827
828#[cfg(test)]
829mod tests;