feat(projects): automate project infrastructure provisioning

app/infra/db/timescaledb/lifecycle.py

@@ -0,0 +1,165 @@
+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}",),
+            )