Skip to main content

laminar_connectors/postgres/cdc/changelog/
mod.rs

1//! CDC envelope construction for `PostgreSQL` row changes.
2//!
3//! `_before` contains the old full row for `REPLICA IDENTITY FULL`, or only
4//! replica-identity fields for a pgoutput key tuple. Unavailable fields are
5//! omitted so they remain distinct from an available SQL `NULL`.
6
7use std::sync::Arc;
8
9use arrow_array::builder::{StringBuilder, TimestampMillisecondBuilder, UInt64Builder};
10use arrow_array::RecordBatch;
11use arrow_schema::SchemaRef;
12
13use crate::error::ConnectorError;
14
15use super::decoder::{ColumnValue, OldTuple, TupleData};
16use super::lsn::Lsn;
17use super::schema::{cdc_envelope_schema, RelationInfo};
18
19/// A CDC operation type.
20#[derive(Debug, Clone, Copy, PartialEq, Eq)]
21pub enum CdcOperation {
22    /// Row inserted.
23    Insert,
24    /// Row updated.
25    Update,
26    /// Row deleted.
27    Delete,
28}
29
30impl CdcOperation {
31    /// Returns the single-character code for the operation.
32    #[must_use]
33    pub fn as_str(self) -> &'static str {
34        match self {
35            CdcOperation::Insert => "I",
36            CdcOperation::Update => "U",
37            CdcOperation::Delete => "D",
38        }
39    }
40}
41
42/// A single change event from CDC.
43#[derive(Debug)]
44pub struct ChangeEvent {
45    /// Fully qualified table name.
46    pub table: String,
47    /// The operation type.
48    pub op: CdcOperation,
49    /// WAL position of this change.
50    pub lsn: Lsn,
51    /// Commit timestamp in milliseconds since Unix epoch.
52    pub ts_ms: i64,
53    /// Available old identity/full-row values as JSON (for UPDATE and DELETE).
54    pub before: Option<String>,
55    /// New row values as JSON (for INSERT and UPDATE).
56    pub after: Option<String>,
57}
58
59/// Converts tuple data to a JSON string using column names from the relation.
60///
61/// Produces a flat JSON object like `{"id": "42", "name": "Alice"}`.
62/// All values are serialized as strings (matching `pgoutput` text format).
63///
64/// Writes JSON directly to a `String` buffer instead of building an
65/// intermediate `HashMap`, avoiding per-row map + key/value cloning.
66pub(crate) fn tuple_json_encoded_len(
67    tuple: &TupleData,
68    relation: &RelationInfo,
69) -> Result<usize, ConnectorError> {
70    tuple_json_encoded_len_inner(tuple, relation, false)
71}
72
73pub(crate) fn old_tuple_json_encoded_len(
74    old_tuple: &OldTuple,
75    relation: &RelationInfo,
76) -> Result<usize, ConnectorError> {
77    let (tuple, identity_only) = old_tuple_parts(old_tuple);
78    tuple_json_encoded_len_inner(tuple, relation, identity_only)
79}
80
81fn tuple_json_encoded_len_inner(
82    tuple: &TupleData,
83    relation: &RelationInfo,
84    identity_only: bool,
85) -> Result<usize, ConnectorError> {
86    validate_tuple_columns(tuple, relation)?;
87    let mut len = 2_usize; // braces
88    let mut first = true;
89    for (index, value) in tuple.columns.iter().enumerate() {
90        let column = &relation.columns[index];
91        if (identity_only && !column.is_key) || matches!(value, ColumnValue::Unchanged) {
92            continue;
93        }
94        if !first {
95            checked_add(&mut len, 1, "JSON separator")?;
96        }
97        first = false;
98        checked_add(&mut len, 3, "JSON key quotes and colon")?;
99        checked_add(&mut len, escaped_json_len(&column.name)?, "JSON key")?;
100        match value {
101            ColumnValue::Text(bytes) => {
102                let text = std::str::from_utf8(bytes).map_err(|_| {
103                    ConnectorError::ReadError("PostgreSQL tuple contains invalid UTF-8".into())
104                })?;
105                checked_add(&mut len, 2, "JSON value quotes")?;
106                checked_add(&mut len, escaped_json_len(text)?, "JSON value")?;
107            }
108            ColumnValue::Null => checked_add(&mut len, 4, "JSON null")?,
109            ColumnValue::Unchanged => unreachable!("unchanged values were skipped"),
110        }
111    }
112    Ok(len)
113}
114
115pub(crate) fn tuple_to_json(
116    tuple: &TupleData,
117    relation: &RelationInfo,
118    encoded_len: usize,
119) -> Result<String, ConnectorError> {
120    tuple_to_json_inner(tuple, relation, encoded_len, false)
121}
122
123pub(crate) fn old_tuple_to_json(
124    old_tuple: &OldTuple,
125    relation: &RelationInfo,
126    encoded_len: usize,
127) -> Result<String, ConnectorError> {
128    let (tuple, identity_only) = old_tuple_parts(old_tuple);
129    tuple_to_json_inner(tuple, relation, encoded_len, identity_only)
130}
131
132fn tuple_to_json_inner(
133    tuple: &TupleData,
134    relation: &RelationInfo,
135    encoded_len: usize,
136    identity_only: bool,
137) -> Result<String, ConnectorError> {
138    validate_tuple_columns(tuple, relation)?;
139    let mut buf = String::new();
140    buf.try_reserve_exact(encoded_len).map_err(|error| {
141        ConnectorError::ReadError(format!(
142            "PostgreSQL CDC could not reserve {encoded_len} JSON bytes: {error}"
143        ))
144    })?;
145    buf.push('{');
146    let mut first = true;
147    for (col_val, col_info) in tuple.columns.iter().zip(&relation.columns) {
148        if identity_only && !col_info.is_key {
149            continue;
150        }
151        let val = match col_val {
152            ColumnValue::Text(bytes) => Some(std::str::from_utf8(bytes).map_err(|_| {
153                ConnectorError::ReadError("PostgreSQL tuple contains invalid UTF-8".into())
154            })?),
155            ColumnValue::Null => None,
156            ColumnValue::Unchanged => continue,
157        };
158        if !first {
159            buf.push(',');
160        }
161        first = false;
162        buf.push('"');
163        escape_json_str(&col_info.name, &mut buf);
164        buf.push('"');
165        buf.push(':');
166        match val {
167            Some(s) => {
168                buf.push('"');
169                escape_json_str(s, &mut buf);
170                buf.push('"');
171            }
172            None => buf.push_str("null"),
173        }
174    }
175    buf.push('}');
176    debug_assert_eq!(buf.len(), encoded_len);
177    Ok(buf)
178}
179
180fn old_tuple_parts(old_tuple: &OldTuple) -> (&TupleData, bool) {
181    match old_tuple {
182        OldTuple::Key(tuple) => (tuple, true),
183        OldTuple::Full(tuple) => (tuple, false),
184    }
185}
186
187fn validate_tuple_columns(
188    tuple: &TupleData,
189    relation: &RelationInfo,
190) -> Result<(), ConnectorError> {
191    if tuple.columns.len() != relation.columns.len() {
192        return Err(ConnectorError::ReadError(format!(
193            "PostgreSQL tuple column count {} does not match relation {} column count {}",
194            tuple.columns.len(),
195            relation.relation_id,
196            relation.columns.len()
197        )));
198    }
199    Ok(())
200}
201
202fn escaped_json_len(value: &str) -> Result<usize, ConnectorError> {
203    let mut len = 0_usize;
204    for character in value.chars() {
205        let bytes = match character {
206            '"' | '\\' | '\n' | '\r' | '\t' => 2,
207            value if value.is_control() => 6,
208            value => value.len_utf8(),
209        };
210        checked_add(&mut len, bytes, "escaped JSON")?;
211    }
212    Ok(len)
213}
214
215/// Escapes a string for JSON output (quotes and control characters).
216fn escape_json_str(s: &str, buf: &mut String) {
217    for ch in s.chars() {
218        match ch {
219            '"' => buf.push_str("\\\""),
220            '\\' => buf.push_str("\\\\"),
221            '\n' => buf.push_str("\\n"),
222            '\r' => buf.push_str("\\r"),
223            '\t' => buf.push_str("\\t"),
224            c if c.is_control() => {
225                const HEX: &[u8; 16] = b"0123456789abcdef";
226                let code = c as u32;
227                buf.push_str("\\u");
228                for shift in [12, 8, 4, 0] {
229                    buf.push(HEX[((code >> shift) & 0x0f) as usize] as char);
230                }
231            }
232            c => buf.push(c),
233        }
234    }
235}
236
237/// Converts a batch of [`ChangeEvent`]s into an Arrow [`RecordBatch`]
238/// using the CDC envelope schema.
239///
240/// # Errors
241///
242/// Returns an error if the Arrow batch construction fails.
243#[derive(Debug)]
244pub(crate) struct ArrowBatchPlan {
245    rows: usize,
246    table_bytes: usize,
247    before_bytes: usize,
248    after_bytes: usize,
249    pub(crate) retained_bytes: usize,
250}
251
252pub(crate) fn plan_record_batch<'a>(
253    events: impl IntoIterator<Item = &'a ChangeEvent>,
254) -> Result<ArrowBatchPlan, ConnectorError> {
255    let mut rows = 0_usize;
256    let mut table_bytes = 0_usize;
257    let mut before_bytes = 0_usize;
258    let mut after_bytes = 0_usize;
259    let mut before_has_null = false;
260    let mut after_has_null = false;
261    for event in events {
262        rows = checked_add_value(rows, 1, "Arrow row count")?;
263        checked_add(&mut table_bytes, event.table.len(), "Arrow table values")?;
264        checked_add(
265            &mut before_bytes,
266            event.before.as_ref().map_or(0, String::len),
267            "Arrow before values",
268        )?;
269        checked_add(
270            &mut after_bytes,
271            event.after.as_ref().map_or(0, String::len),
272            "Arrow after values",
273        )?;
274        before_has_null |= event.before.is_none();
275        after_has_null |= event.after.is_none();
276    }
277
278    let mut retained_bytes = 0_usize;
279    for bytes in [table_bytes, rows, before_bytes, after_bytes] {
280        checked_add(
281            &mut retained_bytes,
282            round_to_arrow_alignment(bytes)?,
283            "Arrow string values",
284        )?;
285    }
286    let offset_bytes = round_to_arrow_alignment(checked_mul(
287        checked_add_value(rows, 1, "Arrow offset rows")?,
288        std::mem::size_of::<i32>(),
289        "Arrow offsets",
290    )?)?;
291    checked_add(
292        &mut retained_bytes,
293        checked_mul(offset_bytes, 4, "Arrow string offsets")?,
294        "Arrow string offsets",
295    )?;
296    let primitive_bytes = round_to_arrow_alignment(checked_mul(
297        rows,
298        std::mem::size_of::<u64>(),
299        "Arrow primitive values",
300    )?)?;
301    checked_add(
302        &mut retained_bytes,
303        checked_mul(primitive_bytes, 2, "Arrow primitive columns")?,
304        "Arrow primitive columns",
305    )?;
306    let nullable_columns = usize::from(before_has_null) + usize::from(after_has_null);
307    let validity_bytes = round_to_arrow_alignment(rows.div_ceil(8))?;
308    checked_add(
309        &mut retained_bytes,
310        checked_mul(validity_bytes, nullable_columns, "Arrow validity buffers")?,
311        "Arrow validity buffers",
312    )?;
313
314    Ok(ArrowBatchPlan {
315        rows,
316        table_bytes,
317        before_bytes,
318        after_bytes,
319        retained_bytes,
320    })
321}
322
323pub(crate) fn events_to_record_batch<I>(
324    events: I,
325    plan: &ArrowBatchPlan,
326) -> Result<RecordBatch, ConnectorError>
327where
328    I: IntoIterator<Item = ChangeEvent>,
329{
330    let schema: SchemaRef = cdc_envelope_schema();
331
332    let mut table_builder = StringBuilder::with_capacity(plan.rows, plan.table_bytes);
333    let mut op_builder = StringBuilder::with_capacity(plan.rows, plan.rows);
334    let mut lsn_builder = UInt64Builder::with_capacity(plan.rows);
335    let mut ts_builder = TimestampMillisecondBuilder::with_capacity(plan.rows);
336    let mut before_builder = StringBuilder::with_capacity(plan.rows, plan.before_bytes);
337    let mut after_builder = StringBuilder::with_capacity(plan.rows, plan.after_bytes);
338
339    let mut rows = 0_usize;
340    for event in events {
341        rows = rows.checked_add(1).ok_or_else(|| {
342            ConnectorError::Internal("PostgreSQL CDC Arrow row-count overflow".into())
343        })?;
344        table_builder.append_value(&event.table);
345        op_builder.append_value(event.op.as_str());
346        lsn_builder.append_value(event.lsn.as_u64());
347        ts_builder.append_value(event.ts_ms);
348
349        match event.before {
350            Some(json) => before_builder.append_value(&json),
351            None => before_builder.append_null(),
352        }
353        match event.after {
354            Some(json) => after_builder.append_value(&json),
355            None => after_builder.append_null(),
356        }
357    }
358    if rows != plan.rows {
359        return Err(ConnectorError::Internal(
360            "PostgreSQL CDC Arrow plan row count changed before construction".into(),
361        ));
362    }
363
364    let batch = RecordBatch::try_new(
365        schema,
366        vec![
367            Arc::new(table_builder.finish()),
368            Arc::new(op_builder.finish()),
369            Arc::new(lsn_builder.finish()),
370            Arc::new(ts_builder.finish()),
371            Arc::new(before_builder.finish()),
372            Arc::new(after_builder.finish()),
373        ],
374    )
375    .map_err(|error| ConnectorError::Internal(format!("Arrow batch build: {error}")))?;
376    let actual_buffer_bytes = batch.columns().iter().try_fold(0_usize, |total, column| {
377        total
378            .checked_add(column.get_buffer_memory_size())
379            .ok_or_else(|| {
380                ConnectorError::Internal(
381                    "PostgreSQL CDC Arrow retained-byte accounting overflow".into(),
382                )
383            })
384    })?;
385    if actual_buffer_bytes > plan.retained_bytes {
386        return Err(ConnectorError::Internal(format!(
387            "PostgreSQL CDC Arrow retained-byte plan was too small: actual={actual_buffer_bytes}, planned={}",
388            plan.retained_bytes
389        )));
390    }
391    Ok(batch)
392}
393
394fn checked_add(total: &mut usize, value: usize, context: &str) -> Result<(), ConnectorError> {
395    *total = checked_add_value(*total, value, context)?;
396    Ok(())
397}
398
399fn checked_add_value(left: usize, right: usize, context: &str) -> Result<usize, ConnectorError> {
400    left.checked_add(right)
401        .ok_or_else(|| ConnectorError::ReadError(format!("PostgreSQL CDC {context} size overflow")))
402}
403
404fn checked_mul(left: usize, right: usize, context: &str) -> Result<usize, ConnectorError> {
405    left.checked_mul(right)
406        .ok_or_else(|| ConnectorError::ReadError(format!("PostgreSQL CDC {context} size overflow")))
407}
408
409fn round_to_arrow_alignment(bytes: usize) -> Result<usize, ConnectorError> {
410    Ok(checked_add_value(bytes, 63, "Arrow alignment")? & !63)
411}
412
413#[cfg(test)]
414mod tests;