From 8acefefc4238b9d32cda2e8670ba53e212837c7c Mon Sep 17 00:00:00 2001 From: zhjwpku Date: Fri, 29 Dec 2023 16:01:46 +0800 Subject: [PATCH] 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 --- patroni/postgresql/mpp/citus.py | 15 +++++---------- patroni/psycopg.py | 5 ++++- tests/__init__.py | 2 ++ tests/test_citus.py | 8 ++++++++ 4 files changed, 19 insertions(+), 11 deletions(-) diff --git a/patroni/postgresql/mpp/citus.py b/patroni/postgresql/mpp/citus.py index b8c205ce8..f3d6394c3 100644 --- a/patroni/postgresql/mpp/citus.py +++ b/patroni/postgresql/mpp/citus.py @@ -8,7 +8,7 @@ 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 @@ def bootstrap(self) -> None: 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() diff --git a/patroni/psycopg.py b/patroni/psycopg.py index 4a92047ca..5d47ad5c5 100644 --- a/patroni/psycopg.py +++ b/patroni/psycopg.py @@ -9,7 +9,8 @@ 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 @@ 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 @@ def quote_literal(value: Any, conn: Optional[Any] = None) -> str: 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``. diff --git a/tests/__init__.py b/tests/__init__.py index 986bd88b0..2f3730f69 100644 --- a/tests/__init__.py +++ b/tests/__init__.py @@ -129,6 +129,8 @@ def execute(self, sql, *params): 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'): diff --git a/tests/test_citus.py b/tests/test_citus.py index dbf6d9cf6..f1d8a020f 100644 --- a/tests/test_citus.py +++ b/tests/test_citus.py @@ -157,3 +157,11 @@ def test_ignore_replication_slot(self): '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'))