From 0ab5b49757acea743be033a861dbd81a12e13304 Mon Sep 17 00:00:00 2001
From: Alexander Kukushkin <cyberdemn@gmail.com>
Date: Tue, 5 Sep 2023 07:26:44 +0200
Subject: [PATCH] Introduce a dedicated postgres connection for REST API
 (#2833)

Sharing a single connection between REST API and the main thread (doing heartbeats) was working mostly fine, except when Postgres becomes so slow that REST API queries start blocking the main loop.

If the dedicated REST API connection isn't available we use the heartbeat connection as a fallback.
---
 patroni/api.py    | 20 +++++++++++++++++---
 tests/test_api.py | 42 ++++++++++++++++++++++++++++++++++--------
 2 files changed, 51 insertions(+), 11 deletions(-)

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')