qail_pg/driver/
replication.rs

1//! Replication helpers.
2//!
3//! Current scope:
4//! - `IDENTIFY_SYSTEM`
5//! - `CREATE_REPLICATION_SLOT ... LOGICAL ...`
6//! - `DROP_REPLICATION_SLOT`
7//! - `START_REPLICATION SLOT ... LOGICAL ...`
8//! - CopyBoth stream message decode (`XLogData`, keepalive)
9//! - Standby status updates (`'r'`)
10
11use super::{
12    PgConnection, PgDriver, PgError, PgResult, PgRow, is_ignorable_session_message,
13    unexpected_backend_message,
14};
15use crate::protocol::{BackendMessage, PgEncoder};
16
17/// Output of `IDENTIFY_SYSTEM`.
18#[derive(Debug, Clone, PartialEq, Eq)]
19pub struct IdentifySystem {
20    /// Cluster system identifier.
21    pub system_id: String,
22    /// Current timeline ID.
23    pub timeline: u32,
24    /// Current WAL/LSN position as text (e.g. `0/16B6C50`).
25    pub xlog_pos: String,
26    /// Database name for logical replication sessions (if provided by server).
27    pub dbname: Option<String>,
28}
29
30/// Output of `CREATE_REPLICATION_SLOT ... LOGICAL ...`.
31#[derive(Debug, Clone, PartialEq, Eq)]
32pub struct ReplicationSlotInfo {
33    /// Created slot name.
34    pub slot_name: String,
35    /// Consistent point at which the slot became valid.
36    pub consistent_point: String,
37    /// Exported snapshot name (if any).
38    pub snapshot_name: Option<String>,
39    /// Output plugin used by this logical slot.
40    pub output_plugin: String,
41}
42
43/// Metadata returned when a logical replication stream starts.
44#[derive(Debug, Clone, PartialEq, Eq)]
45pub struct ReplicationStreamStart {
46    /// Copy format (0=text, 1=binary).
47    pub format: u8,
48    /// Per-column format codes.
49    pub column_formats: Vec<u8>,
50}
51
52/// WAL payload message from `CopyData('w' ...)`.
53#[derive(Debug, Clone, PartialEq, Eq)]
54pub struct ReplicationXLogData {
55    /// WAL start position of this payload.
56    pub wal_start: u64,
57    /// Current WAL end on server.
58    pub wal_end: u64,
59    /// Server clock in microseconds since PostgreSQL epoch (2000-01-01).
60    pub server_time_micros: i64,
61    /// Output-plugin payload bytes.
62    pub data: Vec<u8>,
63}
64
65/// Keepalive message from `CopyData('k' ...)`.
66#[derive(Debug, Clone, PartialEq, Eq)]
67pub struct ReplicationKeepalive {
68    /// Current WAL end on server.
69    pub wal_end: u64,
70    /// Server clock in microseconds since PostgreSQL epoch (2000-01-01).
71    pub server_time_micros: i64,
72    /// Whether server requests an immediate status reply.
73    pub reply_requested: bool,
74}
75
76/// Logical replication stream message decoded from `CopyData`.
77#[derive(Debug, Clone, PartialEq, Eq)]
78pub enum ReplicationStreamMessage {
79    /// XLog data frame (`'w'`).
80    XLogData(ReplicationXLogData),
81    /// Primary keepalive frame (`'k'`).
82    Keepalive(ReplicationKeepalive),
83    /// Unknown CopyData sub-message tag.
84    Raw { tag: u8, payload: Vec<u8> },
85}
86
87/// Plugin options used in `START_REPLICATION ... LOGICAL ... (k 'v', ...)`.
88#[derive(Debug, Clone, PartialEq, Eq)]
89pub struct ReplicationOption {
90    /// Option key (strict identifier).
91    pub key: String,
92    /// Option value (quoted SQL string in command text).
93    pub value: String,
94}
95
96const MAX_REPLICATION_OPTIONS: usize = 64;
97const MAX_REPLICATION_OPTION_VALUE_BYTES: usize = 16 * 1024;
98const MAX_REPLICATION_XLOGDATA_BYTES: usize = 16 * 1024 * 1024;
99
100fn validate_ident(kind: &str, ident: &str) -> PgResult<()> {
101    if ident.is_empty() {
102        return Err(PgError::Query(format!("{} must not be empty", kind)));
103    }
104    if ident.len() > 63 {
105        return Err(PgError::Query(format!(
106            "{} '{}' exceeds PostgreSQL identifier limit (63 bytes)",
107            kind, ident
108        )));
109    }
110    let mut chars = ident.chars();
111    match chars.next() {
112        Some(c) if c == '_' || c.is_ascii_alphabetic() => {}
113        _ => {
114            return Err(PgError::Query(format!(
115                "{} '{}' must start with [A-Za-z_]",
116                kind, ident
117            )));
118        }
119    }
120    if !chars.all(|c| c == '_' || c.is_ascii_alphanumeric()) {
121        return Err(PgError::Query(format!(
122            "{} '{}' contains unsupported characters (allowed: [A-Za-z0-9_])",
123            kind, ident
124        )));
125    }
126    Ok(())
127}
128
129fn sql_single_quote_literal(value: &str) -> PgResult<String> {
130    if value.contains('\0') {
131        return Err(PgError::Query(
132            "replication option value contains NUL byte".to_string(),
133        ));
134    }
135    Ok(value.replace('\'', "''"))
136}
137
138fn parse_lsn_text(lsn: &str) -> PgResult<u64> {
139    let mut parts = lsn.split('/');
140    let high = parts
141        .next()
142        .ok_or_else(|| PgError::Query(format!("Invalid LSN '{}'", lsn)))?;
143    let low = parts
144        .next()
145        .ok_or_else(|| PgError::Query(format!("Invalid LSN '{}'", lsn)))?;
146    if parts.next().is_some() {
147        return Err(PgError::Query(format!("Invalid LSN '{}'", lsn)));
148    }
149    let high = u32::from_str_radix(high, 16)
150        .map_err(|_| PgError::Query(format!("Invalid LSN '{}'", lsn)))?;
151    let low = u32::from_str_radix(low, 16)
152        .map_err(|_| PgError::Query(format!("Invalid LSN '{}'", lsn)))?;
153    Ok(((high as u64) << 32) | (low as u64))
154}
155
156#[cfg(test)]
157fn format_lsn(lsn: u64) -> String {
158    format!("{:X}/{:08X}", (lsn >> 32) as u32, lsn as u32)
159}
160
161fn required_text(row: &PgRow, idx: usize, field: &str) -> PgResult<String> {
162    row.get_string(idx).ok_or_else(|| {
163        PgError::Protocol(format!("Missing or invalid '{}' in replication row", field))
164    })
165}
166
167fn parse_identify_system_row(row: &PgRow) -> PgResult<IdentifySystem> {
168    let system_id = required_text(row, 0, "systemid")?;
169    let timeline = required_text(row, 1, "timeline")?
170        .parse::<u32>()
171        .map_err(|e| PgError::Protocol(format!("Invalid timeline value: {}", e)))?;
172    let xlog_pos = required_text(row, 2, "xlogpos")?;
173    let dbname = row.get_string(3).filter(|v| !v.is_empty());
174
175    Ok(IdentifySystem {
176        system_id,
177        timeline,
178        xlog_pos,
179        dbname,
180    })
181}
182
183fn parse_create_slot_row(row: &PgRow) -> PgResult<ReplicationSlotInfo> {
184    let slot_name = required_text(row, 0, "slot_name")?;
185    let consistent_point = required_text(row, 1, "consistent_point")?;
186    let snapshot_name = row.get_string(2).filter(|v| !v.is_empty());
187    let output_plugin = required_text(row, 3, "output_plugin")?;
188
189    Ok(ReplicationSlotInfo {
190        slot_name,
191        consistent_point,
192        snapshot_name,
193        output_plugin,
194    })
195}
196
197fn build_create_logical_replication_slot_sql(
198    slot_name: &str,
199    output_plugin: &str,
200    temporary: bool,
201    two_phase: bool,
202) -> PgResult<String> {
203    validate_ident("slot_name", slot_name)?;
204    validate_ident("output_plugin", output_plugin)?;
205
206    let mut sql = String::from("CREATE_REPLICATION_SLOT ");
207    sql.push_str(slot_name);
208    if temporary {
209        sql.push_str(" TEMPORARY");
210    }
211    sql.push_str(" LOGICAL ");
212    sql.push_str(output_plugin);
213    if two_phase {
214        sql.push_str(" TWO_PHASE");
215    }
216    Ok(sql)
217}
218
219fn build_drop_replication_slot_sql(slot_name: &str, wait: bool) -> PgResult<String> {
220    validate_ident("slot_name", slot_name)?;
221    let mut sql = String::from("DROP_REPLICATION_SLOT ");
222    sql.push_str(slot_name);
223    if wait {
224        sql.push_str(" WAIT");
225    }
226    Ok(sql)
227}
228
229fn build_start_logical_replication_sql(
230    slot_name: &str,
231    start_lsn: &str,
232    options: &[ReplicationOption],
233) -> PgResult<String> {
234    validate_ident("slot_name", slot_name)?;
235    let _ = parse_lsn_text(start_lsn)?;
236    if options.len() > MAX_REPLICATION_OPTIONS {
237        return Err(PgError::Query(format!(
238            "too many replication options: {} (max {})",
239            options.len(),
240            MAX_REPLICATION_OPTIONS
241        )));
242    }
243
244    let mut sql = format!("START_REPLICATION SLOT {} LOGICAL {}", slot_name, start_lsn);
245    if !options.is_empty() {
246        sql.push_str(" (");
247        for (idx, opt) in options.iter().enumerate() {
248            validate_ident("replication option key", &opt.key)?;
249            if opt.value.len() > MAX_REPLICATION_OPTION_VALUE_BYTES {
250                return Err(PgError::Query(format!(
251                    "replication option '{}' value too large: {} bytes (max {})",
252                    opt.key,
253                    opt.value.len(),
254                    MAX_REPLICATION_OPTION_VALUE_BYTES
255                )));
256            }
257            if idx > 0 {
258                sql.push_str(", ");
259            }
260            sql.push_str(&opt.key);
261            sql.push_str(" '");
262            sql.push_str(&sql_single_quote_literal(&opt.value)?);
263            sql.push('\'');
264        }
265        sql.push(')');
266    }
267    Ok(sql)
268}
269
270fn parse_copy_data_message(payload: &[u8]) -> PgResult<ReplicationStreamMessage> {
271    if payload.is_empty() {
272        return Err(PgError::Protocol(
273            "Replication CopyData payload is empty".to_string(),
274        ));
275    }
276    match payload[0] {
277        b'w' => {
278            if payload.len() < 25 {
279                return Err(PgError::Protocol(format!(
280                    "XLogData payload too short: {} bytes",
281                    payload.len()
282                )));
283            }
284            let wal_start = u64::from_be_bytes(
285                payload[1..9]
286                    .try_into()
287                    .map_err(|_| PgError::Protocol("Invalid wal_start bytes".to_string()))?,
288            );
289            let wal_end = u64::from_be_bytes(
290                payload[9..17]
291                    .try_into()
292                    .map_err(|_| PgError::Protocol("Invalid wal_end bytes".to_string()))?,
293            );
294            let server_time_micros = i64::from_be_bytes(
295                payload[17..25]
296                    .try_into()
297                    .map_err(|_| PgError::Protocol("Invalid server time bytes".to_string()))?,
298            );
299            if wal_end < wal_start {
300                return Err(PgError::Protocol(format!(
301                    "XLogData wal_end {} is behind wal_start {}",
302                    wal_end, wal_start
303                )));
304            }
305            let data_len = payload.len() - 25;
306            if data_len > MAX_REPLICATION_XLOGDATA_BYTES {
307                return Err(PgError::Protocol(format!(
308                    "XLogData payload too large: {} bytes (max {})",
309                    data_len, MAX_REPLICATION_XLOGDATA_BYTES
310                )));
311            }
312            Ok(ReplicationStreamMessage::XLogData(ReplicationXLogData {
313                wal_start,
314                wal_end,
315                server_time_micros,
316                data: payload[25..].to_vec(),
317            }))
318        }
319        b'k' => {
320            if payload.len() != 18 {
321                return Err(PgError::Protocol(format!(
322                    "Keepalive payload must be 18 bytes, got {}",
323                    payload.len()
324                )));
325            }
326            let wal_end =
327                u64::from_be_bytes(payload[1..9].try_into().map_err(|_| {
328                    PgError::Protocol("Invalid keepalive wal_end bytes".to_string())
329                })?);
330            let server_time_micros = i64::from_be_bytes(
331                payload[9..17]
332                    .try_into()
333                    .map_err(|_| PgError::Protocol("Invalid keepalive time bytes".to_string()))?,
334            );
335            let reply_requested = match payload[17] {
336                0 => false,
337                1 => true,
338                other => {
339                    return Err(PgError::Protocol(format!(
340                        "Keepalive reply_requested must be 0 or 1, got {}",
341                        other
342                    )));
343                }
344            };
345            Ok(ReplicationStreamMessage::Keepalive(ReplicationKeepalive {
346                wal_end,
347                server_time_micros,
348                reply_requested,
349            }))
350        }
351        tag => Err(PgError::Protocol(format!(
352            "Unsupported replication CopyData tag '{}'",
353            if tag.is_ascii_graphic() {
354                tag as char
355            } else {
356                '?'
357            }
358        ))),
359    }
360}
361
362fn postgres_epoch_micros_now() -> i64 {
363    const PG_UNIX_EPOCH_DIFF_SECS: i64 = 946_684_800;
364    let now = std::time::SystemTime::now()
365        .duration_since(std::time::UNIX_EPOCH)
366        .unwrap_or_default();
367    let unix_micros = (now.as_secs() as i64) * 1_000_000 + i64::from(now.subsec_micros());
368    unix_micros - PG_UNIX_EPOCH_DIFF_SECS * 1_000_000
369}
370
371fn build_standby_status_update_payload(
372    write_lsn: u64,
373    flush_lsn: u64,
374    apply_lsn: u64,
375    reply_requested: bool,
376) -> Vec<u8> {
377    let mut payload = Vec::with_capacity(1 + 8 + 8 + 8 + 8 + 1);
378    payload.push(b'r');
379    payload.extend_from_slice(&write_lsn.to_be_bytes());
380    payload.extend_from_slice(&flush_lsn.to_be_bytes());
381    payload.extend_from_slice(&apply_lsn.to_be_bytes());
382    payload.extend_from_slice(&postgres_epoch_micros_now().to_be_bytes());
383    payload.push(if reply_requested { 1 } else { 0 });
384    payload
385}
386
387#[inline]
388fn return_with_desync<T>(conn: &mut PgConnection, err: PgError) -> PgResult<T> {
389    if matches!(
390        err,
391        PgError::Protocol(_) | PgError::Connection(_) | PgError::Timeout(_)
392    ) {
393        conn.mark_io_desynced();
394    }
395    Err(err)
396}
397
398impl PgConnection {
399    #[inline]
400    fn ensure_replication_mode(&self, operation: &str) -> PgResult<()> {
401        if self.replication_mode_enabled {
402            return Ok(());
403        }
404        Err(PgError::Protocol(format!(
405            "{} requires connection startup param replication=database",
406            operation
407        )))
408    }
409
410    #[inline]
411    fn ensure_replication_control_idle(&self, operation: &str) -> PgResult<()> {
412        if !self.replication_stream_active {
413            return Ok(());
414        }
415        Err(PgError::Protocol(format!(
416            "{} cannot run while replication stream is active",
417            operation
418        )))
419    }
420
421    #[inline]
422    fn advance_replication_wal_end(&mut self, source: &str, wal_end: u64) -> PgResult<()> {
423        if let Some(prev_wal_end) = self.last_replication_wal_end
424            && wal_end < prev_wal_end
425        {
426            self.replication_stream_active = false;
427            self.last_replication_wal_end = None;
428            return return_with_desync(
429                self,
430                PgError::Protocol(format!(
431                    "Replication {} wal_end regressed: previous {}, current {}",
432                    source, prev_wal_end, wal_end
433                )),
434            );
435        }
436        self.last_replication_wal_end = Some(wal_end);
437        Ok(())
438    }
439
440    /// Run `IDENTIFY_SYSTEM` on a replication connection.
441    pub async fn identify_system(&mut self) -> PgResult<IdentifySystem> {
442        self.ensure_replication_mode("IDENTIFY_SYSTEM")?;
443        self.ensure_replication_control_idle("IDENTIFY_SYSTEM")?;
444        let rows = self.simple_query("IDENTIFY_SYSTEM").await?;
445        let row = rows
446            .first()
447            .ok_or_else(|| PgError::Protocol("IDENTIFY_SYSTEM returned no rows".to_string()))?;
448        parse_identify_system_row(row)
449    }
450
451    /// Create a logical replication slot.
452    ///
453    /// `slot_name` and `output_plugin` are strict SQL identifiers.
454    pub async fn create_logical_replication_slot(
455        &mut self,
456        slot_name: &str,
457        output_plugin: &str,
458        temporary: bool,
459        two_phase: bool,
460    ) -> PgResult<ReplicationSlotInfo> {
461        self.ensure_replication_mode("CREATE_REPLICATION_SLOT")?;
462        self.ensure_replication_control_idle("CREATE_REPLICATION_SLOT")?;
463        let sql = build_create_logical_replication_slot_sql(
464            slot_name,
465            output_plugin,
466            temporary,
467            two_phase,
468        )?;
469        let rows = self.simple_query(&sql).await?;
470        let row = rows.first().ok_or_else(|| {
471            PgError::Protocol("CREATE_REPLICATION_SLOT returned no rows".to_string())
472        })?;
473        parse_create_slot_row(row)
474    }
475
476    /// Drop a replication slot.
477    ///
478    /// `wait=true` uses `DROP_REPLICATION_SLOT <slot> WAIT`.
479    pub async fn drop_replication_slot(&mut self, slot_name: &str, wait: bool) -> PgResult<()> {
480        self.ensure_replication_mode("DROP_REPLICATION_SLOT")?;
481        self.ensure_replication_control_idle("DROP_REPLICATION_SLOT")?;
482        let sql = build_drop_replication_slot_sql(slot_name, wait)?;
483        self.execute_simple(&sql).await
484    }
485
486    /// Start logical replication in CopyBoth mode.
487    ///
488    /// Requires a connection started with `replication=database`.
489    pub async fn start_logical_replication(
490        &mut self,
491        slot_name: &str,
492        start_lsn: &str,
493        options: &[ReplicationOption],
494    ) -> PgResult<ReplicationStreamStart> {
495        self.ensure_replication_mode("START_REPLICATION")?;
496        if self.replication_stream_active {
497            return Err(PgError::Protocol(
498                "logical replication stream already active".to_string(),
499            ));
500        }
501        let sql = build_start_logical_replication_sql(slot_name, start_lsn, options)?;
502        let bytes = PgEncoder::try_encode_query_string(&sql)?;
503        self.send_bytes(&bytes).await?;
504
505        let mut startup_error: Option<PgError> = None;
506        loop {
507            let msg = self.recv().await?;
508            match msg {
509                BackendMessage::CopyBothResponse {
510                    format,
511                    column_formats,
512                } => {
513                    if let Some(err) = startup_error {
514                        return return_with_desync(self, err);
515                    }
516                    if format != 0 {
517                        return return_with_desync(
518                            self,
519                            PgError::Protocol(format!(
520                                "START_REPLICATION returned unsupported CopyBothResponse format {} (expected 0/text)",
521                                format
522                            )),
523                        );
524                    }
525                    if !column_formats.is_empty() {
526                        return return_with_desync(
527                            self,
528                            PgError::Protocol(format!(
529                                "START_REPLICATION returned unexpected CopyBothResponse column formats (expected none, got {})",
530                                column_formats.len()
531                            )),
532                        );
533                    }
534                    self.replication_stream_active = true;
535                    self.last_replication_wal_end = None;
536                    return Ok(ReplicationStreamStart {
537                        format,
538                        column_formats,
539                    });
540                }
541                BackendMessage::ReadyForQuery(_) => {
542                    return return_with_desync(
543                        self,
544                        startup_error.unwrap_or_else(|| {
545                            PgError::Protocol(
546                                "START_REPLICATION ended before CopyBothResponse".to_string(),
547                            )
548                        }),
549                    );
550                }
551                BackendMessage::ErrorResponse(err) => {
552                    if startup_error.is_none() {
553                        startup_error = Some(PgError::QueryServer(err.into()));
554                    }
555                }
556                msg if is_ignorable_session_message(&msg) => {}
557                other => {
558                    return return_with_desync(
559                        self,
560                        unexpected_backend_message("start replication", &other),
561                    );
562                }
563            }
564        }
565    }
566
567    /// Receive the next logical replication stream message.
568    ///
569    /// Uses a no-timeout read path so idle periods do not fail the stream.
570    pub async fn recv_replication_message(&mut self) -> PgResult<ReplicationStreamMessage> {
571        self.ensure_replication_mode("recv_replication_message")?;
572        if !self.replication_stream_active {
573            return Err(PgError::Protocol(
574                "replication stream is not active; call START_REPLICATION first".to_string(),
575            ));
576        }
577        loop {
578            let msg = self.recv_without_timeout().await?;
579            match msg {
580                BackendMessage::CopyData(payload) => match parse_copy_data_message(&payload) {
581                    Ok(ReplicationStreamMessage::XLogData(x)) => {
582                        self.advance_replication_wal_end("XLogData", x.wal_end)?;
583                        return Ok(ReplicationStreamMessage::XLogData(x));
584                    }
585                    Ok(ReplicationStreamMessage::Keepalive(k)) => {
586                        self.advance_replication_wal_end("keepalive", k.wal_end)?;
587                        return Ok(ReplicationStreamMessage::Keepalive(k));
588                    }
589                    Ok(parsed) => return Ok(parsed),
590                    Err(err) => {
591                        self.replication_stream_active = false;
592                        self.last_replication_wal_end = None;
593                        return return_with_desync(self, err);
594                    }
595                },
596                BackendMessage::ErrorResponse(err) => {
597                    self.replication_stream_active = false;
598                    self.last_replication_wal_end = None;
599                    self.mark_io_desynced();
600                    return Err(PgError::QueryServer(err.into()));
601                }
602                BackendMessage::CopyDone => {
603                    self.replication_stream_active = false;
604                    self.last_replication_wal_end = None;
605                    return return_with_desync(
606                        self,
607                        PgError::Protocol("Replication stream ended with CopyDone".to_string()),
608                    );
609                }
610                BackendMessage::ReadyForQuery(_) => {
611                    self.replication_stream_active = false;
612                    self.last_replication_wal_end = None;
613                    return return_with_desync(
614                        self,
615                        PgError::Protocol(
616                            "Replication stream ended with ReadyForQuery".to_string(),
617                        ),
618                    );
619                }
620                msg if is_ignorable_session_message(&msg) => {}
621                other => {
622                    self.replication_stream_active = false;
623                    self.last_replication_wal_end = None;
624                    return return_with_desync(
625                        self,
626                        unexpected_backend_message("replication stream", &other),
627                    );
628                }
629            }
630        }
631    }
632
633    /// Send a standby status update (`CopyData('r' ...)`) to the server.
634    pub async fn send_standby_status_update(
635        &mut self,
636        write_lsn: u64,
637        flush_lsn: u64,
638        apply_lsn: u64,
639        reply_requested: bool,
640    ) -> PgResult<()> {
641        self.ensure_replication_mode("send_standby_status_update")?;
642        if !self.replication_stream_active {
643            return Err(PgError::Protocol(
644                "replication stream is not active; call START_REPLICATION first".to_string(),
645            ));
646        }
647        if flush_lsn > write_lsn {
648            return Err(PgError::Protocol(format!(
649                "Invalid standby status update: flush_lsn {} exceeds write_lsn {}",
650                flush_lsn, write_lsn
651            )));
652        }
653        if apply_lsn > flush_lsn {
654            return Err(PgError::Protocol(format!(
655                "Invalid standby status update: apply_lsn {} exceeds flush_lsn {}",
656                apply_lsn, flush_lsn
657            )));
658        }
659        if let Some(last_wal_end) = self.last_replication_wal_end
660            && write_lsn > last_wal_end
661        {
662            return Err(PgError::Protocol(format!(
663                "Invalid standby status update: write_lsn {} exceeds last seen server wal_end {}",
664                write_lsn, last_wal_end
665            )));
666        }
667        let payload =
668            build_standby_status_update_payload(write_lsn, flush_lsn, apply_lsn, reply_requested);
669        self.send_copy_data(&payload).await
670    }
671}
672
673impl PgDriver {
674    /// Driver wrapper for [`PgConnection::identify_system`].
675    pub async fn identify_system(&mut self) -> PgResult<IdentifySystem> {
676        self.connection.identify_system().await
677    }
678
679    /// Driver wrapper for [`PgConnection::create_logical_replication_slot`].
680    pub async fn create_logical_replication_slot(
681        &mut self,
682        slot_name: &str,
683        output_plugin: &str,
684        temporary: bool,
685        two_phase: bool,
686    ) -> PgResult<ReplicationSlotInfo> {
687        self.connection
688            .create_logical_replication_slot(slot_name, output_plugin, temporary, two_phase)
689            .await
690    }
691
692    /// Driver wrapper for [`PgConnection::drop_replication_slot`].
693    pub async fn drop_replication_slot(&mut self, slot_name: &str, wait: bool) -> PgResult<()> {
694        self.connection.drop_replication_slot(slot_name, wait).await
695    }
696
697    /// Driver wrapper for [`PgConnection::start_logical_replication`].
698    pub async fn start_logical_replication(
699        &mut self,
700        slot_name: &str,
701        start_lsn: &str,
702        options: &[ReplicationOption],
703    ) -> PgResult<ReplicationStreamStart> {
704        self.connection
705            .start_logical_replication(slot_name, start_lsn, options)
706            .await
707    }
708
709    /// Driver wrapper for [`PgConnection::recv_replication_message`].
710    pub async fn recv_replication_message(&mut self) -> PgResult<ReplicationStreamMessage> {
711        self.connection.recv_replication_message().await
712    }
713
714    /// Driver wrapper for [`PgConnection::send_standby_status_update`].
715    pub async fn send_standby_status_update(
716        &mut self,
717        write_lsn: u64,
718        flush_lsn: u64,
719        apply_lsn: u64,
720        reply_requested: bool,
721    ) -> PgResult<()> {
722        self.connection
723            .send_standby_status_update(write_lsn, flush_lsn, apply_lsn, reply_requested)
724            .await
725    }
726}
727
728#[cfg(test)]
729mod tests {
730    use super::*;
731
732    #[cfg(unix)]
733    fn test_conn() -> PgConnection {
734        use crate::driver::connection::StatementCache;
735        use crate::driver::stream::PgStream;
736        use bytes::BytesMut;
737        use std::collections::{HashMap, VecDeque};
738        use std::num::NonZeroUsize;
739        use tokio::net::UnixStream;
740
741        let (unix_stream, _peer) = UnixStream::pair().expect("unix stream pair");
742        PgConnection {
743            stream: PgStream::Unix(unix_stream),
744            buffer: BytesMut::with_capacity(1024),
745            write_buf: BytesMut::with_capacity(1024),
746            sql_buf: BytesMut::with_capacity(256),
747            params_buf: Vec::new(),
748            prepared_statements: HashMap::new(),
749            stmt_cache: StatementCache::new(NonZeroUsize::new(2).expect("non-zero")),
750            column_info_cache: HashMap::new(),
751            process_id: 0,
752            cancel_key_bytes: Vec::new(),
753            requested_protocol_minor: PgConnection::default_protocol_minor(),
754            negotiated_protocol_minor: PgConnection::default_protocol_minor(),
755            notifications: VecDeque::new(),
756            replication_stream_active: false,
757            replication_mode_enabled: true,
758            last_replication_wal_end: None,
759            io_desynced: false,
760            pending_statement_closes: Vec::new(),
761            draining_statement_closes: false,
762        }
763    }
764
765    fn text_row(values: &[Option<&str>]) -> PgRow {
766        PgRow {
767            columns: values
768                .iter()
769                .map(|v| v.map(|s| s.as_bytes().to_vec()))
770                .collect(),
771            column_info: None,
772        }
773    }
774
775    #[cfg(unix)]
776    fn push_backend_frame(conn: &mut PgConnection, msg_type: u8, payload: &[u8]) {
777        conn.buffer.extend_from_slice(&[msg_type]);
778        conn.buffer
779            .extend_from_slice(&((payload.len() + 4) as u32).to_be_bytes());
780        conn.buffer.extend_from_slice(payload);
781    }
782
783    fn error_response_payload(code: &str, message: &str) -> Vec<u8> {
784        let mut payload = Vec::new();
785        payload.push(b'S');
786        payload.extend_from_slice(b"ERROR\0");
787        payload.push(b'C');
788        payload.extend_from_slice(code.as_bytes());
789        payload.push(0);
790        payload.push(b'M');
791        payload.extend_from_slice(message.as_bytes());
792        payload.push(0);
793        payload.push(0);
794        payload
795    }
796
797    #[cfg(unix)]
798    #[tokio::test]
799    async fn replication_wal_regression_marks_connection_desynced() {
800        let mut conn = test_conn();
801        conn.replication_stream_active = true;
802        conn.last_replication_wal_end = Some(20);
803
804        let err = conn
805            .advance_replication_wal_end("XLogData", 10)
806            .expect_err("wal regression must fail");
807
808        assert!(err.to_string().contains("wal_end regressed"));
809        assert!(!conn.replication_stream_active);
810        assert!(conn.last_replication_wal_end.is_none());
811        assert!(conn.is_io_desynced());
812    }
813
814    #[cfg(unix)]
815    #[tokio::test]
816    async fn replication_error_response_marks_connection_desynced() {
817        let mut conn = test_conn();
818        conn.replication_stream_active = true;
819        conn.last_replication_wal_end = Some(20);
820        let payload = error_response_payload("XX000", "replication failed");
821        push_backend_frame(&mut conn, b'E', &payload);
822
823        let err = conn
824            .recv_replication_message()
825            .await
826            .expect_err("server error must fail");
827
828        assert!(matches!(err, PgError::QueryServer(_)));
829        assert!(!conn.replication_stream_active);
830        assert!(conn.last_replication_wal_end.is_none());
831        assert!(conn.is_io_desynced());
832    }
833
834    #[test]
835    fn validate_ident_rejects_bad_names() {
836        assert!(validate_ident("slot_name", "").is_err());
837        assert!(validate_ident("slot_name", "9slot").is_err());
838        assert!(validate_ident("slot_name", "bad-name").is_err());
839        assert!(validate_ident("slot_name", "has space").is_err());
840    }
841
842    #[test]
843    fn validate_ident_accepts_safe_names() {
844        assert!(validate_ident("slot_name", "slot_a1").is_ok());
845        assert!(validate_ident("output_plugin", "pgoutput").is_ok());
846    }
847
848    #[test]
849    fn parse_and_format_lsn_roundtrip() {
850        let lsn = parse_lsn_text("16/B6C50").unwrap();
851        assert_eq!(format_lsn(lsn), "16/000B6C50");
852    }
853
854    #[test]
855    fn build_create_logical_replication_slot_sql_variants() {
856        let sql =
857            build_create_logical_replication_slot_sql("slot_main", "pgoutput", true, true).unwrap();
858        assert_eq!(
859            sql,
860            "CREATE_REPLICATION_SLOT slot_main TEMPORARY LOGICAL pgoutput TWO_PHASE"
861        );
862    }
863
864    #[test]
865    fn build_drop_replication_slot_sql_variants() {
866        let sql = build_drop_replication_slot_sql("slot_main", true).unwrap();
867        assert_eq!(sql, "DROP_REPLICATION_SLOT slot_main WAIT");
868    }
869
870    #[test]
871    fn build_start_logical_replication_sql_with_options() {
872        let sql = build_start_logical_replication_sql(
873            "slot_main",
874            "0/16B6C50",
875            &[
876                ReplicationOption {
877                    key: "proto_version".to_string(),
878                    value: "1".to_string(),
879                },
880                ReplicationOption {
881                    key: "publication_names".to_string(),
882                    value: "pub1,pub2".to_string(),
883                },
884            ],
885        )
886        .unwrap();
887        assert_eq!(
888            sql,
889            "START_REPLICATION SLOT slot_main LOGICAL 0/16B6C50 (proto_version '1', publication_names 'pub1,pub2')"
890        );
891    }
892
893    #[test]
894    fn build_start_logical_replication_sql_rejects_too_many_options() {
895        let options: Vec<ReplicationOption> = (0..=MAX_REPLICATION_OPTIONS)
896            .map(|i| ReplicationOption {
897                key: format!("opt{}", i),
898                value: "x".to_string(),
899            })
900            .collect();
901
902        let err =
903            build_start_logical_replication_sql("slot_main", "0/16B6C50", &options).unwrap_err();
904        assert!(err.to_string().contains("too many replication options"));
905    }
906
907    #[test]
908    fn build_start_logical_replication_sql_rejects_null_value() {
909        let options = vec![ReplicationOption {
910            key: "proto_version".to_string(),
911            value: "1\0oops".to_string(),
912        }];
913        let err =
914            build_start_logical_replication_sql("slot_main", "0/16B6C50", &options).unwrap_err();
915        assert!(err.to_string().contains("contains NUL byte"));
916    }
917
918    #[test]
919    fn build_start_logical_replication_sql_rejects_oversized_value() {
920        let options = vec![ReplicationOption {
921            key: "publication_names".to_string(),
922            value: "x".repeat(MAX_REPLICATION_OPTION_VALUE_BYTES + 1),
923        }];
924        let err =
925            build_start_logical_replication_sql("slot_main", "0/16B6C50", &options).unwrap_err();
926        assert!(err.to_string().contains("value too large"));
927    }
928
929    #[test]
930    fn parse_identify_system_row_happy_path() {
931        let row = text_row(&[
932            Some("7416469842679442267"),
933            Some("1"),
934            Some("0/16B6C50"),
935            Some("app"),
936        ]);
937        let parsed = parse_identify_system_row(&row).unwrap();
938        assert_eq!(parsed.system_id, "7416469842679442267");
939        assert_eq!(parsed.timeline, 1);
940        assert_eq!(parsed.xlog_pos, "0/16B6C50");
941        assert_eq!(parsed.dbname.as_deref(), Some("app"));
942    }
943
944    #[test]
945    fn parse_create_slot_row_happy_path() {
946        let row = text_row(&[
947            Some("slot_main"),
948            Some("0/16B6C88"),
949            Some("00000003-00000041-1"),
950            Some("pgoutput"),
951        ]);
952        let parsed = parse_create_slot_row(&row).unwrap();
953        assert_eq!(parsed.slot_name, "slot_main");
954        assert_eq!(parsed.consistent_point, "0/16B6C88");
955        assert_eq!(parsed.snapshot_name.as_deref(), Some("00000003-00000041-1"));
956        assert_eq!(parsed.output_plugin, "pgoutput");
957    }
958
959    #[test]
960    fn parse_copy_data_xlog_data() {
961        let mut payload = Vec::new();
962        payload.push(b'w');
963        payload.extend_from_slice(&0x10_u64.to_be_bytes());
964        payload.extend_from_slice(&0x20_u64.to_be_bytes());
965        payload.extend_from_slice(&123_i64.to_be_bytes());
966        payload.extend_from_slice(b"hello");
967
968        match parse_copy_data_message(&payload).unwrap() {
969            ReplicationStreamMessage::XLogData(x) => {
970                assert_eq!(x.wal_start, 0x10);
971                assert_eq!(x.wal_end, 0x20);
972                assert_eq!(x.server_time_micros, 123);
973                assert_eq!(x.data, b"hello");
974            }
975            _ => panic!("expected xlog data"),
976        }
977    }
978
979    #[test]
980    fn parse_copy_data_xlog_data_rejects_wal_end_before_start() {
981        let mut payload = Vec::new();
982        payload.push(b'w');
983        payload.extend_from_slice(&0x20_u64.to_be_bytes());
984        payload.extend_from_slice(&0x10_u64.to_be_bytes());
985        payload.extend_from_slice(&123_i64.to_be_bytes());
986        let err = parse_copy_data_message(&payload).unwrap_err();
987        assert!(err.to_string().contains("wal_end"));
988    }
989
990    #[test]
991    fn parse_copy_data_xlog_data_rejects_oversized_data() {
992        let mut payload = Vec::with_capacity(25 + MAX_REPLICATION_XLOGDATA_BYTES + 1);
993        payload.push(b'w');
994        payload.extend_from_slice(&0x10_u64.to_be_bytes());
995        payload.extend_from_slice(&0x20_u64.to_be_bytes());
996        payload.extend_from_slice(&123_i64.to_be_bytes());
997        payload.extend(std::iter::repeat_n(0u8, MAX_REPLICATION_XLOGDATA_BYTES + 1));
998        let err = parse_copy_data_message(&payload).unwrap_err();
999        assert!(err.to_string().contains("payload too large"));
1000    }
1001
1002    #[test]
1003    fn parse_copy_data_keepalive() {
1004        let mut payload = Vec::new();
1005        payload.push(b'k');
1006        payload.extend_from_slice(&0xAB_u64.to_be_bytes());
1007        payload.extend_from_slice(&456_i64.to_be_bytes());
1008        payload.push(1);
1009
1010        match parse_copy_data_message(&payload).unwrap() {
1011            ReplicationStreamMessage::Keepalive(k) => {
1012                assert_eq!(k.wal_end, 0xAB);
1013                assert_eq!(k.server_time_micros, 456);
1014                assert!(k.reply_requested);
1015            }
1016            _ => panic!("expected keepalive"),
1017        }
1018    }
1019
1020    #[test]
1021    fn parse_copy_data_keepalive_rejects_invalid_reply_requested() {
1022        let mut payload = Vec::new();
1023        payload.push(b'k');
1024        payload.extend_from_slice(&0xAB_u64.to_be_bytes());
1025        payload.extend_from_slice(&456_i64.to_be_bytes());
1026        payload.push(2);
1027        let err = parse_copy_data_message(&payload).unwrap_err();
1028        assert!(err.to_string().contains("reply_requested must be 0 or 1"));
1029    }
1030
1031    #[test]
1032    fn parse_copy_data_unknown_tag_rejected() {
1033        let payload = vec![b'x', 1, 2, 3];
1034        let err = parse_copy_data_message(&payload).unwrap_err();
1035        assert!(
1036            err.to_string()
1037                .contains("Unsupported replication CopyData tag")
1038        );
1039    }
1040
1041    #[test]
1042    fn build_standby_status_update_payload_layout() {
1043        let payload = build_standby_status_update_payload(1, 2, 3, true);
1044        assert_eq!(payload.len(), 34);
1045        assert_eq!(payload[0], b'r');
1046        assert_eq!(u64::from_be_bytes(payload[1..9].try_into().unwrap()), 1);
1047        assert_eq!(u64::from_be_bytes(payload[9..17].try_into().unwrap()), 2);
1048        assert_eq!(u64::from_be_bytes(payload[17..25].try_into().unwrap()), 3);
1049        assert_eq!(payload[33], 1);
1050    }
1051}
qail_pg/driver/replication.rs

qail_pg/driver/
replication.rs
