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.
This commit is contained in:
Alexander Kukushkin
2023-12-21 09:25:51 +01:00
committed by GitHub
parent bcfd8438a5
commit dd548c4964
+11 -4
View File
@@ -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}