#![forbid(unsafe_code)]
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use datum::{Keep, Sink};
use datum_cdc::{
CdcOffset, CdcSource, ChangeEvent, ChangeOperation, FileCheckpointStore, PgLsn,
PostgresCdcConfig, SlotLifecycle,
};
use tokio_postgres::{Client, NoTls};
const DEFAULT_PG_URL: &str = "postgresql://datum_cdc@127.0.0.1:55433/datum_cdc";
fn pg_enabled() -> bool {
std::env::var("DATUM_CDC_TEST_PG").as_deref() == Ok("1")
}
fn skip_pg() {
eprintln!("skipping datum-cdc PostgreSQL integration test; set DATUM_CDC_TEST_PG=1 to run");
}
fn pg_url() -> String {
std::env::var("DATUM_CDC_TEST_URL").unwrap_or_else(|_| DEFAULT_PG_URL.to_owned())
}
fn unique_name(prefix: &str) -> String {
let nanos = SystemTime::now()
.duration_since(UNIX_EPOCH)
.expect("system time before unix epoch")
.as_nanos();
format!("{prefix}_{}_{}", std::process::id(), nanos)
}
fn ident(name: &str) -> String {
assert!(
name.bytes()
.all(|byte| byte.is_ascii_alphanumeric() || byte == b'_'),
"test identifiers are generated from safe ascii"
);
format!("\"{name}\"")
}
async fn connect(url: &str) -> Client {
let (client, connection) = tokio_postgres::connect(url, NoTls)
.await
.expect("connect test PostgreSQL");
tokio::spawn(async move {
let _ = connection.await;
});
client
}
async fn setup(client: &Client, table: &str, publication: &str, slot: &str) {
let table_ident = ident(table);
let publication_ident = ident(publication);
client
.execute(
"SELECT pg_drop_replication_slot($1)
WHERE EXISTS (
SELECT 1 FROM pg_replication_slots
WHERE slot_name = $1 AND active = false
)",
&[&slot],
)
.await
.expect("drop stale slot");
client
.batch_execute(&format!(
"
DROP PUBLICATION IF EXISTS {publication_ident};
DROP TABLE IF EXISTS public.{table_ident};
CREATE TABLE public.{table_ident} (
id bigint PRIMARY KEY,
run_id text NOT NULL,
value bigint NOT NULL,
op_seq bigint NOT NULL,
kind text NOT NULL,
commit_ns bigint NOT NULL,
updated_at timestamptz NOT NULL DEFAULT clock_timestamp()
);
ALTER TABLE public.{table_ident} REPLICA IDENTITY FULL;
CREATE PUBLICATION {publication_ident} FOR TABLE public.{table_ident}
WITH (publish = 'insert,update,delete,truncate');
"
))
.await
.expect("create CDC test objects");
}
async fn cleanup(client: &Client, table: &str, publication: &str, slot: &str) {
let table_ident = ident(table);
let publication_ident = ident(publication);
let _ = client
.execute(
"SELECT pg_drop_replication_slot($1)
WHERE EXISTS (
SELECT 1 FROM pg_replication_slots
WHERE slot_name = $1 AND active = false
)",
&[&slot],
)
.await;
let _ = client
.batch_execute(&format!(
"DROP PUBLICATION IF EXISTS {publication_ident}; DROP TABLE IF EXISTS public.{table_ident};"
))
.await;
}
async fn wait_slot_active(client: &Client, slot: &str) {
let deadline = tokio::time::Instant::now() + Duration::from_secs(10);
loop {
let active = client
.query_opt(
"SELECT active FROM pg_replication_slots WHERE slot_name = $1",
&[&slot],
)
.await
.expect("query slot active")
.map(|row| row.get::<_, bool>(0))
.unwrap_or(false);
if active {
return;
}
assert!(
tokio::time::Instant::now() < deadline,
"slot {slot} did not become active"
);
tokio::time::sleep(Duration::from_millis(50)).await;
}
}
async fn wait_confirmed(client: &Client, slot: &str, lsn: PgLsn) {
let deadline = tokio::time::Instant::now() + Duration::from_secs(10);
loop {
let confirmed = client
.query_opt(
"SELECT confirmed_flush_lsn::text FROM pg_replication_slots WHERE slot_name = $1",
&[&slot],
)
.await
.expect("query confirmed lsn")
.and_then(|row| row.get::<_, Option<String>>(0))
.map(|value| PgLsn::parse(&value).expect("parse confirmed lsn"))
.unwrap_or(PgLsn::ZERO);
if confirmed >= lsn {
return;
}
assert!(
tokio::time::Instant::now() < deadline,
"slot {slot} confirmed_flush_lsn {confirmed} did not reach {lsn}"
);
tokio::time::sleep(Duration::from_millis(50)).await;
}
}
fn build_source(
url: &str,
slot: &str,
publication: &str,
lifecycle: SlotLifecycle,
) -> datum_cdc::CdcResult<datum::Source<ChangeEvent, datum_cdc::CdcHandle>> {
CdcSource::postgres()
.connect(PostgresCdcConfig::from_url(url)?)
.slot(slot)
.publication(publication)
.slot_lifecycle(lifecycle)
.status_interval(Duration::from_millis(50))
.idle_wakeup_interval(Duration::from_millis(50))
.disable_reconnect()
.build()
}
fn row_text<'a>(event: &'a ChangeEvent, column: &str) -> &'a str {
let row = match event.op {
ChangeOperation::Insert | ChangeOperation::Update => event.after.as_ref(),
ChangeOperation::Delete => event.before.as_ref(),
ChangeOperation::Truncate => None,
}
.expect("row event has tuple data");
row.get_text(&event.relation, column)
.expect("column exists as text")
}
async fn write_op(client: &Client, table: &str, run_id: &str, id: i64, op_seq: i64, op: &str) {
let table_ident = ident(table);
let value = id * 1_000_000 + op_seq;
let commit_ns = SystemTime::now()
.duration_since(UNIX_EPOCH)
.expect("system time before unix epoch")
.as_nanos() as i64;
match op {
"c" => {
client
.execute(
&format!(
"INSERT INTO public.{table_ident} (id, run_id, value, op_seq, kind, commit_ns)
VALUES ($1, $2, $3, $4, 'insert', $5)"
),
&[&id, &run_id, &value, &op_seq, &commit_ns],
)
.await
.expect("insert test row");
}
"u" => {
client
.execute(
&format!(
"UPDATE public.{table_ident}
SET value = $2, op_seq = $3, kind = 'update', commit_ns = $4
WHERE id = $1"
),
&[&id, &value, &op_seq, &commit_ns],
)
.await
.expect("update test row");
}
"d" => {
client
.execute(
&format!("DELETE FROM public.{table_ident} WHERE id = $1"),
&[&id],
)
.await
.expect("delete test row");
}
other => panic!("unsupported op {other}"),
}
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn insert_update_delete_round_trip_and_per_key_order() {
if !pg_enabled() {
skip_pg();
return;
}
let url = pg_url();
let client = connect(&url).await;
let table = unique_name("datum_cdc_roundtrip");
let publication = unique_name("datum_cdc_pub");
let slot = unique_name("datum_cdc_slot");
setup(&client, &table, &publication, &slot).await;
let source = build_source(&url, &slot, &publication, SlotLifecycle::CreateOwned).unwrap();
let (handle, events_task) = {
let graph = source.take(3).to_mat(Sink::collect(), Keep::both);
let (handle, completion) = graph.run().expect("run CDC graph");
(
handle,
tokio::task::spawn_blocking(move || completion.wait()),
)
};
wait_slot_active(&client, &slot).await;
let run_id = unique_name("run");
write_op(&client, &table, &run_id, 1, 1, "c").await;
write_op(&client, &table, &run_id, 1, 2, "u").await;
write_op(&client, &table, &run_id, 1, 3, "d").await;
let events = events_task
.await
.expect("join completion wait")
.expect("collect events");
assert_eq!(events.len(), 3);
assert_eq!(
events.iter().map(|event| event.op).collect::<Vec<_>>(),
vec![
ChangeOperation::Insert,
ChangeOperation::Update,
ChangeOperation::Delete
]
);
assert_eq!(
events
.iter()
.map(|event| row_text(event, "id").parse::<i64>().unwrap())
.collect::<Vec<_>>(),
vec![1, 1, 1]
);
assert!(events.iter().all(|event| event.schema == "public"));
assert!(events.iter().all(|event| event.table == table));
let _ = handle.stop();
let _ = handle.force_drop_slot().await;
cleanup(&client, &table, &publication, &slot).await;
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn checkpoint_feedback_and_restart_resume_do_not_replay_confirmed_rows() {
if !pg_enabled() {
skip_pg();
return;
}
let url = pg_url();
let client = connect(&url).await;
let table = unique_name("datum_cdc_resume");
let publication = unique_name("datum_cdc_pub");
let slot = unique_name("datum_cdc_slot");
setup(&client, &table, &publication, &slot).await;
let checkpoint_dir = tempfile::tempdir().expect("checkpoint tempdir");
let run_id = unique_name("run");
let first_source = CdcSource::postgres()
.connect(PostgresCdcConfig::from_url(&url).unwrap())
.slot(&slot)
.publication(&publication)
.slot_lifecycle(SlotLifecycle::CreateOwned)
.checkpoint_store(FileCheckpointStore::new(checkpoint_dir.path()))
.status_interval(Duration::from_millis(50))
.idle_wakeup_interval(Duration::from_millis(50))
.disable_reconnect()
.build()
.unwrap();
let (first_handle, first_events_task) = {
let graph = first_source.take(1).to_mat(Sink::collect(), Keep::both);
let (handle, completion) = graph.run().expect("run first CDC graph");
(
handle,
tokio::task::spawn_blocking(move || completion.wait()),
)
};
wait_slot_active(&client, &slot).await;
write_op(&client, &table, &run_id, 1, 1, "c").await;
let first_events = first_events_task
.await
.expect("join first completion")
.expect("collect first event");
assert_eq!(row_text(&first_events[0], "id"), "1");
let first_lsn = first_events[0].lsn.tx_end_lsn;
first_handle
.checkpoint_handle()
.checkpoint(first_events[0].lsn.clone())
.expect("checkpoint first event");
wait_confirmed(&client, &slot, first_lsn).await;
let _ = first_handle.stop();
let second_source = CdcSource::postgres()
.connect(PostgresCdcConfig::from_url(&url).unwrap())
.slot(&slot)
.publication(&publication)
.slot_lifecycle(SlotLifecycle::Existing)
.checkpoint_store(FileCheckpointStore::new(checkpoint_dir.path()))
.status_interval(Duration::from_millis(50))
.idle_wakeup_interval(Duration::from_millis(50))
.disable_reconnect()
.build()
.unwrap();
let (second_handle, second_events_task) = {
let graph = second_source.take(1).to_mat(Sink::collect(), Keep::both);
let (handle, completion) = graph.run().expect("run second CDC graph");
(
handle,
tokio::task::spawn_blocking(move || completion.wait()),
)
};
wait_slot_active(&client, &slot).await;
write_op(&client, &table, &run_id, 2, 2, "c").await;
let second_events = second_events_task
.await
.expect("join second completion")
.expect("collect second event");
assert_eq!(row_text(&second_events[0], "id"), "2");
second_handle
.checkpoint_handle()
.checkpoint(second_events[0].lsn.clone())
.expect("checkpoint second event");
wait_confirmed(&client, &slot, second_events[0].lsn.tx_end_lsn).await;
let _ = second_handle.force_drop_slot().await;
cleanup(&client, &table, &publication, &slot).await;
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn lag_metric_reports_slot_pressure_fields() {
if !pg_enabled() {
skip_pg();
return;
}
let url = pg_url();
let client = connect(&url).await;
let table = unique_name("datum_cdc_lag");
let publication = unique_name("datum_cdc_pub");
let slot = unique_name("datum_cdc_slot");
setup(&client, &table, &publication, &slot).await;
let source = build_source(&url, &slot, &publication, SlotLifecycle::CreateOwned).unwrap();
let (handle, events_task) = {
let graph = source.take(1).to_mat(Sink::collect(), Keep::both);
let (handle, completion) = graph.run().expect("run CDC graph");
(
handle,
tokio::task::spawn_blocking(move || completion.wait()),
)
};
wait_slot_active(&client, &slot).await;
let run_id = unique_name("run");
write_op(&client, &table, &run_id, 1, 1, "c").await;
let events = events_task
.await
.expect("join completion wait")
.expect("collect event");
let lag = handle.lag().await.expect("sample lag");
assert!(lag.retained_wal_bytes >= 0);
assert!(lag.confirmed_lag_bytes >= 0);
handle
.checkpoint_handle()
.checkpoint(CdcOffset {
slot: events[0].lsn.slot.clone(),
tx_end_lsn: events[0].lsn.tx_end_lsn,
commit_lsn: events[0].lsn.commit_lsn,
xid: events[0].lsn.xid,
event_index: events[0].lsn.event_index,
event_count: events[0].lsn.event_count,
})
.expect("checkpoint event");
wait_confirmed(&client, &slot, events[0].lsn.tx_end_lsn).await;
let _ = handle.force_drop_slot().await;
cleanup(&client, &table, &publication, &slot).await;
}