1use tokio_postgres::{Client, NoTls};
2
3use crate::{
4 CdcError, CdcResult, CdcStart, PgLsn, PostgresCdcConfig, SlotLifecycle, config::SourceSettings,
5};
6
7#[derive(Debug, Clone, PartialEq, Eq)]
8pub(crate) struct SlotInfo {
9 pub(crate) plugin: Option<String>,
10 pub(crate) database: Option<String>,
11 pub(crate) active: bool,
12 pub(crate) restart_lsn: Option<PgLsn>,
13 pub(crate) confirmed_flush_lsn: Option<PgLsn>,
14 pub(crate) wal_status: Option<String>,
15}
16
17#[derive(Debug, Clone, Copy, PartialEq, Eq)]
18pub struct CdcLag {
19 pub retained_wal_bytes: i64,
20 pub confirmed_lag_bytes: i64,
21 pub slot_active: bool,
22}
23
24pub(crate) async fn prepare_slot_and_start_lsn(settings: &SourceSettings) -> CdcResult<PgLsn> {
25 let client = connect_admin(&settings.postgres).await?;
26 if settings.validate_publication {
27 validate_publication(&client, &settings.publication).await?;
28 }
29 let mut slot = read_slot(&client, &settings.slot).await?;
30 if slot.is_none() && settings.slot_lifecycle.may_create() {
31 create_slot(&client, &settings.slot, settings.slot_lifecycle.temporary()).await?;
32 slot = read_slot(&client, &settings.slot).await?;
33 }
34 let slot = slot.ok_or_else(|| {
35 CdcError::Slot(format!(
36 "replication slot {:?} does not exist; choose CreateIfMissing/CreateOwned/Temporary or create it first",
37 settings.slot
38 ))
39 })?;
40 validate_slot(settings, &slot)?;
41 choose_start_lsn(settings, &slot)
42}
43
44pub(crate) async fn drop_slot(
45 postgres: &PostgresCdcConfig,
46 slot: &str,
47 lifecycle: SlotLifecycle,
48 force: bool,
49) -> CdcResult<()> {
50 if !force && !lifecycle.may_drop() {
51 return Err(CdcError::Slot(format!(
52 "refusing to drop slot {slot:?}; lifecycle is {lifecycle:?}, not CreateOwned/Temporary"
53 )));
54 }
55 let client = connect_admin(postgres).await?;
56 let Some(info) = read_slot(&client, slot).await? else {
57 return Ok(());
58 };
59 if info.active && !force {
60 return Err(CdcError::Slot(format!(
61 "refusing to drop active slot {slot:?}; stop the source first or pass force"
62 )));
63 }
64 if info.active && force {
65 client
66 .execute(
67 "SELECT pg_terminate_backend(active_pid) FROM pg_replication_slots WHERE slot_name = $1",
68 &[&slot],
69 )
70 .await?;
71 }
72 client
73 .execute("SELECT pg_drop_replication_slot($1)", &[&slot])
74 .await?;
75 Ok(())
76}
77
78pub(crate) async fn sample_lag(postgres: &PostgresCdcConfig, slot: &str) -> CdcResult<CdcLag> {
79 let client = connect_admin(postgres).await?;
80 let row = client
81 .query_opt(
82 "SELECT
83 COALESCE(pg_wal_lsn_diff(pg_current_wal_lsn(), restart_lsn), 0)::bigint,
84 COALESCE(pg_wal_lsn_diff(pg_current_wal_lsn(), confirmed_flush_lsn), 0)::bigint,
85 active
86 FROM pg_replication_slots
87 WHERE slot_name = $1",
88 &[&slot],
89 )
90 .await?
91 .ok_or_else(|| CdcError::Slot(format!("replication slot {slot:?} does not exist")))?;
92 Ok(CdcLag {
93 retained_wal_bytes: row.get(0),
94 confirmed_lag_bytes: row.get(1),
95 slot_active: row.get(2),
96 })
97}
98
99async fn connect_admin(config: &PostgresCdcConfig) -> CdcResult<Client> {
100 let mut pg = tokio_postgres::Config::new();
101 pg.host(&config.host)
102 .port(config.port)
103 .user(&config.user)
104 .dbname(&config.database)
105 .application_name(&config.application_name);
106 if !config.password.is_empty() {
107 pg.password(&config.password);
108 }
109 let (client, connection) = pg.connect(NoTls).await?;
110 tokio::spawn(async move {
111 let _ = connection.await;
112 });
113 Ok(client)
114}
115
116async fn validate_publication(client: &Client, publication: &str) -> CdcResult<()> {
117 let row = client
118 .query_opt(
119 "SELECT pubinsert, pubupdate, pubdelete, pubtruncate
120 FROM pg_publication
121 WHERE pubname = $1",
122 &[&publication],
123 )
124 .await?;
125 let Some(row) = row else {
126 return Err(CdcError::Slot(format!(
127 "publication {publication:?} does not exist"
128 )));
129 };
130 let insert: bool = row.get(0);
131 let update: bool = row.get(1);
132 let delete: bool = row.get(2);
133 let truncate: bool = row.get(3);
134 if !insert || !update || !delete {
135 return Err(CdcError::Slot(format!(
136 "publication {publication:?} must publish insert, update, and delete for datum-cdc MVP (got insert={insert}, update={update}, delete={delete}, truncate={truncate})"
137 )));
138 }
139 Ok(())
140}
141
142async fn read_slot(client: &Client, slot: &str) -> CdcResult<Option<SlotInfo>> {
143 let row = client
144 .query_opt(
145 "SELECT plugin, database, active, restart_lsn::text, confirmed_flush_lsn::text, wal_status
146 FROM pg_replication_slots
147 WHERE slot_name = $1",
148 &[&slot],
149 )
150 .await?;
151 let Some(row) = row else {
152 return Ok(None);
153 };
154 Ok(Some(SlotInfo {
155 plugin: row.get(0),
156 database: row.get(1),
157 active: row.get(2),
158 restart_lsn: parse_optional_lsn(row.get::<_, Option<String>>(3))?,
159 confirmed_flush_lsn: parse_optional_lsn(row.get::<_, Option<String>>(4))?,
160 wal_status: row.get(5),
161 }))
162}
163
164async fn create_slot(client: &Client, slot: &str, temporary: bool) -> CdcResult<()> {
165 client
166 .execute(
167 "SELECT slot_name
168 FROM pg_create_logical_replication_slot($1::name, 'pgoutput', $2::boolean, false)",
169 &[&slot, &temporary],
170 )
171 .await?;
172 Ok(())
173}
174
175fn validate_slot(settings: &SourceSettings, slot: &SlotInfo) -> CdcResult<()> {
176 if slot.plugin.as_deref() != Some("pgoutput") {
177 return Err(CdcError::Slot(format!(
178 "slot {:?} must use pgoutput plugin, got {:?}",
179 settings.slot, slot.plugin
180 )));
181 }
182 if slot.database.as_deref() != Some(settings.postgres.database.as_str()) {
183 return Err(CdcError::Slot(format!(
184 "slot {:?} belongs to database {:?}, not {:?}",
185 settings.slot, slot.database, settings.postgres.database
186 )));
187 }
188 if slot.active {
189 return Err(CdcError::Slot(format!(
190 "slot {:?} is already active; use a dedicated slot per source",
191 settings.slot
192 )));
193 }
194 if slot.wal_status.as_deref() == Some("lost") {
195 return Err(CdcError::Slot(format!(
196 "slot {:?} has wal_status=lost; create a new slot after a fresh snapshot",
197 settings.slot
198 )));
199 }
200 Ok(())
201}
202
203fn choose_start_lsn(settings: &SourceSettings, slot: &SlotInfo) -> CdcResult<PgLsn> {
204 let checkpoint = settings
205 .checkpoint_store
206 .as_ref()
207 .map(|store| store.load(&settings.slot))
208 .transpose()?
209 .flatten();
210 let start = match (&settings.start, checkpoint) {
211 (CdcStart::Lsn(lsn), _) => *lsn,
212 (CdcStart::SlotConfirmed, _) => slot.confirmed_flush_lsn.unwrap_or(PgLsn::ZERO),
213 (CdcStart::CheckpointOrSlot, Some(offset)) => {
214 if let Some(confirmed) = slot.confirmed_flush_lsn
215 && confirmed > offset.tx_end_lsn
216 {
217 return Err(CdcError::Slot(format!(
218 "slot confirmed_flush_lsn {confirmed} is ahead of durable checkpoint {}; possible data loss or competing consumer",
219 offset.tx_end_lsn
220 )));
221 }
222 offset.tx_end_lsn
223 }
224 (CdcStart::CheckpointOrSlot, None) => slot.confirmed_flush_lsn.unwrap_or(PgLsn::ZERO),
225 };
226 if let Some(restart) = slot.restart_lsn
227 && !start.is_zero()
228 && start < restart
229 {
230 return Err(CdcError::Slot(format!(
231 "requested start LSN {start} is older than slot restart_lsn {restart}; WAL has been recycled"
232 )));
233 }
234 Ok(start)
235}
236
237fn parse_optional_lsn(value: Option<String>) -> CdcResult<Option<PgLsn>> {
238 value.as_deref().map(PgLsn::parse).transpose()
239}