hubcio commented on code in PR #3957:
URL: https://github.com/apache/iggy/pull/3957#discussion_r3940086809


##########
core/integration/tests/connectors/postgres/postgres_source.rs:
##########
@@ -127,6 +133,122 @@ async fn json_rows_source_produces_messages_to_iggy(
     }
 }
 
+#[iggy_harness(
+    cluster_nodes = 1,
+    server(connectors_runtime(config_path = 
"tests/connectors/postgres/source.toml")),
+    seed = seeds::connector_stream
+)]
+async fn 
given_delete_after_read_when_iggy_crashes_should_delete_only_after_redelivery(

Review Comment:
   warning: this restarts the whole runtime before checking redelivery, and at 
the 10ms fixture poll the source hits the 5-nack stop long before, so it proves 
redelivery after restart, not nack survival. assert connector status after 
`wait_for_source_errors`, and add a variant with a poll interval longer than 
the outage that restarts only the iggy node.
   
   also postgres_source_cdc.rs:309.



##########
core/connectors/sources/postgres_source/src/lib.rs:
##########
@@ -550,52 +647,86 @@ impl PostgresSource {
 
                 messages.push(processed.message);
                 total_processed += 1;
+                table_processed += 1;
             }
 
-            // Database I/O without holding the lock
-            if !processed_ids.is_empty() {
-                self.mark_or_delete_processed_rows(pool, table, pk_column, 
&processed_ids)
-                    .await?;
+            if self.should_process_rows() && !processed_ids.is_empty() {
+                operations.push(PendingOperation::ProcessRows {
+                    table: table.clone(),
+                    ids: processed_ids,
+                    tracking_boundary: 
self.processing_boundary(max_offset.clone()),
+                });
             }
 
-            // Collect offset update for later
             if let Some(offset) = max_offset {
-                state_updates.push((table.clone(), offset));
+                candidate_state
+                    .tracking_offsets
+                    .insert(table.clone(), offset);
             }
 
             if self.verbose {
-                info!("Fetched {} rows from table '{table}'", messages.len());
+                info!("Fetched {table_processed} rows from table '{table}'");
             } else {
-                debug!("Fetched {} rows from table '{table}'", messages.len());
+                debug!("Fetched {table_processed} rows from table '{table}'");
             }
         }
 
-        // Apply all state updates with a single lock acquisition
-        {
-            let mut state = self.state.lock().await;
-            state.processed_rows += total_processed;
-            for (table, offset) in state_updates {
-                state.tracking_offsets.insert(table, offset);
-            }
-            state.last_poll_time = Utc::now();
-        }
+        let pending = if total_processed > 0 {
+            candidate_state.processed_rows += total_processed;
+            candidate_state.last_poll_time = Utc::now();
+            Some(PendingBatch {
+                state: candidate_state,
+                operations,
+            })
+        } else {
+            None
+        };
 
-        Ok(messages)
+        Ok(PolledBatch { messages, pending })
+    }
+
+    async fn advance_replication_slot(&self, lsn: &str) -> Result<(), Error> {
+        let slot_name = self.replication_slot();
+        let pool = self.get_pool()?;
+        with_retry(
+            || async {
+                match sqlx::query("SELECT pg_replication_slot_advance($1, 
$2::pg_lsn)")
+                    .bind(slot_name)
+                    .bind(lsn)
+                    .execute(pool)
+                    .await
+                {
+                    Err(error) if is_replication_slot_already_advanced(&error) 
=> Ok(()),

Review Comment:
   nit: `22023` is `invalid_parameter_value` in general, not only a target 
below `confirmed_flush`, so a malformed lsn or wrong slot turns into a silent 
success. add a `warn!` when swallowing it.
   
   also at line 1900.



##########
core/connectors/sources/postgres_source/src/lib.rs:
##########
@@ -253,15 +285,71 @@ impl Source for PostgresSource {
             PayloadFormat::JsonDirect | PayloadFormat::Json => Schema::Json,
         };
 
-        let persisted_state = self.serialize_state(&state);
+        let persisted_state = polled
+            .pending
+            .as_ref()
+            .map(|pending| {
+                self.serialize_state(&pending.state).ok_or_else(|| {
+                    Error::Serialization("failed to serialize PostgreSQL 
source state".to_string())
+                })
+            })
+            .transpose()?;
+        *self.pending_batch.lock().await = polled.pending;
 
         Ok(ProducedMessages {
             schema,
-            messages,
+            messages: polled.messages,
             state: persisted_state,
         })
     }
 
+    async fn on_batch_result(&self, result: SourceBatchResult) -> Result<(), 
Error> {
+        let (SourceBatchResult::Ack, Some(pending)) =
+            (result, self.pending_batch.lock().await.take())
+        else {
+            return Ok(());
+        };
+
+        let PendingBatch { state, operations } = pending;
+        let cleanup = async {
+            for operation in operations {
+                match operation {
+                    PendingOperation::ProcessRows {
+                        table,
+                        ids,
+                        tracking_boundary,
+                    } => {
+                        if let Ok(pool) = self.get_pool() {
+                            let _ = self
+                                .mark_or_delete_processed_rows(
+                                    pool,
+                                    &table,
+                                    &ids,
+                                    tracking_boundary.as_deref(),
+                                )
+                                .await;
+                        }
+                    }
+                    PendingOperation::AdvanceReplicationSlot { lsn } => {
+                        let _ = self.advance_replication_slot(&lsn).await;
+                    }
+                }
+            }
+        };
+        if tokio::time::timeout(ACK_CLEANUP_TIMEOUT, cleanup)

Review Comment:
   warning: one budget covers every operation, so with two tables the first 
delete lands, the second gets cancelled mid-query and the cursor commits for 
both. give each operation its own budget, or land the replay todo at line 122.



##########
core/connectors/sources/postgres_source/src/lib.rs:
##########
@@ -253,15 +285,71 @@ impl Source for PostgresSource {
             PayloadFormat::JsonDirect | PayloadFormat::Json => Schema::Json,
         };
 
-        let persisted_state = self.serialize_state(&state);
+        let persisted_state = polled
+            .pending
+            .as_ref()
+            .map(|pending| {
+                self.serialize_state(&pending.state).ok_or_else(|| {
+                    Error::Serialization("failed to serialize PostgreSQL 
source state".to_string())
+                })
+            })
+            .transpose()?;
+        *self.pending_batch.lock().await = polled.pending;
 
         Ok(ProducedMessages {
             schema,
-            messages,
+            messages: polled.messages,
             state: persisted_state,
         })
     }
 
+    async fn on_batch_result(&self, result: SourceBatchResult) -> Result<(), 
Error> {
+        let (SourceBatchResult::Ack, Some(pending)) =
+            (result, self.pending_batch.lock().await.take())
+        else {
+            return Ok(());
+        };
+
+        let PendingBatch { state, operations } = pending;
+        let cleanup = async {
+            for operation in operations {
+                match operation {
+                    PendingOperation::ProcessRows {
+                        table,
+                        ids,
+                        tracking_boundary,
+                    } => {
+                        if let Ok(pool) = self.get_pool() {
+                            let _ = self
+                                .mark_or_delete_processed_rows(
+                                    pool,
+                                    &table,
+                                    &ids,
+                                    tracking_boundary.as_deref(),
+                                )
+                                .await;
+                        }
+                    }
+                    PendingOperation::AdvanceReplicationSlot { lsn } => {
+                        let _ = self.advance_replication_slot(&lsn).await;

Review Comment:
   warning: when the advance keeps failing, line 349 still commits the state, 
so every poll redelivers the same batch, silently. count consecutive advance 
failures and surface them so the operator sees a stuck slot without reading 
`pg_replication_slots`.



##########
Cargo.lock:
##########
@@ -12700,6 +12700,7 @@ version = "0.9.0"
 source = "registry+https://github.com/rust-lang/crates.io-index";
 checksum = "05b44e85bf579a8eeb4ceaa77a3a523baf2bf0e9bac7e40f405d537b5d2d5ccb"
 dependencies = [
+ "bigdecimal",

Review Comment:
   nit: `bigdecimal` sits before `base64 0.22.1` here and at line 12811, which 
is not cargo's order - the lock was hand-edited. regenerate it with `cargo 
build` and commit the result.



##########
core/connectors/sources/postgres_source/src/lib.rs:
##########
@@ -38,13 +41,16 @@ source_connector!(PostgresSource);
 
 const DEFAULT_MAX_RETRIES: u32 = 3;
 const DEFAULT_RETRY_DELAY: &str = "1s";
+const ACK_ACQUIRE_TIMEOUT: Duration = Duration::from_secs(5);
+const ACK_CLEANUP_TIMEOUT: Duration = Duration::from_secs(25);

Review Comment:
   warning: the sdk's 30s batch-result window starts at hand-off, not at the 
ack, so send time and this 25s budget share it. on overrun 
`clear_pending_batch` finds the batch already taken and the sdk returns `Stop`, 
so the source stops for good. cut this to ~10s and derive it from 
`max_retries`/`retry_delay`.



##########
core/connectors/sources/postgres_source/src/lib.rs:
##########
@@ -253,15 +285,71 @@ impl Source for PostgresSource {
             PayloadFormat::JsonDirect | PayloadFormat::Json => Schema::Json,
         };
 
-        let persisted_state = self.serialize_state(&state);
+        let persisted_state = polled
+            .pending
+            .as_ref()
+            .map(|pending| {
+                self.serialize_state(&pending.state).ok_or_else(|| {
+                    Error::Serialization("failed to serialize PostgreSQL 
source state".to_string())
+                })
+            })
+            .transpose()?;
+        *self.pending_batch.lock().await = polled.pending;
 
         Ok(ProducedMessages {
             schema,
-            messages,
+            messages: polled.messages,
             state: persisted_state,
         })
     }
 
+    async fn on_batch_result(&self, result: SourceBatchResult) -> Result<(), 
Error> {
+        let (SourceBatchResult::Ack, Some(pending)) =
+            (result, self.pending_batch.lock().await.take())
+        else {
+            return Ok(());
+        };
+
+        let PendingBatch { state, operations } = pending;
+        let cleanup = async {
+            for operation in operations {
+                match operation {
+                    PendingOperation::ProcessRows {
+                        table,
+                        ids,
+                        tracking_boundary,
+                    } => {
+                        if let Ok(pool) = self.get_pool() {
+                            let _ = self

Review Comment:
   warning: `let _` drops every staged-operation error, including 
`quote_identifier` failures that never reach the `with_retry` log, so a delete 
that always fails is invisible. replace it with `if let Err(error) = ...` and 
an `error!` that names the connector ID.
   
   also at line 334.



##########
core/connectors/sources/postgres_source/src/lib.rs:
##########
@@ -610,26 +741,28 @@ impl PostgresSource {
             .join(", ");
 
         if self.config.delete_after_read.unwrap_or(false) {

Review Comment:
   nit: with both `delete_after_read` and `processed_column` set, delete wins 
and the column is silently ignored while `build_polling_query` still filters on 
`processed = FALSE`. reject the combination in `open()`.
   
   also at line 806.



##########
core/connectors/sources/postgres_source/src/lib.rs:
##########
@@ -2497,44 +2756,132 @@ mod tests {
 
         let src = PostgresSource::new(1, test_config(), Some(connector_state));
 
-        let runtime = tokio::runtime::Runtime::new().unwrap();
-        runtime.block_on(async {
-            let restored = src.state.lock().await;
-            assert_eq!(
-                restored.tracking_offsets.get("users"),
-                Some(&"100".to_string())
-            );
-            assert_eq!(
-                restored.tracking_offsets.get("orders"),
-                Some(&"2024-01-15T10:30:00Z".to_string())
-            );
-            assert_eq!(restored.processed_rows, 500);
-        });
+        let restored = src.state.lock().await;
+        assert_eq!(
+            restored.tracking_offsets.get("users"),
+            Some(&"100".to_string())
+        );
+        assert_eq!(
+            restored.tracking_offsets.get("orders"),
+            Some(&"2024-01-15T10:30:00Z".to_string())
+        );
+        assert_eq!(restored.processed_rows, 500);
+    }
+
+    #[tokio::test]
+    async fn given_no_state_should_start_fresh() {
+        let src = PostgresSource::new(1, test_config(), None);
+
+        let state = src.state.lock().await;
+        assert!(state.tracking_offsets.is_empty());
+        assert_eq!(state.processed_rows, 0);
+    }
+
+    #[tokio::test]
+    async fn 
given_transient_database_errors_when_retrying_should_eventually_succeed() {
+        let attempts = AtomicU32::new(0);
+
+        let result = with_retry(
+            || async {
+                if attempts.fetch_add(1, Ordering::Relaxed) < 2 {
+                    Err(sqlx::Error::PoolTimedOut)
+                } else {
+                    Ok(())
+                }
+            },
+            3,
+            0,
+        )
+        .await;
+
+        assert!(result.is_ok());
+        assert_eq!(attempts.load(Ordering::Relaxed), 3);
     }
 
     #[test]
-    fn given_no_state_should_start_fresh() {
+    fn given_active_replication_slot_sqlstate_should_be_transient() {
+        assert!(is_transient_sqlstate("55006"));
+    }
+
+    #[test]
+    fn 
given_target_below_confirmed_flush_sqlstate_should_be_already_advanced() {
+        assert!(is_replication_slot_already_advanced_sqlstate("22023"));
+    }
+
+    #[tokio::test]
+    async fn given_nack_when_batch_is_staged_should_keep_committed_state() {
         let src = PostgresSource::new(1, test_config(), None);
+        let mut candidate_state = src.state.lock().await.clone();
+        candidate_state
+            .tracking_offsets
+            .insert("users".to_string(), "3".to_string());
+        candidate_state.processed_rows = 3;
+        *src.pending_batch.lock().await = Some(PendingBatch {
+            state: candidate_state,
+            operations: vec![
+                PendingOperation::ProcessRows {
+                    table: "users".to_string(),
+                    ids: vec!["3".to_string()],
+                    tracking_boundary: Some("3".to_string()),
+                },
+                PendingOperation::AdvanceReplicationSlot {
+                    lsn: "0/16D32A0".to_string(),
+                },
+            ],
+        });
 
-        let runtime = tokio::runtime::Runtime::new().unwrap();
-        runtime.block_on(async {
+        src.on_batch_result(SourceBatchResult::Nack)
+            .await
+            .expect("NACK should discard the candidate state");
+
+        {
             let state = src.state.lock().await;
             assert!(state.tracking_offsets.is_empty());
             assert_eq!(state.processed_rows, 0);
+        }
+        assert!(src.pending_batch.lock().await.is_none());
+    }
+
+    #[tokio::test]
+    async fn 
given_ack_when_staged_operation_fails_should_commit_candidate_state() {

Review Comment:
   nit: this builds the source with `pool: None`, so line 322 skips the call 
and `mark_or_delete_processed_rows` never runs - the failed-op path stays 
untested. stage an op whose identifier fails to quote, or point the pool at a 
closed database.



##########
core/connectors/sources/postgres_source/src/lib.rs:
##########
@@ -38,13 +41,16 @@ source_connector!(PostgresSource);
 
 const DEFAULT_MAX_RETRIES: u32 = 3;
 const DEFAULT_RETRY_DELAY: &str = "1s";
+const ACK_ACQUIRE_TIMEOUT: Duration = Duration::from_secs(5);

Review Comment:
   nit: this sets the pool-wide `acquire_timeout` that every poll uses too, not 
just the ack. rename to `POOL_ACQUIRE_TIMEOUT`.



##########
core/connectors/sources/postgres_source/src/lib.rs:
##########
@@ -38,13 +41,16 @@ source_connector!(PostgresSource);
 
 const DEFAULT_MAX_RETRIES: u32 = 3;
 const DEFAULT_RETRY_DELAY: &str = "1s";
+const ACK_ACQUIRE_TIMEOUT: Duration = Duration::from_secs(5);
+const ACK_CLEANUP_TIMEOUT: Duration = Duration::from_secs(25);

Review Comment:
   warning: `max_retries` and `retry_delay` are user config, so 
`max_retries=10` with `retry_delay="5s"` gives a single op ~275s and this cap 
cuts every ack short. validate the retry budget against it in `open()` and 
reject configs that cannot fit.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]

Reply via email to