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)