From 7b1192dbcbcc855dba33622a5dcdf02fcee6b446 Mon Sep 17 00:00:00 2001 From: Rohan Dubey Date: Mon, 24 Aug 2026 03:46:24 +0530 Subject: [PATCH 1/4] 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/4] 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"; From 27db796efa0a6275f7deb2deef368b080729c804 Mon Sep 17 00:00:00 2001 From: Rohan Dubey Date: Tue, 1 Sep 2026 01:05:57 +0530 Subject: [PATCH 3/4] Update stream.rs --- core/connectors/runtime/src/stream.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/core/connectors/runtime/src/stream.rs b/core/connectors/runtime/src/stream.rs index bafd01698b..24c4ff320d 100644 --- a/core/connectors/runtime/src/stream.rs +++ b/core/connectors/runtime/src/stream.rs @@ -142,10 +142,10 @@ fn connection_string_with_token( .filter(|domain| !domain.is_empty()) .map(|domain| format!("&tls_domain={domain}")) .unwrap_or_default(); - append_query_parameters( + Ok(append_query_parameters( &connection_string, &format!("tls=true&tls_ca_file={ca_file}{domain}"), - ) + )) } else { Ok(connection_string) } From 17d9d673b92a026bd894bb76c261f72f358f119d Mon Sep 17 00:00:00 2001 From: Rohan Dubey Date: Thu, 3 Sep 2026 17:48:23 +0530 Subject: [PATCH 4/4] fix(connectors): harden postgres source ack processing --- .claude/skills/connector-source/SKILL.md | 2 + .claude/skills/connector-source/TEMPLATE.md | 16 +- Cargo.lock | 4 + Cargo.toml | 1 + core/connectors/runtime/src/stream.rs | 25 +- .../sources/postgres_source/README.md | 5 + .../sources/postgres_source/src/lib.rs | 437 +++++++++++------- .../tests/connectors/fixtures/mod.rs | 2 +- .../tests/connectors/fixtures/postgres/mod.rs | 2 +- .../connectors/fixtures/postgres/source.rs | 91 ++++ .../connectors/postgres/postgres_source.rs | 60 ++- .../postgres/postgres_source_cdc.rs | 76 ++- 12 files changed, 535 insertions(+), 186 deletions(-) diff --git a/.claude/skills/connector-source/SKILL.md b/.claude/skills/connector-source/SKILL.md index 1d8255108f..1d761462b5 100644 --- a/.claude/skills/connector-source/SKILL.md +++ b/.claude/skills/connector-source/SKILL.md @@ -69,6 +69,8 @@ let persisted = ConnectorState::serialize(&candidate, CONNECTOR_NAME, self.id) 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. +The callback must complete within the SDK's 30-second batch-result window, so +bound connection acquisition and the full retry budget comfortably below 30 seconds. ### Sleep first diff --git a/.claude/skills/connector-source/TEMPLATE.md b/.claude/skills/connector-source/TEMPLATE.md index 301df1a435..04b8c6062e 100644 --- a/.claude/skills/connector-source/TEMPLATE.md +++ b/.claude/skills/connector-source/TEMPLATE.md @@ -163,16 +163,12 @@ impl Source for MySource { } 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 => {} - } + let (SourceBatchResult::Ack, Some(candidate)) = + (result, self.pending.lock().await.take()) + else { + return Ok(()); + }; + *self.state.lock().await = candidate; Ok(()) } diff --git a/Cargo.lock b/Cargo.lock index 251f7f309a..94c41ea573 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -12590,6 +12590,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "05b44e85bf579a8eeb4ceaa77a3a523baf2bf0e9bac7e40f405d537b5d2d5ccb" dependencies = [ "base64", + "bigdecimal", "bytes", "cfg-if", "chrono", @@ -12666,6 +12667,7 @@ version = "0.9.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "90b8020fe17c5f2c245bfa2505d7ef59c5604839527c740266ad2214acebea27" dependencies = [ + "bigdecimal", "bitflags 2.13.1", "byteorder", "bytes", @@ -12696,6 +12698,7 @@ checksum = "87a2bdd6e83f6b3ea525ca9fee568030508b58355a43d0b2c1674d5f79dcd65e" dependencies = [ "atoi", "base64", + "bigdecimal", "bitflags 2.13.1", "byteorder", "chrono", @@ -12712,6 +12715,7 @@ dependencies = [ "log", "md-5 0.11.0", "memchr", + "num-bigint", "rand 0.10.2", "serde", "serde_json", diff --git a/Cargo.toml b/Cargo.toml index 7872f40937..faebf9cf38 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -312,6 +312,7 @@ sqlx = { version = "0.9.0", features = [ "runtime-tokio", "tls-rustls", "postgres", + "bigdecimal", # "mysql": Doris exposes a MySQL-wire frontend; the Doris sink's # integration-test fixture talks to it over this driver. Cargo unifies # features across the workspace, so it is declared once here. diff --git a/core/connectors/runtime/src/stream.rs b/core/connectors/runtime/src/stream.rs index 24c4ff320d..5591444959 100644 --- a/core/connectors/runtime/src/stream.rs +++ b/core/connectors/runtime/src/stream.rs @@ -24,12 +24,8 @@ 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 { - '?' - }; +fn append_query_parameters(connection_string: &str, address: &str, parameters: &str) -> String { + let separator = if address.contains('?') { '&' } else { '?' }; format!("{connection_string}{separator}{parameters}") } @@ -144,6 +140,7 @@ fn connection_string_with_token( .unwrap_or_default(); Ok(append_query_parameters( &connection_string, + &config.address, &format!("tls=true&tls_ca_file={ca_file}{domain}"), )) } else { @@ -205,8 +202,9 @@ mod tests { #[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 address = "127.0.0.1:8090?reconnection_retries=0"; - let result = append_query_parameters(connection_string, "tls=true"); + let result = append_query_parameters(connection_string, address, "tls=true"); assert_eq!( result, @@ -217,12 +215,23 @@ mod tests { #[test] fn given_no_query_when_appending_parameters_should_use_question_mark() { let connection_string = "iggy://user:password@127.0.0.1:8090"; + let address = "127.0.0.1:8090"; - let result = append_query_parameters(connection_string, "tls=true"); + let result = append_query_parameters(connection_string, address, "tls=true"); assert_eq!(result, "iggy://user:password@127.0.0.1:8090?tls=true"); } + #[test] + fn given_question_mark_in_credentials_should_use_address_separator() { + let connection_string = "iggy://user:pass?word@127.0.0.1:8090"; + let address = "127.0.0.1:8090"; + + let result = append_query_parameters(connection_string, address, "tls=true"); + + assert_eq!(result, "iggy://user:pass?word@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 3b1d204e7e..0b142da16a 100644 --- a/core/connectors/sources/postgres_source/README.md +++ b/core/connectors/sources/postgres_source/README.md @@ -242,6 +242,11 @@ 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. +With the generated polling query, a row whose tracking value moves past the +batch boundary between poll and acknowledgement is left unchanged and returns +in a later poll. Custom queries do not apply this boundary because their result +order is not guaranteed. + 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. diff --git a/core/connectors/sources/postgres_source/src/lib.rs b/core/connectors/sources/postgres_source/src/lib.rs index 3c2ab987df..8b738ddd90 100644 --- a/core/connectors/sources/postgres_source/src/lib.rs +++ b/core/connectors/sources/postgres_source/src/lib.rs @@ -31,6 +31,7 @@ use secrecy::{ExposeSecret, SecretString}; use serde::{Deserialize, Serialize}; use sqlx::postgres::PgPoolOptions; use sqlx::postgres::types::{Oid, PgInterval, PgTimeTz}; +use sqlx::types::BigDecimal; use sqlx::{Column, Pool, Postgres, Row, TypeInfo, ValueRef}; use tokio::sync::Mutex; use tracing::{debug, error, info, warn}; @@ -40,6 +41,8 @@ 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); #[derive(Debug)] pub struct PostgresSource { @@ -116,6 +119,7 @@ struct PolledBatch { #[derive(Debug)] struct PendingBatch { state: State, + // TODO: Persist pending operations with the candidate state and replay them during open. operations: Vec, } @@ -124,7 +128,7 @@ enum PendingOperation { ProcessRows { table: String, ids: Vec, - max_offset: String, + tracking_boundary: Option, }, AdvanceReplicationSlot { lsn: String, @@ -300,33 +304,49 @@ impl Source for PostgresSource { } async fn on_batch_result(&self, result: SourceBatchResult) -> Result<(), Error> { - let pending = self.pending_batch.lock().await.take(); - match result { - SourceBatchResult::Ack => {} - SourceBatchResult::Nack => return Ok(()), - } - - let Some(pending) = pending else { + let (SourceBatchResult::Ack, Some(pending)) = + (result, self.pending_batch.lock().await.take()) + else { return Ok(()); }; - for operation in pending.operations { - match operation { - PendingOperation::ProcessRows { - table, - ids, - max_offset, - } => { - self.mark_or_delete_processed_rows(self.get_pool()?, &table, &ids, &max_offset) - .await?; - } - PendingOperation::AdvanceReplicationSlot { lsn } => { - self.advance_replication_slot(&lsn).await?; + 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) + .await + .is_err() + { + warn!( + "PostgreSQL source connector ID: {} exceeded the ACK cleanup budget", + self.id + ); } - *self.state.lock().await = pending.state; + *self.state.lock().await = state; Ok(()) } @@ -357,6 +377,7 @@ impl PostgresSource { let pool = PgPoolOptions::new() .max_connections(max_connections) + .acquire_timeout(ACK_ACQUIRE_TIMEOUT) .connect(self.config.connection_string.expose_secret()) .await .map_err(|e| Error::InitError(format!("Failed to connect to PostgreSQL: {e}")))?; @@ -477,22 +498,33 @@ impl PostgresSource { (!self.config.tables.is_empty()).then_some(self.config.tables.as_slice()); let batch_size = self.config.batch_size.unwrap_or(1000) as i32; + let wal_flush_lsn = with_retry( + || sqlx::query_scalar("SELECT pg_current_wal_flush_lsn()::text").fetch_one(pool), + self.get_max_retries(), + self.retry_delay.as_millis() as u64, + ) + .await + .map_err(|e| Error::Connection(format!("failed to read current WAL flush LSN: {e}")))?; + // Database I/O without holding the lock. upto_nchanges is only // checked at transaction-commit boundaries (a single huge transaction // 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::text AS lsn, data FROM pg_logical_slot_peek_changes($1, NULL, $2)", + let rows = with_retry( + || { + sqlx::query( + "SELECT lsn::text AS lsn, data FROM pg_logical_slot_peek_changes($1, NULL, $2)", + ) + .bind(slot_name) + .bind(batch_size) + .fetch_all(pool) + }, + self.get_max_retries(), + self.retry_delay.as_millis() as u64, ) - .bind(slot_name) - .bind(batch_size) - .fetch_all(pool) .await - .map_err(|e| { - error!("Failed to fetch CDC changes: {e}"); - Error::InvalidRecord - })?; + .map_err(|e| Error::Connection(format!("failed to fetch CDC changes: {e}")))?; let mut messages = Vec::new(); let mut last_lsn = None; @@ -534,7 +566,11 @@ impl PostgresSource { } else { debug!("CDC: Fetched {} change records", messages.len()); } - let pending = if let Some(lsn) = last_lsn { + let lsn = replication_slot_target_lsn(last_lsn, wal_flush_lsn); + let pending = if messages.is_empty() { + self.advance_replication_slot(&lsn).await?; + None + } else { let mut state = self.state.lock().await.clone(); state.processed_rows += messages.len() as u64; state.last_poll_time = Utc::now(); @@ -542,8 +578,6 @@ impl PostgresSource { state, operations: vec![PendingOperation::AdvanceReplicationSlot { lsn }], }) - } else { - None }; Ok(PolledBatch { messages, pending }) @@ -617,15 +651,10 @@ impl PostgresSource { } 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(), ids: processed_ids, - max_offset, + tracking_boundary: self.processing_boundary(max_offset.clone()), }); } @@ -660,18 +689,22 @@ impl PostgresSource { let slot_name = self.replication_slot(); let pool = self.get_pool()?; with_retry( - || { - sqlx::query("SELECT pg_replication_slot_advance($1, $2::pg_lsn)") + || 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(()), + result => result.map(|_| ()), + } }, 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}" )) @@ -684,7 +717,7 @@ impl PostgresSource { pool: &Pool, table: &str, ids: &[String], - max_offset: &str, + tracking_boundary: Option<&str>, ) -> Result<(), Error> { if ids.is_empty() { return Ok(()); @@ -692,8 +725,8 @@ impl PostgresSource { let quoted_table = quote_qualified_identifier(table)?; 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 tracking_condition = + build_tracking_condition(self.tracking_column(), tracking_boundary)?; let ids_list = ids .iter() @@ -709,8 +742,7 @@ impl PostgresSource { if self.config.delete_after_read.unwrap_or(false) { let delete_query = format!( - "DELETE FROM {quoted_table} WHERE {quoted_pk} IN ({ids_list}) \ - AND {quoted_tracking} <= {tracking_boundary}" + "DELETE FROM {quoted_table} WHERE {quoted_pk} IN ({ids_list}){tracking_condition}" ); if self.verbose { @@ -725,15 +757,12 @@ impl PostgresSource { 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}")) - })?; + .map_err(|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}) AND {quoted_tracking} <= {tracking_boundary}" + WHERE {quoted_pk} IN ({ids_list}){tracking_condition}" ); if self.verbose { @@ -748,10 +777,7 @@ impl PostgresSource { 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}")) - })?; + .map_err(|e| Error::Connection(format!("failed to mark rows as processed: {e}")))?; } Ok(()) @@ -780,6 +806,14 @@ impl PostgresSource { self.config.delete_after_read.unwrap_or(false) || self.config.processed_column.is_some() } + fn processing_boundary(&self, max_offset: Option) -> Option { + if self.config.custom_query.is_some() { + None + } else { + max_offset + } + } + fn tracking_column(&self) -> &str { self.config.tracking_column.as_deref().unwrap_or("id") } @@ -868,6 +902,8 @@ impl PostgresSource { let now = Utc::now(); + // TODO: Substitute `$now_unix` before `$now` so the longer placeholder remains intact. + // TODO: Bind or quote `$offset` according to its PostgreSQL type instead of inserting raw data. query .replace("$table", table) .replace("$offset", &offset_value) @@ -966,11 +1002,7 @@ impl PostgresSource { data.insert(column_name.clone(), value.clone()); if column.name() == config.tracking_column { - if let serde_json::Value::String(ref s) = value { - max_offset = Some(s.clone()); - } else if let serde_json::Value::Number(ref n) = value { - max_offset = Some(n.to_string()); - } + max_offset = extract_tracking_value(row, i, &value)?; } if column.name() == config.pk_column { @@ -1124,11 +1156,11 @@ fn extract_column_value( .unwrap_or(serde_json::Value::Null)) } "NUMERIC" => { - let value: Option = row + let value: Option = row .try_get(column_index) .map_err(|_| Error::InvalidRecord)?; Ok(value - .and_then(|s| s.parse::().ok()) + .and_then(|value| value.to_string().parse::().ok()) .map(serde_json::Value::from) .unwrap_or(serde_json::Value::Null)) } @@ -1589,6 +1621,43 @@ fn format_offset_value(value: &str) -> String { } } +fn build_tracking_condition( + tracking_column: &str, + tracking_boundary: Option<&str>, +) -> Result { + let Some(boundary) = tracking_boundary else { + return Ok(String::new()); + }; + let quoted_tracking = quote_identifier(tracking_column)?; + Ok(format!( + " AND ({quoted_tracking} <= {} OR {quoted_tracking} IS NULL)", + format_offset_value(boundary) + )) +} + +fn replication_slot_target_lsn(last_change_lsn: Option, wal_flush_lsn: String) -> String { + last_change_lsn.unwrap_or(wal_flush_lsn) +} + +fn extract_tracking_value( + row: &sqlx::postgres::PgRow, + column_index: usize, + value: &serde_json::Value, +) -> Result, Error> { + if row.columns()[column_index].type_info().name() == "NUMERIC" { + let value: Option = row + .try_get(column_index) + .map_err(|_| Error::InvalidRecord)?; + return Ok(value.map(|value| value.to_string())); + } + + Ok(match value { + serde_json::Value::String(value) => Some(value.clone()), + serde_json::Value::Number(value) => Some(value.to_string()), + _ => None, + }) +} + fn to_snake_case(input: &str) -> String { let mut result = String::new(); let mut prev_was_uppercase = false; @@ -1810,16 +1879,28 @@ fn is_transient_error(e: &sqlx::Error) -> bool { sqlx::Error::PoolTimedOut => true, sqlx::Error::PoolClosed => false, sqlx::Error::Protocol(_) => false, - sqlx::Error::Database(db_err) => db_err.code().is_some_and(|code| { - matches!( - code.as_ref(), - "40001" | "40P01" | "57P01" | "57P02" | "57P03" | "08000" | "08003" | "08006" - ) - }), + sqlx::Error::Database(db_err) => db_err + .code() + .is_some_and(|code| is_transient_sqlstate(code.as_ref())), _ => false, } } +fn is_transient_sqlstate(code: &str) -> bool { + matches!( + code, + "40001" | "40P01" | "55006" | "57P01" | "57P02" | "57P03" | "08000" | "08003" | "08006" + ) +} + +fn is_replication_slot_already_advanced(error: &sqlx::Error) -> bool { + matches!(error, sqlx::Error::Database(database_error) if database_error.code().is_some_and(|code| is_replication_slot_already_advanced_sqlstate(code.as_ref()))) +} + +fn is_replication_slot_already_advanced_sqlstate(code: &str) -> bool { + code == "22023" +} + fn redact_connection_string(conn_str: &str) -> String { if let Some(scheme_end) = conn_str.find("://") { let scheme = &conn_str[..scheme_end + 3]; @@ -1916,6 +1997,49 @@ mod tests { assert!(!query.contains("'42'")); } + #[test] + fn given_exact_numeric_boundary_should_preserve_text_and_include_null_rows() { + let condition = build_tracking_condition("offset", Some("9007199254740993.25")) + .expect("Failed to build tracking condition"); + + assert_eq!( + condition, + " AND (\"offset\" <= 9007199254740993.25 OR \"offset\" IS NULL)" + ); + } + + #[test] + fn given_no_tracking_boundary_should_not_add_tracking_condition() { + let condition = + build_tracking_condition("offset", None).expect("Failed to build tracking condition"); + + assert!(condition.is_empty()); + } + + #[test] + fn given_custom_query_should_not_apply_last_row_as_processing_boundary() { + let mut config = test_config(); + config.custom_query = Some("SELECT id FROM users".to_string()); + let src = PostgresSource::new(1, config, None); + + assert_eq!(src.processing_boundary(Some("42".to_string())), None); + } + + #[test] + fn given_empty_cdc_peek_should_advance_to_pre_peek_wal_flush_lsn() { + let target = replication_slot_target_lsn(None, "0/16D32A0".to_string()); + + assert_eq!(target, "0/16D32A0"); + } + + #[test] + fn given_cdc_changes_should_advance_to_last_change_lsn() { + let target = + replication_slot_target_lsn(Some("0/16D32B0".to_string()), "0/16D32C0".to_string()); + + assert_eq!(target, "0/16D32B0"); + } + #[test] fn given_special_chars_in_identifier_should_escape() { let result = quote_identifier("table\"name").expect("Failed to quote"); @@ -2616,8 +2740,8 @@ mod tests { assert_eq!(redacted, "postgresql://adm***"); } - #[test] - fn given_persisted_state_should_restore_tracking_offsets() { + #[tokio::test] + async fn given_persisted_state_should_restore_tracking_offsets() { let state = State { last_poll_time: Utc::now(), tracking_offsets: HashMap::from([ @@ -2632,31 +2756,25 @@ 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); } - #[test] - fn given_no_state_should_start_fresh() { + #[tokio::test] + async fn given_no_state_should_start_fresh() { let src = PostgresSource::new(1, test_config(), None); - let runtime = tokio::runtime::Runtime::new().unwrap(); - runtime.block_on(async { - let state = src.state.lock().await; - assert!(state.tracking_offsets.is_empty()); - assert_eq!(state.processed_rows, 0); - }); + let state = src.state.lock().await; + assert!(state.tracking_offsets.is_empty()); + assert_eq!(state.processed_rows, 0); } #[tokio::test] @@ -2681,84 +2799,89 @@ 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(), - ids: vec!["3".to_string()], - max_offset: "3".to_string(), - }, - PendingOperation::AdvanceReplicationSlot { - lsn: "0/16D32A0".to_string(), - }, - ], - }); + fn given_active_replication_slot_sqlstate_should_be_transient() { + assert!(is_transient_sqlstate("55006")); + } - src.on_batch_result(SourceBatchResult::Nack) - .await - .expect("NACK should discard the candidate state"); + #[test] + fn given_target_below_confirmed_flush_sqlstate_should_be_already_advanced() { + assert!(is_replication_slot_already_advanced_sqlstate("22023")); + } - { - 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_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(), + }, + ], }); + + 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() { + #[tokio::test] + async fn given_ack_when_staged_operation_fails_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(), - }); + 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()), + }], + }); - src.on_batch_result(SourceBatchResult::Ack) - .await - .expect("ACK should commit the candidate state"); + src.on_batch_result(SourceBatchResult::Ack) + .await + .expect("ACK should commit state even when the staged operation fails"); - { - 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()); - }); + { + 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() { + #[tokio::test] + async fn given_invalid_state_should_start_fresh() { let invalid_state = ConnectorState(b"not valid json".to_vec()); let src = PostgresSource::new(1, test_config(), Some(invalid_state)); - let runtime = tokio::runtime::Runtime::new().unwrap(); - runtime.block_on(async { - let state = src.state.lock().await; - assert!(state.tracking_offsets.is_empty()); - assert_eq!(state.processed_rows, 0); - }); + let state = src.state.lock().await; + assert!(state.tracking_offsets.is_empty()); + assert_eq!(state.processed_rows, 0); } #[test] diff --git a/core/integration/tests/connectors/fixtures/mod.rs b/core/integration/tests/connectors/fixtures/mod.rs index 7eaf6fe510..aa8bf06ed7 100644 --- a/core/integration/tests/connectors/fixtures/mod.rs +++ b/core/integration/tests/connectors/fixtures/mod.rs @@ -79,7 +79,7 @@ pub use postgres::{ PostgresOps, PostgresSinkByteaFixture, PostgresSinkFixture, PostgresSinkJsonFixture, PostgresSourceByteaFixture, PostgresSourceCdcFixture, PostgresSourceDeleteFixture, PostgresSourceJsonFixture, PostgresSourceJsonbFixture, PostgresSourceMarkFixture, - PostgresSourceOps, + PostgresSourceNumericTrackingFixture, PostgresSourceOps, }; pub use quickwit::{QuickwitFixture, QuickwitOps, QuickwitPreCreatedFixture}; pub use redshift::{ diff --git a/core/integration/tests/connectors/fixtures/postgres/mod.rs b/core/integration/tests/connectors/fixtures/postgres/mod.rs index cca22e6e33..ed02c2025f 100644 --- a/core/integration/tests/connectors/fixtures/postgres/mod.rs +++ b/core/integration/tests/connectors/fixtures/postgres/mod.rs @@ -25,5 +25,5 @@ pub use container::{PostgresOps, PostgresSourceOps}; pub use sink::{PostgresSinkByteaFixture, PostgresSinkFixture, PostgresSinkJsonFixture}; pub use source::{ PostgresSourceByteaFixture, PostgresSourceDeleteFixture, PostgresSourceJsonFixture, - PostgresSourceJsonbFixture, PostgresSourceMarkFixture, + PostgresSourceJsonbFixture, PostgresSourceMarkFixture, PostgresSourceNumericTrackingFixture, }; diff --git a/core/integration/tests/connectors/fixtures/postgres/source.rs b/core/integration/tests/connectors/fixtures/postgres/source.rs index 5ad2900c63..9f0b41b891 100644 --- a/core/integration/tests/connectors/fixtures/postgres/source.rs +++ b/core/integration/tests/connectors/fixtures/postgres/source.rs @@ -382,6 +382,97 @@ impl TestFixture for PostgresSourceDeleteFixture { } } +/// PostgreSQL source fixture with an exact NUMERIC tracking column. +pub struct PostgresSourceNumericTrackingFixture { + container: PostgresContainer, +} + +impl PostgresOps for PostgresSourceNumericTrackingFixture { + fn container(&self) -> &PostgresContainer { + &self.container + } +} + +impl PostgresSourceOps for PostgresSourceNumericTrackingFixture { + fn table_name(&self) -> &str { + Self::TABLE + } +} + +impl PostgresSourceNumericTrackingFixture { + const TABLE: &'static str = "test_numeric_tracking"; + + pub async fn create_table(&self, pool: &Pool) { + let query = format!( + "CREATE TABLE IF NOT EXISTS {} ( + id INTEGER PRIMARY KEY, + tracking_value NUMERIC NOT NULL + )", + Self::TABLE + ); + sqlx::query(sqlx::AssertSqlSafe(query)) + .execute(pool) + .await + .unwrap_or_else(|e| panic!("Failed to create table: {e}")); + } + + pub async fn insert_row(&self, pool: &Pool, id: i32, tracking_value: &str) { + let query = format!( + "INSERT INTO {} (id, tracking_value) VALUES ($1, $2::numeric)", + Self::TABLE + ); + sqlx::query(sqlx::AssertSqlSafe(query)) + .bind(id) + .bind(tracking_value) + .execute(pool) + .await + .unwrap_or_else(|e| panic!("Failed to insert row: {e}")); + } + + pub async fn count_rows(&self, pool: &Pool) -> i64 { + PostgresSourceOps::count_rows(self, pool).await + } +} + +#[async_trait] +impl TestFixture for PostgresSourceNumericTrackingFixture { + async fn setup() -> Result { + let container = PostgresContainer::start().await?; + Ok(Self { container }) + } + + fn connectors_runtime_envs(&self) -> HashMap { + let mut envs = HashMap::new(); + envs.insert( + ENV_SOURCE_CONNECTION_STRING.to_string(), + self.container.connection_string.clone(), + ); + envs.insert(ENV_SOURCE_TABLES.to_string(), format!("[{}]", Self::TABLE)); + envs.insert( + ENV_SOURCE_TRACKING_COLUMN.to_string(), + "tracking_value".to_string(), + ); + envs.insert(ENV_SOURCE_PRIMARY_KEY_COLUMN.to_string(), "id".to_string()); + envs.insert(ENV_SOURCE_DELETE_AFTER_READ.to_string(), "true".to_string()); + envs.insert(ENV_SOURCE_INCLUDE_METADATA.to_string(), "true".to_string()); + envs.insert( + ENV_SOURCE_STREAMS_0_STREAM.to_string(), + DEFAULT_TEST_STREAM.to_string(), + ); + envs.insert( + ENV_SOURCE_STREAMS_0_TOPIC.to_string(), + DEFAULT_TEST_TOPIC.to_string(), + ); + envs.insert(ENV_SOURCE_STREAMS_0_SCHEMA.to_string(), "json".to_string()); + envs.insert(ENV_SOURCE_POLL_INTERVAL.to_string(), "10ms".to_string()); + envs.insert( + ENV_SOURCE_PATH.to_string(), + "../../target/debug/libiggy_connector_postgres_source".to_string(), + ); + envs + } +} + /// PostgreSQL source fixture with processed_column marking. pub struct PostgresSourceMarkFixture { container: PostgresContainer, diff --git a/core/integration/tests/connectors/postgres/postgres_source.rs b/core/integration/tests/connectors/postgres/postgres_source.rs index 740295221b..cad53aabc4 100644 --- a/core/integration/tests/connectors/postgres/postgres_source.rs +++ b/core/integration/tests/connectors/postgres/postgres_source.rs @@ -32,7 +32,7 @@ use crate::connectors::create_test_messages; use crate::connectors::fixtures::{ PostgresOps, PostgresSourceByteaFixture, PostgresSourceDeleteFixture, PostgresSourceJsonFixture, PostgresSourceJsonbFixture, PostgresSourceMarkFixture, - PostgresSourceOps, + PostgresSourceNumericTrackingFixture, PostgresSourceOps, }; #[iggy_harness( @@ -449,6 +449,64 @@ async fn delete_after_read_source_removes_rows_after_producing( pool.close().await; } +#[iggy_harness( + cluster_nodes = 1, + server(connectors_runtime(config_path = "tests/connectors/postgres/source.toml")), + seed = seeds::connector_stream +)] +async fn numeric_tracking_source_preserves_exact_ack_boundary( + harness: &TestHarness, + fixture: PostgresSourceNumericTrackingFixture, +) { + const TRACKING_VALUE: &str = "9007199254740993.25"; + + let client = harness.root_client().await.unwrap(); + let pool = fixture.create_pool().await.expect("Failed to create pool"); + fixture.create_table(&pool).await; + fixture.insert_row(&pool, 1, TRACKING_VALUE).await; + + let stream_id: Identifier = seeds::names::STREAM.try_into().unwrap(); + let topic_id: Identifier = seeds::names::TOPIC.try_into().unwrap(); + let consumer_id: Identifier = "numeric_tracking_consumer".try_into().unwrap(); + let mut received = false; + + for _ in 0..POLL_ATTEMPTS { + if let Ok(polled) = client + .poll_messages( + &stream_id, + &topic_id, + None, + &Consumer::new(consumer_id.clone()), + &PollingStrategy::next(), + 1, + true, + ) + .await + && !polled.messages.is_empty() + { + received = true; + break; + } + sleep(Duration::from_millis(POLL_INTERVAL_MS)).await; + } + assert!(received, "NUMERIC tracking row should be delivered"); + + 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, + "Exact NUMERIC tracking boundary should allow ACK cleanup" + ); + + pool.close().await; +} + #[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 c36e36789a..fbc89800f9 100644 --- a/core/integration/tests/connectors/postgres/postgres_source_cdc.rs +++ b/core/integration/tests/connectors/postgres/postgres_source_cdc.rs @@ -75,14 +75,28 @@ async fn poll_cdc_records( } 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)) + for attempt in 0..POLL_ATTEMPTS { + match sqlx::query_scalar::<_, String>( + "SELECT data FROM pg_logical_slot_peek_changes($1, NULL, NULL)", + ) + .bind(DEFAULT_SLOT) + .fetch_all(pool) + .await + { + Ok(changes) => { + return changes.iter().any(|change| change.contains(expected_value)); + } + Err(sqlx::Error::Database(ref database_error)) + if attempt + 1 < POLL_ATTEMPTS + && database_error.code().as_deref() == Some(PG_OBJECT_IN_USE) => + { + sleep(Duration::from_millis(POLL_INTERVAL_MS)).await; + } + Err(error) => panic!("CDC replication slot should be readable: {error}"), + } + } + + panic!("CDC replication slot remained active after {POLL_ATTEMPTS} attempts"); } // End-to-end CDC coverage against a real wal_level=logical container: @@ -241,6 +255,52 @@ 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 idle_cdc_source_advances_slot_to_current_wal( + harness: &TestHarness, + fixture: PostgresSourceCdcFixture, +) { + let pool = fixture.create_pool().await.expect("Failed to create pool"); + fixture.create_table(&pool).await; + + let api_url = harness + .connectors_runtime() + .expect("connectors runtime") + .http_url(); + wait_for_source_status(&Client::new(), &api_url, ConnectorStatus::Running).await; + + sqlx::query("CHECKPOINT") + .execute(&pool) + .await + .expect("Failed to generate WAL without a logical table change"); + let target_lsn: String = sqlx::query_scalar("SELECT pg_current_wal_flush_lsn()::text") + .fetch_one(&pool) + .await + .expect("Failed to read current WAL flush LSN"); + + for _ in 0..POLL_ATTEMPTS { + let reached: bool = sqlx::query_scalar( + "SELECT confirmed_flush_lsn >= $2::pg_lsn FROM pg_replication_slots WHERE slot_name = $1", + ) + .bind(DEFAULT_SLOT) + .bind(&target_lsn) + .fetch_one(&pool) + .await + .expect("Failed to read replication slot position"); + if reached { + pool.close().await; + return; + } + sleep(Duration::from_millis(POLL_INTERVAL_MS)).await; + } + + panic!("Idle CDC slot did not advance to WAL flush LSN {target_lsn}"); +} + #[iggy_harness( cluster_nodes = 1, server(connectors_runtime(config_path = "tests/connectors/postgres/source.toml")),