getDomain/tools/quiet_window_rebuild_release_index.py

#!/usr/bin/env python3
"""Quiet-window helper for rebuilding idx_detect_job_items_release_node_job.

Stops node-agent and worker services, keeps clearing blockers for the release-index
DROP INDEX backend until it disappears, waits briefly for CREATE INDEX to appear,
then restores services.
"""

from __future__ import annotations

import os
import subprocess
import sys
import time
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
DOMAINCHECK_ROOT = ROOT / "domainCheck"
if str(DOMAINCHECK_ROOT) not in sys.path:
    sys.path.insert(0, str(DOMAINCHECK_ROOT))

from app.utils.database import Database  # noqa: E402


ENV_FILE = Path("/etc/default/domaincheck-worker")
INDEX_NAME = "idx_detect_job_items_release_node_job"
POLL_SECONDS = 5
MAX_DROP_LOOPS = 72  # about 6 minutes
MAX_RESTORE_LOOPS = 24


def load_env(path: Path) -> None:
    if not path.exists():
        return
    for raw in path.read_text(encoding="utf-8").splitlines():
        line = raw.strip()
        if not line or line.startswith("#") or "=" not in line:
            continue
        key, value = line.split("=", 1)
        os.environ[key] = value


def log(message: str) -> None:
    stamp = time.strftime("%Y-%m-%d %H:%M:%S")
    print(f"[{stamp}] {message}", flush=True)


def run(cmd: str) -> tuple[int, str, str]:
    proc = subprocess.run(
        cmd,
        shell=True,
        text=True,
        capture_output=True,
        check=False,
    )
    return proc.returncode, proc.stdout.strip(), proc.stderr.strip()


def service_state(name: str) -> str:
    _, out, _ = run(f"systemctl is-active {name} || true")
    return out.strip() or "unknown"


def templated_running_count() -> int:
    _, out, _ = run(
        "systemctl list-units --type=service --state=running 'domaincheck-worker@*.service' --no-legend | wc -l"
    )
    try:
        return int((out or "0").strip())
    except Exception:
        return -1


def stop_services() -> None:
    for cmd in [
        "systemctl stop domaincheck-node-agent.service",
        "systemctl stop 'domaincheck-worker@*.service'",
        "systemctl stop domaincheck-worker.service",
    ]:
        code, out, err = run(cmd)
        log(f"stop cmd={cmd!r} code={code} out={out!r} err={err!r}")


def start_services() -> None:
    for cmd in [
        "systemctl start domaincheck-node-agent.service",
        "systemctl start domaincheck-worker.service",
        "systemctl start --all 'domaincheck-worker@*.service'",
    ]:
        code, out, err = run(cmd)
        log(f"start cmd={cmd!r} code={code} out={out!r} err={err!r}")


def wait_for_quiet() -> None:
    for loop in range(MAX_RESTORE_LOOPS):
        templated = templated_running_count()
        base = service_state("domaincheck-worker.service")
        node_agent = service_state("domaincheck-node-agent.service")
        log(
            f"quiet_check loop={loop} templated={templated} "
            f"base={base} node_agent={node_agent}"
        )
        if templated == 0 and base != "active" and node_agent != "active":
            return
        time.sleep(POLL_SECONDS)


def wait_for_restore() -> None:
    for loop in range(MAX_RESTORE_LOOPS):
        templated = templated_running_count()
        base = service_state("domaincheck-worker.service")
        node_agent = service_state("domaincheck-node-agent.service")
        log(
            f"restore_check loop={loop} templated={templated} "
            f"base={base} node_agent={node_agent}"
        )
        if templated >= 190 and base == "active" and node_agent == "active":
            return
        time.sleep(POLL_SECONDS)


def fetch_drop_pids(cur) -> list[int]:
    cur.execute(
        """
        SELECT pid
        FROM pg_stat_activity
        WHERE query LIKE %s
        ORDER BY pid
        """,
        (f"DROP INDEX CONCURRENTLY IF EXISTS {INDEX_NAME}%",),
    )
    return [int(row[0]) for row in cur.fetchall()]


def fetch_create_progress(cur) -> list[tuple]:
    cur.execute(
        """
        SELECT pid, phase, blocks_done, blocks_total, current_locker_pid
        FROM pg_stat_progress_create_index
        WHERE index_relid = (SELECT oid FROM pg_class WHERE relname = %s)
        ORDER BY pid
        """,
        (INDEX_NAME,),
    )
    return list(cur.fetchall())


def fetch_index_state(cur) -> str:
    cur.execute(
        """
        SELECT
            c.relname,
            i.indisvalid,
            i.indisready
        FROM pg_class c
        JOIN pg_index i ON i.indexrelid = c.oid
        WHERE c.relname = %s
        """,
        (INDEX_NAME,),
    )
    row = cur.fetchone()
    if not row:
        return "missing"
    return f"{row[0]}|{int(bool(row[1]))}|{int(bool(row[2]))}"


def main() -> int:
    load_env(ENV_FILE)
    db = Database()
    conn = db.get_connection()
    if not conn:
        log("db_connect_failed")
        return 1
    conn.autocommit = True
    cur = conn.cursor()
    try:
        log(f"quiet-window start index={INDEX_NAME}")
        stop_services()
        wait_for_quiet()

        for loop in range(MAX_DROP_LOOPS):
            drop_pids = fetch_drop_pids(cur)
            log(f"drop_loop={loop} drop_pids={drop_pids or ['none']}")
            if not drop_pids:
                break
            for pid in drop_pids:
                cur.execute("SELECT pg_blocking_pids(%s)", (pid,))
                row = cur.fetchone()
                blockers = [int(v) for v in (row[0] or [])] if row and row[0] else []
                log(f"drop_pid={pid} blockers={blockers or ['none']}")
                for blocker in blockers:
                    cur.execute("SELECT pg_terminate_backend(%s)", (blocker,))
                    result = cur.fetchone()
                    log(f"terminate blocker={blocker} result={result!r}")
            time.sleep(POLL_SECONDS)

        for loop in range(12):
            progress = fetch_create_progress(cur)
            state = fetch_index_state(cur)
            log(f"create_wait loop={loop} index_state={state} progress={progress or ['none']}")
            if progress:
                break
            time.sleep(POLL_SECONDS)
    finally:
        try:
            start_services()
            wait_for_restore()
        finally:
            cur.close()
            db.close(conn=conn)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())