feat(projects): automate project infrastructure provisioning

app/native/wndb/core/project_templates.py

@@ -0,0 +1,191 @@
+from __future__ import annotations
+
+import time
+
+from psycopg import sql
+from psycopg.conninfo import make_conninfo
+
+from app.core.config import settings
+from app.infra.db.project_routing import get_project_template_database_name
+
+from .connection import admin_connection, close_project_pool, project_connection
+from .model_replace import copy_network_tables
+from .projects import (
+    _ensure_project_model_template_ready,
+    copy_project,
+    delete_project,
+    have_project,
+)
+
+
+PUBLICATION_NAME = "wndb_network_pub"
+SUBSCRIPTION_NAME = "wndb_network_sub"
+
+
+def _slot_name(project: str) -> str:
+    return f"{project}_network_slot"
+
+
+def _publisher_conninfo(project: str) -> str:
+    return make_conninfo(
+        dbname=project,
+        host=settings.DB_HOST,
+        port=settings.DB_PORT,
+        user=settings.DB_USER,
+        password=settings.DB_PASSWORD,
+    )
+
+
+def ensure_replication_worker_capacity() -> None:
+    """Keep one worker slot spare while admitting one new project subscription."""
+    with admin_connection() as conn, conn.cursor() as cur:
+        cur.execute(
+            """
+            select current_setting('max_worker_processes')::integer as maximum,
+                   count(*) filter (
+                       where backend_type in (
+                           'logical replication launcher',
+                           'logical replication worker'
+                       )
+                   ) as used
+            from pg_stat_activity
+            """
+        )
+        row = cur.fetchone()
+    maximum = int(row["maximum"])
+    used = int(row["used"])
+    if used + 1 >= maximum:
+        raise RuntimeError(
+            "PostgreSQL has no reserved logical-replication worker capacity "
+            f"(used={used}, max_worker_processes={maximum}); increase "
+            "max_worker_processes before provisioning another project"
+        )
+
+
+def _create_publication_and_slot(project: str) -> None:
+    slot_name = _slot_name(project)
+    with project_connection(project) as conn, conn.cursor() as cur:
+        cur.execute(
+            "select n.nspname as schema_name, c.relname as table_name "
+            "from pg_class c join pg_namespace n on n.oid = c.relnamespace "
+            "where n.nspname = 'network' and c.relkind in ('r', 'p') "
+            "order by c.relname"
+        )
+        tables = [(row["schema_name"], row["table_name"]) for row in cur.fetchall()]
+        if not tables:
+            raise RuntimeError("Project database has no network tables to publish")
+        cur.execute(
+            sql.SQL("create publication {} for table {}").format(
+                sql.Identifier(PUBLICATION_NAME),
+                sql.SQL(", ").join(
+                    sql.Identifier(schema_name, table_name)
+                    for schema_name, table_name in tables
+                ),
+            )
+        )
+        cur.execute(
+            "select slot_name from pg_create_logical_replication_slot(%s, 'pgoutput')",
+            (slot_name,),
+        )
+
+
+def _create_subscription(project: str, template: str) -> None:
+    with project_connection(template) as conn, conn.cursor() as cur:
+        cur.execute(
+            sql.SQL(
+                "create subscription {} connection {} publication {} "
+                "with (create_slot = false, copy_data = false, slot_name = {})"
+            ).format(
+                sql.Identifier(SUBSCRIPTION_NAME),
+                sql.Literal(_publisher_conninfo(project)),
+                sql.Identifier(PUBLICATION_NAME),
+                sql.Literal(_slot_name(project)),
+            )
+        )
+
+
+def _wait_until_ready(template: str) -> None:
+    last_error: RuntimeError | None = None
+    for _ in range(50):
+        try:
+            _ensure_project_model_template_ready(template)
+            return
+        except RuntimeError as exc:
+            last_error = exc
+            time.sleep(0.1)
+    raise last_error or RuntimeError(
+        f"Project model template {template!r} did not become ready"
+    )
+
+
+def create_project_model_template(project: str) -> str:
+    template = get_project_template_database_name(project)
+    if have_project(template):
+        raise ValueError(f"Project model template {template!r} already exists")
+    template_created = False
+    try:
+        _create_publication_and_slot(project)
+        copy_project(
+            settings.WNDB_SCHEMA_TEMPLATE_DB_NAME,
+            template,
+            allow_template_source=True,
+            allow_template_target=True,
+        )
+        template_created = True
+        copy_network_tables(project, template)
+        _create_subscription(project, template)
+        _wait_until_ready(template)
+        return template
+    except Exception:
+        if template_created:
+            delete_project_model_template(project)
+        else:
+            _drop_publisher_objects(project)
+        raise
+
+
+def _drop_subscription(template: str) -> None:
+    if not have_project(template):
+        return
+    try:
+        with project_connection(template) as conn, conn.cursor() as cur:
+            cur.execute(
+                sql.SQL("drop subscription if exists {}").format(
+                    sql.Identifier(SUBSCRIPTION_NAME)
+                )
+            )
+    finally:
+        close_project_pool(template)
+
+
+def _drop_publisher_objects(
+    project: str,
+    *,
+    drop_publication: bool = True,
+    drop_slot: bool = True,
+) -> None:
+    if not have_project(project):
+        return
+    with project_connection(project) as conn, conn.cursor() as cur:
+        if drop_slot:
+            cur.execute(
+                "select pg_drop_replication_slot(slot_name) "
+                "from pg_replication_slots where slot_name = %s and active = false",
+                (_slot_name(project),),
+            )
+        if drop_publication:
+            cur.execute(
+                sql.SQL("drop publication if exists {}").format(
+                    sql.Identifier(PUBLICATION_NAME)
+                )
+            )
+
+
+def delete_project_model_template(project: str) -> None:
+    template = get_project_template_database_name(project)
+    try:
+        _drop_subscription(template)
+    finally:
+        _drop_publisher_objects(project)
+        if have_project(template):
+            delete_project(template, allow_template=True)