From 366829e3791097a233aa194c51b5f149d0163107 Mon Sep 17 00:00:00 2001 From: Alexander Kukushkin Date: Thu, 17 Aug 2023 15:42:11 +0200 Subject: [PATCH] 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 | 25 +++----- patroni/postgresql/__init__.py | 101 ++++++++++++++++--------------- patroni/postgresql/citus.py | 20 +++--- patroni/postgresql/config.py | 4 +- patroni/postgresql/connection.py | 63 +++++++++++++++---- patroni/postgresql/slots.py | 31 ++++------ tests/__init__.py | 10 +-- tests/test_api.py | 19 +++--- tests/test_postgresql.py | 15 +++-- tests/test_rewind.py | 5 +- tests/test_slots.py | 12 ++-- 11 files changed, 165 insertions(+), 140 deletions(-) diff --git a/patroni/api.py b/patroni/api.py index 18170334..b01c24d8 100644 --- a/patroni/api.py +++ b/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: diff --git a/patroni/postgresql/__init__.py b/patroni/postgresql/__init__.py index af87c3e2..7f3bdccb 100644 --- a/patroni/postgresql/__init__.py +++ b/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 diff --git a/patroni/postgresql/citus.py b/patroni/postgresql/citus.py index a3a8fe45..f659e325 100644 --- a/patroni/postgresql/citus.py +++ b/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. diff --git a/patroni/postgresql/config.py b/patroni/postgresql/config.py index 6d1fe384..51bbf9af 100644 --- a/patroni/postgresql/config.py +++ b/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) diff --git a/patroni/postgresql/connection.py b/patroni/postgresql/connection.py index 277a4889..5bb97a1f 100644 --- a/patroni/postgresql/connection.py +++ b/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 diff --git a/patroni/postgresql/slots.py b/patroni/postgresql/slots.py index 681aed00..319119ce 100644 --- a/patroni/postgresql/slots.py +++ b/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 diff --git a/tests/__init__.py b/tests/__init__.py index 01d74616..79ef0229 100644 --- a/tests/__init__.py +++ b/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] diff --git a/tests/test_api.py b/tests/test_api.py index b5aeff8a..b1e56fdf 100644 --- a/tests/test_api.py +++ b/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()) diff --git a/tests/test_postgresql.py b/tests/test_postgresql.py index 4f1e36da..510c2a90 100644 --- a/tests/test_postgresql.py +++ b/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() diff --git a/tests/test_rewind.py b/tests/test_rewind.py index f40c46af..a54c27e9 100644 --- a/tests/test_rewind.py +++ b/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) diff --git a/tests/test_slots.py b/tests/test_slots.py index 0210d701..df8c6955 100644 --- a/tests/test_slots.py +++ b/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)