Repository: zalando/patroni
Pull request #2994 (2023-12-26 06:18:08)

CREATE DATABASE cannot be executed from a function
This was introduced by #2990, I came across this when I tried the Citus on K8s demo, 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
```

So I kind of revert part of #2990 and add a warning when Exception.

coveralls: ## Pull Request Test Coverage Report for [Build 7327177771](https://coveralls.io/builds/64758257)


* **2** of **4**   **(50.0%)**  changed or added relevant lines in **1** file are covered.
* No unchanged relevant lines lost coverage.
* Overall coverage decreased (**-0.02%**) to **99.844%**

---

|  Changes Missing Coverage | Covered Lines | Changed/Added Lines | % |
| :-----|--------------|--------|---: |
| [patroni/postgresql/mpp/citus.py](https://coveralls.io/builds/64758257/source?filename=patroni%2Fpostgresql%2Fmpp%2Fcitus.py#L395) | 2 | 4 | 50.0%
<!-- | **Total:** | **2** | **4** | **50.0%** | -->


|  Totals | [![Coverage Status](https://coveralls.io/builds/64758257/badge)](https://coveralls.io/builds/64758257) |
| :-- | --: |
| Change from base [Build 7285720031](https://coveralls.io/builds/64703697): |  -0.02% |
| Covered Lines: | 13455 |
| Relevant Lines: | 13476 |

---
##### 💛  - [Coveralls](https://coveralls.io)

zhjwpku: 


> What is the reason for modifying psycopg2 stubs? Please also add a unit-test.

I was hesitating to add these psycopg2 stubs, just searched `ProgrammingError` and thought may be somewhere need the same for `DuplicateDatabase`, I'll delete them and add a unittest, thanks for the review.
zhjwpku: @CyberDem0n The behave test on macos seems not happy, but I doubt this is a random failure, can you take a look at it? Thanks.
hughcapet: 👍 
CyberDem0n: 👍

Patch:
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``.


Test patch:
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'))