diff --git a/patroni/postgresql/mpp/citus.py b/patroni/postgresql/mpp/citus.py index b8c205ce..f3d6394c 100644 --- a/patroni/postgresql/mpp/citus.py +++ b/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() diff --git a/patroni/psycopg.py b/patroni/psycopg.py index 4a92047c..5d47ad5c 100644 --- a/patroni/psycopg.py +++ b/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``. diff --git a/tests/__init__.py b/tests/__init__.py index 986bd88b..2f3730f6 100644 --- a/tests/__init__.py +++ b/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'): diff --git a/tests/test_citus.py b/tests/test_citus.py index dbf6d9cf..f1d8a020 100644 --- a/tests/test_citus.py +++ b/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'))