Refactor drop_replication_slot() and _drop_incorrect_slots() (#2534)

Use CTE to avoid running the second query if pg_drop_replication_slot() failed

patroni/postgresql/slots.py

@@ -179,31 +179,33 @@ class SlotsHandler(object):
         return False
 
     def drop_replication_slot(self, name):
-        cursor = self._query(('SELECT pg_catalog.pg_drop_replication_slot(%s) WHERE EXISTS (SELECT 1 ' +
-                              'FROM pg_catalog.pg_replication_slots WHERE slot_name = %s AND NOT active)'), name, name)
-        # In normal situation rowcount should be 1, otherwise either slot doesn't exists or it is still active
-        return cursor.rowcount == 1
+        """Returns a tuple(active, dropped)"""
+        cursor = self._query(('WITH slots AS (SELECT slot_name, active' +
+                              ' FROM pg_catalog.pg_replication_slots WHERE slot_name = %s),' +
+                              ' dropped AS (SELECT pg_catalog.pg_drop_replication_slot(slot_name),' +
+                              ' true AS dropped FROM slots WHERE not active) ' +
+                              'SELECT active, COALESCE(dropped, false) FROM slots' +
+                              ' FULL OUTER JOIN dropped ON true'), name)
+        return cursor.fetchone() if cursor.rowcount == 1 else (False, False)
 
     def _drop_incorrect_slots(self, cluster, slots, paused):
         # drop old replication slots which are not presented in desired slots
         for name in set(self._replication_slots) - set(slots):
             if not paused and not self.ignore_replication_slot(cluster, name):
-                if not self.drop_replication_slot(name):
+                active, dropped = self.drop_replication_slot(name)
+                if dropped:
+                    logger.info("Dropped unknown replication slot '%s'", name)
+                else:
                     self._schedule_load_slots = True
-                    # Check if slot is still active, this is not considered an error
-                    cursor = self._query(('SELECT 1 FROM pg_catalog.pg_replication_slots WHERE ' +
-                                          'slot_name = %s AND active'), name)
-                    if (cursor.rowcount == 1):
+                    if active:
                         logger.debug("Unable to drop unknown replication slot '%s', slot is still active", name)
                     else:
                         logger.error("Failed to drop replication slot '%s'", name)
-                else:
-                    logger.info("Dropped unknown replication slot '%s'", name)
         for name, value in slots.items():
             if name in self._replication_slots and not compare_slots(value, self._replication_slots[name]):
                 logger.info("Trying to drop replication slot '%s' because value is changing from %s to %s",
                             name, self._replication_slots[name], value)
-                if self.drop_replication_slot(name):
+                if self.drop_replication_slot(name) == (False, True):
                     self._replication_slots.pop(name)
                 else:
                     logger.error("Failed to drop replication slot '%s'", name)

tests/__init__.py

@@ -98,6 +98,8 @@ class MockCursor(object):
             self.results = [('ls', 'logical', 'a', 'b', 100, 500, b'123456')]
         elif sql.startswith('SELECT slot_name'):
             self.results = [('blabla', 'physical'), ('foobar', 'physical'), ('ls', 'logical', 'a', 'b', 5, 100, 500)]
+        elif sql.startswith('WITH slots AS (SELECT slot_name, active'):
+            self.results = [(False, True)]
         elif sql.startswith('SELECT CASE WHEN pg_catalog.pg_is_in_recovery()'):
             self.results = [(1, 2, 1, 0, False, 1, 1, None, None, [{"slot_name": "ls", "confirmed_flush_lsn": 12345}])]
         elif sql.startswith('SELECT pg_catalog.pg_is_in_recovery()'):

tests/test_slots.py

@@ -43,7 +43,10 @@ class TestSlotsHandler(BaseTestPostgresql):
         with mock.patch('patroni.postgresql.Postgresql._query', Mock(side_effect=psycopg.OperationalError)):
             self.s.sync_replication_slots(cluster, False)
         self.p.set_role('standby_leader')
-        self.s.sync_replication_slots(cluster, False)
+        with patch.object(SlotsHandler, 'drop_replication_slot', Mock(return_value=(True, False))),\
+                patch('patroni.postgresql.slots.logger.debug') as mock_debug:
+            self.s.sync_replication_slots(cluster, False)
+            mock_debug.assert_called_once()
         self.p.set_role('replica')
         with patch.object(Postgresql, 'is_leader', Mock(return_value=False)),\
                 patch.object(SlotsHandler, 'drop_replication_slot') as mock_drop:
@@ -52,8 +55,7 @@ class TestSlotsHandler(BaseTestPostgresql):
         self.p.set_role('master')
         with mock.patch('patroni.postgresql.Postgresql.role', new_callable=PropertyMock(return_value='replica')):
             self.s.sync_replication_slots(cluster, False)
-        with patch.object(SlotsHandler, 'drop_replication_slot', Mock(return_value=True)),\
-                patch('patroni.dcs.logger.error', new_callable=Mock()) as errorlog_mock:
+        with patch('patroni.dcs.logger.error', new_callable=Mock()) as errorlog_mock:
             alias1 = Member(0, 'test-3', 28, {'conn_url': 'postgres://replicator:[email protected]:5436/postgres'})
             alias2 = Member(0, 'test.3', 28, {'conn_url': 'postgres://replicator:[email protected]:5436/postgres'})
             cluster.members.extend([alias1, alias2])