laminar_connectors/postgres/cdc/source/
mod.rs

1//! `PostgreSQL` CDC source connector implementation.
2//!
3//! Implements [`SourceConnector`] for streaming logical replication changes
4//! from `PostgreSQL` into `LaminarDB` as Arrow `RecordBatch`es.
5//!
6use arrow_array::RecordBatch;
7use arrow_schema::SchemaRef;
8use async_trait::async_trait;
9use bytes::Bytes;
10use std::collections::{BTreeMap, VecDeque};
11use std::sync::Arc;
12use tokio::sync::Notify;
13use tokio::sync::{OwnedSemaphorePermit, Semaphore};
14
15use crate::checkpoint::SourceCheckpoint;
16use crate::config::{ConnectorConfig, ConnectorState};
17use crate::connector::{
18    ConnectorTaskOwner, ConnectorTaskTracker, SourceBatch, SourceConnector, SourceContract,
19    SourcePosition, SourceStart,
20};
21use crate::error::ConnectorError;
22
23use super::changelog::{
24    events_to_record_batch, old_tuple_json_encoded_len, old_tuple_to_json, plan_record_batch,
25    tuple_json_encoded_len, tuple_to_json, CdcOperation, ChangeEvent,
26};
27use super::config::PostgresCdcConfig;
28use super::decoder::{decode_message, OldTuple, WalMessage};
29use super::lsn::Lsn;
30use super::metrics::PostgresCdcMetrics;
31use super::postgres_io::{source_config_digest, PostgresCheckpointBinding};
32use super::schema::{cdc_envelope_schema, RelationCache, RelationInfo};
33
34#[cfg(not(test))]
35const PGWIRE_IN_FLIGHT_EVENTS: usize = 1;
36#[cfg(not(test))]
37const RAW_WAL_QUEUE_CAPACITY: usize = 4_096;
38const INITIAL_BOOTSTRAP_NOT_ADMITTED: &str = "[LDB-5060] PostgreSQL CDC initial startup is not admitted until a complete, certified snapshot-to-WAL handoff is implemented";
39const CHECKPOINT_CONNECTOR: &str = "postgres-cdc";
40const CHECKPOINT_VERSION: &str = "3";
41const SYSTEM_IDENTIFIER_METADATA: &str = "system_identifier";
42const TIMELINE_ID_METADATA: &str = "timeline_id";
43const DATABASE_OID_METADATA: &str = "database_oid";
44const PUBLICATION_OID_METADATA: &str = "publication_oid";
45const PUBLICATION_DEFINITION_METADATA: &str = "publication_definition_sha256";
46const SOURCE_CONFIG_METADATA: &str = "source_config_sha256";
47const SLOT_PLUGIN_METADATA: &str = "slot_plugin";
48const SLOT_TWO_PHASE_METADATA: &str = "slot_two_phase";
49const SLOT_FAILOVER_METADATA: &str = "slot_failover";
50
51/// `PostgreSQL` CDC source connector.
52///
53/// Streams row-level changes from `PostgreSQL` using logical replication
54/// (`pgoutput` plugin). Changes are emitted as Arrow `RecordBatch`es
55/// in the CDC envelope format.
56///
57/// # Envelope Schema
58///
59/// | Column   | Type          | Nullable | Description                              |
60/// |----------|---------------|----------|------------------------------------------|
61/// | `_table` | Utf8          | no       | Schema-qualified table name              |
62/// | `_op`    | Utf8          | no       | Operation: I, U, D                       |
63/// | `_lsn`   | UInt64        | no       | WAL position                             |
64/// | `_ts_ms` | Timestamp(ms) | no       | Commit timestamp                         |
65/// | `_before`| Utf8          | yes      | Available old identity/full-row JSON     |
66/// | `_after` | Utf8          | yes      | New row JSON (for I, U)                  |
67pub struct PostgresCdcSource {
68    /// Connector configuration.
69    config: PostgresCdcConfig,
70
71    /// Current lifecycle state.
72    state: ConnectorState,
73
74    /// Output schema (CDC envelope).
75    schema: SchemaRef,
76
77    /// Lock-free metrics.
78    metrics: Arc<PostgresCdcMetrics>,
79
80    /// Cached relation (table) schemas from Relation messages.
81    relation_cache: RelationCache,
82
83    /// Committed transactions awaiting `poll_batch()` in WAL order.
84    committed_transactions: VecDeque<CommittedTransaction>,
85
86    /// Number of decoded events retained across the current and committed transactions.
87    buffered_event_count: usize,
88
89    /// Variable-width bytes retained by decoded events across current and committed transactions.
90    /// Event container capacities are measured separately when enforcing the decoded-stage limit.
91    buffered_event_bytes: usize,
92
93    /// Current transaction state.
94    current_txn: Option<TransactionState>,
95
96    /// Confirmed flush LSN (last acknowledged position).
97    confirmed_flush_lsn: Lsn,
98
99    /// Write LSN (latest position received from server).
100    write_lsn: Lsn,
101
102    /// Polled LSN — tracks the latest position drained into a batch.
103    /// Decoupled from `confirmed_flush_lsn` so the PG replication slot
104    /// is only advanced when the pipeline actually checkpoints.
105    polled_lsn: Lsn,
106
107    /// Exact database, publication, slot, and filter identity bound to checkpoints.
108    checkpoint_binding: Option<PostgresCheckpointBinding>,
109
110    /// Pending WAL messages used only by deterministic decoder tests.
111    #[cfg(test)]
112    pending_messages: VecDeque<Vec<u8>>,
113
114    /// Notification handle signalled when WAL data arrives from the reader task.
115    data_ready: Arc<Notify>,
116
117    /// Channel receiver for WAL events from the background reader task.
118    wal_rx: Option<WalPayloadRx>,
119
120    /// Background WAL reader task handle.
121    reader_handle: Option<tokio::task::JoinHandle<()>>,
122
123    /// Shutdown signal for the background reader task.
124    reader_shutdown: Option<tokio::sync::watch::Sender<bool>>,
125
126    /// Sender for feeding confirmed flush LSN back to the reader task.
127    /// The reader uses this to call `update_applied_lsn` only for
128    /// durably-checkpointed positions (prevents at-least-once violation).
129    confirmed_lsn_tx: Option<tokio::sync::watch::Sender<u64>>,
130
131    /// One-item readiness lookahead retained across bounded polls.
132    pending_payloads: VecDeque<OwnedWalPayload>,
133
134    /// Weighted budget shared by the reader queue and payloads deferred between polls.
135    wal_byte_budget: Option<Arc<Semaphore>>,
136
137    /// Fatal reader error delivered out of band so a full WAL queue cannot hide it.
138    wal_terminal_error: Option<WalTerminalError>,
139
140    /// Admission authority and terminal observer for this connector generation.
141    task_owner: ConnectorTaskOwner,
142    task_tracker: ConnectorTaskTracker,
143}
144
145impl Drop for PostgresCdcSource {
146    fn drop(&mut self) {
147        if let Some(shutdown) = self.reader_shutdown.take() {
148            shutdown.send_replace(true);
149        }
150        if let Some(handle) = self.reader_handle.take() {
151            reap_postgres_reader(handle, &self.task_owner);
152        }
153    }
154}
155
156fn reap_postgres_reader(handle: tokio::task::JoinHandle<()>, task_owner: &ConnectorTaskOwner) {
157    let Some(reaper_guard) = task_owner.track() else {
158        tracing::warn!("PostgreSQL CDC task generation was sealed before reader reaping");
159        return;
160    };
161    let Ok(runtime) = tokio::runtime::Handle::try_current() else {
162        // The reader's own guard remains authoritative. Dropping its runtime
163        // destroys the future and therefore resolves the generation tracker.
164        drop(reaper_guard);
165        return;
166    };
167    drop(runtime.spawn(async move {
168        let _reaper_guard = reaper_guard;
169        if let Err(error) = handle.await {
170            tracing::debug!(%error, "PostgreSQL CDC retired reader task reaped");
171        }
172    }));
173}
174
175/// In-progress transaction state.
176#[derive(Debug)]
177struct TransactionState {
178    /// Final LSN of the transaction.
179    final_lsn: Lsn,
180    /// Commit timestamp in milliseconds.
181    commit_ts_ms: i64,
182    /// Change events accumulated in this transaction.
183    events: VecDeque<ChangeEvent>,
184}
185
186/// A transaction is resumable only after every one of its events has been emitted.
187#[derive(Debug)]
188struct CommittedTransaction {
189    end_lsn: Lsn,
190    events: VecDeque<ChangeEvent>,
191}
192
193mod accounting;
194mod checkpoint;
195mod decoding;
196mod drain;
197mod lifecycle;
198mod reader;
199#[cfg(not(test))]
200mod startup;
201
202use accounting::{conservative_deque_growth_bytes, planned_event_bytes, retained_event_bytes};
203use checkpoint::{validate_checkpoint_identity, validate_live_binding, write_checkpoint_binding};
204#[cfg(not(test))]
205use reader::run_wal_reader;
206use reader::{
207    logical_wal_payload_bytes, OwnedWalPayload, WalPayload, WalPayloadRx, WalTerminalError,
208};
209#[cfg(test)]
210use reader::{
211    publish_terminal_wal_error, retained_wal_payload_bytes, send_wal_or_shutdown,
212    take_confirmed_lsn,
213};
214
215impl PostgresCdcSource {
216    /// Creates a new `PostgreSQL` CDC source with the given configuration.
217    #[must_use]
218    pub fn new(mut config: PostgresCdcConfig, registry: Option<&prometheus::Registry>) -> Self {
219        config.normalize_table_filters();
220        let (task_owner, task_tracker) = ConnectorTaskOwner::new();
221        Self {
222            config,
223            state: ConnectorState::Created,
224            schema: cdc_envelope_schema(),
225            metrics: Arc::new(PostgresCdcMetrics::new(registry)),
226            relation_cache: RelationCache::new(),
227            committed_transactions: VecDeque::new(),
228            buffered_event_count: 0,
229            buffered_event_bytes: 0,
230            current_txn: None,
231            confirmed_flush_lsn: Lsn::ZERO,
232            write_lsn: Lsn::ZERO,
233            polled_lsn: Lsn::ZERO,
234            checkpoint_binding: None,
235            #[cfg(test)]
236            pending_messages: VecDeque::new(),
237            data_ready: Arc::new(Notify::new()),
238            wal_rx: None,
239            reader_handle: None,
240            reader_shutdown: None,
241            confirmed_lsn_tx: None,
242            pending_payloads: VecDeque::new(),
243            wal_byte_budget: None,
244            wal_terminal_error: None,
245            task_owner,
246            task_tracker,
247        }
248    }
249
250    /// Creates a new source from a generic [`ConnectorConfig`].
251    ///
252    /// # Errors
253    ///
254    /// Returns `ConnectorError` if the configuration is invalid.
255    pub fn from_config(config: &ConnectorConfig) -> Result<Self, ConnectorError> {
256        let pg_config = PostgresCdcConfig::from_config(config)?;
257        Ok(Self::new(pg_config, None))
258    }
259
260    /// Returns a reference to the CDC configuration.
261    #[must_use]
262    pub fn config(&self) -> &PostgresCdcConfig {
263        &self.config
264    }
265
266    /// Returns the current confirmed flush LSN.
267    #[must_use]
268    pub fn confirmed_flush_lsn(&self) -> Lsn {
269        self.confirmed_flush_lsn
270    }
271
272    /// Returns the current write LSN.
273    #[must_use]
274    pub fn write_lsn(&self) -> Lsn {
275        self.write_lsn
276    }
277
278    /// Returns the current replication lag in bytes.
279    #[must_use]
280    pub fn replication_lag_bytes(&self) -> u64 {
281        self.write_lsn.diff(self.confirmed_flush_lsn)
282    }
283
284    /// Returns a reference to the relation cache.
285    #[must_use]
286    pub fn relation_cache(&self) -> &RelationCache {
287        &self.relation_cache
288    }
289
290    /// Returns the number of buffered events.
291    #[must_use]
292    pub fn buffered_events(&self) -> usize {
293        self.buffered_event_count
294    }
295}
296
297// ── Test helpers ──
298
299#[cfg(test)]
300impl PostgresCdcSource {
301    /// Injects a pre-built change event directly into the event buffer.
302    fn inject_event(&mut self, event: ChangeEvent) {
303        let end_lsn = event.lsn;
304        self.buffered_event_bytes = self
305            .buffered_event_bytes
306            .checked_add(retained_event_bytes(&event).expect("test event size must be valid"))
307            .expect("test buffered-event bytes must be valid");
308        self.committed_transactions.push_back(CommittedTransaction {
309            end_lsn,
310            events: VecDeque::from([event]),
311        });
312        self.buffered_event_count += 1;
313    }
314
315    /// Builds a binary pgoutput Relation message for testing.
316    fn build_relation_message(
317        relation_id: u32,
318        namespace: &str,
319        name: &str,
320        columns: &[(u8, &str, u32, i32)], // (flags, name, oid, modifier)
321    ) -> Vec<u8> {
322        let mut buf = vec![b'R'];
323        buf.extend_from_slice(&relation_id.to_be_bytes());
324        buf.extend_from_slice(namespace.as_bytes());
325        buf.push(0);
326        buf.extend_from_slice(name.as_bytes());
327        buf.push(0);
328        buf.push(b'd'); // replica identity = default
329        buf.extend_from_slice(&(columns.len() as i16).to_be_bytes());
330        for (flags, col_name, oid, modifier) in columns {
331            buf.push(*flags);
332            buf.extend_from_slice(col_name.as_bytes());
333            buf.push(0);
334            buf.extend_from_slice(&oid.to_be_bytes());
335            buf.extend_from_slice(&modifier.to_be_bytes());
336        }
337        buf
338    }
339
340    /// Builds a binary pgoutput Begin message for testing.
341    fn build_begin_message(final_lsn: u64, commit_ts_us: i64, xid: u32) -> Vec<u8> {
342        let mut buf = vec![b'B'];
343        buf.extend_from_slice(&final_lsn.to_be_bytes());
344        buf.extend_from_slice(&commit_ts_us.to_be_bytes());
345        buf.extend_from_slice(&xid.to_be_bytes());
346        buf
347    }
348
349    /// Builds a binary pgoutput Commit message for testing.
350    fn build_commit_message(commit_lsn: u64, end_lsn: u64, commit_ts_us: i64) -> Vec<u8> {
351        let mut buf = vec![b'C'];
352        buf.push(0); // flags
353        buf.extend_from_slice(&commit_lsn.to_be_bytes());
354        buf.extend_from_slice(&end_lsn.to_be_bytes());
355        buf.extend_from_slice(&commit_ts_us.to_be_bytes());
356        buf
357    }
358
359    /// Builds a binary pgoutput Insert message for testing.
360    fn build_insert_message(relation_id: u32, values: &[Option<&str>]) -> Vec<u8> {
361        let mut buf = vec![b'I'];
362        buf.extend_from_slice(&relation_id.to_be_bytes());
363        buf.push(b'N');
364        buf.extend_from_slice(&(values.len() as i16).to_be_bytes());
365        for val in values {
366            match val {
367                Some(s) => {
368                    buf.push(b't');
369                    buf.extend_from_slice(&(s.len() as i32).to_be_bytes());
370                    buf.extend_from_slice(s.as_bytes());
371                }
372                None => buf.push(b'n'),
373            }
374        }
375        buf
376    }
377
378    /// Builds a binary pgoutput Delete message for testing.
379    fn build_delete_message(relation_id: u32, values: &[Option<&str>]) -> Vec<u8> {
380        let mut buf = vec![b'D'];
381        buf.extend_from_slice(&relation_id.to_be_bytes());
382        buf.push(b'K'); // key identity
383        buf.extend_from_slice(&(values.len() as i16).to_be_bytes());
384        for val in values {
385            match val {
386                Some(s) => {
387                    buf.push(b't');
388                    buf.extend_from_slice(&(s.len() as i32).to_be_bytes());
389                    buf.extend_from_slice(s.as_bytes());
390                }
391                None => buf.push(b'n'),
392            }
393        }
394        buf
395    }
396
397    /// Builds a binary pgoutput Truncate message for testing.
398    fn build_truncate_message(relation_ids: &[u32], options: u8) -> Vec<u8> {
399        let mut buf = vec![b'T'];
400        buf.extend_from_slice(&(relation_ids.len() as i32).to_be_bytes());
401        buf.push(options);
402        for id in relation_ids {
403            buf.extend_from_slice(&id.to_be_bytes());
404        }
405        buf
406    }
407
408    /// Builds a binary pgoutput Update message (with old tuple) for testing.
409    fn build_update_message(
410        relation_id: u32,
411        old_tuple_tag: u8,
412        old_values: &[Option<&str>],
413        new_values: &[Option<&str>],
414    ) -> Vec<u8> {
415        assert!(matches!(old_tuple_tag, b'K' | b'O'));
416        let mut buf = vec![b'U'];
417        buf.extend_from_slice(&relation_id.to_be_bytes());
418        buf.push(old_tuple_tag);
419        buf.extend_from_slice(&(old_values.len() as i16).to_be_bytes());
420        for val in old_values {
421            match val {
422                Some(s) => {
423                    buf.push(b't');
424                    buf.extend_from_slice(&(s.len() as i32).to_be_bytes());
425                    buf.extend_from_slice(s.as_bytes());
426                }
427                None => buf.push(b'n'),
428            }
429        }
430        // New tuple
431        buf.push(b'N');
432        buf.extend_from_slice(&(new_values.len() as i16).to_be_bytes());
433        for val in new_values {
434            match val {
435                Some(s) => {
436                    buf.push(b't');
437                    buf.extend_from_slice(&(s.len() as i32).to_be_bytes());
438                    buf.extend_from_slice(s.as_bytes());
439                }
440                None => buf.push(b'n'),
441            }
442        }
443        buf
444    }
445}
446
447#[cfg(test)]
448mod tests;
laminar_connectors/postgres/cdc/source/mod.rs

laminar_connectors/postgres/cdc/source/
mod.rs
