From dd548c49645c78d0b85cc723bf3bec500528c3da Mon Sep 17 00:00:00 2001 From: Alexander Kukushkin Date: Thu, 21 Dec 2023 09:25:51 +0100 Subject: [PATCH] Create citus database and extension idempotently (#2990) Consider a task: we want to create an extension _before_ citus in a database. Currently `post_bootstrab` script is executed before `CitusHandler.bootstrap()` method, which seems to allow doing that, but in fact `CitusHandler.bootstrap()` will fail to create already existing database and as a result the whole bootstrap will fail. Changing the order of execution of `post_bootstrab` hook and `CitusHandler.bootstrap()` seems to be useless, because it will not allow creating another extension _before_ citus. Therefore the only way of solving it is making CREATE DATABASE and CREATE EXTENSION idempotent. It will allow to create citus database and all dependencies from the `post_bootstrab` hook. --- patroni/postgresql/mpp/citus.py | 15 +++++++++++---- 1 file changed, 11 insertions(+), 4 deletions(-) diff --git a/patroni/postgresql/mpp/citus.py b/patroni/postgresql/mpp/citus.py index a17d1435..b8c205ce 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 +from ...psycopg import connect, quote_ident, quote_literal from ...utils import parse_int if TYPE_CHECKING: # pragma: no cover @@ -389,9 +389,16 @@ 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('CREATE DATABASE {0}'.format( - quote_ident(self._config['database'], conn)).encode('utf-8')) + cur.execute(sql.encode('utf-8')) finally: conn.close() @@ -399,7 +406,7 @@ class CitusHandler(Citus, AbstractMPPHandler, 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}