diff --git a/patroni/postgresql/citus.py b/patroni/postgresql/citus.py index b50dc1d0..944a8945 100644 --- a/patroni/postgresql/citus.py +++ b/patroni/postgresql/citus.py @@ -7,7 +7,7 @@ from urllib.parse import urlparse from typing import Any, Dict, List, Optional, Union, Tuple, TYPE_CHECKING from ..dcs import CITUS_COORDINATOR_GROUP_ID, Cluster -from ..psycopg import connect, quote_ident +from ..psycopg import connect, quote_ident, quote_literal if TYPE_CHECKING: # pragma: no cover from . import Postgresql @@ -361,9 +361,16 @@ class CitusHandler(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('CREATE DATABASE {0}'.format( - quote_ident(self._config['database'], conn)).encode('utf-8')) + cur.execute(sql.encode('utf-8')) finally: conn.close() @@ -371,7 +378,7 @@ class CitusHandler(Thread): conn = connect(**conn_kwargs) try: with conn.cursor() as cur: - cur.execute('CREATE EXTENSION citus') + cur.execute('CREATE EXTENSION IF NOT EXISTS citus') superuser = self._postgresql.config.superuser params = {k: superuser[k] for k in ('password', 'sslcert', 'sslkey') if k in superuser}