TJWaterServerBinary/app/infra/db/timescaledb/lifecycle.py

from __future__ import annotations

from contextlib import contextmanager
import re
from threading import RLock
from typing import Iterator

from psycopg import Connection, sql
from psycopg.rows import dict_row
from psycopg_pool import ConnectionPool

from app.core.config import get_timescaledb_pgconn_string, settings

from .sync_pool import close_timescale_pool


_DATABASE_NAME = re.compile(r"^[a-z][a-z0-9_]{0,49}$")
_SERVER_DATABASES = frozenset({"template0", "template1", "postgres"})
_admin_pool: ConnectionPool | None = None
_admin_conninfo: str | None = None
_lock = RLock()


def validate_timescale_database_name(name: str, *, allow_template: bool = False) -> str:
    if not _DATABASE_NAME.fullmatch(name):
        raise ValueError(
            "TimescaleDB database name must start with a lowercase letter and "
            "contain only lowercase letters, digits, and underscores"
        )
    protected = {*_SERVER_DATABASES, settings.TIMESCALEDB_SCHEMA_TEMPLATE_DB_NAME}
    if name in protected and not (
        allow_template and name == settings.TIMESCALEDB_SCHEMA_TEMPLATE_DB_NAME
    ):
        raise ValueError(f"TimescaleDB database {name!r} is protected")
    return name


def _get_admin_pool() -> ConnectionPool:
    global _admin_pool, _admin_conninfo
    conninfo = get_timescaledb_pgconn_string(db_name="postgres")
    with _lock:
        if (
            _admin_pool is not None
            and not _admin_pool.closed
            and _admin_conninfo == conninfo
        ):
            return _admin_pool
        if _admin_pool is not None and not _admin_pool.closed:
            _admin_pool.close()
        _admin_pool = ConnectionPool(
            conninfo=conninfo,
            min_size=0,
            max_size=2,
            kwargs={"autocommit": True, "row_factory": dict_row},
            check=ConnectionPool.check_connection,
            open=True,
        )
        _admin_conninfo = conninfo
        return _admin_pool


@contextmanager
def timescale_admin_connection() -> Iterator[Connection]:
    with _get_admin_pool().connection() as conn:
        yield conn


def timescale_database_exists(name: str) -> bool:
    with timescale_admin_connection() as conn, conn.cursor() as cur:
        cur.execute("select 1 from pg_database where datname = %s", (name,))
        return cur.fetchone() is not None


def require_timescale_schema_template() -> str:
    template = settings.TIMESCALEDB_SCHEMA_TEMPLATE_DB_NAME
    validate_timescale_database_name(template, allow_template=True)
    if not timescale_database_exists(template):
        raise RuntimeError(
            f"TimescaleDB schema template {template!r} does not exist"
        )
    return template


def create_timescale_database(name: str) -> None:
    validate_timescale_database_name(name)
    template = require_timescale_schema_template()
    close_timescale_pool(name)
    with timescale_admin_connection() as conn, conn.cursor() as cur:
        cur.execute(
            "select pg_advisory_lock(hashtextextended(%s, 0))",
            (f"tjwater:timescaledb:{name}",),
        )
        try:
            cur.execute("select 1 from pg_database where datname = %s", (name,))
            if cur.fetchone() is not None:
                raise ValueError(f"TimescaleDB database {name!r} already exists")
            cur.execute(
                "select datallowconn from pg_database where datname = %s",
                (template,),
            )
            row = cur.fetchone()
            if row is None:
                raise RuntimeError(
                    f"TimescaleDB schema template {template!r} does not exist"
                )
            template_allowed = bool(row["datallowconn"])
            if template_allowed:
                cur.execute(
                    "update pg_database set datallowconn = false where datname = %s",
                    (template,),
                )
            try:
                cur.execute(
                    "select pg_terminate_backend(pid) from pg_stat_activity "
                    "where datname = %s and pid <> pg_backend_pid()",
                    (template,),
                )
                cur.execute(
                    sql.SQL("create database {} with template = {}").format(
                        sql.Identifier(name),
                        sql.Identifier(template),
                    )
                )
            finally:
                if template_allowed:
                    cur.execute(
                        "update pg_database set datallowconn = true where datname = %s",
                        (template,),
                    )
        finally:
            cur.execute(
                "select pg_advisory_unlock(hashtextextended(%s, 0))",
                (f"tjwater:timescaledb:{name}",),
            )


def delete_timescale_database(name: str) -> None:
    validate_timescale_database_name(name)
    close_timescale_pool(name)
    with timescale_admin_connection() as conn, conn.cursor() as cur:
        cur.execute(
            "select pg_advisory_lock(hashtextextended(%s, 0))",
            (f"tjwater:timescaledb:{name}",),
        )
        try:
            cur.execute("select 1 from pg_database where datname = %s", (name,))
            if cur.fetchone() is None:
                return
            cur.execute(
                "update pg_database set datallowconn = false where datname = %s",
                (name,),
            )
            cur.execute(
                "select pg_terminate_backend(pid) from pg_stat_activity "
                "where datname = %s and pid <> pg_backend_pid()",
                (name,),
            )
            cur.execute(
                sql.SQL("drop database {}").format(sql.Identifier(name))
            )
        finally:
            cur.execute(
                "select pg_advisory_unlock(hashtextextended(%s, 0))",
                (f"tjwater:timescaledb:{name}",),
            )