TJWaterServerBinary/app/native/wndb/core/project_templates.py

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)