laminar_connectors/postgres/cdc/source/
mod.rs

1//! `PostgreSQL` CDC source connector implementation.
2//!
3//! Streams the changes of one primary-keyed table from `PostgreSQL` logical replication
4//! (`pgoutput`) as typed Arrow rows of the declared source schema. A fresh source claims a new
5//! slot name and creates the slot once a checkpoint commits the claim; it then either copies the
6//! table in chunks while streaming from the slot's consistent point (`snapshot.mode=initial`) or
7//! starts at that point (`snapshot.mode=never`). Later starts resume from the engine checkpoint.
8
9use arrow_array::RecordBatch;
10use arrow_schema::{Schema, SchemaRef};
11use bytes::Bytes;
12use std::collections::VecDeque;
13use std::sync::Arc;
14use tokio::sync::Notify;
15use tokio::sync::{OwnedSemaphorePermit, Semaphore};
16
17use crate::config::ConnectorState;
18use crate::connector::{ConnectorTaskOwner, ConnectorTaskTracker, SourceMutation};
19use crate::error::ConnectorError;
20
21use super::config::PostgresCdcConfig;
22use super::lsn::Lsn;
23use super::metrics::PostgresCdcMetrics;
24use super::postgres_io::PostgresCheckpointBinding;
25use super::schema::RelationInfo;
26use super::typed_rows::{RowBuilder, RowLayout};
27
28mod checkpoint;
29mod claim;
30mod decoding;
31mod drain;
32mod lifecycle;
33mod reader;
34mod snapshot;
35mod startup;
36
37use claim::{ClaimTask, SlotClaim};
38use reader::{OwnedWalPayload, WalPayload, WalPayloadRx, WalTerminalError};
39use snapshot::{ChunkedSnapshot, CopyPlan};
40
41const PGWIRE_IN_FLIGHT_EVENTS: usize = 1;
42const RAW_WAL_QUEUE_CAPACITY: usize = 4_096;
43
44/// Where the source currently reads from.
45enum Phase {
46    /// Not started, or closed.
47    Idle,
48    /// Intake held until a checkpoint commits the claim and its slot exists.
49    Claiming(ClaimTask),
50    /// `snapshot.mode=never`: the slot exists, and the reader starts at its consistent point once
51    /// a committed `streaming` cursor names it (`released`). A restart from the claim may stream
52    /// from a later point on another slot, so nothing may be emitted under it.
53    AwaitingStream { released: bool },
54    /// Copying the table in chunks while streaming, once a cursor naming the slot committed.
55    Snapshot(Box<ChunkedSnapshot>),
56    /// Streaming committed transactions from the slot.
57    Streaming,
58}
59
60/// One emitted row: which image it reads and how it is applied.
61#[derive(Debug, Clone, Copy, PartialEq, Eq)]
62enum RowImage {
63    Put,
64    Tombstone,
65    Weighted(i64),
66}
67
68impl RowImage {
69    /// Whether only key columns are emitted, the `__weight`, and the upsert mutation.
70    fn parts(self) -> (bool, Option<i64>, Option<SourceMutation>) {
71        match self {
72            Self::Put => (false, None, Some(SourceMutation::Put)),
73            Self::Tombstone => (true, None, Some(SourceMutation::Tombstone)),
74            Self::Weighted(weight) => (false, Some(weight), None),
75        }
76    }
77}
78
79/// Declared-column texts of one row; `None` is SQL `NULL`.
80type RowValues = Box<[Option<Box<[u8]>>]>;
81
82/// A row decoded during the copy, kept as texts until the drain knows where its key lies.
83struct StagedRow {
84    image: RowImage,
85    values: RowValues,
86    /// Planned retained bytes, charged to the decoded stage.
87    bytes: usize,
88}
89
90/// The staged rows of the open transaction.
91#[derive(Default)]
92struct StagedRows {
93    rows: Vec<StagedRow>,
94    bytes: usize,
95}
96
97/// The transaction being decoded; its rows become visible only at COMMIT.
98struct OpenTransaction {
99    final_lsn: Lsn,
100    xid: u32,
101    /// Rows of a transaction decoded during the copy.
102    staged: Option<StagedRows>,
103}
104
105/// A decoded transaction, resumable only after every one of its rows has been emitted.
106struct CommittedTransaction {
107    end_lsn: Lsn,
108    xid: u32,
109    /// `None` when the transaction changed no captured rows, or was decoded during the copy.
110    records: Option<RecordBatch>,
111    /// Row mutations in upsert mode; empty in changelog mode.
112    mutations: Box<[SourceMutation]>,
113    /// Rows of a transaction decoded during the copy, not built until drained.
114    staged: Option<Box<[StagedRow]>>,
115    retained_bytes: usize,
116}
117
118impl CommittedTransaction {
119    fn rows(&self) -> usize {
120        self.records.as_ref().map_or(0, RecordBatch::num_rows)
121            + self.staged.as_ref().map_or(0, |staged| staged.len())
122    }
123}
124
125/// `PostgreSQL` CDC source connector.
126///
127/// In `output.mode=upsert` every row is a [`SourceMutation::Put`] of the complete current row or
128/// a key-only [`SourceMutation::Tombstone`]; a primary-key change is a tombstone of the old key
129/// followed by a put of the new row. In `output.mode=changelog` rows carry a trailing `__weight`:
130/// `+1` for inserted and updated-to images, `-1` for deleted and updated-from images.
131pub struct PostgresCdcSource {
132    config: PostgresCdcConfig,
133    state: ConnectorState,
134    /// The error that failed the source, repeated by every later poll.
135    failure: Option<String>,
136    /// Declared output schema; empty until started.
137    schema: SchemaRef,
138    metrics: Arc<PostgresCdcMetrics>,
139    phase: Phase,
140
141    /// Declared-schema binding fixed at startup.
142    layout: Option<RowLayout>,
143    /// The bound relation every `pgoutput` Relation message must match exactly.
144    relation: Option<RelationInfo>,
145    relation_announced: bool,
146
147    open_transaction: Option<OpenTransaction>,
148    open_rows: Option<RowBuilder>,
149    open_mutations: Vec<SourceMutation>,
150    committed: VecDeque<CommittedTransaction>,
151    /// Bytes retained by `committed` batches.
152    committed_bytes: usize,
153
154    /// Last position whose durable commit was handed to the replication worker.
155    confirmed_flush_lsn: Lsn,
156    /// Latest position received from the server.
157    write_lsn: Lsn,
158    /// The resumable cursor: the end LSN of the last transaction drained into a batch, or a later
159    /// keepalive position received while no transaction was open or undrained.
160    polled_lsn: Lsn,
161
162    /// Exact database and publication identity bound to checkpoints.
163    checkpoint_binding: Option<PostgresCheckpointBinding>,
164    /// The claim whose slot this source creates or reads.
165    claim: Option<SlotClaim>,
166    /// Random per `start()`, in this incarnation's `application_name`.
167    incarnation: String,
168    /// The owned slot's consistent point, once the slot exists.
169    consistent_point: Option<Lsn>,
170    /// SQL of the initial copy, read at startup when the source copies.
171    copy_plan: Option<Arc<CopyPlan>>,
172
173    #[cfg(test)]
174    pending_messages: VecDeque<Vec<u8>>,
175
176    data_ready: Arc<Notify>,
177    wal_rx: Option<WalPayloadRx>,
178    reader_handle: Option<tokio::task::JoinHandle<()>>,
179    reader_shutdown: Option<tokio::sync::watch::Sender<bool>>,
180    /// Durable progress reported by the replication worker on every status update,
181    /// independently of the reader task and of a full event queue.
182    applied_lsn: Option<pgwire_replication::AppliedLsnHandle>,
183    /// Payloads taken from the reader queue but not yet decoded, in WAL order.
184    pending_payloads: VecDeque<OwnedWalPayload>,
185    wal_byte_budget: Option<Arc<Semaphore>>,
186    /// Fatal reader error delivered out of band so a full WAL queue cannot hide it.
187    wal_terminal_error: Option<WalTerminalError>,
188    /// Earliest instant of the next live publication/table revalidation.
189    next_contract_check: Option<tokio::time::Instant>,
190
191    task_owner: ConnectorTaskOwner,
192    task_tracker: ConnectorTaskTracker,
193}
194
195impl Drop for PostgresCdcSource {
196    fn drop(&mut self) {
197        if let Some(shutdown) = self.reader_shutdown.take() {
198            shutdown.send_replace(true);
199        }
200        if let Some(handle) = self.reader_handle.take() {
201            reap_postgres_reader(handle, &self.task_owner);
202        }
203    }
204}
205
206fn reap_postgres_reader(handle: tokio::task::JoinHandle<()>, task_owner: &ConnectorTaskOwner) {
207    let Some(reaper_guard) = task_owner.track() else {
208        tracing::warn!("PostgreSQL CDC task generation was sealed before reader reaping");
209        return;
210    };
211    let Ok(runtime) = tokio::runtime::Handle::try_current() else {
212        // The reader's own guard remains authoritative. Dropping its runtime
213        // destroys the future and therefore resolves the generation tracker.
214        drop(reaper_guard);
215        return;
216    };
217    drop(runtime.spawn(async move {
218        let _reaper_guard = reaper_guard;
219        if let Err(error) = handle.await {
220            tracing::debug!(%error, "PostgreSQL CDC retired reader task reaped");
221        }
222    }));
223}
224
225impl PostgresCdcSource {
226    /// Creates a new `PostgreSQL` CDC source with the given configuration.
227    #[must_use]
228    pub fn new(config: PostgresCdcConfig, registry: Option<&prometheus::Registry>) -> Self {
229        let (task_owner, task_tracker) = ConnectorTaskOwner::new();
230        Self {
231            config,
232            state: ConnectorState::Created,
233            failure: None,
234            schema: Arc::new(Schema::empty()),
235            metrics: Arc::new(PostgresCdcMetrics::new(registry)),
236            phase: Phase::Idle,
237            layout: None,
238            relation: None,
239            relation_announced: false,
240            open_transaction: None,
241            open_rows: None,
242            open_mutations: Vec::new(),
243            committed: VecDeque::new(),
244            committed_bytes: 0,
245            confirmed_flush_lsn: Lsn::ZERO,
246            write_lsn: Lsn::ZERO,
247            polled_lsn: Lsn::ZERO,
248            checkpoint_binding: None,
249            claim: None,
250            incarnation: String::new(),
251            consistent_point: None,
252            copy_plan: None,
253            #[cfg(test)]
254            pending_messages: VecDeque::new(),
255            data_ready: Arc::new(Notify::new()),
256            wal_rx: None,
257            reader_handle: None,
258            reader_shutdown: None,
259            applied_lsn: None,
260            pending_payloads: VecDeque::new(),
261            wal_byte_budget: None,
262            wal_terminal_error: None,
263            next_contract_check: None,
264            task_owner,
265            task_tracker,
266        }
267    }
268
269    /// Returns a reference to the CDC configuration.
270    #[must_use]
271    pub fn config(&self) -> &PostgresCdcConfig {
272        &self.config
273    }
274
275    /// Returns the last durably committed position handed to the replication worker.
276    #[must_use]
277    pub fn confirmed_flush_lsn(&self) -> Lsn {
278        self.confirmed_flush_lsn
279    }
280
281    /// Returns the latest position received from the server.
282    #[must_use]
283    pub fn write_lsn(&self) -> Lsn {
284        self.write_lsn
285    }
286
287    /// Returns the current replication lag in bytes.
288    #[must_use]
289    pub fn replication_lag_bytes(&self) -> u64 {
290        self.write_lsn.diff(self.confirmed_flush_lsn)
291    }
292
293    /// Rows decoded but not yet emitted, across the open and committed transactions.
294    #[must_use]
295    pub fn buffered_rows(&self) -> usize {
296        let committed: usize = self.committed.iter().map(CommittedTransaction::rows).sum();
297        let staged = self
298            .open_transaction
299            .as_ref()
300            .and_then(|transaction| transaction.staged.as_ref())
301            .map_or(0, |staged| staged.rows.len());
302        committed + staged + self.open_rows.as_ref().map_or(0, RowBuilder::len)
303    }
304
305    fn fail(&mut self, error: ConnectorError) -> ConnectorError {
306        self.state = ConnectorState::Failed;
307        self.failure.get_or_insert_with(|| error.to_string());
308        error
309    }
310}
311
312#[cfg(test)]
313mod tests;
laminar_connectors/postgres/cdc/source/mod.rs

laminar_connectors/postgres/cdc/source/
mod.rs
