elefant_client/postgres_client/replication/
client_ext.rs

1use super::messages::Lsn;
2use super::stream::ReplicationStream;
3use crate::pool::ConnectionFactory;
4use crate::postgres_client::PostgresClient;
5use crate::protocol::FrontendMessage;
6use crate::ElefantClientError;
7use std::borrow::Cow;
8
9/// Quotes an identifier for use in replication protocol commands using
10/// PostgreSQL's standard double-quoting rules: wrap in double quotes and
11/// escape any internal double quotes by doubling them.
12fn quote_identifier(name: &str) -> String {
13    let escaped = name.replace('"', "\"\"");
14    format!("\"{escaped}\"")
15}
16
17impl<F: ConnectionFactory> PostgresClient<F> {
18    pub async fn create_replication_slot(
19        &mut self,
20        slot_name: &str,
21        output_plugin: &str,
22    ) -> Result<(String, Lsn), ElefantClientError> {
23        let slot_quoted = quote_identifier(slot_name);
24        let plugin_quoted = quote_identifier(output_plugin);
25        let query = format!("CREATE_REPLICATION_SLOT {slot_quoted} LOGICAL {plugin_quoted}");
26        let mut result = self.query_simple(&query).await?;
27        let mut slot = String::new();
28        let mut lsn = Lsn(0);
29
30        loop {
31            match result.next_result_set().await? {
32                crate::postgres_client::QueryResultSet::QueryProcessingComplete => break,
33                crate::postgres_client::QueryResultSet::RowDescriptionReceived(mut reader) => {
34                    if let Some(row) = reader.next_row().await? {
35                        slot = row.get_text::<String>(0)?;
36                        lsn = row.get_text(1)?;
37                    }
38                }
39            }
40        }
41
42        Ok((slot, lsn))
43    }
44
45    pub async fn drop_replication_slot(
46        &mut self,
47        slot_name: &str,
48    ) -> Result<(), ElefantClientError> {
49        let slot_quoted = quote_identifier(slot_name);
50        let query = format!("DROP_REPLICATION_SLOT {slot_quoted}");
51        self.execute_non_query_simple(&query).await
52    }
53
54    pub async fn start_replication(
55        &mut self,
56        slot_name: &str,
57        lsn: Lsn,
58        options: &str,
59    ) -> Result<ReplicationStream<'_, F>, ElefantClientError> {
60        let slot_quoted = quote_identifier(slot_name);
61        let query = format!("START_REPLICATION SLOT {slot_quoted} LOGICAL {lsn} ({options})");
62
63        self.start_new_query().await?;
64        self.connection
65            .write_frontend_message(&FrontendMessage::Query(crate::protocol::Query {
66                query: Cow::Borrowed(&query),
67            }))
68            .await?;
69        self.connection.flush().await?;
70
71        let msg = self.read_next_backend_message().await?;
72        match msg {
73            crate::protocol::BackendMessage::CopyBothResponse(_) => {
74                Ok(ReplicationStream::new(self))
75            }
76            _ => Err(ElefantClientError::UnexpectedBackendMessage(format!(
77                "Expected CopyBothResponse, got {msg:?}"
78            ))),
79        }
80    }
81}
82
83#[cfg(test)]
84mod tests {
85    use super::*;
86
87    #[test]
88    fn quote_identifier_simple_name() {
89        assert_eq!(quote_identifier("my_slot"), "\"my_slot\"");
90    }
91
92    #[test]
93    fn quote_identifier_escapes_double_quotes() {
94        assert_eq!(quote_identifier(r#"my"slot"#), r#""my""slot""#);
95    }
96
97    #[test]
98    fn quote_identifier_handles_spaces_and_special_chars() {
99        assert_eq!(
100            quote_identifier("my replication slot"),
101            "\"my replication slot\""
102        );
103        assert_eq!(quote_identifier("UPPER_CASE"), "\"UPPER_CASE\"");
104    }
105
106    #[test]
107    fn quote_identifier_handles_empty_string() {
108        assert_eq!(quote_identifier(""), "\"\"");
109    }
110}
elefant_client/postgres_client/replication/client_ext.rs

elefant_client/postgres_client/replication/
client_ext.rs
