1use 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#[derive(Debug, Clone, Copy, PartialEq, Eq)]
21pub enum CdcOperation {
22 Insert,
24 Update,
26 Delete,
28}
29
30impl CdcOperation {
31 #[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#[derive(Debug)]
44pub struct ChangeEvent {
45 pub table: String,
47 pub op: CdcOperation,
49 pub lsn: Lsn,
51 pub ts_ms: i64,
53 pub before: Option<String>,
55 pub after: Option<String>,
57}
58
59pub(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; 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
215fn 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#[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 {
415 use super::*;
416 use crate::postgres::cdc::schema::RelationInfo;
417 use crate::postgres::cdc::types::PgColumn;
418 use crate::postgres::cdc::types::{INT8_OID, TEXT_OID};
419 use bytes::Bytes;
420
421 fn sample_relation() -> RelationInfo {
422 RelationInfo {
423 relation_id: 16384,
424 namespace: "public".to_string(),
425 name: "users".to_string(),
426 replica_identity: 'd',
427 columns: vec![
428 PgColumn::new("id".to_string(), INT8_OID, -1, true),
429 PgColumn::new("name".to_string(), TEXT_OID, -1, false),
430 ],
431 }
432 }
433
434 #[test]
435 fn test_tuple_to_json() {
436 let relation = sample_relation();
437 let tuple = TupleData {
438 columns: vec![
439 ColumnValue::Text(Bytes::from_static(b"42")),
440 ColumnValue::Text(Bytes::from_static(b"Alice")),
441 ],
442 };
443
444 let encoded_len = tuple_json_encoded_len(&tuple, &relation).unwrap();
445 let json = tuple_to_json(&tuple, &relation, encoded_len).unwrap();
446 assert_eq!(json.len(), encoded_len);
447 let parsed: serde_json::Value = serde_json::from_str(&json).unwrap();
448 assert_eq!(parsed["id"], "42");
449 assert_eq!(parsed["name"], "Alice");
450 }
451
452 #[test]
453 fn test_tuple_to_json_with_null() {
454 let relation = sample_relation();
455 let tuple = TupleData {
456 columns: vec![
457 ColumnValue::Text(Bytes::from_static(b"42")),
458 ColumnValue::Null,
459 ],
460 };
461
462 let encoded_len = tuple_json_encoded_len(&tuple, &relation).unwrap();
463 let json = tuple_to_json(&tuple, &relation, encoded_len).unwrap();
464 let parsed: serde_json::Value = serde_json::from_str(&json).unwrap();
465 assert_eq!(parsed["id"], "42");
466 assert!(parsed["name"].is_null());
467 }
468
469 #[test]
470 fn test_tuple_to_json_unchanged_omitted() {
471 let relation = sample_relation();
472 let tuple = TupleData {
473 columns: vec![
474 ColumnValue::Text(Bytes::from_static(b"42")),
475 ColumnValue::Unchanged,
476 ],
477 };
478
479 let encoded_len = tuple_json_encoded_len(&tuple, &relation).unwrap();
480 let json = tuple_to_json(&tuple, &relation, encoded_len).unwrap();
481 let parsed: serde_json::Value = serde_json::from_str(&json).unwrap();
482 assert_eq!(parsed["id"], "42");
483 assert!(parsed.get("name").is_none());
485 }
486
487 #[test]
488 fn key_old_tuple_omits_unavailable_non_identity_fields() {
489 let relation = RelationInfo {
490 columns: vec![
491 PgColumn::new("id".to_string(), INT8_OID, -1, true),
492 PgColumn::new("name".to_string(), TEXT_OID, -1, false),
493 PgColumn::new("note".to_string(), TEXT_OID, -1, false),
494 ],
495 ..sample_relation()
496 };
497 let old_tuple = OldTuple::Key(TupleData {
498 columns: vec![
499 ColumnValue::Text(Bytes::from_static(b"42")),
500 ColumnValue::Text(Bytes::from_static(b"unavailable")),
501 ColumnValue::Null,
502 ],
503 });
504
505 let encoded_len = old_tuple_json_encoded_len(&old_tuple, &relation).unwrap();
506 let json = old_tuple_to_json(&old_tuple, &relation, encoded_len).unwrap();
507
508 assert_eq!(json, r#"{"id":"42"}"#);
509 assert_eq!(json.len(), encoded_len);
510 }
511
512 #[test]
513 fn full_old_tuple_retains_non_key_fields_and_explicit_null() {
514 let relation = sample_relation();
515 let old_tuple = OldTuple::Full(TupleData {
516 columns: vec![
517 ColumnValue::Text(Bytes::from_static(b"42")),
518 ColumnValue::Null,
519 ],
520 });
521
522 let encoded_len = old_tuple_json_encoded_len(&old_tuple, &relation).unwrap();
523 let json = old_tuple_to_json(&old_tuple, &relation, encoded_len).unwrap();
524
525 assert_eq!(json, r#"{"id":"42","name":null}"#);
526 assert_eq!(json.len(), encoded_len);
527 }
528
529 #[test]
530 fn key_old_tuple_requires_full_relation_cardinality() {
531 let relation = sample_relation();
532 let old_tuple = OldTuple::Key(TupleData {
533 columns: vec![ColumnValue::Text(Bytes::from_static(b"42"))],
534 });
535
536 let error = old_tuple_json_encoded_len(&old_tuple, &relation).unwrap_err();
537 assert!(error.to_string().contains("column count"), "{error}");
538 }
539
540 #[test]
541 fn test_events_to_record_batch_insert() {
542 let events = vec![ChangeEvent {
543 table: "users".to_string(),
544 op: CdcOperation::Insert,
545 lsn: Lsn::new(0x100),
546 ts_ms: 1_700_000_000_000,
547 before: None,
548 after: Some(r#"{"id":"1","name":"Alice"}"#.to_string()),
549 }];
550
551 let plan = plan_record_batch(&events).unwrap();
552 let batch = events_to_record_batch(events, &plan).unwrap();
553 assert_eq!(batch.num_rows(), 1);
554 assert_eq!(batch.num_columns(), 6);
555 }
556
557 #[test]
558 fn test_events_to_record_batch_mixed() {
559 let events = vec![
560 ChangeEvent {
561 table: "users".to_string(),
562 op: CdcOperation::Insert,
563 lsn: Lsn::new(0x100),
564 ts_ms: 1_700_000_000_000,
565 before: None,
566 after: Some(r#"{"id":"1"}"#.to_string()),
567 },
568 ChangeEvent {
569 table: "users".to_string(),
570 op: CdcOperation::Update,
571 lsn: Lsn::new(0x200),
572 ts_ms: 1_700_000_000_001,
573 before: Some(r#"{"id":"1","name":"Alice"}"#.to_string()),
574 after: Some(r#"{"id":"1","name":"Bob"}"#.to_string()),
575 },
576 ChangeEvent {
577 table: "users".to_string(),
578 op: CdcOperation::Delete,
579 lsn: Lsn::new(0x300),
580 ts_ms: 1_700_000_000_002,
581 before: Some(r#"{"id":"1"}"#.to_string()),
582 after: None,
583 },
584 ];
585
586 let plan = plan_record_batch(&events).unwrap();
587 let batch = events_to_record_batch(events, &plan).unwrap();
588 assert_eq!(batch.num_rows(), 3);
589 }
590
591 #[test]
592 fn test_events_to_record_batch_empty() {
593 let events: Vec<ChangeEvent> = vec![];
594 let plan = plan_record_batch(&events).unwrap();
595 let batch = events_to_record_batch(events, &plan).unwrap();
596 assert_eq!(batch.num_rows(), 0);
597 assert_eq!(batch.num_columns(), 6);
598 }
599
600 #[test]
601 fn test_cdc_operation_as_str() {
602 assert_eq!(CdcOperation::Insert.as_str(), "I");
603 assert_eq!(CdcOperation::Update.as_str(), "U");
604 assert_eq!(CdcOperation::Delete.as_str(), "D");
605 }
606
607 #[test]
608 fn json_preflight_matches_all_escape_classes() {
609 let relation = RelationInfo {
610 columns: vec![PgColumn::new(
611 "control\nkey".to_string(),
612 TEXT_OID,
613 -1,
614 false,
615 )],
616 ..sample_relation()
617 };
618 let tuple = TupleData {
619 columns: vec![ColumnValue::Text(Bytes::from_static(
620 b"quote\" slash\\ newline\n tab\t",
621 ))],
622 };
623 let encoded_len = tuple_json_encoded_len(&tuple, &relation).unwrap();
624 let json = tuple_to_json(&tuple, &relation, encoded_len).unwrap();
625 assert_eq!(json.len(), encoded_len);
626 serde_json::from_str::<serde_json::Value>(&json).unwrap();
627 }
628
629 #[test]
630 fn json_preflight_rejects_invalid_text_and_column_count_drift() {
631 let relation = sample_relation();
632 let invalid_text = TupleData {
633 columns: vec![
634 ColumnValue::Text(Bytes::from_static(&[0xff])),
635 ColumnValue::Null,
636 ],
637 };
638 assert!(tuple_json_encoded_len(&invalid_text, &relation)
639 .unwrap_err()
640 .to_string()
641 .contains("UTF-8"));
642
643 let truncated = TupleData {
644 columns: vec![ColumnValue::Null],
645 };
646 assert!(tuple_json_encoded_len(&truncated, &relation)
647 .unwrap_err()
648 .to_string()
649 .contains("column count"));
650 }
651
652 #[test]
653 fn arrow_plan_covers_actual_retained_buffers() {
654 let events = vec![
655 ChangeEvent {
656 table: "public.users".into(),
657 op: CdcOperation::Insert,
658 lsn: Lsn::new(1),
659 ts_ms: 1,
660 before: None,
661 after: Some("{\"id\":\"1\"}".into()),
662 },
663 ChangeEvent {
664 table: "public.users".into(),
665 op: CdcOperation::Update,
666 lsn: Lsn::new(2),
667 ts_ms: 2,
668 before: Some("{\"id\":\"1\"}".into()),
669 after: Some("{\"id\":\"2\"}".into()),
670 },
671 ];
672 let plan = plan_record_batch(&events).unwrap();
673 let planned = plan.retained_bytes;
674 let batch = events_to_record_batch(events, &plan).unwrap();
675 let actual = batch
676 .columns()
677 .iter()
678 .map(|column| column.get_buffer_memory_size())
679 .sum::<usize>();
680 assert!(actual <= planned, "{actual} > {planned}");
681 }
682}