diff --git a/patroni/api.py b/patroni/api.py index a7a754a3..4647e7c2 100644 --- a/patroni/api.py +++ b/patroni/api.py @@ -1375,6 +1375,9 @@ class RestApiServer(ThreadingMixIn, HTTPServer, Thread): def query(self, sql: str, *params: Any) -> List[Tuple[Any, ...]]: """Execute *sql* query with *params* and optionally return results. + .. note:: + Prefer to use own connection to postgres and fallback to ``heartbeat`` when own isn't available. + :param sql: the SQL statement to be run. :param params: positional arguments to be used as parameters for *sql*. @@ -1384,10 +1387,21 @@ 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. """ + # We first try to get a heartbeat connection because it is always required for the main thread. try: - return self.patroni.postgresql.query(sql, *params, retry=False) - except RetryFailedError as e: - raise PostgresConnectionException(str(e)) + heartbeat_connection = self.patroni.postgresql.connection_pool.get('heartbeat') + heartbeat_connection.get() # try to open psycopg connection to postgres + except psycopg.Error as exc: + raise PostgresConnectionException('connection problems') from exc + + try: + connection = self.patroni.postgresql.connection_pool.get('restapi') + connection.get() # try to open psycopg connection to postgres + except psycopg.Error: + logger.debug('restapi connection to postgres is not available') + connection = heartbeat_connection + + return connection.query(sql, *params) @staticmethod def _set_fd_cloexec(fd: socket.socket) -> None: diff --git a/tests/test_api.py b/tests/test_api.py index f433ca18..fa9a6280 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -11,9 +11,12 @@ from socketserver import ThreadingMixIn from patroni.api import RestApiHandler, RestApiServer from patroni.config import GlobalConfig from patroni.dcs import ClusterConfig, Member +from patroni.exceptions import PostgresConnectionException from patroni.ha import _MemberStatus +from patroni.psycopg import OperationalError from patroni.utils import RetryFailedError, tzutc +from . import MockConnect, psycopg_connect from .test_ha import get_cluster_initialized_without_leader @@ -21,8 +24,29 @@ future_restart_time = datetime.datetime.now(tzutc) + datetime.timedelta(days=5) postmaster_start_time = datetime.datetime.now(tzutc) -class MockPostgresql(object): +class MockConnection: + @staticmethod + def get(*args): + return psycopg_connect() + + @staticmethod + def query(sql, *params): + 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 MockConnectionPool: + + @staticmethod + def get(*args): + return MockConnection() + + +class MockPostgresql: + + connection_pool = MockConnectionPool() name = 'test' state = 'running' role = 'primary' @@ -54,12 +78,6 @@ 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 @@ -487,7 +505,7 @@ class TestRestApiHandler(unittest.TestCase): @patch('time.sleep', Mock()) def test_RestApiServer_query(self): - with patch.object(MockPostgresql, 'query', Mock(side_effect=RetryFailedError('bla'))): + with patch.object(MockConnection, 'query', Mock(side_effect=RetryFailedError('bla'))): self.assertIsNotNone(MockRestApiServer(RestApiHandler, 'GET /patroni')) @patch('time.sleep', Mock()) @@ -659,3 +677,11 @@ class TestRestApiServer(unittest.TestCase): def test_get_certificate_serial_number(self): self.assertIsNone(self.srv.get_certificate_serial_number()) + + def test_query(self): + with patch.object(MockConnection, 'get', Mock(side_effect=OperationalError)): + self.assertRaises(PostgresConnectionException, self.srv.query, 'SELECT 1') + with patch.object(MockConnection, 'get', Mock(side_effect=[MockConnect(), OperationalError])), \ + patch.object(MockConnection, 'query') as mock_query: + self.srv.query('SELECT 1') + mock_query.assert_called_once_with('SELECT 1')