1use 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#[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
145pub struct PostgresSink {
150 config: PostgresSinkConfig,
152 schema: SchemaRef,
154 user_schema: SchemaRef,
156 state: ConnectorState,
158 buffer: Vec<RecordBatch>,
160 buffered_rows: usize,
162 buffered_retained_bytes: usize,
164 metrics: PostgresSinkMetrics,
166 upsert_sql: Option<String>,
168 copy_sql: Option<String>,
170 create_table_sql: Option<String>,
172 delete_sql: Option<String>,
174 #[cfg(feature = "postgres-sink")]
176 pool: Option<Pool>,
177}
178
179impl PostgresSink {
180 #[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 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 #[must_use]
246 pub fn state(&self) -> ConnectorState {
247 self.state
248 }
249
250 #[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 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 #[must_use]
273 pub fn sink_metrics(&self) -> &PostgresSinkMetrics {
274 &self.metrics
275 }
276
277 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 #[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 #[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 #[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 let mut encode_buf = BytesMut::with_capacity(COPY_ENCODE_INITIAL_CAPACITY);
342 encoder.write_header(&mut encode_buf);
343 encoder
344 .write_batch(©_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 #[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 #[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 for batch in buffer {
431 validate_changelog_input(batch)?;
432 }
433
434 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 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 #[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 #[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 #[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#[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 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 self.config.validate()?;
694 validate_sink_schema(&self.schema, &self.config)?;
695 } else {
696 self.apply_connector_config(config)?;
697 }
698
699 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 let client = pool.get().await.map_err(|e| {
712 ConnectorError::ConnectionFailed(format!("initial connection failed: {e}"))
713 })?;
714
715 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;