From 7b1192dbcbcc855dba33622a5dcdf02fcee6b446 Mon Sep 17 00:00:00 2001 From: Rohan Dubey Date: Mon, 24 Aug 2026 03:46:24 +0530 Subject: [PATCH 1/2] fix(connectors): defer postgres source progress until ack --- .../sources/postgres_source/README.md | 13 +- .../sources/postgres_source/src/lib.rs | 300 ++++++++++++++---- .../src/harness/handle/connectors_runtime.rs | 12 +- .../connectors/postgres/postgres_source.rs | 180 ++++++++++- .../postgres/postgres_source_cdc.rs | 4 +- 5 files changed, 433 insertions(+), 76 deletions(-) diff --git a/core/connectors/sources/postgres_source/README.md b/core/connectors/sources/postgres_source/README.md index d0619bdeab..76278d8130 100644 --- a/core/connectors/sources/postgres_source/README.md +++ b/core/connectors/sources/postgres_source/README.md @@ -209,7 +209,7 @@ LIMIT $limit ### Delete After Read -Deletes rows from the source table after successful processing: +Deletes rows from the source table only after Iggy acknowledges the batch: ```toml [plugin_config] @@ -219,7 +219,7 @@ primary_key_column = "id" ### Mark as Processed -Updates a boolean column instead of deleting: +Updates a boolean column after Iggy acknowledges the batch instead of deleting: ```toml [plugin_config] @@ -267,6 +267,10 @@ tables = ["users", "orders"] capture_operations = ["INSERT", "UPDATE", "DELETE"] ``` +The connector peeks at logical changes and advances the replication slot only +after Iggy acknowledges the batch. A failed delivery leaves the slot unchanged +so the next poll can read the same changes again. + The `pg_replicate` backend requires the `cdc_pg_replicate` feature flag at build time. ### Slot Naming @@ -274,8 +278,9 @@ The `pg_replicate` backend requires the `cdc_pg_replicate` feature flag at build Each CDC connector must use a unique `replication_slot`. Setup accepts any pre-existing `test_decoding` slot, so two connectors pointed at the same database with the default `replication_slot = "iggy_slot"` will silently -share one slot. `pg_logical_slot_get_changes` consumes changes on read, so -each connector only sees a subset of the other's changes instead of erroring. +share one slot. Each connector peeks from and advances the same slot after +delivery, so one connector can move the shared position past changes that the +other has not processed. Set an explicit, distinct `replication_slot` per connector instance. ### Decommissioning diff --git a/core/connectors/sources/postgres_source/src/lib.rs b/core/connectors/sources/postgres_source/src/lib.rs index 8deb355169..16b648085c 100644 --- a/core/connectors/sources/postgres_source/src/lib.rs +++ b/core/connectors/sources/postgres_source/src/lib.rs @@ -15,21 +15,23 @@ // specific language governing permissions and limitations // under the License. +use std::collections::HashMap; +use std::str::FromStr; +use std::time::Duration; + use async_trait::async_trait; use chrono::{NaiveDate, NaiveDateTime, NaiveTime}; use humantime::Duration as HumanDuration; use iggy_common::{DateTime, Utc}; use iggy_connector_sdk::{ - ConnectorState, Error, ProducedMessage, ProducedMessages, Schema, Source, source_connector, + ConnectorState, Error, ProducedMessage, ProducedMessages, Schema, Source, + source::SourceBatchResult, source_connector, }; use secrecy::{ExposeSecret, SecretString}; use serde::{Deserialize, Serialize}; use sqlx::postgres::PgPoolOptions; use sqlx::postgres::types::{Oid, PgInterval, PgTimeTz}; use sqlx::{Column, Pool, Postgres, Row, TypeInfo, ValueRef}; -use std::collections::HashMap; -use std::str::FromStr; -use std::time::Duration; use tokio::sync::Mutex; use tracing::{debug, error, info, warn}; use uuid::Uuid; @@ -45,6 +47,7 @@ pub struct PostgresSource { pool: Option>, config: PostgresSourceConfig, state: Mutex, + pending_batch: Mutex>, verbose: bool, retry_delay: Duration, poll_interval: Duration, @@ -97,13 +100,38 @@ impl PayloadFormat { } } -#[derive(Debug, Serialize, Deserialize)] +#[derive(Debug, Clone, Serialize, Deserialize)] struct State { last_poll_time: DateTime, tracking_offsets: HashMap, processed_rows: u64, } +#[derive(Debug)] +struct PolledBatch { + messages: Vec, + pending: Option, +} + +#[derive(Debug)] +struct PendingBatch { + state: State, + operations: Vec, +} + +#[derive(Debug)] +enum PendingOperation { + ProcessRows { + table: String, + primary_key_column: String, + ids: Vec, + }, + AdvanceReplicationSlot { + slot_name: String, + lsn: String, + }, +} + #[derive(Debug, Serialize, Deserialize)] pub struct DatabaseRecord { pub table_name: String, @@ -162,6 +190,7 @@ impl PostgresSource { tracking_offsets: HashMap::new(), processed_rows: 0, })), + pending_batch: Mutex::new(None), verbose, retry_delay, poll_interval, @@ -221,7 +250,7 @@ impl Source for PostgresSource { let poll_interval = self.poll_interval; tokio::time::sleep(poll_interval).await; - let messages = match self.config.mode.as_str() { + let polled = match self.config.mode.as_str() { "polling" => self.poll_tables().await?, "cdc" => self.poll_cdc().await?, _ => { @@ -230,20 +259,23 @@ impl Source for PostgresSource { } }; - let state = self.state.lock().await; + let processed_rows = match polled.pending.as_ref() { + Some(pending) => pending.state.processed_rows, + None => self.state.lock().await.processed_rows, + }; if self.verbose { info!( "PostgreSQL source connector ID: {} produced {} messages. Total processed: {}", self.id, - messages.len(), - state.processed_rows + polled.messages.len(), + processed_rows ); } else { debug!( "PostgreSQL source connector ID: {} produced {} messages. Total processed: {}", self.id, - messages.len(), - state.processed_rows + polled.messages.len(), + processed_rows ); } @@ -253,15 +285,59 @@ 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 pending = self.pending_batch.lock().await.take(); + if result == SourceBatchResult::Nack { + return Ok(()); + } + + let Some(pending) = pending else { + return Ok(()); + }; + + for operation in pending.operations { + match operation { + PendingOperation::ProcessRows { + table, + primary_key_column, + ids, + } => { + self.mark_or_delete_processed_rows( + self.get_pool()?, + &table, + &primary_key_column, + &ids, + ) + .await?; + } + PendingOperation::AdvanceReplicationSlot { slot_name, lsn } => { + self.advance_replication_slot(&slot_name, &lsn).await?; + } + } + } + + *self.state.lock().await = pending.state; + Ok(()) + } + async fn close(&mut self) -> Result<(), Error> { if let Some(pool) = self.pool.take() { pool.close().await; @@ -389,7 +465,7 @@ impl PostgresSource { Ok(()) } - async fn poll_cdc(&self) -> Result, Error> { + async fn poll_cdc(&self) -> Result { match self.config.cdc_backend.as_deref().unwrap_or("builtin") { "builtin" => self.poll_cdc_builtin().await, "pg_replicate" => Err(Error::InitError( @@ -399,7 +475,7 @@ impl PostgresSource { } } - async fn poll_cdc_builtin(&self) -> Result, Error> { + async fn poll_cdc_builtin(&self) -> Result { let pool = self.get_pool()?; let slot_name = self @@ -422,38 +498,39 @@ impl PostgresSource { // can still exceed it), so this isn't a hard per-call cap - but it // stops the backlog from growing unbounded across many transactions // the way NULL (no limit at all) did. - let rows = - sqlx::query("SELECT lsn, xid, data FROM pg_logical_slot_get_changes($1, NULL, $2)") - .bind(slot_name) - .bind(batch_size) - .fetch_all(pool) - .await - .map_err(|e| { - error!("Failed to fetch CDC changes: {e}"); - Error::InvalidRecord - })?; + let rows = sqlx::query( + "SELECT lsn::text AS lsn, xid, data FROM pg_logical_slot_peek_changes($1, NULL, $2)", + ) + .bind(slot_name) + .bind(batch_size) + .fetch_all(pool) + .await + .map_err(|e| { + error!("Failed to fetch CDC changes: {e}"); + Error::InvalidRecord + })?; let mut messages = Vec::new(); + let mut last_lsn = None; for row in rows { - let data: String = match row.try_get("data") { - Ok(data) => data, - Err(e) => { - error!("Skipping CDC row with unreadable data column: {e}"); - continue; - } - }; + let lsn: String = row.try_get("lsn").map_err(|e| { + error!("Failed to read CDC row LSN: {e}"); + Error::InvalidRecord + })?; + let data: String = row.try_get("data").map_err(|e| { + error!("Failed to read CDC row data: {e}"); + Error::InvalidRecord + })?; + last_lsn = Some(lsn); if let Some(change_record) = self.parse_logical_replication_message(&data, &capture_ops, captured_tables) { - let payload = match simd_json::to_vec(&change_record) { - Ok(payload) => payload, - Err(e) => { - error!("Skipping CDC row that failed to serialize: {e}"); - continue; - } - }; + let payload = simd_json::to_vec(&change_record).map_err(|e| { + error!("Failed to serialize CDC row: {e}"); + Error::InvalidRecord + })?; let message = ProducedMessage { id: Some(Uuid::new_v4().as_u128()), @@ -468,23 +545,34 @@ impl PostgresSource { } } - // Update state with minimal lock time - if !messages.is_empty() { - let mut state = self.state.lock().await; - state.processed_rows += messages.len() as u64; - } - if self.verbose { info!("CDC: Fetched {} change records", messages.len()); } else { debug!("CDC: Fetched {} change records", messages.len()); } - Ok(messages) + let pending = if let Some(lsn) = last_lsn { + let mut state = self.state.lock().await.clone(); + state.processed_rows += messages.len() as u64; + state.last_poll_time = Utc::now(); + Some(PendingBatch { + state, + operations: vec![PendingOperation::AdvanceReplicationSlot { + slot_name: slot_name.to_string(), + lsn, + }], + }) + } else { + None + }; + + Ok(PolledBatch { messages, pending }) } - async fn poll_tables(&self) -> Result, Error> { + async fn poll_tables(&self) -> Result { let pool = self.get_pool()?; let mut messages = Vec::new(); + let mut operations = Vec::new(); + let mut candidate_state = self.state.lock().await.clone(); let batch_size = self.config.batch_size.unwrap_or(1000); let tracking_column = self.config.tracking_column.as_deref().unwrap_or("id"); @@ -514,11 +602,7 @@ impl PostgresSource { ..row_config }; - // Get last offset with minimal lock time - let last_offset = { - let state = self.state.lock().await; - state.tracking_offsets.get(table).cloned() - }; + let last_offset = candidate_state.tracking_offsets.get(table).cloned(); let query = if let Some(custom_query) = &self.config.custom_query { self.validate_custom_query(custom_query)?; @@ -552,10 +636,12 @@ impl PostgresSource { total_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?; + operations.push(PendingOperation::ProcessRows { + table: table.clone(), + primary_key_column: pk_column.to_string(), + ids: processed_ids, + }); } // Collect offset update for later @@ -570,17 +656,36 @@ impl PostgresSource { } } - // Apply all state updates with a single lock acquisition - { - let mut state = self.state.lock().await; - state.processed_rows += total_processed; + let pending = if total_processed > 0 { + candidate_state.processed_rows += total_processed; for (table, offset) in state_updates { - state.tracking_offsets.insert(table, offset); + candidate_state.tracking_offsets.insert(table, offset); } - state.last_poll_time = Utc::now(); - } + 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, slot_name: &str, lsn: &str) -> Result<(), Error> { + sqlx::query("SELECT pg_replication_slot_advance($1, $2::pg_lsn)") + .bind(slot_name) + .bind(lsn) + .execute(self.get_pool()?) + .await + .map_err(|e| { + error!("Failed to advance replication slot '{slot_name}' to {lsn}: {e}"); + Error::Connection(format!( + "failed to advance replication slot '{slot_name}' to {lsn}: {e}" + )) + })?; + Ok(()) } async fn mark_or_delete_processed_rows( @@ -2524,6 +2629,75 @@ mod tests { }); } + #[test] + fn given_nack_when_batch_is_staged_should_keep_committed_state() { + let src = PostgresSource::new(1, test_config(), None); + let runtime = tokio::runtime::Runtime::new().expect("failed to create test runtime"); + runtime.block_on(async { + 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(), + primary_key_column: "id".to_string(), + ids: vec!["3".to_string()], + }, + PendingOperation::AdvanceReplicationSlot { + slot_name: "iggy_slot".to_string(), + lsn: "0/16D32A0".to_string(), + }, + ], + }); + + 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()); + }); + } + + #[test] + fn given_ack_when_batch_is_staged_should_commit_candidate_state() { + let src = PostgresSource::new(1, test_config(), None); + let runtime = tokio::runtime::Runtime::new().expect("failed to create test runtime"); + runtime.block_on(async { + 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::new(), + }); + + src.on_batch_result(SourceBatchResult::Ack) + .await + .expect("ACK should commit the candidate state"); + + { + let state = src.state.lock().await; + assert_eq!( + state.tracking_offsets.get("users").map(String::as_str), + Some("3") + ); + assert_eq!(state.processed_rows, 3); + } + assert!(src.pending_batch.lock().await.is_none()); + }); + } + #[test] fn given_invalid_state_should_start_fresh() { let invalid_state = ConnectorState(b"not valid json".to_vec()); diff --git a/core/integration/src/harness/handle/connectors_runtime.rs b/core/integration/src/harness/handle/connectors_runtime.rs index f86eb91436..1718972ab7 100644 --- a/core/integration/src/harness/handle/connectors_runtime.rs +++ b/core/integration/src/harness/handle/connectors_runtime.rs @@ -42,6 +42,7 @@ pub struct ConnectorsRuntimeHandle { child_handle: Option, server_address: SocketAddr, iggy_address: Option, + iggy_connection_options: Option, stdout_path: Option, stderr_path: Option, _port_reserver: SinglePortReserver, @@ -78,6 +79,10 @@ impl ConnectorsRuntimeHandle { common::collect_logs(&self.stdout_path, &self.stderr_path) } + pub fn set_iggy_connection_options(&mut self, options: impl Into) { + self.iggy_connection_options = Some(options.into()); + } + fn build_envs(&mut self) { let state_path = self.context.connectors_runtime_state_path(self.server_id); self.envs.insert( @@ -90,8 +95,12 @@ impl ConnectorsRuntimeHandle { ); if let Some(addr) = self.iggy_address { + let address = self + .iggy_connection_options + .as_ref() + .map_or_else(|| addr.to_string(), |options| format!("{addr}?{options}")); self.envs - .insert("IGGY_CONNECTORS_IGGY_ADDRESS".to_string(), addr.to_string()); + .insert("IGGY_CONNECTORS_IGGY_ADDRESS".to_string(), address); } if let Some(ref config_path) = self.config.config_path { @@ -125,6 +134,7 @@ impl ConnectorsRuntimeHandle { child_handle: None, server_address, iggy_address: None, + iggy_connection_options: None, stdout_path: None, stderr_path: None, _port_reserver: reserver, diff --git a/core/integration/tests/connectors/postgres/postgres_source.rs b/core/integration/tests/connectors/postgres/postgres_source.rs index a66a3c5e67..34214aeaa7 100644 --- a/core/integration/tests/connectors/postgres/postgres_source.rs +++ b/core/integration/tests/connectors/postgres/postgres_source.rs @@ -15,6 +15,16 @@ // specific language governing permissions and limitations // under the License. +use std::time::Duration; + +use iggy_common::MessageClient; +use iggy_common::{Consumer, Identifier, PollingStrategy}; +use iggy_connector_sdk::api::{ConnectorRuntimeStats, ConnectorStatus}; +use integration::harness::seeds; +use integration::iggy_harness; +use reqwest::Client; +use tokio::time::{sleep, timeout}; + use super::{DatabaseRecord, POLL_ATTEMPTS, POLL_INTERVAL_MS, TEST_MESSAGE_COUNT}; use crate::connectors::create_test_messages; use crate::connectors::fixtures::{ @@ -22,12 +32,10 @@ use crate::connectors::fixtures::{ PostgresSourceJsonFixture, PostgresSourceJsonbFixture, PostgresSourceMarkFixture, PostgresSourceOps, }; -use iggy_common::MessageClient; -use iggy_common::{Consumer, Identifier, PollingStrategy}; -use integration::harness::seeds; -use integration::iggy_harness; -use std::time::Duration; -use tokio::time::sleep; + +const API_KEY: &str = "test-api-key"; +const SOURCE_KEY: &str = "postgres"; +const SEND_FAILURE_TIMEOUT: Duration = Duration::from_secs(25); #[iggy_harness( server(connectors_runtime(config_path = "tests/connectors/postgres/source.toml")), @@ -127,6 +135,166 @@ 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_rows_in_postgres_when_iggy_server_crashes_should_redeliver_after_restart( + harness: &mut TestHarness, + fixture: PostgresSourceJsonFixture, +) { + let pool = fixture.create_pool().await.expect("Failed to create pool"); + fixture.create_table(&pool).await; + + harness + .server_mut() + .stop_dependents() + .expect("Failed to stop connectors runtime"); + harness + .server_mut() + .connectors_runtime_mut() + .expect("connectors runtime") + // Keep a failed send bounded instead of waiting indefinitely for Iggy to return. + .set_iggy_connection_options("reconnection_retries=0"); + harness + .server_mut() + .start_dependents() + .await + .expect("Failed to restart connectors runtime"); + + let api_url = harness + .connectors_runtime() + .expect("connectors runtime") + .http_url(); + let http = Client::new(); + let errors_before_failure = source_errors(&http, &api_url).await; + + harness.kill_node(0).expect("Failed to kill Iggy server"); + + let expected = create_test_messages(TEST_MESSAGE_COUNT); + let mut transaction = pool.begin().await.expect("Failed to begin transaction"); + let insert = format!( + "INSERT INTO {} (id, name, count, amount, active, timestamp, tag) \ + VALUES ($1, $2, $3, $4, $5, $6, $7)", + fixture.table_name() + ); + for message in &expected { + let tag = format!("{:<10}", format!("tag_{}", message.id)); + sqlx::query(sqlx::AssertSqlSafe(insert.as_str())) + .bind(message.id as i32) + .bind(&message.name) + .bind(message.count as i32) + .bind(message.amount) + .bind(message.active) + .bind(message.timestamp) + .bind(tag) + .execute(&mut *transaction) + .await + .expect("Failed to insert source row"); + } + transaction + .commit() + .await + .expect("Failed to commit source rows"); + + // The second error proves that NACK made the same rows eligible for another poll. + wait_for_source_errors(&http, &api_url, errors_before_failure + 2).await; + harness + .server_mut() + .stop_dependents() + .expect("Failed to stop connectors runtime"); + harness + .restart_node(0) + .expect("Failed to restart Iggy server"); + harness + .server_mut() + .start_dependents() + .await + .expect("Failed to restart connectors runtime"); + + let client = harness.root_client().await.unwrap(); + let stream_id: Identifier = seeds::names::STREAM.try_into().unwrap(); + let topic_id: Identifier = seeds::names::TOPIC.try_into().unwrap(); + let consumer_id: Identifier = "send_failure_consumer".try_into().unwrap(); + let mut received = Vec::new(); + + for _ in 0..POLL_ATTEMPTS { + if let Ok(polled) = client + .poll_messages( + &stream_id, + &topic_id, + None, + &Consumer::new(consumer_id.clone()), + &PollingStrategy::next(), + 10, + true, + ) + .await + { + received.extend(polled.messages.into_iter().filter_map(|message| { + serde_json::from_slice::(&message.payload).ok() + })); + if received.len() >= TEST_MESSAGE_COUNT { + break; + } + } + sleep(Duration::from_millis(POLL_INTERVAL_MS)).await; + } + + assert_eq!( + received.len(), + TEST_MESSAGE_COUNT, + "Rows polled during the failed send should be delivered after restart" + ); + for (record, expected) in received.iter().zip(expected) { + assert_eq!(record.data.id, expected.id); + } + + pool.close().await; +} + +async fn source_errors(http: &Client, api_url: &str) -> u64 { + http.get(format!("{api_url}/stats")) + .header("api-key", API_KEY) + .send() + .await + .expect("runtime stats should be available") + .json::() + .await + .expect("runtime stats should be valid") + .connectors + .into_iter() + .find(|connector| connector.key == SOURCE_KEY) + .expect("PostgreSQL source stats should be present") + .errors +} + +async fn wait_for_source_errors(http: &Client, api_url: &str, minimum_errors: u64) { + timeout(SEND_FAILURE_TIMEOUT, async { + loop { + if let Ok(response) = http + .get(format!("{api_url}/stats")) + .header("api-key", API_KEY) + .send() + .await + && let Ok(stats) = response.json::().await + && let Some(source) = stats + .connectors + .iter() + .find(|connector| connector.key == SOURCE_KEY) + && source.status == ConnectorStatus::Error + && source.errors >= minimum_errors + { + break; + } + sleep(Duration::from_millis(POLL_INTERVAL_MS)).await; + } + }) + .await + .expect("PostgreSQL source did not retry the NACKed batch"); +} + #[iggy_harness( server(connectors_runtime(config_path = "tests/connectors/postgres/source.toml")), seed = seeds::connector_stream diff --git a/core/integration/tests/connectors/postgres/postgres_source_cdc.rs b/core/integration/tests/connectors/postgres/postgres_source_cdc.rs index b9a485e832..efefa8bd60 100644 --- a/core/integration/tests/connectors/postgres/postgres_source_cdc.rs +++ b/core/integration/tests/connectors/postgres/postgres_source_cdc.rs @@ -340,8 +340,8 @@ async fn delete_source_config_version(http: &Client, api_url: &str, version: u64 ); } -// The connector calls pg_logical_slot_get_changes on a fixed poll interval and -// briefly holds the slot active during each call. A drop landing in that window +// The connector peeks changes and advances the slot on a fixed poll interval. +// Both operations briefly hold the slot active. A drop landing in that window // gets ERROR 55006 (slot is active for PID ...), so retry past transient hits // instead of dropping while the poller is guaranteed stopped. const PG_OBJECT_IN_USE: &str = "55006"; From 045a6809d484d4ac96045a89cd9c5a606aa0d7ea Mon Sep 17 00:00:00 2001 From: Rohan Dubey Date: Wed, 26 Aug 2026 00:52:27 +0530 Subject: [PATCH 2/2] Addressing comments --- .claude/skills/connector-source/SKILL.md | 34 +-- .claude/skills/connector-source/TEMPLATE.md | 53 +++-- core/connectors/runtime/src/stream.rs | 35 ++- .../sources/postgres_source/README.md | 22 +- .../sources/postgres_source/src/lib.rs | 222 +++++++++++------- .../src/harness/handle/connectors_runtime.rs | 4 + .../tests/connectors/postgres/mod.rs | 42 +++- .../connectors/postgres/postgres_source.rs | 124 +++------- .../postgres/postgres_source_cdc.rs | 130 +++++++++- .../tests/connectors/postgres/restart.rs | 3 +- 10 files changed, 452 insertions(+), 217 deletions(-) diff --git a/.claude/skills/connector-source/SKILL.md b/.claude/skills/connector-source/SKILL.md index 5627b1812d..1d8255108f 100644 --- a/.claude/skills/connector-source/SKILL.md +++ b/.claude/skills/connector-source/SKILL.md @@ -45,25 +45,31 @@ The macro shares the source as `Arc` across the FFI callback and forwarding l ### Lock discipline -Never hold the state `Mutex` across upstream I/O. Canonical pattern (matches `sources/postgres_source/src/lib.rs::poll_tables`): +Never hold the state `Mutex` across upstream I/O. Build a candidate from committed +state, then stage it until the runtime reports the batch result: ```rust -let cursor = { self.state.lock().await.cursor.clone() }; // brief read -let rows = client.query(&sql, &[&cursor]).await?; // no lock held -let persisted = { // brief write - let mut state = self.state.lock().await; - state.cursor = Some(new_cursor); - ConnectorState::serialize(&*state, CONNECTOR_NAME, self.id) -}; +let mut candidate = self.state.lock().await.clone(); +let rows = client.query(&sql, &[&candidate.cursor]).await?; +candidate.cursor = Some(new_cursor); +let persisted = ConnectorState::serialize(&candidate, CONNECTOR_NAME, self.id) + .ok_or_else(|| Error::Serialization("failed to serialize source state".into()))?; +*self.pending.lock().await = Some(candidate); ``` ### State persistence -- `ConnectorState` is `Vec` via MessagePack (`rmp_serde`). Use `ConnectorState::serialize(&state, NAME, id)` + `ConnectorState::deserialize::(NAME, id)`. Both return `Option` and log on failure (non-fatal). -- Runtime saves to `{state_path}/source_{key}.state` only after a successful Iggy send. Between `poll()` returning and the runtime persisting the save, a crash leaves the same cursor for the next poll - downstream must tolerate at-least-once. -- **Always return state in every `ProducedMessages`**, including empty polls. Empty results still need to advance watermarks (timestamp sources) or affirm "nothing new." +- `ConnectorState` is `Vec` via MessagePack (`rmp_serde`). Use `ConnectorState::serialize(&state, NAME, id)` + `ConnectorState::deserialize::(NAME, id)`. +- `poll()` must not commit cursors or destructive work. Return messages with candidate state and keep the corresponding work staged. +- The runtime sends the batch, persists its candidate state, then calls `on_batch_result(Ack)`. Commit staged in-memory state and external delete/mark operations only on ACK. A NACK discards the candidate so the same data can be polled again. +- Return `state: None` for an empty poll when no watermark changed. If an empty poll advances a watermark, stage and return the new state through the same ACK handshake. +- Treat candidate-state serialization failure as a poll error. Do not send messages without the state needed to resume them safely. - Keep `State` small - rewritten every batch. No unbounded vecs. +The SDK allows one in-flight batch. Five consecutive NACKs stop the source and +require a manual restart. Returning `Err` from `on_batch_result` is fatal, so +retry transient backend failures inside the callback before returning an error. + ### Sleep first `poll()` must `sleep(self.poll_interval).await` before any work. Without it, an empty source spins a CPU. @@ -97,7 +103,7 @@ Match `ProducedMessages.schema` to the bytes in `messages[i].payload`: | Transient fetch failure (retry-worthy) | `Error::Connection` or `Error::HttpRequestFailed` | | Permanent fetch failure (auth, schema gone) | `Error::PermanentHttpError` | | Row failed to serialize | `Error::Serialization(...)` | -| State serialization failed | log + skip (non-fatal) | +| State serialization failed | `Error::Serialization(...)` | Returning `Err` from `poll()` is only logged by the SDK's FFI bridge (`sdk/src/source.rs::handle_messages`) - the loop continues, the next @@ -127,7 +133,7 @@ Iggy consumer-loop labels use literal API names (`offset=`, `current_offset=`). 1. `async fn poll(&mut self)` - won't compile. Use `&self` + `Mutex`. 2. Holding `state.lock()` across the fetch I/O - blocks `close()`, causes shutdown timeouts. 3. Forgetting to sleep - 100% CPU on idle source. -4. Returning state only on success - state should advance on empty polls too. +4. Committing a cursor or deleting source data in `poll()` - stage it and wait for ACK. 5. Unbounded data in `State` - rewritten every batch. keep O(constant). 6. `std::sync::Mutex` - blocks the executor. Use `tokio::sync::Mutex`. 7. Not setting `ProducedMessage.id` when a stable ID exists - loses idempotency. @@ -137,7 +143,7 @@ Iggy consumer-loop labels use literal API names (`offset=`, `current_offset=`). Mandatory four canonical source state tests (see [connector-testing](../connector-testing/SKILL.md) for the full pattern). Copy from `sources/random_source/src/lib.rs::tests`. Plus config defaults, payload building, schema selection. -Integration tests under `core/integration/tests/connectors//` for any source backed by external infra. Use `#[iggy_harness]` + a `TestFixture` backed by `testcontainers-modules`. Reference: `core/integration/tests/connectors/postgres/postgres_source.rs` (multi-mode tests) + `restart.rs` (state survives restart). +Integration tests under `core/integration/tests/connectors//` for any source backed by external infra. Use `#[iggy_harness]` + a `TestFixture` backed by `testcontainers-modules`. Reference: `core/integration/tests/connectors/postgres/postgres_source.rs` (multi-mode tests) + `restart.rs` (state survives restart). Exercise both ACK and NACK paths when the source stages cursors or destructive work. ## Before declaring done diff --git a/.claude/skills/connector-source/TEMPLATE.md b/.claude/skills/connector-source/TEMPLATE.md index 7a27953654..301df1a435 100644 --- a/.claude/skills/connector-source/TEMPLATE.md +++ b/.claude/skills/connector-source/TEMPLATE.md @@ -16,13 +16,15 @@ helpers below. ```rust /* Apache 2.0 header */ +use std::str::FromStr; +use std::time::Duration; + use async_trait::async_trait; use iggy_connector_sdk::{ - ConnectorState, Error, ProducedMessage, ProducedMessages, Schema, Source, source_connector, + ConnectorState, Error, ProducedMessage, ProducedMessages, Schema, Source, + source::SourceBatchResult, source_connector, }; use serde::{Deserialize, Serialize}; -use std::str::FromStr; -use std::time::Duration; use tokio::sync::Mutex; use tokio::time::sleep; use tracing::{debug, error, info, warn}; @@ -41,7 +43,7 @@ pub struct MySourceConfig { pub verbose_logging: Option, } -#[derive(Debug, Serialize, Deserialize)] +#[derive(Debug, Clone, Serialize, Deserialize)] struct State { cursor: Option, // WAL LSN, scroll id, timestamp, ... last_offset: u64, @@ -56,6 +58,7 @@ pub struct MySource { verbose: bool, client: Option, state: Mutex, + pending: Mutex>, } impl MySource { @@ -88,6 +91,7 @@ impl MySource { last_offset: 0, messages_produced: 0, })), + pending: Mutex::new(None), } } } @@ -109,9 +113,9 @@ impl Source for MySource { async fn poll(&self) -> Result { sleep(self.poll_interval).await; // sleep first - backpressure - let cursor = { self.state.lock().await.cursor.clone() }; // brief read + let mut candidate = self.state.lock().await.clone(); - let fetched = self.fetch_since(cursor.as_deref()).await?; // no lock held + let fetched = self.fetch_since(candidate.cursor.as_deref()).await?; let mut messages = Vec::with_capacity(fetched.len()); let mut next_cursor = None; @@ -137,22 +141,41 @@ impl Source for MySource { ); } - let persisted = { // brief write - let mut state = self.state.lock().await; - state.messages_produced += messages.len() as u64; - if let Some(c) = next_cursor { - state.cursor = Some(c); - } - ConnectorState::serialize(&*state, CONNECTOR_NAME, self.id) - }; + if messages.is_empty() { + return Ok(ProducedMessages { + schema: Schema::Json, + messages, + state: None, + }); + } + + candidate.messages_produced += messages.len() as u64; + candidate.cursor = next_cursor; + let persisted = ConnectorState::serialize(&candidate, CONNECTOR_NAME, self.id) + .ok_or_else(|| Error::Serialization("failed to serialize source state".into()))?; + *self.pending.lock().await = Some(candidate); Ok(ProducedMessages { schema: Schema::Json, messages, - state: persisted, + state: Some(persisted), }) } + async fn on_batch_result(&self, result: SourceBatchResult) -> Result<(), Error> { + let pending = self.pending.lock().await.take(); + match result { + SourceBatchResult::Ack => { + let Some(candidate) = pending else { + return Ok(()); + }; + *self.state.lock().await = candidate; + } + SourceBatchResult::Nack => {} + } + Ok(()) + } + async fn close(&mut self) -> Result<(), Error> { if let Some(client) = self.client.take() { let _ = client; // or `client.close().await;` for sqlx pools diff --git a/core/connectors/runtime/src/stream.rs b/core/connectors/runtime/src/stream.rs index 574fdd34c5..d1b5c41d92 100644 --- a/core/connectors/runtime/src/stream.rs +++ b/core/connectors/runtime/src/stream.rs @@ -24,6 +24,15 @@ use crate::error::RuntimeError; const TOKEN_FILE_PREFIX: &str = "file:"; +fn append_query_parameters(connection_string: &str, parameters: &str) -> String { + let separator = if connection_string.contains('?') { + '&' + } else { + '?' + }; + format!("{connection_string}{separator}{parameters}") +} + fn expand_home(path: &str) -> PathBuf { if let Some(rest) = path.strip_prefix("~/") { if let Some(home) = dirs::home_dir() { @@ -127,7 +136,10 @@ async fn create_client( .filter(|domain| !domain.is_empty()) .map(|domain| format!("&tls_domain={domain}")) .unwrap_or_default(); - format!("{connection_string}?tls=true&tls_ca_file={ca_file}{domain}") + append_query_parameters( + &connection_string, + &format!("tls=true&tls_ca_file={ca_file}{domain}"), + ) } else { connection_string }; @@ -181,6 +193,27 @@ mod tests { assert_eq!(result, PathBuf::from("relative/path")); } + #[test] + fn given_existing_query_when_appending_parameters_should_use_ampersand() { + let connection_string = "iggy://user:password@127.0.0.1:8090?reconnection_retries=0"; + + let result = append_query_parameters(connection_string, "tls=true"); + + assert_eq!( + result, + "iggy://user:password@127.0.0.1:8090?reconnection_retries=0&tls=true" + ); + } + + #[test] + fn given_no_query_when_appending_parameters_should_use_question_mark() { + let connection_string = "iggy://user:password@127.0.0.1:8090"; + + let result = append_query_parameters(connection_string, "tls=true"); + + assert_eq!(result, "iggy://user:password@127.0.0.1:8090?tls=true"); + } + #[test] fn test_resolve_token_direct_value() { let token = "my-secret-token"; diff --git a/core/connectors/sources/postgres_source/README.md b/core/connectors/sources/postgres_source/README.md index 76278d8130..3b1d204e7e 100644 --- a/core/connectors/sources/postgres_source/README.md +++ b/core/connectors/sources/postgres_source/README.md @@ -12,7 +12,7 @@ The PostgreSQL source connector fetches data from PostgreSQL databases and strea - **Mark as Processed**: Mark rows as processed using a boolean column - **Multiple Tables**: Monitor multiple tables simultaneously - **Batch Processing**: Fetch data in configurable batch sizes -- **Offset Tracking**: Keep track of processed records to avoid duplicates +- **Offset Tracking**: Resume incremental polling from the last acknowledged offset ## Configuration @@ -54,7 +54,7 @@ cdc_backend = "builtin" | `connection_string` | string | required | PostgreSQL connection string | | `mode` | string | required | `polling` or `cdc` | | `tables` | array | required | List of tables to monitor | -| `poll_interval` | string | `1s` | How often to poll (e.g., `1s`, `5m`) | +| `poll_interval` | string | `10s` | How often to poll (e.g., `1s`, `5m`) | | `batch_size` | u32 | `1000` | Max rows per poll | | `tracking_column` | string | `id` | Column for incremental updates | | `initial_offset` | string | none | Starting value for tracking column | @@ -74,6 +74,13 @@ cdc_backend = "builtin" | `max_retries` | u32 | `3` | Max retry attempts for transient errors | | `retry_delay` | string | `1s` | Base delay between retries (e.g., `500ms`, `2s`) | +## Delivery Failures + +Delivery is at-least-once, so consumers must tolerate duplicates. A failed send +NACKs the batch and leaves its database progress uncommitted for redelivery. +After five consecutive NACKs, the source stops and requires a manual connector +restart. + ## Output Modes ### JSON Mode (Default) @@ -235,6 +242,10 @@ ALTER TABLE users ADD COLUMN is_processed BOOLEAN DEFAULT false; When `processed_column` is set, the connector automatically adds a `WHERE is_processed = FALSE` filter to the polling query, so only unprocessed rows are fetched. This improves polling efficiency as the table grows. +The connector persists the acknowledged offset before deleting or marking rows. +If it stops in between, the rows have been delivered but may remain unchanged in +PostgreSQL. The persisted offset prevents those rows from being selected again. + ## Supported Column Types The connector handles these PostgreSQL types in JSON mode: @@ -254,7 +265,7 @@ The connector handles these PostgreSQL types in JSON mode: ## CDC Mode -CDC requires PostgreSQL logical replication setup: +CDC requires PostgreSQL 11 or newer and logical replication setup: 1. Set `wal_level = logical` in `postgresql.conf` 2. Restart PostgreSQL @@ -271,6 +282,11 @@ The connector peeks at logical changes and advances the replication slot only after Iggy acknowledges the batch. A failed delivery leaves the slot unchanged so the next poll can read the same changes again. +Advancing the slot fast-forwards through the WAL range that was just peeked, so +each acknowledged batch is decoded twice. Poll and decode errors do not change +the connector's runtime status. Monitor `confirmed_flush_lsn`, retained WAL, and +replication slot lag in PostgreSQL to detect a stuck CDC poller. + The `pg_replicate` backend requires the `cdc_pg_replicate` feature flag at build time. ### Slot Naming diff --git a/core/connectors/sources/postgres_source/src/lib.rs b/core/connectors/sources/postgres_source/src/lib.rs index 16b648085c..3c2ab987df 100644 --- a/core/connectors/sources/postgres_source/src/lib.rs +++ b/core/connectors/sources/postgres_source/src/lib.rs @@ -123,11 +123,10 @@ struct PendingBatch { enum PendingOperation { ProcessRows { table: String, - primary_key_column: String, ids: Vec, + max_offset: String, }, AdvanceReplicationSlot { - slot_name: String, lsn: String, }, } @@ -259,10 +258,7 @@ impl Source for PostgresSource { } }; - let processed_rows = match polled.pending.as_ref() { - Some(pending) => pending.state.processed_rows, - None => self.state.lock().await.processed_rows, - }; + let processed_rows = self.state.lock().await.processed_rows; if self.verbose { info!( "PostgreSQL source connector ID: {} produced {} messages. Total processed: {}", @@ -305,8 +301,9 @@ impl Source for PostgresSource { async fn on_batch_result(&self, result: SourceBatchResult) -> Result<(), Error> { let pending = self.pending_batch.lock().await.take(); - if result == SourceBatchResult::Nack { - return Ok(()); + match result { + SourceBatchResult::Ack => {} + SourceBatchResult::Nack => return Ok(()), } let Some(pending) = pending else { @@ -317,19 +314,14 @@ impl Source for PostgresSource { match operation { PendingOperation::ProcessRows { table, - primary_key_column, ids, + max_offset, } => { - self.mark_or_delete_processed_rows( - self.get_pool()?, - &table, - &primary_key_column, - &ids, - ) - .await?; + self.mark_or_delete_processed_rows(self.get_pool()?, &table, &ids, &max_offset) + .await?; } - PendingOperation::AdvanceReplicationSlot { slot_name, lsn } => { - self.advance_replication_slot(&slot_name, &lsn).await?; + PendingOperation::AdvanceReplicationSlot { lsn } => { + self.advance_replication_slot(&lsn).await?; } } } @@ -423,11 +415,7 @@ impl PostgresSource { } } - let slot_name = self - .config - .replication_slot - .as_deref() - .unwrap_or("iggy_slot"); + let slot_name = self.replication_slot(); let existing_plugin: Option = sqlx::query_scalar("SELECT plugin FROM pg_replication_slots WHERE slot_name = $1") @@ -478,11 +466,7 @@ impl PostgresSource { async fn poll_cdc_builtin(&self) -> Result { let pool = self.get_pool()?; - let slot_name = self - .config - .replication_slot - .as_deref() - .unwrap_or("iggy_slot"); + let slot_name = self.replication_slot(); let capture_ops = self .config .capture_operations @@ -499,7 +483,7 @@ impl PostgresSource { // stops the backlog from growing unbounded across many transactions // the way NULL (no limit at all) did. let rows = sqlx::query( - "SELECT lsn::text AS lsn, xid, data FROM pg_logical_slot_peek_changes($1, NULL, $2)", + "SELECT lsn::text AS lsn, data FROM pg_logical_slot_peek_changes($1, NULL, $2)", ) .bind(slot_name) .bind(batch_size) @@ -556,10 +540,7 @@ impl PostgresSource { state.last_poll_time = Utc::now(); Some(PendingBatch { state, - operations: vec![PendingOperation::AdvanceReplicationSlot { - slot_name: slot_name.to_string(), - lsn, - }], + operations: vec![PendingOperation::AdvanceReplicationSlot { lsn }], }) } else { None @@ -575,12 +556,8 @@ impl PostgresSource { let mut candidate_state = self.state.lock().await.clone(); let batch_size = self.config.batch_size.unwrap_or(1000); - let tracking_column = self.config.tracking_column.as_deref().unwrap_or("id"); - let pk_column = self - .config - .primary_key_column - .as_deref() - .unwrap_or(tracking_column); + let tracking_column = self.tracking_column(); + let pk_column = self.primary_key_column(); let row_config = RowProcessingConfig { table: "", @@ -592,8 +569,6 @@ impl PostgresSource { include_metadata: self.config.include_metadata.unwrap_or(true), }; - // Collect state updates to apply after processing - let mut state_updates: Vec<(String, String)> = Vec::new(); let mut total_processed: u64 = 0; for table in &self.config.tables { @@ -617,10 +592,14 @@ impl PostgresSource { self.get_max_retries(), self.retry_delay.as_millis() as u64, ) - .await?; + .await + .map_err(|e| { + Error::Connection(format!("failed to poll PostgreSQL table '{table}': {e}")) + })?; let mut max_offset: Option = None; let mut processed_ids: Vec = Vec::new(); + let mut table_processed = 0; for row in rows { let processed = self.process_row(&row, &table_config)?; @@ -634,33 +613,37 @@ impl PostgresSource { messages.push(processed.message); total_processed += 1; + table_processed += 1; } - if !processed_ids.is_empty() { + if self.should_process_rows() && !processed_ids.is_empty() { + let max_offset = max_offset.clone().ok_or_else(|| { + Error::InvalidRecordValue(format!( + "tracking column '{tracking_column}' is missing from rows read from '{table}'" + )) + })?; operations.push(PendingOperation::ProcessRows { table: table.clone(), - primary_key_column: pk_column.to_string(), ids: processed_ids, + max_offset, }); } - // 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}'"); } } let pending = if total_processed > 0 { candidate_state.processed_rows += total_processed; - for (table, offset) in state_updates { - candidate_state.tracking_offsets.insert(table, offset); - } candidate_state.last_poll_time = Utc::now(); Some(PendingBatch { state: candidate_state, @@ -673,18 +656,26 @@ impl PostgresSource { Ok(PolledBatch { messages, pending }) } - async fn advance_replication_slot(&self, slot_name: &str, lsn: &str) -> Result<(), Error> { - sqlx::query("SELECT pg_replication_slot_advance($1, $2::pg_lsn)") - .bind(slot_name) - .bind(lsn) - .execute(self.get_pool()?) - .await - .map_err(|e| { - error!("Failed to advance replication slot '{slot_name}' to {lsn}: {e}"); - Error::Connection(format!( - "failed to advance replication slot '{slot_name}' to {lsn}: {e}" - )) - })?; + async fn advance_replication_slot(&self, lsn: &str) -> Result<(), Error> { + let slot_name = self.replication_slot(); + let pool = self.get_pool()?; + with_retry( + || { + sqlx::query("SELECT pg_replication_slot_advance($1, $2::pg_lsn)") + .bind(slot_name) + .bind(lsn) + .execute(pool) + }, + self.get_max_retries(), + self.retry_delay.as_millis() as u64, + ) + .await + .map_err(|e| { + error!("Failed to advance replication slot '{slot_name}' to {lsn}: {e}"); + Error::Connection(format!( + "failed to advance replication slot '{slot_name}' to {lsn}: {e}" + )) + })?; Ok(()) } @@ -692,15 +683,17 @@ impl PostgresSource { &self, pool: &Pool, table: &str, - pk_column: &str, ids: &[String], + max_offset: &str, ) -> Result<(), Error> { if ids.is_empty() { return Ok(()); } let quoted_table = quote_qualified_identifier(table)?; - let quoted_pk = quote_identifier(pk_column)?; + let quoted_pk = quote_identifier(self.primary_key_column())?; + let quoted_tracking = quote_identifier(self.tracking_column())?; + let tracking_boundary = format_offset_value(max_offset); let ids_list = ids .iter() @@ -715,8 +708,10 @@ impl PostgresSource { .join(", "); if self.config.delete_after_read.unwrap_or(false) { - let delete_query = - format!("DELETE FROM {quoted_table} WHERE {quoted_pk} IN ({ids_list})"); + let delete_query = format!( + "DELETE FROM {quoted_table} WHERE {quoted_pk} IN ({ids_list}) \ + AND {quoted_tracking} <= {tracking_boundary}" + ); if self.verbose { info!("Deleting {} processed rows from '{table}'", ids.len()); @@ -724,17 +719,21 @@ impl PostgresSource { debug!("Deleting {} processed rows from '{table}'", ids.len()); } - sqlx::query(sqlx::AssertSqlSafe(delete_query)) - .execute(pool) - .await - .map_err(|e| { - error!("Failed to delete processed rows: {e}"); - Error::InvalidRecord - })?; + with_retry( + || sqlx::query(sqlx::AssertSqlSafe(delete_query.as_str())).execute(pool), + self.get_max_retries(), + self.retry_delay.as_millis() as u64, + ) + .await + .map_err(|e| { + error!("Failed to delete processed rows: {e}"); + Error::Connection(format!("failed to delete processed rows: {e}")) + })?; } else if let Some(processed_col) = &self.config.processed_column { let quoted_processed = quote_identifier(processed_col)?; let update_query = format!( - "UPDATE {quoted_table} SET {quoted_processed} = TRUE WHERE {quoted_pk} IN ({ids_list})" + "UPDATE {quoted_table} SET {quoted_processed} = TRUE \ + WHERE {quoted_pk} IN ({ids_list}) AND {quoted_tracking} <= {tracking_boundary}" ); if self.verbose { @@ -743,13 +742,16 @@ impl PostgresSource { debug!("Marking {} rows as processed in '{table}'", ids.len()); } - sqlx::query(sqlx::AssertSqlSafe(update_query)) - .execute(pool) - .await - .map_err(|e| { - error!("Failed to mark rows as processed: {e}"); - Error::InvalidRecord - })?; + with_retry( + || sqlx::query(sqlx::AssertSqlSafe(update_query.as_str())).execute(pool), + self.get_max_retries(), + self.retry_delay.as_millis() as u64, + ) + .await + .map_err(|e| { + error!("Failed to mark rows as processed: {e}"); + Error::Connection(format!("failed to mark rows as processed: {e}")) + })?; } Ok(()) @@ -774,6 +776,28 @@ impl PostgresSource { self.config.max_retries.unwrap_or(DEFAULT_MAX_RETRIES) } + fn should_process_rows(&self) -> bool { + self.config.delete_after_read.unwrap_or(false) || self.config.processed_column.is_some() + } + + fn tracking_column(&self) -> &str { + self.config.tracking_column.as_deref().unwrap_or("id") + } + + fn primary_key_column(&self) -> &str { + self.config + .primary_key_column + .as_deref() + .unwrap_or_else(|| self.tracking_column()) + } + + fn replication_slot(&self) -> &str { + self.config + .replication_slot + .as_deref() + .unwrap_or("iggy_slot") + } + fn build_polling_query( &self, table: &str, @@ -1752,7 +1776,11 @@ fn parse_bare_scalar(token: &str) -> serde_json::Value { } } -async fn with_retry(operation: F, max_retries: u32, delay_ms: u64) -> Result +async fn with_retry( + operation: F, + max_retries: u32, + delay_ms: u64, +) -> Result where F: Fn() -> Fut, Fut: std::future::Future>, @@ -1765,7 +1793,7 @@ where attempts += 1; if attempts >= max_retries || !is_transient_error(&e) { error!("Database operation failed after {attempts} attempts: {e}"); - return Err(Error::InvalidRecord); + return Err(e); } warn!( "Transient database error (attempt {attempts}/{max_retries}): {e}. Retrying in {delay_ms}ms..." @@ -1809,6 +1837,8 @@ mod cdc_fixtures; #[cfg(test)] mod tests { + use std::sync::atomic::{AtomicU32, Ordering}; + use super::*; fn test_config() -> PostgresSourceConfig { @@ -2629,6 +2659,27 @@ mod tests { }); } + #[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_nack_when_batch_is_staged_should_keep_committed_state() { let src = PostgresSource::new(1, test_config(), None); @@ -2644,11 +2695,10 @@ mod tests { operations: vec![ PendingOperation::ProcessRows { table: "users".to_string(), - primary_key_column: "id".to_string(), ids: vec!["3".to_string()], + max_offset: "3".to_string(), }, PendingOperation::AdvanceReplicationSlot { - slot_name: "iggy_slot".to_string(), lsn: "0/16D32A0".to_string(), }, ], diff --git a/core/integration/src/harness/handle/connectors_runtime.rs b/core/integration/src/harness/handle/connectors_runtime.rs index 1718972ab7..4693f04896 100644 --- a/core/integration/src/harness/handle/connectors_runtime.rs +++ b/core/integration/src/harness/handle/connectors_runtime.rs @@ -83,6 +83,10 @@ impl ConnectorsRuntimeHandle { self.iggy_connection_options = Some(options.into()); } + pub fn clear_iggy_connection_options(&mut self) { + self.iggy_connection_options = None; + } + fn build_envs(&mut self) { let state_path = self.context.connectors_runtime_state_path(self.server_id); self.envs.insert( diff --git a/core/integration/tests/connectors/postgres/mod.rs b/core/integration/tests/connectors/postgres/mod.rs index ee992b36b8..3d673f3789 100644 --- a/core/integration/tests/connectors/postgres/mod.rs +++ b/core/integration/tests/connectors/postgres/mod.rs @@ -20,12 +20,22 @@ mod postgres_source; mod postgres_source_cdc; mod restart; -use crate::connectors::TestMessage; +use std::time::Duration; + +use iggy_connector_sdk::api::{ConnectorRuntimeStats, ConnectorStats, ConnectorStatus}; +use reqwest::Client; use serde::Deserialize; +use tokio::time::{sleep, timeout}; + +use crate::connectors::TestMessage; +const API_KEY: &str = "test-api-key"; +const SOURCE_KEY: &str = "postgres"; +const DEFAULT_SLOT: &str = "iggy_slot"; const TEST_MESSAGE_COUNT: usize = 3; const POLL_ATTEMPTS: usize = 100; const POLL_INTERVAL_MS: u64 = 50; +const SEND_FAILURE_TIMEOUT: Duration = Duration::from_secs(25); #[derive(Debug, Deserialize)] struct DatabaseRecord { @@ -33,3 +43,33 @@ struct DatabaseRecord { operation_type: String, data: TestMessage, } + +async fn source_stats(http: &Client, api_url: &str) -> Option { + let response = http + .get(format!("{api_url}/stats")) + .header("api-key", API_KEY) + .send() + .await + .ok()?; + let stats = response.json::().await.ok()?; + stats + .connectors + .into_iter() + .find(|connector| connector.key == SOURCE_KEY) +} + +async fn wait_for_source_errors(http: &Client, api_url: &str, minimum_errors: u64) { + timeout(SEND_FAILURE_TIMEOUT, async { + loop { + if let Some(source) = source_stats(http, api_url).await + && source.status == ConnectorStatus::Error + && source.errors >= minimum_errors + { + break; + } + sleep(Duration::from_millis(POLL_INTERVAL_MS)).await; + } + }) + .await + .expect("PostgreSQL source did not retry the NACKed batch"); +} diff --git a/core/integration/tests/connectors/postgres/postgres_source.rs b/core/integration/tests/connectors/postgres/postgres_source.rs index 34214aeaa7..740295221b 100644 --- a/core/integration/tests/connectors/postgres/postgres_source.rs +++ b/core/integration/tests/connectors/postgres/postgres_source.rs @@ -19,13 +19,15 @@ use std::time::Duration; use iggy_common::MessageClient; use iggy_common::{Consumer, Identifier, PollingStrategy}; -use iggy_connector_sdk::api::{ConnectorRuntimeStats, ConnectorStatus}; use integration::harness::seeds; use integration::iggy_harness; use reqwest::Client; -use tokio::time::{sleep, timeout}; +use tokio::time::sleep; -use super::{DatabaseRecord, POLL_ATTEMPTS, POLL_INTERVAL_MS, TEST_MESSAGE_COUNT}; +use super::{ + DatabaseRecord, POLL_ATTEMPTS, POLL_INTERVAL_MS, TEST_MESSAGE_COUNT, source_stats, + wait_for_source_errors, +}; use crate::connectors::create_test_messages; use crate::connectors::fixtures::{ PostgresOps, PostgresSourceByteaFixture, PostgresSourceDeleteFixture, @@ -33,10 +35,6 @@ use crate::connectors::fixtures::{ PostgresSourceOps, }; -const API_KEY: &str = "test-api-key"; -const SOURCE_KEY: &str = "postgres"; -const SEND_FAILURE_TIMEOUT: Duration = Duration::from_secs(25); - #[iggy_harness( server(connectors_runtime(config_path = "tests/connectors/postgres/source.toml")), seed = seeds::connector_stream @@ -140,9 +138,9 @@ async fn json_rows_source_produces_messages_to_iggy( server(connectors_runtime(config_path = "tests/connectors/postgres/source.toml")), seed = seeds::connector_stream )] -async fn given_rows_in_postgres_when_iggy_server_crashes_should_redeliver_after_restart( +async fn given_delete_after_read_when_iggy_crashes_should_delete_only_after_redelivery( harness: &mut TestHarness, - fixture: PostgresSourceJsonFixture, + fixture: PostgresSourceDeleteFixture, ) { let pool = fixture.create_pool().await.expect("Failed to create pool"); fixture.create_table(&pool).await; @@ -168,38 +166,26 @@ async fn given_rows_in_postgres_when_iggy_server_crashes_should_redeliver_after_ .expect("connectors runtime") .http_url(); let http = Client::new(); - let errors_before_failure = source_errors(&http, &api_url).await; + let errors_before_failure = source_stats(&http, &api_url) + .await + .expect("PostgreSQL source stats should be present") + .errors; harness.kill_node(0).expect("Failed to kill Iggy server"); - let expected = create_test_messages(TEST_MESSAGE_COUNT); - let mut transaction = pool.begin().await.expect("Failed to begin transaction"); - let insert = format!( - "INSERT INTO {} (id, name, count, amount, active, timestamp, tag) \ - VALUES ($1, $2, $3, $4, $5, $6, $7)", - fixture.table_name() - ); - for message in &expected { - let tag = format!("{:<10}", format!("tag_{}", message.id)); - sqlx::query(sqlx::AssertSqlSafe(insert.as_str())) - .bind(message.id as i32) - .bind(&message.name) - .bind(message.count as i32) - .bind(message.amount) - .bind(message.active) - .bind(message.timestamp) - .bind(tag) - .execute(&mut *transaction) - .await - .expect("Failed to insert source row"); + for index in 0..TEST_MESSAGE_COUNT { + fixture + .insert_row(&pool, &format!("row_{index}"), index as i32) + .await; } - transaction - .commit() - .await - .expect("Failed to commit source rows"); - // The second error proves that NACK made the same rows eligible for another poll. wait_for_source_errors(&http, &api_url, errors_before_failure + 2).await; + assert_eq!( + fixture.count_rows(&pool).await, + TEST_MESSAGE_COUNT as i64, + "NACKed rows must not be deleted" + ); + harness .server_mut() .stop_dependents() @@ -207,6 +193,11 @@ async fn given_rows_in_postgres_when_iggy_server_crashes_should_redeliver_after_ harness .restart_node(0) .expect("Failed to restart Iggy server"); + harness + .server_mut() + .connectors_runtime_mut() + .expect("connectors runtime") + .clear_iggy_connection_options(); harness .server_mut() .start_dependents() @@ -217,7 +208,7 @@ async fn given_rows_in_postgres_when_iggy_server_crashes_should_redeliver_after_ let stream_id: Identifier = seeds::names::STREAM.try_into().unwrap(); let topic_id: Identifier = seeds::names::TOPIC.try_into().unwrap(); let consumer_id: Identifier = "send_failure_consumer".try_into().unwrap(); - let mut received = Vec::new(); + let mut received = 0; for _ in 0..POLL_ATTEMPTS { if let Ok(polled) = client @@ -232,10 +223,8 @@ async fn given_rows_in_postgres_when_iggy_server_crashes_should_redeliver_after_ ) .await { - received.extend(polled.messages.into_iter().filter_map(|message| { - serde_json::from_slice::(&message.payload).ok() - })); - if received.len() >= TEST_MESSAGE_COUNT { + received += polled.messages.len(); + if received >= TEST_MESSAGE_COUNT { break; } } @@ -243,58 +232,23 @@ async fn given_rows_in_postgres_when_iggy_server_crashes_should_redeliver_after_ } assert_eq!( - received.len(), - TEST_MESSAGE_COUNT, + received, TEST_MESSAGE_COUNT, "Rows polled during the failed send should be delivered after restart" ); - for (record, expected) in received.iter().zip(expected) { - assert_eq!(record.data.id, expected.id); + + let mut remaining_rows = fixture.count_rows(&pool).await; + for _ in 0..POLL_ATTEMPTS { + if remaining_rows == 0 { + break; + } + sleep(Duration::from_millis(POLL_INTERVAL_MS)).await; + remaining_rows = fixture.count_rows(&pool).await; } + assert_eq!(remaining_rows, 0, "ACKed rows should be deleted"); pool.close().await; } -async fn source_errors(http: &Client, api_url: &str) -> u64 { - http.get(format!("{api_url}/stats")) - .header("api-key", API_KEY) - .send() - .await - .expect("runtime stats should be available") - .json::() - .await - .expect("runtime stats should be valid") - .connectors - .into_iter() - .find(|connector| connector.key == SOURCE_KEY) - .expect("PostgreSQL source stats should be present") - .errors -} - -async fn wait_for_source_errors(http: &Client, api_url: &str, minimum_errors: u64) { - timeout(SEND_FAILURE_TIMEOUT, async { - loop { - if let Ok(response) = http - .get(format!("{api_url}/stats")) - .header("api-key", API_KEY) - .send() - .await - && let Ok(stats) = response.json::().await - && let Some(source) = stats - .connectors - .iter() - .find(|connector| connector.key == SOURCE_KEY) - && source.status == ConnectorStatus::Error - && source.errors >= minimum_errors - { - break; - } - sleep(Duration::from_millis(POLL_INTERVAL_MS)).await; - } - }) - .await - .expect("PostgreSQL source did not retry the NACKed batch"); -} - #[iggy_harness( server(connectors_runtime(config_path = "tests/connectors/postgres/source.toml")), seed = seeds::connector_stream diff --git a/core/integration/tests/connectors/postgres/postgres_source_cdc.rs b/core/integration/tests/connectors/postgres/postgres_source_cdc.rs index efefa8bd60..c36e36789a 100644 --- a/core/integration/tests/connectors/postgres/postgres_source_cdc.rs +++ b/core/integration/tests/connectors/postgres/postgres_source_cdc.rs @@ -15,7 +15,10 @@ // specific language governing permissions and limitations // under the License. -use super::{POLL_ATTEMPTS, POLL_INTERVAL_MS}; +use super::{ + API_KEY, DEFAULT_SLOT, POLL_ATTEMPTS, POLL_INTERVAL_MS, SOURCE_KEY, source_stats, + wait_for_source_errors, +}; use crate::connectors::create_test_messages; use crate::connectors::fixtures::{PostgresOps, PostgresSourceCdcFixture, PostgresSourceOps}; use iggy::prelude::IggyClient; @@ -29,10 +32,6 @@ use serde::Deserialize; use std::time::Duration; use tokio::time::sleep; -const API_KEY: &str = "test-api-key"; -const SOURCE_KEY: &str = "postgres"; -const DEFAULT_SLOT: &str = "iggy_slot"; - #[derive(Debug, Deserialize)] struct CdcRecord { table_name: String, @@ -75,6 +74,17 @@ async fn poll_cdc_records( received } +async fn slot_contains_change(pool: &sqlx::PgPool, expected_value: &str) -> bool { + let changes = sqlx::query_scalar::<_, String>( + "SELECT data FROM pg_logical_slot_peek_changes($1, NULL, NULL)", + ) + .bind(DEFAULT_SLOT) + .fetch_all(pool) + .await + .expect("CDC replication slot should be readable"); + changes.iter().any(|change| change.contains(expected_value)) +} + // End-to-end CDC coverage against a real wal_level=logical container: // INSERT, UPDATE, PK-changing UPDATE, DELETE, a rolled-back transaction // (must produce nothing), an untracked table (must be filtered out), a @@ -231,6 +241,108 @@ async fn cdc_source_captures_insert_update_delete( pool.close().await; } +#[iggy_harness( + cluster_nodes = 1, + server(connectors_runtime(config_path = "tests/connectors/postgres/source.toml")), + seed = seeds::connector_stream +)] +async fn given_cdc_change_when_iggy_crashes_should_advance_slot_only_after_redelivery( + harness: &mut TestHarness, + fixture: PostgresSourceCdcFixture, +) { + let pool = fixture.create_pool().await.expect("Failed to create pool"); + fixture.create_table(&pool).await; + + harness + .server_mut() + .stop_dependents() + .expect("Failed to stop connectors runtime"); + harness + .server_mut() + .connectors_runtime_mut() + .expect("connectors runtime") + .set_iggy_connection_options("reconnection_retries=0"); + harness + .server_mut() + .start_dependents() + .await + .expect("Failed to restart connectors runtime"); + + let api_url = harness + .connectors_runtime() + .expect("connectors runtime") + .http_url(); + let http = Client::new(); + wait_for_source_status(&http, &api_url, ConnectorStatus::Running).await; + let errors_before_failure = source_stats(&http, &api_url) + .await + .expect("PostgreSQL source stats should be present") + .errors; + harness.kill_node(0).expect("Failed to kill Iggy server"); + + let [expected] = create_test_messages(1).try_into().unwrap(); + fixture + .insert_row( + &pool, + expected.id as i32, + &expected.name, + expected.count as i32, + expected.amount, + expected.active, + expected.timestamp, + ) + .await; + + wait_for_source_errors(&http, &api_url, errors_before_failure + 2).await; + assert!( + slot_contains_change(&pool, &expected.name).await, + "NACKed CDC change must remain available in the replication slot" + ); + + harness + .server_mut() + .stop_dependents() + .expect("Failed to stop connectors runtime"); + harness + .restart_node(0) + .expect("Failed to restart Iggy server"); + harness + .server_mut() + .connectors_runtime_mut() + .expect("connectors runtime") + .clear_iggy_connection_options(); + harness + .server_mut() + .start_dependents() + .await + .expect("Failed to restart connectors runtime"); + + let client = harness.root_client().await.unwrap(); + let stream_id: Identifier = seeds::names::STREAM.try_into().unwrap(); + let topic_id: Identifier = seeds::names::TOPIC.try_into().unwrap(); + let consumer_id: Identifier = "cdc_send_failure_consumer".try_into().unwrap(); + let received = poll_cdc_records(&client, &stream_id, &topic_id, &consumer_id, 1).await; + + assert_eq!(received.len(), 1, "CDC change should be redelivered"); + assert_eq!(received[0].operation_type, "INSERT"); + assert_eq!(received[0].data["id"], serde_json::json!(expected.id)); + + let mut change_remains = slot_contains_change(&pool, &expected.name).await; + for _ in 0..POLL_ATTEMPTS { + if !change_remains { + break; + } + sleep(Duration::from_millis(POLL_INTERVAL_MS)).await; + change_remains = slot_contains_change(&pool, &expected.name).await; + } + assert!( + !change_remains, + "ACKed CDC change should be consumed from the replication slot" + ); + + pool.close().await; +} + async fn wait_for_source_status( http: &Client, api_url: &str, @@ -376,11 +488,9 @@ async fn drop_replication_slot_retrying(pool: &sqlx::PgPool, slot: &str) { // connector that silently drops every change or emits wrong data - the // same silent-death shape as the slot mismatch above. Config is fixed // one field at a time until restart succeeds and CDC resumes. -// 3. Changes written while the connector is down (the slot retains WAL -// regardless of consumer state) - not the at-least-once crash window -// where the slot has already been consumed but send/state-persist -// hasn't happened yet. That gap remains open until the slot-peek/LSN -// work lands. +// 3. Changes written while the connector is down. The slot retains WAL +// regardless of consumer state, then the connector peeks and advances +// it only after Iggy acknowledges the recovered batch. #[iggy_harness( server(connectors_runtime(config_path = "tests/connectors/postgres/source_cdc_restart.toml")), seed = seeds::connector_stream diff --git a/core/integration/tests/connectors/postgres/restart.rs b/core/integration/tests/connectors/postgres/restart.rs index fd226785a5..bd58d3547c 100644 --- a/core/integration/tests/connectors/postgres/restart.rs +++ b/core/integration/tests/connectors/postgres/restart.rs @@ -15,7 +15,7 @@ // specific language governing permissions and limitations // under the License. -use super::{POLL_ATTEMPTS, POLL_INTERVAL_MS, TEST_MESSAGE_COUNT}; +use super::{API_KEY, POLL_ATTEMPTS, POLL_INTERVAL_MS, TEST_MESSAGE_COUNT}; use crate::connectors::fixtures::{PostgresOps, PostgresSinkFixture}; use crate::connectors::{TestMessage, create_test_messages}; use bytes::Bytes; @@ -29,7 +29,6 @@ use reqwest::Client; use std::time::Duration; use tokio::time::sleep; -const API_KEY: &str = "test-api-key"; const SINK_TABLE: &str = "iggy_messages"; const SINK_KEY: &str = "postgres";