1use std::cmp::Ordering;
16use std::sync::LazyLock;
17
18use anyhow::Context;
19use futures::stream::BoxStream;
20use futures::{StreamExt, pin_mut};
21use futures_async_stream::{for_await, try_stream};
22use itertools::Itertools;
23use risingwave_common::catalog::{Field, Schema};
24use risingwave_common::log::LogSuppressor;
25use risingwave_common::row::{OwnedRow, Row};
26use risingwave_common::types::{DataType, Datum, ScalarImpl, ToOwnedDatum};
27use risingwave_common::util::iter_util::ZipEqFast;
28use serde::{Deserialize, Serialize};
29use thiserror_ext::AsReport;
30use tokio_postgres::types::{PgLsn, Type as PgType};
31
32use crate::connector_common::create_pg_client;
33use crate::connector_common::postgres::postgres_point_type;
34use crate::error::{ConnectorError, ConnectorResult};
35use crate::parser::scalar_adapter::ScalarAdapter;
36use crate::parser::{
37 postgres_cell_to_scalar_impl_strict, postgres_row_to_owned_row_with_strict_pk,
38};
39use crate::source::CdcTableSnapshotSplit;
40use crate::source::cdc::external::{
41 CDC_TABLE_SPLIT_ID_START, CdcOffset, CdcOffsetParseFunc, CdcTableSnapshotSplitOption,
42 DebeziumOffset, ExternalTableConfig, ExternalTableReader, SchemaTableName,
43};
44
45#[derive(Debug, Clone, Default, Serialize, Deserialize)]
46pub struct PostgresOffset {
47 pub txid: i64,
48 pub lsn: u64,
51 #[serde(default)]
53 pub lsn_commit: Option<u64>,
54 #[serde(default)]
55 pub lsn_proc: Option<u64>,
56}
57
58impl PartialOrd for PostgresOffset {
59 fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
60 Some(self.cmp(other))
61 }
62}
63
64impl Eq for PostgresOffset {}
65impl PartialEq for PostgresOffset {
66 fn eq(&self, other: &Self) -> bool {
67 match (
68 self.lsn_commit,
69 self.lsn_proc,
70 other.lsn_commit,
71 other.lsn_proc,
72 ) {
73 (_, Some(_), _, Some(_)) => {
74 self.lsn_commit == other.lsn_commit && self.lsn_proc == other.lsn_proc
75 }
76 _ => self.lsn == other.lsn,
77 }
78 }
79}
80
81impl Ord for PostgresOffset {
83 fn cmp(&self, other: &Self) -> Ordering {
84 match (
85 self.lsn_commit,
86 self.lsn_proc,
87 other.lsn_commit,
88 other.lsn_proc,
89 ) {
90 (_, Some(self_proc), _, Some(other_proc)) => {
91 match self.lsn_commit.cmp(&other.lsn_commit) {
94 Ordering::Equal => self_proc.cmp(&other_proc),
95 other_result => other_result,
96 }
97 }
98 _ => {
99 static LOG_SUPPRESSOR: LazyLock<LogSuppressor> =
101 LazyLock::new(LogSuppressor::default);
102 if let Ok(suppressed_count) = LOG_SUPPRESSOR.check() {
103 tracing::warn!(
104 suppressed_count,
105 self_lsn = self.lsn,
106 other_lsn = other.lsn,
107 "lsn_commit and lsn_proc are missing, fall back to lsn comparison"
108 );
109 }
110 self.lsn.cmp(&other.lsn)
111 }
112 }
113 }
114}
115
116impl PostgresOffset {
117 pub fn parse_debezium_offset(offset: &str) -> ConnectorResult<Self> {
118 let dbz_offset: DebeziumOffset = serde_json::from_str(offset)
119 .with_context(|| format!("invalid upstream offset: {}", offset))?;
120
121 let lsn = dbz_offset
122 .source_offset
123 .lsn
124 .context("invalid postgres lsn")?;
125
126 let lsn_commit = dbz_offset.source_offset.lsn_commit;
128
129 let lsn_proc = dbz_offset
130 .source_offset
131 .lsn_proc
132 .context("invalid postgres lsn_proc")?;
133
134 Ok(Self {
135 txid: dbz_offset
136 .source_offset
137 .txid
138 .context("invalid postgres txid")?,
139 lsn,
140 lsn_commit,
141 lsn_proc: Some(lsn_proc),
142 })
143 }
144}
145
146pub struct PostgresExternalTableReader {
147 rw_schema: Schema,
148 field_names: String,
149 pk_indices: Vec<usize>,
150 client: tokio::sync::Mutex<tokio_postgres::Client>,
151 schema_table_name: SchemaTableName,
152}
153
154impl ExternalTableReader for PostgresExternalTableReader {
155 async fn current_cdc_offset(&self) -> ConnectorResult<CdcOffset> {
156 let mut client = self.client.lock().await;
157 let trxn = client.transaction().await?;
159 let row = trxn.query_one("SELECT pg_current_wal_lsn()", &[]).await?;
160 let mut pg_offset = PostgresOffset::default();
161 let pg_lsn = row.get::<_, PgLsn>(0);
162 tracing::debug!("current lsn: {}", pg_lsn);
163 pg_offset.lsn = pg_lsn.into();
164
165 let txid_row = trxn.query_one("SELECT txid_current()", &[]).await?;
166 let txid: i64 = txid_row.get::<_, i64>(0);
167 pg_offset.txid = txid;
168
169 trxn.commit().await?;
171
172 Ok(CdcOffset::Postgres(pg_offset))
173 }
174
175 fn snapshot_read(
176 &self,
177 table_name: SchemaTableName,
178 start_pk: Option<OwnedRow>,
179 primary_keys: Vec<String>,
180 limit: u32,
181 ) -> BoxStream<'_, ConnectorResult<OwnedRow>> {
182 assert_eq!(table_name, self.schema_table_name);
183 self.snapshot_read_inner(table_name, start_pk, primary_keys, limit)
184 }
185
186 #[try_stream(boxed, ok = CdcTableSnapshotSplit, error = ConnectorError)]
187 async fn get_parallel_cdc_splits(&self, options: CdcTableSnapshotSplitOption) {
188 let backfill_num_rows_per_split = options.backfill_num_rows_per_split;
189 if backfill_num_rows_per_split == 0 {
190 return Err(anyhow::anyhow!(
191 "invalid backfill_num_rows_per_split, must be greater than 0"
192 )
193 .into());
194 }
195 if options.backfill_split_pk_column_index as usize >= self.pk_indices.len() {
196 return Err(anyhow::anyhow!(format!(
197 "invalid backfill_split_pk_column_index {}, out of bound",
198 options.backfill_split_pk_column_index
199 ))
200 .into());
201 }
202 let split_column = self.split_column(&options);
203 let row_stream = if options.backfill_as_even_splits
204 && is_supported_even_split_data_type(&split_column.data_type)
205 {
206 tracing::info!(?self.schema_table_name, ?self.rw_schema, ?self.pk_indices, ?split_column, "Get parallel cdc table snapshot even splits.");
208 self.as_even_splits(options)
209 } else {
210 tracing::info!(?self.schema_table_name, ?self.rw_schema, ?self.pk_indices, ?split_column, "Get parallel cdc table snapshot uneven splits.");
211 self.as_uneven_splits(options)
212 };
213 pin_mut!(row_stream);
214 #[for_await]
215 for row in row_stream {
216 let row = row?;
217 yield row;
218 }
219 }
220
221 fn split_snapshot_read(
222 &self,
223 table_name: SchemaTableName,
224 left: OwnedRow,
225 right: OwnedRow,
226 split_columns: Vec<Field>,
227 ) -> BoxStream<'_, ConnectorResult<OwnedRow>> {
228 assert_eq!(table_name, self.schema_table_name);
229 self.split_snapshot_read_inner(table_name, left, right, split_columns)
230 }
231}
232
233impl PostgresExternalTableReader {
234 pub async fn new(
235 config: ExternalTableConfig,
236 rw_schema: Schema,
237 pk_indices: Vec<usize>,
238 schema_table_name: SchemaTableName,
239 table_id: u32,
240 ) -> ConnectorResult<Self> {
241 tracing::info!(
242 ?rw_schema,
243 ?pk_indices,
244 "create postgres external table reader"
245 );
246 let application_name = format!(
248 "risingwave-postgres-source-reader-{}-{}.{}",
249 table_id, schema_table_name.schema_name, schema_table_name.table_name
250 );
251 let client = create_pg_client(
252 &config.pg_connection_config()?,
253 None,
254 Some(&application_name),
255 )
256 .await?;
257
258 let composite_columns = client
266 .query(
267 "SELECT a.attname, (t.typcategory = 'A') AS is_array \
268 FROM pg_attribute a \
269 JOIN pg_class c ON a.attrelid = c.oid \
270 JOIN pg_namespace n ON c.relnamespace = n.oid \
271 JOIN pg_type t ON a.atttypid = t.oid \
272 LEFT JOIN pg_type t_elem ON t.typelem = t_elem.oid \
273 WHERE n.nspname = $1 \
274 AND c.relname = $2 \
275 AND a.attnum > 0 \
276 AND NOT a.attisdropped \
277 AND (t.typtype = 'c' \
278 OR (t.typcategory = 'A' AND t_elem.typtype = 'c'))",
279 &[
280 &schema_table_name.schema_name,
281 &schema_table_name.table_name,
282 ],
283 )
284 .await
285 .map(|rows| {
286 rows.into_iter()
287 .map(|row| (row.get::<_, String>(0), row.get::<_, bool>(1)))
288 .collect::<std::collections::HashMap<_, _>>()
289 })
290 .unwrap_or_else(|err| {
291 tracing::warn!(
292 error = %err.as_report(),
293 schema = %schema_table_name.schema_name,
294 table = %schema_table_name.table_name,
295 "failed to discover postgres composite columns; falling back to no text cast"
296 );
297 std::collections::HashMap::new()
298 });
299
300 let field_names = rw_schema
301 .fields
302 .iter()
303 .map(|f| {
304 let quoted = Self::quote_column(&f.name);
305 match composite_columns.get(&f.name) {
306 Some(false) if matches!(f.data_type, DataType::Varchar) => {
307 format!("{quoted}::text AS {quoted}")
308 }
309 Some(true) if matches!(f.data_type, DataType::List(_)) => {
310 format!(
311 "CASE WHEN {quoted} IS NULL THEN NULL \
312 ELSE ARRAY(SELECT t::text FROM unnest({quoted}) t)::text[] \
313 END AS {quoted}"
314 )
315 }
316 _ => quoted,
317 }
318 })
319 .join(",");
320
321 Ok(Self {
322 rw_schema,
323 field_names,
324 pk_indices,
325 client: tokio::sync::Mutex::new(client),
326 schema_table_name,
327 })
328 }
329
330 pub fn get_normalized_table_name(table_name: &SchemaTableName) -> String {
331 format!(
332 "\"{}\".\"{}\"",
333 table_name.schema_name, table_name.table_name
334 )
335 }
336
337 pub fn get_cdc_offset_parser() -> CdcOffsetParseFunc {
338 Box::new(move |offset| {
339 Ok(CdcOffset::Postgres(PostgresOffset::parse_debezium_offset(
340 offset,
341 )?))
342 })
343 }
344
345 #[try_stream(boxed, ok = OwnedRow, error = ConnectorError)]
346 async fn snapshot_read_inner(
347 &self,
348 table_name: SchemaTableName,
349 start_pk_row: Option<OwnedRow>,
350 primary_keys: Vec<String>,
351 scan_limit: u32,
352 ) {
353 let order_key = Self::get_order_key(&primary_keys);
354 let client = self.client.lock().await;
355 client.execute("set time zone '+00:00'", &[]).await?;
356
357 let stream = match start_pk_row {
358 Some(ref pk_row) => {
359 let prepared_scan_stmt = {
362 let primary_keys = self
363 .pk_indices
364 .iter()
365 .map(|i| self.rw_schema.fields[*i].name.clone())
366 .collect_vec();
367
368 let order_key = Self::get_order_key(&primary_keys);
369 let scan_sql = format!(
370 "SELECT {} FROM {} WHERE {} ORDER BY {} LIMIT {scan_limit}",
371 self.field_names,
372 Self::get_normalized_table_name(&table_name),
373 Self::filter_expression(&primary_keys),
374 order_key,
375 );
376 client.prepare(&scan_sql).await?
377 };
378
379 let params: Vec<Option<ScalarAdapter>> = pk_row
380 .iter()
381 .zip_eq_fast(prepared_scan_stmt.params())
382 .map(|(datum, ty)| {
383 datum
384 .map(|scalar| ScalarAdapter::from_scalar(scalar, ty))
385 .transpose()
386 })
387 .try_collect()?;
388
389 client.query_raw(&prepared_scan_stmt, ¶ms).await?
390 }
391 None => {
392 let sql = format!(
393 "SELECT {} FROM {} ORDER BY {} LIMIT {scan_limit}",
394 self.field_names,
395 Self::get_normalized_table_name(&table_name),
396 order_key,
397 );
398 let params: Vec<Option<ScalarAdapter>> = vec![];
399 client.query_raw(&sql, ¶ms).await?
400 }
401 };
402
403 let row_stream = stream.map(|row| {
404 let row = row?;
405 postgres_row_to_owned_row_with_strict_pk(row, &self.rw_schema, &self.pk_indices)
406 .map_err(ConnectorError::from)
407 });
408
409 pin_mut!(row_stream);
410 #[for_await]
411 for row in row_stream {
412 let row = row?;
413 yield row;
414 }
415 }
416
417 fn filter_expression(columns: &[String]) -> String {
419 let mut col_expr = String::new();
420 let mut arg_expr = String::new();
421 for (i, column) in columns.iter().enumerate() {
422 if i > 0 {
423 col_expr.push_str(", ");
424 arg_expr.push_str(", ");
425 }
426 col_expr.push_str(&Self::quote_column(column));
427 arg_expr.push_str(format!("${}", i + 1).as_str());
428 }
429 format!("({}) > ({})", col_expr, arg_expr)
430 }
431
432 fn split_filter_expression(
434 columns: &[String],
435 is_first_split: bool,
436 is_last_split: bool,
437 ) -> String {
438 let mut left_col_expr = String::new();
439 let mut left_arg_expr = String::new();
440 let mut right_col_expr = String::new();
441 let mut right_arg_expr = String::new();
442 let mut c = 1;
443 if !is_first_split {
444 for (i, column) in columns.iter().enumerate() {
445 if i > 0 {
446 left_col_expr.push_str(", ");
447 left_arg_expr.push_str(", ");
448 }
449 left_col_expr.push_str(&Self::quote_column(column));
450 left_arg_expr.push_str(format!("${}", c).as_str());
451 c += 1;
452 }
453 }
454 if !is_last_split {
455 for (i, column) in columns.iter().enumerate() {
456 if i > 0 {
457 right_col_expr.push_str(", ");
458 right_arg_expr.push_str(", ");
459 }
460 right_col_expr.push_str(&Self::quote_column(column));
461 right_arg_expr.push_str(format!("${}", c).as_str());
462 c += 1;
463 }
464 }
465 if is_first_split && is_last_split {
466 "1 = 1".to_owned()
467 } else if is_first_split {
468 format!("({}) < ({})", right_col_expr, right_arg_expr,)
469 } else if is_last_split {
470 format!("({}) >= ({})", left_col_expr, left_arg_expr,)
471 } else {
472 format!(
473 "({}) >= ({}) AND ({}) < ({})",
474 left_col_expr, left_arg_expr, right_col_expr, right_arg_expr,
475 )
476 }
477 }
478
479 fn get_order_key(primary_keys: &Vec<String>) -> String {
480 primary_keys
481 .iter()
482 .map(|col| Self::quote_column(col))
483 .join(",")
484 }
485
486 fn quote_column(column: &str) -> String {
487 format!("\"{}\"", column)
488 }
489
490 async fn min_and_max(
491 &self,
492 split_column: &Field,
493 ) -> ConnectorResult<Option<(ScalarImpl, ScalarImpl)>> {
494 let sql = format!(
495 "SELECT MIN({}), MAX({}) FROM {}",
496 split_column.name,
497 split_column.name,
498 Self::get_normalized_table_name(&self.schema_table_name),
499 );
500 let client = self.client.lock().await;
501 let rows = client.query(&sql, &[]).await?;
502 if rows.is_empty() {
503 Ok(None)
504 } else {
505 let row = &rows[0];
506 let min = postgres_cell_to_scalar_impl_strict(
507 row,
508 &split_column.data_type,
509 0,
510 &split_column.name,
511 )?;
512 let max = postgres_cell_to_scalar_impl_strict(
513 row,
514 &split_column.data_type,
515 1,
516 &split_column.name,
517 )?;
518 match (min, max) {
519 (Some(min), Some(max)) => Ok(Some((min, max))),
520 _ => Ok(None),
521 }
522 }
523 }
524
525 async fn next_split_right_bound_exclusive(
526 &self,
527 left_value: &ScalarImpl,
528 max_value: &ScalarImpl,
529 max_split_size: u64,
530 split_column: &Field,
531 ) -> ConnectorResult<Option<Datum>> {
532 let sql = format!(
533 "WITH t as (SELECT {} FROM {} WHERE {} >= $1 ORDER BY {} ASC LIMIT {}) SELECT CASE WHEN MAX({}) < $2 THEN MAX({}) ELSE NULL END FROM t",
534 Self::quote_column(&split_column.name),
535 Self::get_normalized_table_name(&self.schema_table_name),
536 Self::quote_column(&split_column.name),
537 Self::quote_column(&split_column.name),
538 max_split_size,
539 Self::quote_column(&split_column.name),
540 Self::quote_column(&split_column.name),
541 );
542 let client = self.client.lock().await;
543 let prepared_stmt = client.prepare(&sql).await?;
544 let params: Vec<Option<ScalarAdapter>> = vec![
545 Some(ScalarAdapter::from_scalar(
546 left_value.as_scalar_ref_impl(),
547 &prepared_stmt.params()[0],
548 )?),
549 Some(ScalarAdapter::from_scalar(
550 max_value.as_scalar_ref_impl(),
551 &prepared_stmt.params()[1],
552 )?),
553 ];
554 let stream = client.query_raw(&prepared_stmt, ¶ms).await?;
555 let datum_stream = stream.map(|row| {
556 let row = row?;
557 Ok::<_, ConnectorError>(postgres_cell_to_scalar_impl_strict(
558 &row,
559 &split_column.data_type,
560 0,
561 &split_column.name,
562 )?)
563 });
564 pin_mut!(datum_stream);
565 #[for_await]
566 for datum in datum_stream {
567 let right = datum?;
568 return Ok(Some(right.to_owned_datum()));
569 }
570 Ok(None)
571 }
572
573 async fn next_greater_bound(
574 &self,
575 start_offset: &ScalarImpl,
576 max_value: &ScalarImpl,
577 split_column: &Field,
578 ) -> ConnectorResult<Option<Datum>> {
579 let sql = format!(
580 "SELECT MIN({}) FROM {} WHERE {} > $1 AND {} <$2",
581 Self::quote_column(&split_column.name),
582 Self::get_normalized_table_name(&self.schema_table_name),
583 Self::quote_column(&split_column.name),
584 Self::quote_column(&split_column.name),
585 );
586 let client = self.client.lock().await;
587 let prepared_stmt = client.prepare(&sql).await?;
588 let params: Vec<Option<ScalarAdapter>> = vec![
589 Some(ScalarAdapter::from_scalar(
590 start_offset.as_scalar_ref_impl(),
591 &prepared_stmt.params()[0],
592 )?),
593 Some(ScalarAdapter::from_scalar(
594 max_value.as_scalar_ref_impl(),
595 &prepared_stmt.params()[1],
596 )?),
597 ];
598 let stream = client.query_raw(&prepared_stmt, ¶ms).await?;
599 let datum_stream = stream.map(|row| {
600 let row = row?;
601 Ok::<_, ConnectorError>(postgres_cell_to_scalar_impl_strict(
602 &row,
603 &split_column.data_type,
604 0,
605 &split_column.name,
606 )?)
607 });
608 pin_mut!(datum_stream);
609 #[for_await]
610 for datum in datum_stream {
611 let right = datum?;
612 return Ok(Some(right));
613 }
614 Ok(None)
615 }
616
617 #[try_stream(boxed, ok = OwnedRow, error = ConnectorError)]
618 async fn split_snapshot_read_inner(
619 &self,
620 table_name: SchemaTableName,
621 left: OwnedRow,
622 right: OwnedRow,
623 split_columns: Vec<Field>,
624 ) {
625 assert_eq!(
640 split_columns.len(),
641 1,
642 "multiple split columns is not supported yet"
643 );
644 assert_eq!(left.len(), 1, "multiple split columns is not supported yet");
645 assert_eq!(
646 right.len(),
647 1,
648 "multiple split columns is not supported yet"
649 );
650 let is_first_split = left[0].is_none();
651 let is_last_split = right[0].is_none();
652 let split_column_names = split_columns.iter().map(|c| c.name.clone()).collect_vec();
653 let client = self.client.lock().await;
654 client.execute("set time zone '+00:00'", &[]).await?;
655 let prepared_scan_stmt = {
658 let scan_sql = format!(
659 "SELECT {} FROM {} WHERE {}",
660 self.field_names,
661 Self::get_normalized_table_name(&table_name),
662 Self::split_filter_expression(&split_column_names, is_first_split, is_last_split),
663 );
664 client.prepare(&scan_sql).await?
665 };
666
667 let mut params: Vec<Option<ScalarAdapter>> = vec![];
668 if !is_first_split {
669 let left_params: Vec<Option<ScalarAdapter>> = left
670 .iter()
671 .zip_eq_fast(prepared_scan_stmt.params().iter().take(left.len()))
672 .map(|(datum, ty)| {
673 datum
674 .map(|scalar| ScalarAdapter::from_scalar(scalar, ty))
675 .transpose()
676 })
677 .try_collect()?;
678 params.extend(left_params);
679 }
680 if !is_last_split {
681 let right_params: Vec<Option<ScalarAdapter>> = right
682 .iter()
683 .zip_eq_fast(prepared_scan_stmt.params().iter().skip(params.len()))
684 .map(|(datum, ty)| {
685 datum
686 .map(|scalar| ScalarAdapter::from_scalar(scalar, ty))
687 .transpose()
688 })
689 .try_collect()?;
690 params.extend(right_params);
691 }
692
693 let stream = client.query_raw(&prepared_scan_stmt, ¶ms).await?;
694 let row_stream = stream.map(|row| {
695 let row = row?;
696 postgres_row_to_owned_row_with_strict_pk(row, &self.rw_schema, &self.pk_indices)
697 .map_err(ConnectorError::from)
698 });
699
700 pin_mut!(row_stream);
701 #[for_await]
702 for row in row_stream {
703 let row = row?;
704 yield row;
705 }
706 }
707
708 #[try_stream(boxed, ok = CdcTableSnapshotSplit, error = ConnectorError)]
709 async fn as_uneven_splits(&self, options: CdcTableSnapshotSplitOption) {
710 let split_column = self.split_column(&options);
711 let mut split_id = CDC_TABLE_SPLIT_ID_START;
712 let Some((min_value, max_value)) = self.min_and_max(&split_column).await? else {
713 let left_bound_row = OwnedRow::new(vec![None]);
714 let right_bound_row = OwnedRow::new(vec![None]);
715 let split = CdcTableSnapshotSplit {
716 split_id,
717 left_bound_inclusive: left_bound_row,
718 right_bound_exclusive: right_bound_row,
719 };
720 yield split;
721 return Ok(());
722 };
723 let mut next_left_bound_inclusive = min_value.clone();
725 loop {
726 let left_bound_inclusive: Datum = if next_left_bound_inclusive == min_value {
727 None
728 } else {
729 Some(next_left_bound_inclusive.clone())
730 };
731 let right_bound_exclusive;
732 let mut next_right = self
733 .next_split_right_bound_exclusive(
734 &next_left_bound_inclusive,
735 &max_value,
736 options.backfill_num_rows_per_split,
737 &split_column,
738 )
739 .await?;
740 if let Some(Some(ref inner)) = next_right
741 && *inner == next_left_bound_inclusive
742 {
743 next_right = self
744 .next_greater_bound(&next_left_bound_inclusive, &max_value, &split_column)
745 .await?;
746 }
747 if let Some(next_right) = next_right {
748 match next_right {
749 None => {
750 right_bound_exclusive = None;
752 }
753 Some(next_right) => {
754 next_left_bound_inclusive = next_right.clone();
755 right_bound_exclusive = Some(next_right);
756 }
757 }
758 } else {
759 right_bound_exclusive = None;
761 };
762 let is_completed = right_bound_exclusive.is_none();
763 if is_completed && left_bound_inclusive.is_none() {
764 assert_eq!(split_id, 1);
765 }
766 tracing::info!(
767 split_id,
768 ?left_bound_inclusive,
769 ?right_bound_exclusive,
770 "New CDC table snapshot split."
771 );
772 let left_bound_row = OwnedRow::new(vec![left_bound_inclusive]);
773 let right_bound_row = OwnedRow::new(vec![right_bound_exclusive]);
774 let split = CdcTableSnapshotSplit {
775 split_id,
776 left_bound_inclusive: left_bound_row,
777 right_bound_exclusive: right_bound_row,
778 };
779 try_increase_split_id(&mut split_id)?;
780 yield split;
781 if is_completed {
782 break;
783 }
784 }
785 }
786
787 #[try_stream(boxed, ok = CdcTableSnapshotSplit, error = ConnectorError)]
788 async fn as_even_splits(&self, options: CdcTableSnapshotSplitOption) {
789 let split_column = self.split_column(&options);
790 let mut split_id = 1;
791 let Some((min_value, max_value)) = self.min_and_max(&split_column).await? else {
792 let left_bound_row = OwnedRow::new(vec![None]);
793 let right_bound_row = OwnedRow::new(vec![None]);
794 let split = CdcTableSnapshotSplit {
795 split_id,
796 left_bound_inclusive: left_bound_row,
797 right_bound_exclusive: right_bound_row,
798 };
799 yield split;
800 return Ok(());
801 };
802 let min_value = min_value.as_integral();
803 let max_value = max_value.as_integral();
804 let saturated_split_max_size = options
805 .backfill_num_rows_per_split
806 .try_into()
807 .unwrap_or(i64::MAX);
808 let mut left = None;
809 let mut right = Some(min_value.saturating_add(saturated_split_max_size));
810 loop {
811 let mut is_completed = false;
812 if right.as_ref().map(|r| *r >= max_value).unwrap_or(true) {
813 right = None;
814 is_completed = true;
815 }
816 let split = CdcTableSnapshotSplit {
817 split_id,
818 left_bound_inclusive: OwnedRow::new(vec![
819 left.map(|l| to_int_scalar(l, &split_column.data_type)),
820 ]),
821 right_bound_exclusive: OwnedRow::new(vec![
822 right.map(|r| to_int_scalar(r, &split_column.data_type)),
823 ]),
824 };
825 try_increase_split_id(&mut split_id)?;
826 yield split;
827 if is_completed {
828 break;
829 }
830 left = right;
831 right = left.map(|l| l.saturating_add(saturated_split_max_size));
832 }
833 }
834
835 fn split_column(&self, options: &CdcTableSnapshotSplitOption) -> Field {
836 self.rw_schema.fields[self.pk_indices[options.backfill_split_pk_column_index as usize]]
837 .clone()
838 }
839}
840
841fn to_int_scalar(i: i64, data_type: &DataType) -> ScalarImpl {
842 match data_type {
843 DataType::Int16 => ScalarImpl::Int16(i.try_into().unwrap()),
844 DataType::Int32 => ScalarImpl::Int32(i.try_into().unwrap()),
845 DataType::Int64 => ScalarImpl::Int64(i),
846 _ => {
847 panic!("Can't convert int {} to ScalarImpl::{}", i, data_type)
848 }
849 }
850}
851
852fn try_increase_split_id(split_id: &mut i64) -> ConnectorResult<()> {
853 match split_id.checked_add(1) {
854 Some(s) => {
855 *split_id = s;
856 Ok(())
857 }
858 None => Err(anyhow::anyhow!("too many CDC snapshot splits").into()),
859 }
860}
861
862fn is_supported_even_split_data_type(data_type: &DataType) -> bool {
864 matches!(
865 data_type,
866 DataType::Int16 | DataType::Int32 | DataType::Int64
867 )
868}
869
870pub fn type_name_to_pg_type(ty_name: &str) -> Option<PgType> {
871 let ty_name_lower = ty_name.to_lowercase();
872 if let Some(base_type) = ty_name_lower.strip_prefix('_') {
874 match base_type {
875 "int2" => Some(PgType::INT2_ARRAY),
876 "int4" => Some(PgType::INT4_ARRAY),
877 "int8" => Some(PgType::INT8_ARRAY),
878 "bit" => Some(PgType::BIT_ARRAY),
879 "float4" => Some(PgType::FLOAT4_ARRAY),
880 "float8" => Some(PgType::FLOAT8_ARRAY),
881 "numeric" => Some(PgType::NUMERIC_ARRAY),
882 "bool" => Some(PgType::BOOL_ARRAY),
883 "xml" | "macaddr" | "macaddr8" | "cidr" | "inet" | "int4range" | "int8range"
884 | "numrange" | "tsrange" | "tstzrange" | "daterange" | "citext" => {
885 Some(PgType::VARCHAR_ARRAY)
886 }
887 "varchar" => Some(PgType::VARCHAR_ARRAY),
888 "text" => Some(PgType::TEXT_ARRAY),
889 "bytea" => Some(PgType::BYTEA_ARRAY),
890 "geometry" | "geography" => Some(PgType::BYTEA_ARRAY), "date" => Some(PgType::DATE_ARRAY),
892 "time" => Some(PgType::TIME_ARRAY),
893 "timetz" => Some(PgType::TIMETZ_ARRAY),
894 "timestamp" => Some(PgType::TIMESTAMP_ARRAY),
895 "timestamptz" => Some(PgType::TIMESTAMPTZ_ARRAY),
896 "interval" => Some(PgType::INTERVAL_ARRAY),
897 "json" => Some(PgType::JSON_ARRAY),
898 "jsonb" => Some(PgType::JSONB_ARRAY),
899 "uuid" => Some(PgType::UUID_ARRAY),
900 "point" => Some(PgType::POINT_ARRAY),
901 "oid" => Some(PgType::OID_ARRAY),
902 "money" => Some(PgType::MONEY_ARRAY),
903 _ => None,
904 }
905 } else {
906 match ty_name_lower.as_str() {
908 "int2" => Some(PgType::INT2),
909 "bit" => Some(PgType::BIT),
910 "int" | "int4" => Some(PgType::INT4),
911 "int8" => Some(PgType::INT8),
912 "float4" => Some(PgType::FLOAT4),
913 "float8" => Some(PgType::FLOAT8),
914 "numeric" => Some(PgType::NUMERIC),
915 "money" => Some(PgType::MONEY),
916 "boolean" | "bool" => Some(PgType::BOOL),
917 "inet" | "xml" | "varchar" | "character varying" | "int4range" | "int8range"
918 | "numrange" | "tsrange" | "tstzrange" | "daterange" | "macaddr" | "macaddr8"
919 | "cidr" => Some(PgType::VARCHAR),
920 "char" | "character" | "bpchar" => Some(PgType::BPCHAR),
921 "citext" | "text" => Some(PgType::TEXT),
922 "bytea" => Some(PgType::BYTEA),
923 "geometry" | "geography" => Some(PgType::BYTEA), "date" => Some(PgType::DATE),
925 "time" => Some(PgType::TIME),
926 "timetz" => Some(PgType::TIMETZ),
927 "timestamp" => Some(PgType::TIMESTAMP),
928 "timestamptz" => Some(PgType::TIMESTAMPTZ),
929 "interval" => Some(PgType::INTERVAL),
930 "json" => Some(PgType::JSON),
931 "jsonb" => Some(PgType::JSONB),
932 "uuid" => Some(PgType::UUID),
933 "point" => Some(PgType::POINT),
934 "oid" => Some(PgType::OID),
935 _ => None,
936 }
937 }
938}
939
940pub fn pg_type_to_rw_type(pg_type: &PgType) -> ConnectorResult<DataType> {
944 let data_type = match *pg_type {
945 PgType::BOOL => DataType::Boolean,
946 PgType::BIT => DataType::Boolean,
947 PgType::INT2 => DataType::Int16,
948 PgType::INT4 => DataType::Int32,
949 PgType::INT8 => DataType::Int64,
950 PgType::FLOAT4 => DataType::Float32,
951 PgType::FLOAT8 => DataType::Float64,
952 PgType::NUMERIC | PgType::MONEY => DataType::Decimal,
953 PgType::DATE => DataType::Date,
954 PgType::TIME => DataType::Time,
955 PgType::TIMETZ => DataType::Time,
956 PgType::POINT => postgres_point_type(),
957 PgType::TIMESTAMP => DataType::Timestamp,
958 PgType::TIMESTAMPTZ => DataType::Timestamptz,
959 PgType::INTERVAL => DataType::Interval,
960 PgType::VARCHAR | PgType::TEXT | PgType::BPCHAR | PgType::UUID => DataType::Varchar,
961 PgType::BYTEA => DataType::Bytea,
962 PgType::JSON | PgType::JSONB => DataType::Jsonb,
963 PgType::BOOL_ARRAY => DataType::Boolean.list(),
965 PgType::BIT_ARRAY => DataType::Boolean.list(),
966 PgType::INT2_ARRAY => DataType::Int16.list(),
967 PgType::INT4_ARRAY => DataType::Int32.list(),
968 PgType::INT8_ARRAY => DataType::Int64.list(),
969 PgType::FLOAT4_ARRAY => DataType::Float32.list(),
970 PgType::FLOAT8_ARRAY => DataType::Float64.list(),
971 PgType::NUMERIC_ARRAY => DataType::Decimal.list(),
972 PgType::VARCHAR_ARRAY => DataType::Varchar.list(),
973 PgType::TEXT_ARRAY => DataType::Varchar.list(),
974 PgType::BYTEA_ARRAY => DataType::Bytea.list(),
975 PgType::DATE_ARRAY => DataType::Date.list(),
976 PgType::TIME_ARRAY => DataType::Time.list(),
977 PgType::TIMESTAMP_ARRAY => DataType::Timestamp.list(),
978 PgType::TIMESTAMPTZ_ARRAY => DataType::Timestamptz.list(),
979 PgType::INTERVAL_ARRAY => DataType::Interval.list(),
980 PgType::JSON_ARRAY => DataType::Jsonb.list(),
981 PgType::JSONB_ARRAY => DataType::Jsonb.list(),
982 PgType::UUID_ARRAY => DataType::Varchar.list(),
983 PgType::OID => DataType::Int64,
984 PgType::OID_ARRAY => DataType::Int64.list(),
985 PgType::MONEY_ARRAY => DataType::Decimal.list(),
986 PgType::POINT_ARRAY => {
989 return Err(anyhow::anyhow!("unsupported postgres type: {}", pg_type).into());
990 }
991 _ => {
992 return Err(anyhow::anyhow!("unsupported postgres type: {}", pg_type).into());
993 }
994 };
995 Ok(data_type)
996}
997
998#[cfg(test)]
999mod tests {
1000 use std::cmp::Ordering;
1001 use std::collections::HashMap;
1002
1003 use futures::pin_mut;
1004 use futures_async_stream::for_await;
1005 use maplit::{convert_args, hashmap};
1006 use risingwave_common::catalog::{ColumnDesc, ColumnId, Field, Schema};
1007 use risingwave_common::row::OwnedRow;
1008 use risingwave_common::types::{DataType, ScalarImpl};
1009
1010 use crate::connector_common::PostgresExternalTable;
1011 use crate::source::cdc::external::postgres::{PostgresExternalTableReader, PostgresOffset};
1012 use crate::source::cdc::external::{ExternalTableConfig, ExternalTableReader, SchemaTableName};
1013
1014 #[ignore]
1015 #[tokio::test]
1016 async fn test_postgres_schema() {
1017 let config = ExternalTableConfig {
1018 connector: "postgres-cdc".to_owned(),
1019 host: "localhost".to_owned(),
1020 port: "8432".to_owned(),
1021 username: "myuser".to_owned(),
1022 password: "123456".to_owned(),
1023 database: "mydb".to_owned(),
1024 schema: "public".to_owned(),
1025 table: "mytest".to_owned(),
1026 ssl_mode: Default::default(),
1027 ssl_root_cert: None,
1028 encrypt: "false".to_owned(),
1029 };
1030
1031 let table = PostgresExternalTable::connect(
1032 &config.pg_connection_config().unwrap(),
1033 &config.schema,
1034 &config.table,
1035 false,
1036 Some("SELECT"),
1037 )
1038 .await
1039 .unwrap();
1040
1041 println!("columns: {:?}", table.column_descs());
1042 println!("primary keys: {:?}", table.pk_names());
1043 }
1044
1045 #[test]
1046 fn test_postgres_offset() {
1047 let off1 = PostgresOffset {
1048 txid: 4,
1049 lsn: 2,
1050 ..Default::default()
1051 };
1052 let off2 = PostgresOffset {
1053 txid: 1,
1054 lsn: 3,
1055 ..Default::default()
1056 };
1057 let off3 = PostgresOffset {
1058 txid: 5,
1059 lsn: 1,
1060 ..Default::default()
1061 };
1062
1063 assert!(off1 < off2);
1064 assert!(off3 < off1);
1065 assert!(off2 > off3);
1066 }
1067
1068 #[test]
1069 fn test_postgres_offset_partial_ord_with_lsn_commit() {
1070 let off1 = PostgresOffset {
1072 txid: 1,
1073 lsn: 100,
1074 lsn_commit: Some(200),
1075 lsn_proc: Some(150),
1076 };
1077 let off2 = PostgresOffset {
1078 txid: 2,
1079 lsn: 300,
1080 lsn_commit: Some(250),
1081 lsn_proc: Some(200),
1082 };
1083
1084 assert!(off1 < off2);
1086
1087 let off3 = PostgresOffset {
1089 txid: 3,
1090 lsn: 500,
1091 lsn_commit: Some(200), lsn_proc: Some(160), };
1094
1095 assert!(off1 < off3);
1097
1098 let off4 = PostgresOffset {
1100 txid: 4,
1101 lsn: 400,
1102 lsn_commit: Some(100), lsn_proc: None, };
1105
1106 assert!(off1 < off4);
1108
1109 let off5 = PostgresOffset {
1111 txid: 5,
1112 lsn: 50, lsn_commit: None, lsn_proc: Some(300), };
1116
1117 assert!(off5 < off1);
1119
1120 let off6 = PostgresOffset {
1122 txid: 6,
1123 lsn: 600,
1124 lsn_commit: Some(500),
1125 lsn_proc: Some(300),
1126 };
1127 let off7 = PostgresOffset {
1128 txid: 7,
1129 lsn: 700,
1130 lsn_commit: Some(500), lsn_proc: Some(400), };
1133
1134 assert!(off6 < off7);
1136
1137 let off8 = PostgresOffset {
1139 txid: 8,
1140 lsn: 800,
1141 lsn_commit: Some(500), lsn_proc: Some(200), };
1144
1145 assert!(off8 < off6);
1146 assert!(off8 < off7);
1147
1148 let off9 = PostgresOffset {
1150 txid: 9,
1151 lsn: 900,
1152 lsn_commit: Some(500), lsn_proc: Some(300), };
1155
1156 assert_eq!(off6.partial_cmp(&off9), Some(Ordering::Equal));
1158 }
1159
1160 #[test]
1161 fn test_debezium_offset_parsing() {
1162 let debezium_offset_with_fields = r#"{
1164 "sourcePartition": {"server": "RW_CDC_1004"},
1165 "sourceOffset": {
1166 "last_snapshot_record": false,
1167 "lsn": 29973552,
1168 "txId": 1046,
1169 "ts_usec": 1670826189008456,
1170 "snapshot": true,
1171 "lsn_commit": 29973600,
1172 "lsn_proc": 29973580
1173 },
1174 "isHeartbeat": false
1175 }"#;
1176
1177 let offset = PostgresOffset::parse_debezium_offset(debezium_offset_with_fields).unwrap();
1178 assert_eq!(offset.txid, 1046);
1179 assert_eq!(offset.lsn, 29973552);
1180 assert_eq!(offset.lsn_commit, Some(29973600));
1181 assert_eq!(offset.lsn_proc, Some(29973580));
1182
1183 let debezium_offset_missing_fields = r#"{
1185 "sourcePartition": {"server": "RW_CDC_1004"},
1186 "sourceOffset": {
1187 "last_snapshot_record": false,
1188 "lsn": 29973552,
1189 "txId": 1046,
1190 "ts_usec": 1670826189008456,
1191 "snapshot": true
1192 },
1193 "isHeartbeat": false
1194 }"#;
1195
1196 let result = PostgresOffset::parse_debezium_offset(debezium_offset_missing_fields);
1197 assert!(result.is_err());
1198 let error_msg = result.unwrap_err().to_string();
1199 assert!(error_msg.contains("invalid postgres lsn_proc"));
1200 }
1201
1202 #[test]
1203 fn test_filter_expression() {
1204 let cols = vec!["v1".to_owned()];
1205 let expr = PostgresExternalTableReader::filter_expression(&cols);
1206 assert_eq!(expr, "(\"v1\") > ($1)");
1207
1208 let cols = vec!["v1".to_owned(), "v2".to_owned()];
1209 let expr = PostgresExternalTableReader::filter_expression(&cols);
1210 assert_eq!(expr, "(\"v1\", \"v2\") > ($1, $2)");
1211
1212 let cols = vec!["v1".to_owned(), "v2".to_owned(), "v3".to_owned()];
1213 let expr = PostgresExternalTableReader::filter_expression(&cols);
1214 assert_eq!(expr, "(\"v1\", \"v2\", \"v3\") > ($1, $2, $3)");
1215 }
1216
1217 #[test]
1218 fn test_split_filter_expression() {
1219 let cols = vec!["v1".to_owned()];
1220 let expr = PostgresExternalTableReader::split_filter_expression(&cols, true, true);
1221 assert_eq!(expr, "1 = 1");
1222
1223 let expr = PostgresExternalTableReader::split_filter_expression(&cols, true, false);
1224 assert_eq!(expr, "(\"v1\") < ($1)");
1225
1226 let expr = PostgresExternalTableReader::split_filter_expression(&cols, false, true);
1227 assert_eq!(expr, "(\"v1\") >= ($1)");
1228
1229 let expr = PostgresExternalTableReader::split_filter_expression(&cols, false, false);
1230 assert_eq!(expr, "(\"v1\") >= ($1) AND (\"v1\") < ($2)");
1231 }
1232
1233 #[ignore]
1235 #[tokio::test]
1236 async fn test_pg_table_reader() {
1237 let columns = [
1238 ColumnDesc::named("v1", ColumnId::new(1), DataType::Int32),
1239 ColumnDesc::named("v2", ColumnId::new(2), DataType::Varchar),
1240 ColumnDesc::named("v3", ColumnId::new(3), DataType::Decimal),
1241 ColumnDesc::named("v4", ColumnId::new(4), DataType::Date),
1242 ];
1243 let rw_schema = Schema {
1244 fields: columns.iter().map(Field::from).collect(),
1245 };
1246
1247 let props: HashMap<String, String> = convert_args!(hashmap!(
1248 "hostname" => "localhost",
1249 "port" => "8432",
1250 "username" => "myuser",
1251 "password" => "123456",
1252 "database.name" => "mydb",
1253 "schema.name" => "public",
1254 "table.name" => "t1"));
1255
1256 let config =
1257 serde_json::from_value::<ExternalTableConfig>(serde_json::to_value(props).unwrap())
1258 .unwrap();
1259 let schema_table_name = SchemaTableName {
1260 schema_name: "public".to_owned(),
1261 table_name: "t1".to_owned(),
1262 };
1263 let reader = PostgresExternalTableReader::new(
1264 config,
1265 rw_schema,
1266 vec![0, 1],
1267 schema_table_name.clone(),
1268 233,
1269 )
1270 .await
1271 .unwrap();
1272
1273 let offset = reader.current_cdc_offset().await.unwrap();
1274 println!("CdcOffset: {:?}", offset);
1275
1276 let start_pk = OwnedRow::new(vec![Some(ScalarImpl::from(3)), Some(ScalarImpl::from("c"))]);
1277 let stream = reader.snapshot_read(
1278 schema_table_name,
1279 Some(start_pk),
1280 vec!["v1".to_owned(), "v2".to_owned()],
1281 1000,
1282 );
1283
1284 pin_mut!(stream);
1285 #[for_await]
1286 for row in stream {
1287 println!("OwnedRow: {:?}", row);
1288 }
1289 }
1290}