inline _create_replication_slot() to _create_recovery_conf()

Author: funbringer

testgres/backup.py

@@ -154,12 +154,13 @@ def spawn_primary(self, name=None, destroy=True):
 
         return node
 
-    def spawn_replica(self, name=None, destroy=True, slot_name=None):
+    def spawn_replica(self, name=None, destroy=True, slot=None):
         """
         Create a replica of the original node from a backup.
 
         Args:
             name: replica's application name.
+            slot: create a replication slot with the specified name.
             destroy: should we convert this backup into a node?
 
         Returns:
@@ -171,11 +172,11 @@ def spawn_replica(self, name=None, destroy=True, slot_name=None):
 
         # Assign it a master and a recovery file (private magic)
         node._assign_master(self.original_node)
-        node._create_recovery_conf(username=self.username, slot_name=slot_name)
+        node._create_recovery_conf(username=self.username, slot=slot)
 
         return node
 
     def cleanup(self):
         if self._available:
-            rmtree(self.base_dir, ignore_errors=True)
             self._available = False
+            rmtree(self.base_dir, ignore_errors=True)

testgres/node.py

@@ -278,7 +278,7 @@ def _assign_master(self, master):
         # now this node has a master
         self._master = master
 
-    def _create_recovery_conf(self, username, slot_name=None):
+    def _create_recovery_conf(self, username, slot=None):
         """NOTE: this is a private method!"""
 
         # fetch master of this node
@@ -306,8 +306,26 @@ def _create_recovery_conf(self, username, slot_name=None):
             "standby_mode=on\n"
         ).format(conninfo)
 
-        if slot_name:
-            line += "primary_slot_name={}\n".format(slot_name)
+        if slot:
+            # Connect to master for some additional actions
+            with master.connect(username=username) as con:
+                # check if slot already exists
+                res = con.execute("""
+                    select exists (
+                        select from pg_catalog.pg_replication_slots
+                        where slot_name = $1
+                    )
+                """, slot)
+
+                if res[0][0]:
+                    raise TestgresException("Slot '{}' already exists".format(slot))
+
+                # TODO: we should drop this slot after replica's cleanup()
+                con.execute("""
+                    select pg_catalog.pg_create_physical_replication_slot($1)
+                """, slot)
+
+            line += "primary_slot_name={}\n".format(slot)
 
         self.append_conf(RECOVERY_CONF_FILE, line)
 
@@ -352,28 +370,6 @@ def _collect_special_files(self):
 
         return result
 
-    def _create_replication_slot(self, slot_name, dbname=None, username=None):
-        """
-        Create a physical replication slot.
-
-        Args:
-            slot_name: slot name
-            dbname: database name
-            username: database user name
-        """
-        rs = self.execute("select exists (select * from pg_replication_slots "
-                          "where slot_name = '{}')".format(slot_name),
-                          dbname=dbname, username=username)
-
-        if rs[0][0]:
-            raise TestgresException("Slot '{}' already exists".format(slot_name))
-
-        query = (
-            "select pg_create_physical_replication_slot('{}')"
-        ).format(slot_name)
-
-        self.execute(query=query, dbname=dbname, username=username)
-
     def init(self, initdb_params=None, **kwargs):
         """
         Perform initdb for this node.
@@ -969,26 +965,24 @@ def backup(self, **kwargs):
 
         return NodeBackup(node=self, **kwargs)
 
-    def replicate(self, name=None, slot_name=None, **kwargs):
+    def replicate(self, name=None, slot=None, **kwargs):
         """
         Create a binary replica of this node.
 
         Args:
             name: replica's application name.
+            slot: create a replication slot with the specified name.
             username: database user name.
             xlog_method: a method for collecting the logs ('fetch' | 'stream').
             base_dir: the base directory for data files and logs
         """
 
-        if slot_name:
-            self._create_replication_slot(slot_name, **kwargs)
-
         backup = self.backup(**kwargs)
 
         # transform backup into a replica
         return backup.spawn_replica(name=name,
                                     destroy=True,
-                                    slot_name=slot_name)
+                                    slot=slot)
 
     def catchup(self, dbname=None, username=None):
         """

tests/test_simple.py

@@ -384,19 +384,15 @@ def test_replicate(self):
                 self.assertListEqual(res, [])
 
     def test_replication_slots(self):
-        query_create = 'create table test as select generate_series(1, 2) as val'
-
         with get_new_node() as node:
             node.init(allow_streaming=True).start()
-            node.execute(query_create)
 
-            with node.replicate(slot_name='slot1').start() as replica:
-                res = replica.execute('select * from test')
-                self.assertListEqual(res, [(1, ), (2, )])
+            with node.replicate(slot='slot1').start() as replica:
+                replica.execute('select 1')
 
                 # cannot create new slot with the same name
                 with self.assertRaises(TestgresException):
-                    node._create_replication_slot('slot1')
+                    node.replicate(slot='slot1')
 
     def test_incorrect_catchup(self):
         with get_new_node() as node: