create/drop replication slot

Author: jeff-davis

tokio-postgres/src/replication_client.rs

@@ -57,6 +57,13 @@ impl TimelineHistory {
     }
 }
 
+#[derive(Debug)]
+pub enum SnapshotMode {
+    ExportSnapshot,
+    NoExportSnapshot,
+    UseSnapshot,
+}
+
 /// Replication client connection.
 ///
 /// A replication client is used to issue replication commands, begin
@@ -204,6 +211,62 @@ impl ReplicationClient {
         })
     }
 
+    /// Create physical replication slot
+    pub async fn create_physical_replication_slot(
+        &self,
+        slot_name: &str,
+        temporary: bool,
+        reserve_wal: bool
+    ) -> Result<(), Error> {
+        let iclient = self.0.inner();
+        let temporary_str = if temporary { " TEMPORARY" } else { "" };
+        let reserve_wal_str = if reserve_wal { " RESERVE_WAL" } else { "" };
+        let command = format!("CREATE_REPLICATION_SLOT {}{} PHYSICAL{}",
+                              escape_identifier(slot_name),
+                              temporary_str,
+                              reserve_wal_str);
+        let buf = simple_query::encode(iclient, &command)?;
+        let _responses = iclient.send(RequestMessages::Single(FrontendMessage::Raw(buf)))?;
+        Ok(())
+    }
+
+    /// Create logical replication slot.
+    pub async fn create_logical_replication_slot(
+        &self,
+        slot_name: &str,
+        temporary: bool,
+        plugin_name: &str,
+        snapshot_mode: Option<SnapshotMode>,
+    ) -> Result<(), Error> {
+        let iclient = self.0.inner();
+        let temporary_str = if temporary { " TEMPORARY" } else { "" };
+        let snapshot_str =  snapshot_mode.map_or("", |mode| {
+            match mode {
+                SnapshotMode::ExportSnapshot => " EXPORT_SNAPSHOT",
+                SnapshotMode::NoExportSnapshot => " NOEXPORT_SNAPSHOT",
+                SnapshotMode::UseSnapshot => " USE_SNAPSHOT",
+            }
+        });
+        let command = format!("CREATE_REPLICATION_SLOT {}{} LOGICAL {}{}",
+                              escape_identifier(slot_name),
+                              temporary_str,
+                              escape_identifier(plugin_name),
+                              snapshot_str);
+        let buf = simple_query::encode(iclient, &command)?;
+        let _responses = iclient.send(RequestMessages::Single(FrontendMessage::Raw(buf)))?;
+        Ok(())
+    }
+
+    /// Drop replication slot
+    pub async fn drop_replication_slot(&self, slot_name: &str, wait: bool) -> Result<(), Error> {
+        let iclient = self.0.inner();
+        let wait_str = if wait { " WAIT" } else { "" };
+        let command = format!("DROP_REPLICATION_SLOT {}{}", escape_identifier(slot_name), wait_str);
+        let buf = simple_query::encode(iclient, &command)?;
+        let _responses = iclient.send(RequestMessages::Single(FrontendMessage::Raw(buf)))?;
+        Ok(())
+    }
+
     /// Begin physical replication, consuming the replication client and producing a replication stream.
     ///
     /// Replication begins starting with the given Log Sequence Number

tokio-postgres/tests/test/replication.rs

@@ -15,6 +15,7 @@ const LOGICAL_INSERT_TAG: u8 = b'I';
 // test for:
 //   - identify_system
 //   - show
+//   - slot create/drop
 //   - physical replication
 #[tokio::test]
 async fn physical_replication() {
@@ -25,6 +26,10 @@ async fn physical_replication() {
     let show_port = rclient.show("port").await.unwrap();
     assert_eq!(show_port, "5433");
 
+    let slot = "test_physical_slot";
+    let _ = rclient.drop_replication_slot(slot, false);
+    let _ = rclient.create_physical_replication_slot(slot, false, false);
+
     let physical_stream = rclient
         .start_physical_replication(None, identify_system.xlogpos(), None)
         .await
@@ -51,7 +56,7 @@ async fn physical_replication() {
 }
 
 // test for:
-//   X create/drop slot
+//   - create/drop slot
 //   X standby_status_update
 //   - logical replication
 #[tokio::test]
@@ -63,15 +68,8 @@ async fn logical_replication() {
 
     let slot = "test_logical_slot";
     let plugin = "pgoutput";
-    let _nrows = sclient
-        .execute("select pg_drop_replication_slot($1)", &[&slot.to_string()])
-        .await
-        .unwrap();
-
-    let _nrows = sclient
-        .execute("select pg_create_logical_replication_slot($1, $2)", &[&slot.to_string(), &plugin.to_string()])
-        .await
-        .unwrap();
+    let _ = rclient.drop_replication_slot(slot, false);
+    let _ = rclient.create_logical_replication_slot(slot, false, plugin, None);
 
     let xlog_start = identify_system.xlogpos();
     let options = &vec![("proto_version","1"), ("publication_names", "test_logical_pub")];