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";