From d00f5a645b235a52c74212ed7afa816ee40a8975 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/citus.py | 15 +++++++++++---- 1 file changed, 11 insertions(+), 4 deletions(-) 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}