use tokio_postgres::{Client, NoTls};
use crate::{
CdcError, CdcResult, CdcStart, PgLsn, PostgresCdcConfig, SlotLifecycle, config::SourceSettings,
};
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) struct SlotInfo {
pub(crate) plugin: Option<String>,
pub(crate) database: Option<String>,
pub(crate) active: bool,
pub(crate) restart_lsn: Option<PgLsn>,
pub(crate) confirmed_flush_lsn: Option<PgLsn>,
pub(crate) wal_status: Option<String>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct CdcLag {
pub retained_wal_bytes: i64,
pub confirmed_lag_bytes: i64,
pub slot_active: bool,
}
pub(crate) async fn prepare_slot_and_start_lsn(settings: &SourceSettings) -> CdcResult<PgLsn> {
let client = connect_admin(&settings.postgres).await?;
if settings.validate_publication {
validate_publication(&client, &settings.publication).await?;
}
let mut slot = read_slot(&client, &settings.slot).await?;
if slot.is_none() && settings.slot_lifecycle.may_create() {
create_slot(&client, &settings.slot, settings.slot_lifecycle.temporary()).await?;
slot = read_slot(&client, &settings.slot).await?;
}
let slot = slot.ok_or_else(|| {
CdcError::Slot(format!(
"replication slot {:?} does not exist; choose CreateIfMissing/CreateOwned/Temporary or create it first",
settings.slot
))
})?;
validate_slot(settings, &slot)?;
choose_start_lsn(settings, &slot)
}
pub(crate) async fn drop_slot(
postgres: &PostgresCdcConfig,
slot: &str,
lifecycle: SlotLifecycle,
force: bool,
) -> CdcResult<()> {
if !force && !lifecycle.may_drop() {
return Err(CdcError::Slot(format!(
"refusing to drop slot {slot:?}; lifecycle is {lifecycle:?}, not CreateOwned/Temporary"
)));
}
let client = connect_admin(postgres).await?;
let Some(info) = read_slot(&client, slot).await? else {
return Ok(());
};
if info.active && !force {
return Err(CdcError::Slot(format!(
"refusing to drop active slot {slot:?}; stop the source first or pass force"
)));
}
if info.active && force {
client
.execute(
"SELECT pg_terminate_backend(active_pid) FROM pg_replication_slots WHERE slot_name = $1",
&[&slot],
)
.await?;
}
client
.execute("SELECT pg_drop_replication_slot($1)", &[&slot])
.await?;
Ok(())
}
pub(crate) async fn sample_lag(postgres: &PostgresCdcConfig, slot: &str) -> CdcResult<CdcLag> {
let client = connect_admin(postgres).await?;
let row = client
.query_opt(
"SELECT
COALESCE(pg_wal_lsn_diff(pg_current_wal_lsn(), restart_lsn), 0)::bigint,
COALESCE(pg_wal_lsn_diff(pg_current_wal_lsn(), confirmed_flush_lsn), 0)::bigint,
active
FROM pg_replication_slots
WHERE slot_name = $1",
&[&slot],
)
.await?
.ok_or_else(|| CdcError::Slot(format!("replication slot {slot:?} does not exist")))?;
Ok(CdcLag {
retained_wal_bytes: row.get(0),
confirmed_lag_bytes: row.get(1),
slot_active: row.get(2),
})
}
async fn connect_admin(config: &PostgresCdcConfig) -> CdcResult<Client> {
let mut pg = tokio_postgres::Config::new();
pg.host(&config.host)
.port(config.port)
.user(&config.user)
.dbname(&config.database)
.application_name(&config.application_name);
if !config.password.is_empty() {
pg.password(&config.password);
}
let (client, connection) = pg.connect(NoTls).await?;
tokio::spawn(async move {
let _ = connection.await;
});
Ok(client)
}
async fn validate_publication(client: &Client, publication: &str) -> CdcResult<()> {
let row = client
.query_opt(
"SELECT pubinsert, pubupdate, pubdelete, pubtruncate
FROM pg_publication
WHERE pubname = $1",
&[&publication],
)
.await?;
let Some(row) = row else {
return Err(CdcError::Slot(format!(
"publication {publication:?} does not exist"
)));
};
let insert: bool = row.get(0);
let update: bool = row.get(1);
let delete: bool = row.get(2);
let truncate: bool = row.get(3);
if !insert || !update || !delete {
return Err(CdcError::Slot(format!(
"publication {publication:?} must publish insert, update, and delete for datum-cdc MVP (got insert={insert}, update={update}, delete={delete}, truncate={truncate})"
)));
}
Ok(())
}
async fn read_slot(client: &Client, slot: &str) -> CdcResult<Option<SlotInfo>> {
let row = client
.query_opt(
"SELECT plugin, database, active, restart_lsn::text, confirmed_flush_lsn::text, wal_status
FROM pg_replication_slots
WHERE slot_name = $1",
&[&slot],
)
.await?;
let Some(row) = row else {
return Ok(None);
};
Ok(Some(SlotInfo {
plugin: row.get(0),
database: row.get(1),
active: row.get(2),
restart_lsn: parse_optional_lsn(row.get::<_, Option<String>>(3))?,
confirmed_flush_lsn: parse_optional_lsn(row.get::<_, Option<String>>(4))?,
wal_status: row.get(5),
}))
}
async fn create_slot(client: &Client, slot: &str, temporary: bool) -> CdcResult<()> {
client
.execute(
"SELECT slot_name
FROM pg_create_logical_replication_slot($1::name, 'pgoutput', $2::boolean, false)",
&[&slot, &temporary],
)
.await?;
Ok(())
}
fn validate_slot(settings: &SourceSettings, slot: &SlotInfo) -> CdcResult<()> {
if slot.plugin.as_deref() != Some("pgoutput") {
return Err(CdcError::Slot(format!(
"slot {:?} must use pgoutput plugin, got {:?}",
settings.slot, slot.plugin
)));
}
if slot.database.as_deref() != Some(settings.postgres.database.as_str()) {
return Err(CdcError::Slot(format!(
"slot {:?} belongs to database {:?}, not {:?}",
settings.slot, slot.database, settings.postgres.database
)));
}
if slot.active {
return Err(CdcError::Slot(format!(
"slot {:?} is already active; use a dedicated slot per source",
settings.slot
)));
}
if slot.wal_status.as_deref() == Some("lost") {
return Err(CdcError::Slot(format!(
"slot {:?} has wal_status=lost; create a new slot after a fresh snapshot",
settings.slot
)));
}
Ok(())
}
fn choose_start_lsn(settings: &SourceSettings, slot: &SlotInfo) -> CdcResult<PgLsn> {
let checkpoint = settings
.checkpoint_store
.as_ref()
.map(|store| store.load(&settings.slot))
.transpose()?
.flatten();
let start = match (&settings.start, checkpoint) {
(CdcStart::Lsn(lsn), _) => *lsn,
(CdcStart::SlotConfirmed, _) => slot.confirmed_flush_lsn.unwrap_or(PgLsn::ZERO),
(CdcStart::CheckpointOrSlot, Some(offset)) => {
if let Some(confirmed) = slot.confirmed_flush_lsn
&& confirmed > offset.tx_end_lsn
{
return Err(CdcError::Slot(format!(
"slot confirmed_flush_lsn {confirmed} is ahead of durable checkpoint {}; possible data loss or competing consumer",
offset.tx_end_lsn
)));
}
offset.tx_end_lsn
}
(CdcStart::CheckpointOrSlot, None) => slot.confirmed_flush_lsn.unwrap_or(PgLsn::ZERO),
};
if let Some(restart) = slot.restart_lsn
&& !start.is_zero()
&& start < restart
{
return Err(CdcError::Slot(format!(
"requested start LSN {start} is older than slot restart_lsn {restart}; WAL has been recycled"
)));
}
Ok(start)
}
fn parse_optional_lsn(value: Option<String>) -> CdcResult<Option<PgLsn>> {
value.as_deref().map(PgLsn::parse).transpose()
}