Pass errors to the sink trait instead of dropping the batch

pg_replicate/src/pipeline/batching/data_pipeline.rs

@@ -74,14 +74,10 @@ impl<Src: Source, Snk: BatchSink> BatchDataPipeline<Src, Snk> {
 
             while let Some(batch) = batch_timeout_stream.next().await {
                 info!("got {} table copy events in a batch", batch.len());
-                //TODO: Avoid a vec copy
-                let mut rows = Vec::with_capacity(batch.len());
-                for row in batch {
-                    rows.push(row.map_err(SourceError::TableCopyStream)?);
-                }
                 self.sink
-                    .write_table_rows(rows, table_schema.table_id)
-                    .await?;
+                    .write_table_rows(batch, table_schema.table_id)
+                    .await
+                    .map_err(PipelineError::SinkError)?;
             }
 
             self.sink.table_copied(table_schema.table_id).await?;

pg_replicate/src/pipeline/data_pipeline.rs

@@ -56,7 +56,7 @@ impl<Src: Source, Snk: Sink> DataPipeline<Src, Snk> {
             pin!(table_rows);
 
             while let Some(row) = table_rows.next().await {
-                let row = row.map_err(SourceError::TableCopyStream)?;
+                let row = row.map_err(SourceError::TableCopyStream);
                 self.sink
                     .write_table_row(row, table_schema.table_id)
                     .await?;
@@ -77,8 +77,8 @@ impl<Src: Source, Snk: Sink> DataPipeline<Src, Snk> {
         pin!(cdc_events);
 
         while let Some(cdc_event) = cdc_events.next().await {
-            let cdc_event = cdc_event.map_err(SourceError::CdcStream)?;
-            let send_status_update = if let CdcEvent::KeepAliveRequested { reply } = cdc_event {
+            let cdc_event = cdc_event.map_err(SourceError::CdcStream);
+            let send_status_update = if let Ok(CdcEvent::KeepAliveRequested { reply }) = cdc_event {
                 reply
             } else {
                 false

pg_replicate/src/pipeline/sinks/bigquery.rs

@@ -9,7 +9,7 @@ use tracing::info;
 use crate::{
     clients::bigquery::BigQueryClient,
     conversions::{cdc_event::CdcEvent, table_row::TableRow, Cell},
-    pipeline::PipelineResumptionState,
+    pipeline::{PipelineResumptionState, sources::postgres::TableCopyStreamError},
     table::{ColumnSchema, TableId, TableName, TableSchema},
 };
 
@@ -159,19 +159,24 @@ impl BatchSink for BigQueryBatchSink {
 
     async fn write_table_rows(
         &mut self,
-        mut table_rows: Vec<TableRow>,
+        mut table_rows: Vec<Result<TableRow, TableCopyStreamError>>,
         table_id: TableId,
     ) -> Result<(), SinkError> {
         let table_schema = self.get_table_schema(table_id)?;
         let table_name = Self::table_name_in_bq(&table_schema.table_name);
         let table_descriptor = table_schema.into();
 
-        for table_row in &mut table_rows {
-            table_row.values.push(Cell::String("UPSERT".to_string()));
-        }
+        let new_rows = table_rows
+            .drain(..)
+            .filter_map(|row| row.ok())
+            .map(|mut row| {
+                row.values.push(Cell::String("UPSERT".to_string()));
+                row
+            })
+            .collect::<Vec<TableRow>>();
 
         self.client
-            .stream_rows(&self.dataset_id, table_name, &table_descriptor, &table_rows)
+            .stream_rows(&self.dataset_id, table_name, &table_descriptor, &new_rows)
             .await?;
 
         Ok(())

pg_replicate/src/pipeline/sinks/mod.rs

@@ -16,7 +16,10 @@ use crate::{
     table::{TableId, TableSchema},
 };
 
-use super::PipelineResumptionState;
+use super::{
+    sources::{postgres::TableCopyStreamError, SourceError},
+    PipelineResumptionState,
+};
 
 #[cfg(feature = "duckdb")]
 use self::duckdb::{DuckDbExecutorError, DuckDbRequest};
@@ -58,8 +61,15 @@ pub trait Sink {
         &mut self,
         table_schemas: HashMap<TableId, TableSchema>,
     ) -> Result<(), SinkError>;
-    async fn write_table_row(&mut self, row: TableRow, table_id: TableId) -> Result<(), SinkError>;
-    async fn write_cdc_event(&mut self, event: CdcEvent) -> Result<PgLsn, SinkError>;
+    async fn write_table_row(
+        &mut self,
+        row: Result<TableRow, SourceError>,
+        table_id: TableId,
+    ) -> Result<(), SinkError>;
+    async fn write_cdc_event(
+        &mut self,
+        event: Result<CdcEvent, SourceError>,
+    ) -> Result<PgLsn, SinkError>;
     async fn table_copied(&mut self, table_id: TableId) -> Result<(), SinkError>;
     async fn truncate_table(&mut self, table_id: TableId) -> Result<(), SinkError>;
 }
@@ -73,7 +83,7 @@ pub trait BatchSink {
     ) -> Result<(), SinkError>;
     async fn write_table_rows(
         &mut self,
-        rows: Vec<TableRow>,
+        rows: Vec<Result<TableRow, TableCopyStreamError>>,
         table_id: TableId,
     ) -> Result<(), SinkError>;
     async fn write_cdc_events(&mut self, events: Vec<CdcEvent>) -> Result<PgLsn, SinkError>;