Refactor Connection class (#2815)

1. stop using the same cursor all the time, it creates problems when not carefully used from different threads.
2. introduce query() method in the Connection class and make it return a result set when it is possible.
3. refactor most of the code that is relying (directly or indirectly) on the Connection object to use the query() method as much as possible.

This refactoring helps with reducing code complexity and will help with future introduction of a separate database connection for the REST API thread. The last one will help to improve reliability when system is under significant stress when simple monitoring queries are taking seconds to execute and the REST API starts blocking the main thread.

patroni/api.py

@@ -1183,20 +1183,18 @@ class RestApiHandler(BaseHTTPRequestHandler):
                 self.command = mname
         return ret
 
-    def query(self, sql: str, *params: Any, **kwargs: Any) -> List[Tuple[Any, ...]]:
-        """Execute *sql* query with *params*.
+    def query(self, sql: str, *params: Any, retry: bool = False) -> List[Tuple[Any, ...]]:
+        """Execute *sql* query with *params* and optionally return results.
 
         :param sql: the SQL statement to be run.
         :param params: positional arguments to call :func:`RestApiServer.query` with.
-        :param kwargs: can contain the key ``retry``. If the key is present its value should be a :class:`bool` which
-            indicates whether the query should be retried upon failure or given up immediately.
+        :param retry: whether the query should be retried upon failure or given up immediately.
 
         :returns: a list of rows that were fetched from the database.
         """
-        if not kwargs.get('retry', False):
+        if not retry:
             return self.server.query(sql, *params)
-        retry = Retry(delay=1, retry_exceptions=PostgresConnectionException)
-        return retry(self.server.query, sql, *params)
+        return Retry(delay=1, retry_exceptions=PostgresConnectionException)(self.server.query, sql, *params)
 
     def get_postgresql_status(self, retry: bool = False) -> Dict[str, Any]:
         """Builds an object representing a status of "postgres".
@@ -1368,7 +1366,7 @@ class RestApiServer(ThreadingMixIn, HTTPServer, Thread):
         self.daemon = True
 
     def query(self, sql: str, *params: Any) -> List[Tuple[Any, ...]]:
-        """Execute *sql* query with *params*.
+        """Execute *sql* query with *params* and optionally return results.
 
         :param sql: the SQL statement to be run.
         :param params: positional arguments to be used as parameters for *sql*.
@@ -1379,15 +1377,10 @@ class RestApiServer(ThreadingMixIn, HTTPServer, Thread):
             :class:`psycopg.Error`: if had issues while executing *sql*.
             :class:`~patroni.exceptions.PostgresConnectionException`: if had issues while connecting to the database.
         """
-        cursor = None
         try:
-            with self.patroni.postgresql.connection().cursor() as cursor:
-                cursor.execute(sql.encode('utf-8'), params)
-                return [r for r in cursor]
-        except psycopg.Error as e:
-            if cursor and cursor.connection.closed == 0:
-                raise e
-            raise PostgresConnectionException('connection problems')
+            return self.patroni.postgresql.query(sql, *params, retry=False)
+        except RetryFailedError as e:
+            raise PostgresConnectionException(str(e))
 
     @staticmethod
     def _set_fd_cloexec(fd: socket.socket) -> None:

patroni/postgresql/__init__.py

@@ -332,36 +332,50 @@ class Postgresql(object):
         self._connection.set_conn_kwargs(kwargs.copy())
         self.citus_handler.set_conn_kwargs(kwargs.copy())
 
-    def _query(self, sql: str, *params: Any) -> Union['Cursor[Any]', 'cursor']:
-        """We are always using the same cursor, therefore this method is not thread-safe!!!
-        You can call it from different threads only if you are holding explicit `AsyncExecutor` lock,
-        because the main thread is always holding this lock when running HA cycle."""
-        cursor = None
-        try:
-            cursor = self._connection.cursor()
-            cursor.execute(sql.encode('utf-8'), params or None)
-            return cursor
-        except psycopg.Error as e:
-            if cursor and cursor.connection.closed == 0:
-                # When connected via unix socket, psycopg2 can't recoginze 'connection lost'
-                # and leaves `_cursor_holder.connection.closed == 0`, but psycopg2.OperationalError
-                # is still raised (what is correct). It doesn't make sense to continiue with existing
-                # connection and we will close it, to avoid its reuse by the `cursor` method.
-                if isinstance(e, psycopg.OperationalError):
-                    self._connection.close()
-                else:
-                    raise e
-            if self.state == 'restarting':
-                raise RetryFailedError('cluster is being restarted')
-            raise PostgresConnectionException('connection problems')
+    def _query(self, sql: str, *params: Any) -> List[Tuple[Any, ...]]:
+        """Execute *sql* query with *params* and optionally return results.
 
-    def query(self, sql: str, *args: Any, **kwargs: Any) -> Union['Cursor[Any]', 'cursor']:
-        if not kwargs.get('retry', True):
-            return self._query(sql, *args)
+        :param sql: SQL statement to execute.
+        :param params: parameters to pass.
+
+        :returns: a query response as a list of tuples if there is any.
+        :raises:
+            :exc:`~psycopg.Error` if had issues while executing *sql*.
+
+            :exc:`~patroni.exceptions.PostgresConnectionException`: if had issues while connecting to the database.
+
+            :exc:`~patroni.utils.RetryFailedError`: if it was detected that connection/query failed due to PostgreSQL
+            restart.
+        """
         try:
-            return self.retry(self._query, sql, *args)
-        except RetryFailedError as e:
-            raise PostgresConnectionException(str(e))
+            return self._connection.query(sql, *params)
+        except PostgresConnectionException as exc:
+            if self.state == 'restarting':
+                raise RetryFailedError('cluster is being restarted') from exc
+            raise
+
+    def query(self, sql: str, *params: Any, retry: bool = True) -> List[Tuple[Any, ...]]:
+        """Execute *sql* query with *params* and optionally return results.
+
+        :param sql: SQL statement to execute.
+        :param params: parameters to pass.
+        :param retry: whether the query should be retried upon failure or given up immediately.
+
+        :returns: a query response as a list of tuples if there is any.
+        :raises:
+            :exc:`~psycopg.Error` if had issues while executing *sql*.
+
+            :exc:`~patroni.exceptions.PostgresConnectionException`: if had issues while connecting to the database.
+
+            :exc:`~patroni.utils.RetryFailedError`: if it was detected that connection/query failed due to PostgreSQL
+            restart or if retry deadline was exceeded.
+        """
+        if not retry:
+            return self._query(sql, *params)
+        try:
+            return self.retry(self._query, sql, *params)
+        except RetryFailedError as exc:
+            raise PostgresConnectionException(str(exc)) from exc
 
     def pg_control_exists(self) -> bool:
         return os.path.isfile(self._pg_control)
@@ -431,7 +445,7 @@ class Postgresql(object):
     def _cluster_info_state_get(self, name: str) -> Optional[Any]:
         if not self._cluster_info_state:
             try:
-                result = self._is_leader_retry(self._query, self.cluster_info_query).fetchone()
+                result = self._is_leader_retry(self._query, self.cluster_info_query)[0]
                 cluster_info_state = dict(zip(['timeline', 'wal_position', 'replayed_location',
                                                'received_location', 'replay_paused', 'pg_control_timeline',
                                                'received_tli', 'slot_name', 'conninfo', 'receiver_state',
@@ -873,11 +887,10 @@ class Postgresql(object):
 
     def _wait_for_connection_close(self, postmaster: PostmasterProcess) -> None:
         try:
-            with self.connection().cursor() as cur:
-                while postmaster.is_running():  # Need a timeout here?
-                    cur.execute("SELECT 1")
-                    time.sleep(STOP_POLLING_INTERVAL)
-        except psycopg.Error:
+            while postmaster.is_running():  # Need a timeout here?
+                self._connection.query("SELECT 1")
+                time.sleep(STOP_POLLING_INTERVAL)
+        except (psycopg.Error, PostgresConnectionException):
             pass
 
     def reload(self, block_callbacks: bool = False) -> bool:
@@ -1178,26 +1191,16 @@ class Postgresql(object):
             received_location = self.received_location()
             pg_control_timeline = self._cluster_info_state_get('pg_control_timeline')
         else:
-            with self.connection().cursor() as cursor:
-                cursor.execute(self.cluster_info_query.encode('utf-8'))
-                row = cursor.fetchone()
-                if TYPE_CHECKING:  # pragma: no cover
-                    assert row is not None
-                (timeline, wal_position, replayed_location, received_location, _, pg_control_timeline) = row[:6]
+            timeline, wal_position, replayed_location, received_location, _, pg_control_timeline = \
+                self._query(self.cluster_info_query)[0][:6]
 
         wal_position = self._wal_position(bool(timeline), wal_position, received_location, replayed_location)
-        return (timeline, wal_position, pg_control_timeline)
+        return timeline, wal_position, pg_control_timeline
 
     def postmaster_start_time(self) -> Optional[str]:
         try:
-            query = "SELECT " + self.POSTMASTER_START_TIME
-            if current_thread().ident == self.__thread_ident:
-                row = self.query(query).fetchone()
-            else:
-                with self.connection().cursor() as cursor:
-                    cursor.execute(query)
-                    row = cursor.fetchone()
-            return row[0].isoformat(sep=' ') if row else None
+            sql = "SELECT " + self.POSTMASTER_START_TIME
+            return self.query(sql, retry=current_thread().ident == self.__thread_ident)[0][0].isoformat(sep=' ')
         except psycopg.Error:
             return None
 

patroni/postgresql/citus.py

@@ -11,8 +11,6 @@ from ..dcs import CITUS_COORDINATOR_GROUP_ID, Cluster
 from ..psycopg import connect, quote_ident
 
 if TYPE_CHECKING:  # pragma: no cover
-    from psycopg import Cursor
-    from psycopg2 import cursor
     from . import Postgresql
 
 CITUS_SLOT_NAME_RE = re.compile(r'^citus_shard_(move|split)_slot(_[1-9][0-9]*){2,3}$')
@@ -109,12 +107,10 @@ class CitusHandler(Thread):
             self._tasks[:] = []
             self._in_flight = None
 
-    def query(self, sql: str, *params: Any) -> Union['Cursor[Any]', 'cursor']:
+    def query(self, sql: str, *params: Any) -> List[Tuple[Any, ...]]:
         try:
             logger.debug('query(%s, %s)', sql, params)
-            cursor = self._connection.cursor()
-            cursor.execute(sql.encode('utf-8'), params or None)
-            return cursor
+            return self._connection.query(sql, *params)
         except Exception as e:
             logger.error('Exception when executing query "%s", (%s): %r', sql, params, e)
             self._connection.close()
@@ -132,13 +128,13 @@ class CitusHandler(Thread):
             self._schedule_load_pg_dist_node = False
 
         try:
-            cursor = self.query("SELECT nodeid, groupid, nodename, nodeport, noderole"
-                                " FROM pg_catalog.pg_dist_node WHERE noderole = 'primary'")
+            rows = self.query("SELECT nodeid, groupid, nodename, nodeport, noderole"
+                              " FROM pg_catalog.pg_dist_node WHERE noderole = 'primary'")
         except Exception:
             return False
 
         with self._condition:
-            self._pg_dist_node = {r[1]: PgDistNode(r[1], r[2], r[3], 'after_promote', r[0]) for r in cursor}
+            self._pg_dist_node = {r[1]: PgDistNode(r[1], r[2], r[3], 'after_promote', r[0]) for r in rows}
         return True
 
     def sync_pg_dist_node(self, cluster: Cluster) -> None:
@@ -211,10 +207,8 @@ class CitusHandler(Thread):
             self.query('SELECT pg_catalog.citus_update_node(%s, %s, %s, true, %s)',
                        task.nodeid, task.host, task.port, task.cooldown)
         elif task.event != 'before_demote':
-            row = self.query("SELECT pg_catalog.citus_add_node(%s, %s, %s, 'primary', 'default')",
-                             task.host, task.port, task.group).fetchone()
-            if row is not None:
-                task.nodeid = row[0]
+            task.nodeid = self.query("SELECT pg_catalog.citus_add_node(%s, %s, %s, 'primary', 'default')",
+                                     task.host, task.port, task.group)[0][0]
 
     def process_task(self, task: PgDistNode) -> bool:
         """Updates a single row in `pg_dist_node` table, optionally in a transaction.

patroni/postgresql/config.py

@@ -1093,10 +1093,10 @@ class ConfigHandler(object):
             if self._postgresql.major_version >= 90500:
                 time.sleep(1)
                 try:
-                    pending_restart = (self._postgresql.query(
+                    pending_restart = self._postgresql.query(
                         'SELECT COUNT(*) FROM pg_catalog.pg_settings'
                         ' WHERE pg_catalog.lower(name) != ALL(%s) AND pending_restart',
-                        [n.lower() for n in self._RECOVERY_PARAMETERS]).fetchone() or (0,))[0] > 0
+                        [n.lower() for n in self._RECOVERY_PARAMETERS])[0][0] > 0
                     self._postgresql.set_pending_restart(pending_restart)
                 except Exception as e:
                     logger.warning('Exception %r when running query', e)

patroni/postgresql/connection.py

@@ -2,45 +2,86 @@ import logging
 
 from contextlib import contextmanager
 from threading import Lock
-from typing import Any, Dict, Iterator, Union, TYPE_CHECKING
+from typing import Any, Dict, Iterator, List, Union, Tuple, TYPE_CHECKING
 if TYPE_CHECKING:  # pragma: no cover
     from psycopg import Connection as Connection3, Cursor
     from psycopg2 import connection, cursor
 
 from .. import psycopg
+from ..exceptions import PostgresConnectionException
 
 logger = logging.getLogger(__name__)
 
 
-class Connection(object):
+class Connection:
+    """Helper class to manage connections from Patroni to PostgreSQL.
+
+    :ivar server_version: PostgreSQL version in integer format where we are connected to.
+    """
+
     server_version: int
 
     def __init__(self) -> None:
-        self._lock = Lock()
+        """Create an instance of :class:`Connection` class."""
+        self._lock = Lock()  # used to make sure that only one connection to postgres is established
         self._connection = None
-        self._cursor_holder = None
 
     def set_conn_kwargs(self, conn_kwargs: Dict[str, Any]) -> None:
+        """Set connection parameters, like user, password, host, port and so on.
+
+        :param conn_kwargs: connection parameters as a dictionary.
+        """
         self._conn_kwargs = conn_kwargs
 
     def get(self) -> Union['connection', 'Connection3[Any]']:
+        """Get ``psycopg``/``psycopg2`` connection object.
+
+        .. note::
+            Opens a new connection if necessary.
+
+        :returns: ``psycopg`` or ``psycopg2`` connection object.
+        """
         with self._lock:
             if not self._connection or self._connection.closed != 0:
+                logger.info("establishing a new patroni connection to postgres")
                 self._connection = psycopg.connect(**self._conn_kwargs)
                 self.server_version = getattr(self._connection, 'server_version', 0)
         return self._connection
 
-    def cursor(self) -> Union['cursor', 'Cursor[Any]']:
-        if not self._cursor_holder or self._cursor_holder.closed or self._cursor_holder.connection.closed != 0:
-            logger.info("establishing a new patroni connection to the postgres cluster")
-            self._cursor_holder = self.get().cursor()
-        return self._cursor_holder
+    def query(self, sql: str, *params: Any) -> List[Tuple[Any, ...]]:
+        """Execute a query with parameters and optionally returns a response.
+
+        :param sql: SQL statement to execute.
+        :param params: parameters to pass.
+
+        :returns: a query response as a list of tuples if there is any.
+        :raises:
+            :exc:`~psycopg.Error` if had issues while executing *sql*.
+
+            :exc:`~patroni.exceptions.PostgresConnectionException`: if had issues while connecting to the database.
+        """
+        cursor = None
+        try:
+            with self.get().cursor() as cursor:
+                cursor.execute(sql.encode('utf-8'), params or None)
+                return cursor.fetchall() if cursor.rowcount and cursor.rowcount > 0 else []
+        except psycopg.Error as exc:
+            if cursor and cursor.connection.closed == 0:
+                # When connected via unix socket, psycopg2 can't recoginze 'connection lost' and leaves
+                # `self._connection.closed == 0`, but the generic exception is raised. It doesn't make
+                # sense to continue with existing connection and we will close it, to avoid its reuse.
+                if type(exc) in (psycopg.DatabaseError, psycopg.OperationalError):
+                    self.close()
+                else:
+                    raise exc
+            raise PostgresConnectionException('connection problems') from exc
 
     def close(self) -> None:
+        """Close the psycopg connection to postgres."""
         if self._connection and self._connection.closed == 0:
             self._connection.close()
-            logger.info("closed patroni connection to the postgresql cluster")
-        self._cursor_holder = self._connection = None
+            logger.info("closed patroni connection to postgres")
+        self._connection = None
 
 
 @contextmanager

patroni/postgresql/slots.py

@@ -186,7 +186,7 @@ class SlotsHandler:
         self.pg_replslot_dir = os.path.join(self._postgresql.data_dir, 'pg_replslot')
         self.schedule()
 
-    def _query(self, sql: str, *params: Any) -> Union['cursor', 'Cursor[Any]']:
+    def _query(self, sql: str, *params: Any) -> List[Tuple[Any, ...]]:
         """Helper method for :meth:`Postgresql.query`.
 
         :param sql: SQL statement to execute.
@@ -263,9 +263,8 @@ class SlotsHandler:
             extra = ", catalog_xmin, pg_catalog.pg_wal_lsn_diff(confirmed_flush_lsn, '0/0')::bigint" \
                 if self._postgresql.major_version >= 100000 else ""
             skip_temp_slots = ' WHERE NOT temporary' if self._postgresql.major_version >= 100000 else ''
-            cursor = self._query(f'SELECT slot_name, slot_type, plugin, database, datoid'
-                                 f'{extra} FROM pg_catalog.pg_replication_slots{skip_temp_slots}')
-            for r in cursor:
+            for r in self._query('SELECT slot_name, slot_type, plugin, database, datoid'
+                                 f'{extra} FROM pg_catalog.pg_replication_slots{skip_temp_slots}'):
                 value = {'type': r[1]}
                 if r[1] == 'logical':
                     value.update(plugin=r[2], database=r[3], datoid=r[4])
@@ -308,16 +307,13 @@ class SlotsHandler:
                   ``dropped`` is ``True`` if the slot was successfully dropped. If the slot was not found return
                   ``False`` for both.
         """
-        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)
-        row = cursor.fetchone()
-        if not row:
-            row = (False, False)
-        return row
+        rows = 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 rows[0] if rows else (False, False)
 
     def _drop_incorrect_slots(self, cluster: Cluster, slots: Dict[str, Any], paused: bool) -> None:
         """Compare required slots and configured as permanent slots with those found, dropping extraneous ones.
@@ -595,11 +591,8 @@ class SlotsHandler:
 
         # Replica isn't streaming or the hot_standby_feedback isn't enabled
         try:
-            cur = self._query("SELECT pg_catalog.current_setting('hot_standby_feedback')::boolean")
-            row = cur.fetchone()
-            if row and not row[0]:
-                logger.error('Logical slot failover requires "hot_standby_feedback".'
-                             ' Please check postgresql.auto.conf')
+            if not self._query("SELECT pg_catalog.current_setting('hot_standby_feedback')::boolean")[0][0]:
+                logger.error('Logical slot failover requires "hot_standby_feedback". Please check postgresql.auto.conf')
         except Exception as e:
             logger.error('Failed to check the hot_standby_feedback setting: %r', e)
         return False

tests/__init__.py

@@ -106,7 +106,7 @@ class MockCursor(object):
         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)] if self.rowcount == 1 else [None]
+            self.results = [(False, True)] if self.rowcount == 1 else []
         elif sql.startswith('SELECT CASE WHEN pg_catalog.pg_is_in_recovery()'):
             self.results = [(1, 2, 1, 0, False, 1, 1, None, None, 'streaming', '',
                              [{"slot_name": "ls", "confirmed_flush_lsn": 12345}],
@@ -114,10 +114,7 @@ class MockCursor(object):
         elif sql.startswith('SELECT pg_catalog.pg_is_in_recovery()'):
             self.results = [(False, 2)]
         elif sql.startswith('SELECT pg_catalog.pg_postmaster_start_time'):
-            replication_info = '[{"application_name":"walreceiver","client_addr":"1.2.3.4",' +\
-                               '"state":"streaming","sync_state":"async","sync_priority":0}]'
-            now = datetime.datetime.now(tzutc)
-            self.results = [(now, 0, '', 0, '', False, now, 'streaming', None, replication_info)]
+            self.results = [(datetime.datetime.now(tzutc),)]
         elif sql.startswith('SELECT name, current_setting(name) FROM pg_settings'):
             self.results = [('data_directory', 'data'),
                             ('hba_file', os.path.join('data', 'pg_hba.conf')),
@@ -143,6 +140,8 @@ class MockCursor(object):
                             ('listen_addresses', '*', None, 'string', 'postmaster'),
                             ('autovacuum', 'on', None, 'bool', 'sighup'),
                             ('unix_socket_directories', '/tmp', None, 'string', 'postmaster')]
+        elif sql.startswith('SELECT COUNT(*) FROM pg_catalog.pg_settings'):
+            self.results = [(1,)]
         elif sql.startswith('IDENTIFY_SYSTEM'):
             self.results = [('1', 3, '0/402EEC0', '')]
         elif sql.startswith('TIMELINE_HISTORY '):
@@ -156,6 +155,7 @@ class MockCursor(object):
             self.results = [(1, 0, 'host1', 5432, 'primary'), (2, 1, 'host2', 5432, 'primary')]
         else:
             self.results = [(None, None, None, None, None, None, None, None, None, None)]
+        self.rowcount = len(self.results)
 
     def fetchone(self):
         return self.results[0]

tests/test_api.py

@@ -3,8 +3,6 @@ import json
 import unittest
 import socket
 
-import patroni.psycopg as psycopg
-
 from http.server import HTTPServer
 from io import BytesIO as IO
 from mock import Mock, PropertyMock, patch
@@ -14,9 +12,8 @@ from patroni.api import RestApiHandler, RestApiServer
 from patroni.config import GlobalConfig
 from patroni.dcs import ClusterConfig, Member
 from patroni.ha import _MemberStatus
-from patroni.utils import tzutc
+from patroni.utils import RetryFailedError, tzutc
 
-from . import psycopg_connect, MockCursor
 from .test_ha import get_cluster_initialized_without_leader
 
 
@@ -41,10 +38,6 @@ class MockPostgresql(object):
     TL_LSN = 'CASE WHEN pg_catalog.pg_is_in_recovery()'
     citus_handler = Mock()
 
-    @staticmethod
-    def connection():
-        return psycopg_connect()
-
     @staticmethod
     def postmaster_start_time():
         return postmaster_start_time
@@ -61,6 +54,12 @@ class MockPostgresql(object):
     def replication_state_from_parameters(*args):
         return 'streaming'
 
+    @staticmethod
+    def query(sql, *params, retry=False):
+        return [(postmaster_start_time, 0, '', 0, '', False, postmaster_start_time, 'streaming', None,
+                 '[{"application_name":"walreceiver","client_addr":"1.2.3.4",'
+                 + '"state":"streaming","sync_state":"async","sync_priority":0}]')]
+
 
 class MockWatchdog(object):
     is_healthy = False
@@ -488,9 +487,7 @@ class TestRestApiHandler(unittest.TestCase):
 
     @patch('time.sleep', Mock())
     def test_RestApiServer_query(self):
-        with patch.object(MockCursor, 'execute', Mock(side_effect=psycopg.OperationalError)):
-            self.assertIsNotNone(MockRestApiServer(RestApiHandler, 'GET /patroni'))
-        with patch.object(MockPostgresql, 'connection', Mock(side_effect=psycopg.OperationalError)):
+        with patch.object(MockPostgresql, 'query', Mock(side_effect=RetryFailedError('bla'))):
             self.assertIsNotNone(MockRestApiServer(RestApiHandler, 'GET /patroni'))
 
     @patch('time.sleep', Mock())

tests/test_postgresql.py

@@ -545,9 +545,7 @@ class TestPostgresql(BaseTestPostgresql):
 
     @patch('time.sleep', Mock())
     @patch.object(Postgresql, 'is_running', Mock(return_value=True))
-    @patch.object(MockCursor, 'fetchone')
-    def test_reload_config(self, mock_fetchone):
-        mock_fetchone.return_value = (1,)
+    def test_reload_config(self):
         parameters = self._PARAMETERS.copy()
         parameters.pop('f.oo')
         parameters['wal_buffers'] = '512'
@@ -555,9 +553,14 @@ class TestPostgresql(BaseTestPostgresql):
                   'authentication': {},
                   'retry_timeout': 10, 'listen': '*', 'krbsrvname': 'postgres', 'parameters': parameters}
         self.p.reload_config(config)
-        mock_fetchone.side_effect = Exception
         parameters['b.ar'] = 'bar'
-        self.p.reload_config(config)
+        with patch.object(MockCursor, 'fetchall',
+                          Mock(side_effect=[[('wal_block_size', '8191', None, 'integer', 'internal'),
+                                             ('wal_segment_size', '2048', '8kB', 'integer', 'internal'),
+                                             ('shared_buffers', '16384', '8kB', 'integer', 'postmaster'),
+                                             ('wal_buffers', '-1', '8kB', 'integer', 'postmaster'),
+                                             ('port', '5433', None, 'integer', 'postmaster')], Exception])):
+            self.p.reload_config(config)
         parameters['autovacuum'] = 'on'
         self.p.reload_config(config)
         parameters['autovacuum'] = 'off'
@@ -585,7 +588,7 @@ class TestPostgresql(BaseTestPostgresql):
 
     def test_postmaster_start_time(self):
         now = datetime.datetime.now()
-        with patch.object(MockCursor, "fetchone", Mock(return_value=(now, True, '', '', '', '', False))):
+        with patch.object(MockCursor, "fetchall", Mock(return_value=[(now, True, '', '', '', '', False)])):
             self.assertEqual(self.p.postmaster_start_time(), now.isoformat(sep=' '))
             t = Thread(target=self.p.postmaster_start_time)
             t.start()

tests/test_rewind.py

@@ -92,8 +92,9 @@ class TestRewind(BaseTestPostgresql):
             self.r.rewind_or_reinitialize_needed_and_possible(self.leader)
 
         with patch.object(Postgresql, 'is_running', Mock(return_value=True)), \
-                patch.object(MockCursor, 'fetchone',
-                             Mock(side_effect=[(0, 0, 1, 1, 0, 0, 0, 0, 0, None, None, None), Exception])):
+                patch.object(MockCursor, 'fetchone', Mock(side_effect=Exception)), \
+                patch.object(MockCursor, 'fetchall',
+                             Mock(return_value=[(0, 0, 1, 1, 0, 0, 0, 0, 0, None, None, None)])):
             self.r.rewind_or_reinitialize_needed_and_possible(self.leader)
 
     @patch.object(CancellableSubprocess, 'call', mock_cancellable_call)

tests/test_slots.py

@@ -93,17 +93,17 @@ class TestSlotsHandler(BaseTestPostgresql):
         self.s.sync_replication_slots(cluster, False)
         with patch.object(Postgresql, '_query') as mock_query:
             self.p.reset_cluster_info_state(None)
-            mock_query.return_value.fetchone.return_value = (
+            mock_query.return_value = [(
                 1, 0, 0, 0, 0, 0, 0, 0, 0, None, None,
                 [{"slot_name": "ls", "type": "logical", "datoid": 5, "plugin": "b",
-                  "confirmed_flush_lsn": 12345, "catalog_xmin": 105}])
+                  "confirmed_flush_lsn": 12345, "catalog_xmin": 105}])]
             self.assertEqual(self.p.slots(), {'ls': 12345})
 
             self.p.reset_cluster_info_state(None)
-            mock_query.return_value.fetchone.return_value = (
+            mock_query.return_value = [(
                 1, 0, 0, 0, 0, 0, 0, 0, 0, None, None,
                 [{"slot_name": "ls", "type": "logical", "datoid": 6, "plugin": "b",
-                  "confirmed_flush_lsn": 12345, "catalog_xmin": 105}])
+                  "confirmed_flush_lsn": 12345, "catalog_xmin": 105}])]
             self.assertEqual(self.p.slots(), {})
 
     @patch.object(Postgresql, 'is_primary', Mock(return_value=False))
@@ -137,10 +137,10 @@ class TestSlotsHandler(BaseTestPostgresql):
     def test_check_logical_slots_readiness(self):
         self.s.copy_logical_slots(self.cluster, ['ls'])
         with patch.object(MockCursor, '__iter__', Mock(return_value=iter([('postgresql0', None)]))), \
-                patch.object(MockCursor, 'fetchone', Mock(side_effect=Exception)):
+                patch.object(MockCursor, 'fetchall', Mock(side_effect=Exception)):
             self.assertFalse(self.s.check_logical_slots_readiness(self.cluster, None))
         with patch.object(MockCursor, '__iter__', Mock(return_value=iter([('postgresql0', None)]))), \
-                patch.object(MockCursor, 'fetchone', Mock(return_value=(False,))):
+                patch.object(MockCursor, 'fetchall', Mock(return_value=[(False,)])):
             self.assertFalse(self.s.check_logical_slots_readiness(self.cluster, None))
         with patch.object(MockCursor, '__iter__', Mock(return_value=iter([('ls', 100)]))):
             self.s.check_logical_slots_readiness(self.cluster, None)