Fix Citus bootstrap - CREATE DATABASE cannot be executed from a function (#2994)

This was introduced by #2990: pod cannot be started and show the
following logs:

```
2023-12-26 03:29:25.569 UTC [47] CONTEXT:  SQL statement "CREATE DATABASE "citus""
        PL/pgSQL function inline_code_block line 5 at SQL statement
2023-12-26 03:29:25.569 UTC [47] STATEMENT:  DO $$
        BEGIN
            PERFORM * FROM pg_catalog.pg_database WHERE datname = 'citus';
            IF NOT FOUND THEN
                CREATE DATABASE "citus";
            END IF;
        END;$$
2023-12-26 03:29:25,570 ERROR: post_bootstrap
Traceback (most recent call last):
  File "/usr/local/lib/python3.11/dist-packages/patroni/postgresql/bootstrap.py", line 474, in post_bootstrap
    self._postgresql.citus_handler.bootstrap()
  File "/usr/local/lib/python3.11/dist-packages/patroni/postgresql/mpp/citus.py", line 401, in bootstrap
    cur.execute(sql.encode('utf-8'))
psycopg2.errors.ActiveSqlTransaction: CREATE DATABASE cannot be executed from a function
CONTEXT:  SQL statement "CREATE DATABASE "citus""
PL/pgSQL function inline_code_block line 5 at SQL statement
```
---------

Signed-off-by: Zhao Junwang <[email protected]>

patroni/postgresql/mpp/citus.py

@@ -8,7 +8,7 @@ from typing import Any, Dict, List, Optional, Union, Tuple, TYPE_CHECKING
 
 from . import AbstractMPP, AbstractMPPHandler
 from ...dcs import Cluster
-from ...psycopg import connect, quote_ident, quote_literal
+from ...psycopg import connect, quote_ident, DuplicateDatabase
 from ...utils import parse_int
 
 if TYPE_CHECKING:  # pragma: no cover
@@ -389,16 +389,11 @@ class CitusHandler(Citus, AbstractMPPHandler, Thread):
         if self._config['database'] != self._postgresql.database:
             conn = connect(**conn_kwargs)
             try:
-                database = self._config['database']
-                sql = """DO $$
-BEGIN
-    PERFORM * FROM pg_catalog.pg_database WHERE datname = {0};
-    IF NOT FOUND THEN
-        CREATE DATABASE {1};
-    END IF;
-END;$$""".format(quote_literal(database), quote_ident(database, conn))
                 with conn.cursor() as cur:
-                    cur.execute(sql.encode('utf-8'))
+                    cur.execute('CREATE DATABASE {0}'.format(
+                        quote_ident(self._config['database'], conn)).encode('utf-8'))
+            except DuplicateDatabase as e:
+                logger.debug('Exception when creating database: %r', e)
             finally:
                 conn.close()
 

patroni/psycopg.py

@@ -9,7 +9,8 @@ if TYPE_CHECKING:  # pragma: no cover
     from psycopg import Connection
     from psycopg2 import connection, cursor
 
-__all__ = ['connect', 'quote_ident', 'quote_literal', 'DatabaseError', 'Error', 'OperationalError', 'ProgrammingError']
+__all__ = ['connect', 'quote_ident', 'quote_literal', 'DatabaseError', 'Error', 'OperationalError', 'ProgrammingError',
+           'DuplicateDatabase']
 
 _legacy = False
 try:
@@ -18,6 +19,7 @@ try:
     if parse_version(__version__) < MIN_PSYCOPG2:
         raise ImportError
     from psycopg2 import connect as _connect, Error, DatabaseError, OperationalError, ProgrammingError
+    from psycopg2.errors import DuplicateDatabase
     from psycopg2.extensions import adapt
 
     try:
@@ -43,6 +45,7 @@ try:
         return value.getquoted().decode('utf-8')
 except ImportError:
     from psycopg import connect as __connect, sql, Error, DatabaseError, OperationalError, ProgrammingError
+    from psycopg.errors import DuplicateDatabase
 
     def _connect(dsn: Optional[str] = None, **kwargs: Any) -> 'Connection[Any]':
         """Call :func:`psycopg.connect` with *dsn* and ``**kwargs``.

tests/__init__.py

@@ -129,6 +129,8 @@ class MockCursor(object):
             sql = sql.decode('utf-8')
         if sql.startswith('blabla'):
             raise psycopg.ProgrammingError()
+        if sql.startswith('CREATE DATABASE'):
+            raise psycopg.DuplicateDatabase()
         elif sql == 'CHECKPOINT' or sql.startswith('SELECT pg_catalog.pg_create_'):
             raise psycopg.OperationalError()
         elif sql.startswith('RetryFailedError'):

tests/test_citus.py

@@ -157,3 +157,11 @@ class TestCitus(BaseTestPostgresql):
                                                          'type': 'logical', 'database': 'citus', 'plugin': 'pgoutput'}))
         self.assertTrue(self.c.ignore_replication_slot({'name': 'citus_shard_split_slot_1_2_3',
                                                         'type': 'logical', 'database': 'citus', 'plugin': 'citus'}))
+
+    @patch('patroni.postgresql.mpp.citus.logger.debug')
+    @patch('patroni.postgresql.mpp.citus.connect', psycopg_connect)
+    @patch('patroni.postgresql.mpp.citus.quote_ident', Mock())
+    def test_bootstrap_duplicate_database(self, mock_logger):
+        self.c.bootstrap()
+        mock_logger.assert_called_once()
+        self.assertTrue(mock_logger.call_args[0][0].startswith('Exception when creating database'))