Repository: zalando/patroni
Pull request #2158 (2022-01-05 16:11:48)

Compatibility with legacy psycopg2
Psycopg2 version installed from Ubuntu 18.04 packages doesn't have `UndefinedFile` exception yet.

sdudoladov: :+1: 
CyberDem0n: :+1:

Patch:
diff --git a/patroni/postgresql/slots.py b/patroni/postgresql/slots.py
index 09d62463..8e6a2d87 100644
--- a/patroni/postgresql/slots.py
+++ b/patroni/postgresql/slots.py
@@ -8,7 +8,7 @@ from contextlib import contextmanager
 
 from .connection import get_connection_cursor
 from .misc import format_lsn
-from ..psycopg import UndefinedFile
+from ..psycopg import OperationalError
 
 logger = logging.getLogger(__name__)
 
@@ -201,7 +201,7 @@ class SlotsHandler(object):
                                     (name, format_lsn(int(cluster.slots[name]))))
                     except Exception as e:
                         logger.error("Failed to advance logical replication slot '%s': %r", name, e)
-                        if isinstance(e, UndefinedFile):
+                        if isinstance(e, OperationalError) and e.diag.sqlstate == '58P01':  # WAL file is gone
                             create_slots.append(name)
                         self._schedule_load_slots = True
         return create_slots
diff --git a/patroni/psycopg.py b/patroni/psycopg.py
index 487c2161..c1f23e69 100644
--- a/patroni/psycopg.py
+++ b/patroni/psycopg.py
@@ -1,5 +1,4 @@
-__all__ = ['connect', 'quote_ident', 'quote_literal', 'DatabaseError',
-           'Error', 'OperationalError', 'ProgrammingError', 'UndefinedFile']
+__all__ = ['connect', 'quote_ident', 'quote_literal', 'DatabaseError', 'Error', 'OperationalError', 'ProgrammingError']
 
 _legacy = False
 try:
@@ -8,7 +7,6 @@ try:
     if parse_version(__version__) < MIN_PSYCOPG2:
         raise ImportError
     from psycopg2 import connect, Error, DatabaseError, OperationalError, ProgrammingError
-    from psycopg2.errors import UndefinedFile
     from psycopg2.extensions import adapt
 
     try:
@@ -23,7 +21,6 @@ try:
         return value.getquoted().decode('utf-8')
 except ImportError:
     from psycopg import connect as _connect, sql, Error, DatabaseError, OperationalError, ProgrammingError
-    from psycopg.errors import UndefinedFile
 
     def connect(*args, **kwargs):
         ret = _connect(*args, **kwargs)


Test patch:
diff --git a/tests/test_slots.py b/tests/test_slots.py
index 9eb1cdec..6ab66a82 100644
--- a/tests/test_slots.py
+++ b/tests/test_slots.py
@@ -86,7 +86,9 @@ class TestSlotsHandler(BaseTestPostgresql):
                           [self.me, self.other, self.leadermem], None, None, None, {'ls': 12346})
         self.assertEqual(self.s.sync_replication_slots(cluster, False), [])
         self.s._schedule_load_slots = False
-        with patch.object(MockCursor, 'execute', Mock(side_effect=psycopg.UndefinedFile)):
+        with patch.object(MockCursor, 'execute', Mock(side_effect=psycopg.OperationalError)),\
+                patch.object(psycopg.OperationalError, 'diag') as mock_diag:
+            type(mock_diag).sqlstate = PropertyMock(return_value='58P01')
             self.assertEqual(self.s.sync_replication_slots(cluster, False), ['ls'])
         cluster.slots['ls'] = 'a'
         self.assertEqual(self.s.sync_replication_slots(cluster, False), [])