Merge branch 'master' of github.com:zalando/patroni into feature/citus-secondaries

README.rst

@@ -74,9 +74,8 @@ There are a few options available:
 
 ::
 
-    sudo apt-get install python-psycopg2   # install python2 psycopg2 module on Debian/Ubuntu
-    sudo apt-get install python3-psycopg2  # install python3 psycopg2 module on Debian/Ubuntu
-    sudo yum install python-psycopg2       # install python2 psycopg2 on RedHat/Fedora/CentOS
+    sudo apt-get install python3-psycopg2  # install psycopg2 module on Debian/Ubuntu
+    sudo yum install python3-psycopg2      # install psycopg2 on RedHat/Fedora/CentOS
 
 2. Install psycopg2 from the binary package
 
@@ -94,7 +93,7 @@ There are a few options available:
 
 ::
 
-    pip install psycopg[binary]
+    pip install psycopg[binary]>=3.0.0
 
 **General installation for pip**
 

docs/README.rst

@@ -46,9 +46,8 @@ There are a few options available:
 
 ::
 
-    sudo apt-get install python-psycopg2   # install python2 psycopg2 module on Debian/Ubuntu
-    sudo apt-get install python3-psycopg2  # install python3 psycopg2 module on Debian/Ubuntu
-    sudo yum install python-psycopg2       # install python2 psycopg2 on RedHat/Fedora/CentOS
+    sudo apt-get install python3-psycopg2  # install psycopg2 module on Debian/Ubuntu
+    sudo yum install python3-psycopg2      # install psycopg2 on RedHat/Fedora/CentOS
 
 2. Install psycopg2 from the binary package
 
@@ -165,10 +164,6 @@ Applications Should Not Use Superusers
 
 When connecting from an application, always use a non-superuser. Patroni requires access to the database to function properly. By using a superuser from an application, you can potentially use the entire connection pool, including the connections reserved for superusers, with the ``superuser_reserved_connections`` setting. If Patroni cannot access the Primary because the connection pool is full, behavior will be undesirable.
 
-.. |Build Status| image:: https://travis-ci.org/zalando/patroni.svg?branch=master
-   :target: https://travis-ci.org/zalando/patroni
-.. |Coverage Status| image:: https://coveralls.io/repos/zalando/patroni/badge.svg?branch=master
-   :target: https://coveralls.io/r/zalando/patroni?branch=master
 
 Testing Your HA Solution
 --------------------------------------

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}$')
@@ -384,12 +382,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()
@@ -407,13 +403,13 @@ class CitusHandler(Thread):
             self._schedule_load_pg_dist_group = False
 
         try:
-            cursor = self.query('SELECT groupid, nodename, nodeport, noderole, nodeid FROM pg_catalog.pg_dist_node')
+            rows = self.query('SELECT groupid, nodename, nodeport, noderole, nodeid FROM pg_catalog.pg_dist_node')
         except Exception:
             return False
 
         pg_dist_group: Dict[int, PgDistTask] = {}
 
-        for row in cursor:
+        for row in rows:
             if row[0] not in pg_dist_group:
                 pg_dist_group[row[0]] = PgDistTask(row[0], nodes=set(), event='after_promote')
             pg_dist_group[row[0]].add(PgDistNode(*row[1:]))
@@ -491,10 +487,8 @@ class CitusHandler(Thread):
             self.query('SELECT pg_catalog.citus_update_node(%s, %s, %s, true, %s)',
                        node.nodeid, host, node.port, cooldown)
         elif node.role != 'demoted':
-            row = self.query("SELECT pg_catalog.citus_add_node(%s, %s, %s, %s, 'default')",
-                             node.host, node.port, group, node.role).fetchone()
-            if row is not None:
-                node.nodeid = row[0]
+            node.nodeid = self.query("SELECT pg_catalog.citus_add_node(%s, %s, %s, %s, 'default')",
+                                     node.host, node.port, group, node.role)[0][0]
 
     def update_group(self, task: PgDistTask, transaction: bool) -> None:
         current_state = self._in_flight\

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 '):
@@ -160,6 +159,7 @@ class MockCursor(object):
                             (1, '127.0.0.1', 5438, 'secondary', 5)]
         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)