diff --git a/.claude/skills/connector-source/SKILL.md b/.claude/skills/connector-source/SKILL.md index 5627b1812d..1d761462b5 100644 --- a/.claude/skills/connector-source/SKILL.md +++ b/.claude/skills/connector-source/SKILL.md @@ -45,25 +45,33 @@ 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. +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 `poll()` must `sleep(self.poll_interval).await` before any work. Without it, an empty source spins a CPU. @@ -97,7 +105,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 +135,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 +145,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..04b8c6062e 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,37 @@ 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 (SourceBatchResult::Ack, Some(candidate)) = + (result, self.pending.lock().await.take()) + else { + return Ok(()); + }; + *self.state.lock().await = candidate; + 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/Cargo.lock b/Cargo.lock index 56ecb71dd9..ed0a67c77b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -12700,6 +12700,7 @@ version = "0.9.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "05b44e85bf579a8eeb4ceaa77a3a523baf2bf0e9bac7e40f405d537b5d2d5ccb" dependencies = [ + "bigdecimal", "base64 0.22.1", "bytes", "cfg-if", @@ -12777,6 +12778,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", @@ -12806,6 +12808,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "87a2bdd6e83f6b3ea525ca9fee568030508b58355a43d0b2c1674d5f79dcd65e" dependencies = [ "atoi", + "bigdecimal", "base64 0.22.1", "bitflags 2.13.1", "byteorder", @@ -12823,6 +12826,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 af7a697ca9..2c72dbc96f 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -306,7 +306,10 @@ sqlx = { version = "0.9.0", features = [ "runtime-tokio", "tls-rustls", "postgres", - # "mysql": the Doris integration fixture talks to Doris over its MySQL frontend. + "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. "mysql", "chrono", "uuid", diff --git a/core/connectors/runtime/src/stream.rs b/core/connectors/runtime/src/stream.rs index 4082ffc8ae..5591444959 100644 --- a/core/connectors/runtime/src/stream.rs +++ b/core/connectors/runtime/src/stream.rs @@ -24,6 +24,11 @@ use crate::error::RuntimeError; const TOKEN_FILE_PREFIX: &str = "file:"; +fn append_query_parameters(connection_string: &str, address: &str, parameters: &str) -> String { + let separator = if address.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() { @@ -133,8 +138,10 @@ fn connection_string_with_token( .filter(|domain| !domain.is_empty()) .map(|domain| format!("&tls_domain={domain}")) .unwrap_or_default(); - Ok(format!( - "{connection_string}?tls=true&tls_ca_file={ca_file}{domain}" + Ok(append_query_parameters( + &connection_string, + &config.address, + &format!("tls=true&tls_ca_file={ca_file}{domain}"), )) } else { Ok(connection_string) @@ -192,6 +199,39 @@ 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 address = "127.0.0.1:8090?reconnection_retries=0"; + + let result = append_query_parameters(connection_string, address, "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 address = "127.0.0.1:8090"; + + 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 d0619bdeab..0b142da16a 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) @@ -209,7 +216,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 +226,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] @@ -235,6 +242,15 @@ 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. + ## Supported Column Types The connector handles these PostgreSQL types in JSON mode: @@ -254,7 +270,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 @@ -267,6 +283,15 @@ 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. + +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 @@ -274,8 +299,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..8b738ddd90 100644 --- a/core/connectors/sources/postgres_source/src/lib.rs +++ b/core/connectors/sources/postgres_source/src/lib.rs @@ -15,21 +15,24 @@ // 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::types::BigDecimal; 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; @@ -38,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 { @@ -45,6 +50,7 @@ pub struct PostgresSource { pool: Option>, config: PostgresSourceConfig, state: Mutex, + pending_batch: Mutex>, verbose: bool, retry_delay: Duration, poll_interval: Duration, @@ -97,13 +103,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, + // TODO: Persist pending operations with the candidate state and replay them during open. + operations: Vec, +} + +#[derive(Debug)] +enum PendingOperation { + ProcessRows { + table: String, + ids: Vec, + tracking_boundary: Option, + }, + AdvanceReplicationSlot { + lsn: String, + }, +} + #[derive(Debug, Serialize, Deserialize)] pub struct DatabaseRecord { pub table_name: String, @@ -162,6 +193,7 @@ impl PostgresSource { tracking_offsets: HashMap::new(), processed_rows: 0, })), + pending_batch: Mutex::new(None), verbose, retry_delay, poll_interval, @@ -221,7 +253,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 +262,20 @@ impl Source for PostgresSource { } }; - let state = self.state.lock().await; + let processed_rows = 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,71 @@ impl Source for PostgresSource { PayloadFormat::JsonDirect | PayloadFormat::Json => Schema::Json, }; - let persisted_state = self.serialize_state(&state); + let persisted_state = polled + .pending + .as_ref() + .map(|pending| { + self.serialize_state(&pending.state).ok_or_else(|| { + Error::Serialization("failed to serialize PostgreSQL source state".to_string()) + }) + }) + .transpose()?; + *self.pending_batch.lock().await = polled.pending; Ok(ProducedMessages { schema, - messages, + messages: polled.messages, state: persisted_state, }) } + async fn on_batch_result(&self, result: SourceBatchResult) -> Result<(), Error> { + let (SourceBatchResult::Ack, Some(pending)) = + (result, self.pending_batch.lock().await.take()) + else { + return Ok(()); + }; + + let PendingBatch { state, operations } = pending; + let cleanup = async { + for operation in operations { + match operation { + PendingOperation::ProcessRows { + table, + ids, + tracking_boundary, + } => { + if let Ok(pool) = self.get_pool() { + let _ = self + .mark_or_delete_processed_rows( + pool, + &table, + &ids, + tracking_boundary.as_deref(), + ) + .await; + } + } + PendingOperation::AdvanceReplicationSlot { lsn } => { + let _ = self.advance_replication_slot(&lsn).await; + } + } + } + }; + if tokio::time::timeout(ACK_CLEANUP_TIMEOUT, cleanup) + .await + .is_err() + { + warn!( + "PostgreSQL source connector ID: {} exceeded the ACK cleanup budget", + self.id + ); + } + + *self.state.lock().await = state; + Ok(()) + } + async fn close(&mut self) -> Result<(), Error> { if let Some(pool) = self.pool.take() { pool.close().await; @@ -289,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}")))?; @@ -347,11 +436,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") @@ -389,7 +474,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,14 +484,10 @@ 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 - .config - .replication_slot - .as_deref() - .unwrap_or("iggy_slot"); + let slot_name = self.replication_slot(); let capture_ops = self .config .capture_operations @@ -417,43 +498,55 @@ 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, xid, data FROM pg_logical_slot_get_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) - .await - .map_err(|e| { - error!("Failed to fetch CDC changes: {e}"); - Error::InvalidRecord - })?; + }, + self.get_max_retries(), + self.retry_delay.as_millis() as u64, + ) + .await + .map_err(|e| Error::Connection(format!("failed to fetch CDC changes: {e}")))?; 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,31 +561,37 @@ 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 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(); + Some(PendingBatch { + state, + operations: vec![PendingOperation::AdvanceReplicationSlot { lsn }], + }) + }; + + 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"); - 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: "", @@ -504,8 +603,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 { @@ -514,11 +611,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)?; @@ -533,10 +626,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)?; @@ -550,52 +647,86 @@ impl PostgresSource { messages.push(processed.message); total_processed += 1; + table_processed += 1; } - // Database I/O without holding the lock - if !processed_ids.is_empty() { - self.mark_or_delete_processed_rows(pool, table, pk_column, &processed_ids) - .await?; + if self.should_process_rows() && !processed_ids.is_empty() { + operations.push(PendingOperation::ProcessRows { + table: table.clone(), + ids: processed_ids, + tracking_boundary: self.processing_boundary(max_offset.clone()), + }); } - // Collect offset update for later if let Some(offset) = max_offset { - state_updates.push((table.clone(), offset)); + candidate_state + .tracking_offsets + .insert(table.clone(), offset); } if self.verbose { - info!("Fetched {} rows from table '{table}'", messages.len()); + info!("Fetched {table_processed} rows from table '{table}'"); } else { - debug!("Fetched {} rows from table '{table}'", messages.len()); + debug!("Fetched {table_processed} rows from table '{table}'"); } } - // Apply all state updates with a single lock acquisition - { - let mut state = self.state.lock().await; - state.processed_rows += total_processed; - for (table, offset) in state_updates { - state.tracking_offsets.insert(table, offset); - } - state.last_poll_time = Utc::now(); - } + let pending = if total_processed > 0 { + candidate_state.processed_rows += total_processed; + candidate_state.last_poll_time = Utc::now(); + Some(PendingBatch { + state: candidate_state, + operations, + }) + } else { + None + }; - Ok(messages) + Ok(PolledBatch { messages, pending }) + } + + async fn advance_replication_slot(&self, lsn: &str) -> Result<(), Error> { + let slot_name = self.replication_slot(); + let pool = self.get_pool()?; + with_retry( + || async { + match sqlx::query("SELECT pg_replication_slot_advance($1, $2::pg_lsn)") + .bind(slot_name) + .bind(lsn) + .execute(pool) + .await + { + Err(error) if is_replication_slot_already_advanced(&error) => Ok(()), + result => result.map(|_| ()), + } + }, + self.get_max_retries(), + self.retry_delay.as_millis() as u64, + ) + .await + .map_err(|e| { + Error::Connection(format!( + "failed to advance replication slot '{slot_name}' to {lsn}: {e}" + )) + })?; + Ok(()) } async fn mark_or_delete_processed_rows( &self, pool: &Pool, table: &str, - pk_column: &str, ids: &[String], + tracking_boundary: Option<&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 tracking_condition = + build_tracking_condition(self.tracking_column(), tracking_boundary)?; let ids_list = ids .iter() @@ -610,8 +741,9 @@ 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}){tracking_condition}" + ); if self.verbose { info!("Deleting {} processed rows from '{table}'", ids.len()); @@ -619,17 +751,18 @@ 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::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}){tracking_condition}" ); if self.verbose { @@ -638,13 +771,13 @@ 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::Connection(format!("failed to mark rows as processed: {e}")))?; } Ok(()) @@ -669,6 +802,36 @@ 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 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") + } + + 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, @@ -739,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) @@ -837,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 { @@ -995,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)) } @@ -1460,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; @@ -1647,7 +1845,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>, @@ -1660,7 +1862,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..." @@ -1677,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]; @@ -1704,6 +1918,8 @@ mod cdc_fixtures; #[cfg(test)] mod tests { + use std::sync::atomic::{AtomicU32, Ordering}; + use super::*; fn test_config() -> PostgresSourceConfig { @@ -1781,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"); @@ -2481,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([ @@ -2497,44 +2756,132 @@ mod tests { let src = PostgresSource::new(1, test_config(), Some(connector_state)); - let runtime = tokio::runtime::Runtime::new().unwrap(); - runtime.block_on(async { - let restored = src.state.lock().await; - assert_eq!( - restored.tracking_offsets.get("users"), - Some(&"100".to_string()) - ); - assert_eq!( - restored.tracking_offsets.get("orders"), - Some(&"2024-01-15T10:30:00Z".to_string()) - ); - assert_eq!(restored.processed_rows, 500); - }); + let restored = src.state.lock().await; + assert_eq!( + restored.tracking_offsets.get("users"), + Some(&"100".to_string()) + ); + assert_eq!( + restored.tracking_offsets.get("orders"), + Some(&"2024-01-15T10:30:00Z".to_string()) + ); + assert_eq!(restored.processed_rows, 500); + } + + #[tokio::test] + async fn given_no_state_should_start_fresh() { + let src = PostgresSource::new(1, test_config(), None); + + let state = src.state.lock().await; + assert!(state.tracking_offsets.is_empty()); + assert_eq!(state.processed_rows, 0); + } + + #[tokio::test] + async fn given_transient_database_errors_when_retrying_should_eventually_succeed() { + let attempts = AtomicU32::new(0); + + let result = with_retry( + || async { + if attempts.fetch_add(1, Ordering::Relaxed) < 2 { + Err(sqlx::Error::PoolTimedOut) + } else { + Ok(()) + } + }, + 3, + 0, + ) + .await; + + assert!(result.is_ok()); + assert_eq!(attempts.load(Ordering::Relaxed), 3); } #[test] - fn given_no_state_should_start_fresh() { + fn given_active_replication_slot_sqlstate_should_be_transient() { + assert!(is_transient_sqlstate("55006")); + } + + #[test] + fn given_target_below_confirmed_flush_sqlstate_should_be_already_advanced() { + assert!(is_replication_slot_already_advanced_sqlstate("22023")); + } + + #[tokio::test] + async fn given_nack_when_batch_is_staged_should_keep_committed_state() { let src = PostgresSource::new(1, test_config(), None); + let mut candidate_state = src.state.lock().await.clone(); + candidate_state + .tracking_offsets + .insert("users".to_string(), "3".to_string()); + candidate_state.processed_rows = 3; + *src.pending_batch.lock().await = Some(PendingBatch { + state: candidate_state, + operations: vec![ + PendingOperation::ProcessRows { + table: "users".to_string(), + ids: vec!["3".to_string()], + tracking_boundary: Some("3".to_string()), + }, + PendingOperation::AdvanceReplicationSlot { + lsn: "0/16D32A0".to_string(), + }, + ], + }); - let runtime = tokio::runtime::Runtime::new().unwrap(); - runtime.block_on(async { + src.on_batch_result(SourceBatchResult::Nack) + .await + .expect("NACK should discard the candidate state"); + + { let state = src.state.lock().await; assert!(state.tracking_offsets.is_empty()); assert_eq!(state.processed_rows, 0); + } + assert!(src.pending_batch.lock().await.is_none()); + } + + #[tokio::test] + async fn given_ack_when_staged_operation_fails_should_commit_candidate_state() { + 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()), + }], }); + + 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()); } - #[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/src/harness/handle/connectors_runtime.rs b/core/integration/src/harness/handle/connectors_runtime.rs index 8a581c4813..66d2a90faf 100644 --- a/core/integration/src/harness/handle/connectors_runtime.rs +++ b/core/integration/src/harness/handle/connectors_runtime.rs @@ -44,6 +44,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, @@ -80,6 +81,14 @@ 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()); + } + + 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( @@ -92,8 +101,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 { @@ -127,6 +140,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/fixtures/mod.rs b/core/integration/tests/connectors/fixtures/mod.rs index e9d7a1332e..60d2de82c7 100644 --- a/core/integration/tests/connectors/fixtures/mod.rs +++ b/core/integration/tests/connectors/fixtures/mod.rs @@ -80,7 +80,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/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 a66a3c5e67..cad53aabc4 100644 --- a/core/integration/tests/connectors/postgres/postgres_source.rs +++ b/core/integration/tests/connectors/postgres/postgres_source.rs @@ -15,20 +15,26 @@ // specific language governing permissions and limitations // under the License. -use super::{DatabaseRecord, POLL_ATTEMPTS, POLL_INTERVAL_MS, TEST_MESSAGE_COUNT}; -use crate::connectors::create_test_messages; -use crate::connectors::fixtures::{ - PostgresOps, PostgresSourceByteaFixture, PostgresSourceDeleteFixture, - PostgresSourceJsonFixture, PostgresSourceJsonbFixture, PostgresSourceMarkFixture, - PostgresSourceOps, -}; +use std::time::Duration; + use iggy_common::MessageClient; use iggy_common::{Consumer, Identifier, PollingStrategy}; use integration::harness::seeds; use integration::iggy_harness; -use std::time::Duration; +use reqwest::Client; use tokio::time::sleep; +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, + PostgresSourceJsonFixture, PostgresSourceJsonbFixture, PostgresSourceMarkFixture, + PostgresSourceNumericTrackingFixture, PostgresSourceOps, +}; + #[iggy_harness( server(connectors_runtime(config_path = "tests/connectors/postgres/source.toml")), seed = seeds::connector_stream @@ -127,6 +133,122 @@ async fn json_rows_source_produces_messages_to_iggy( } } +#[iggy_harness( + cluster_nodes = 1, + server(connectors_runtime(config_path = "tests/connectors/postgres/source.toml")), + seed = seeds::connector_stream +)] +async fn given_delete_after_read_when_iggy_crashes_should_delete_only_after_redelivery( + harness: &mut TestHarness, + fixture: PostgresSourceDeleteFixture, +) { + 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_stats(&http, &api_url) + .await + .expect("PostgreSQL source stats should be present") + .errors; + + harness.kill_node(0).expect("Failed to kill Iggy server"); + + for index in 0..TEST_MESSAGE_COUNT { + fixture + .insert_row(&pool, &format!("row_{index}"), index as i32) + .await; + } + + 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() + .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 = "send_failure_consumer".try_into().unwrap(); + let mut received = 0; + + 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 += polled.messages.len(); + if received >= TEST_MESSAGE_COUNT { + break; + } + } + sleep(Duration::from_millis(POLL_INTERVAL_MS)).await; + } + + assert_eq!( + received, TEST_MESSAGE_COUNT, + "Rows polled during the failed send should be delivered after restart" + ); + + 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; +} + #[iggy_harness( server(connectors_runtime(config_path = "tests/connectors/postgres/source.toml")), seed = seeds::connector_stream @@ -327,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 b9a485e832..fbc89800f9 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,31 @@ async fn poll_cdc_records( received } +async fn slot_contains_change(pool: &sqlx::PgPool, expected_value: &str) -> bool { + 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: // INSERT, UPDATE, PK-changing UPDATE, DELETE, a rolled-back transaction // (must produce nothing), an untracked table (must be filtered out), a @@ -231,6 +255,154 @@ 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")), + 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, @@ -340,8 +512,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"; @@ -376,11 +548,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";