first

domain-api/deploy/multi-region/prune_cluster_nodes.py

@@ -0,0 +1,82 @@
+#!/usr/bin/env python3
+from __future__ import annotations
+
+import argparse
+import sys
+from datetime import datetime, timedelta
+from pathlib import Path
+
+API_ROOT = Path(__file__).resolve().parents[2]
+if str(API_ROOT) not in sys.path:
+    sys.path.insert(0, str(API_ROOT))
+
+from app.core.db import get_db
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Prune stale/offline cluster node records from detect_worker_nodes.")
+    parser.add_argument("--minutes", type=int, default=30, help="delete nodes whose heartbeat is older than this many minutes")
+    parser.add_argument("--node-code", action="append", default=[], help="delete a specific node code; can be passed multiple times")
+    parser.add_argument("--dry-run", action="store_true")
+    args = parser.parse_args()
+
+    cutoff = datetime.now() - timedelta(minutes=max(1, int(args.minutes or 30)))
+    removed: list[dict] = []
+    with get_db() as conn:
+        with conn.cursor() as cur:
+            if args.node_code:
+                cur.execute(
+                    """
+                    SELECT node_code, region, role, status, last_heartbeat_at
+                    FROM detect_worker_nodes
+                    WHERE node_code = ANY(%s)
+                    ORDER BY node_code ASC
+                    """,
+                    (list(args.node_code),),
+                )
+            else:
+                cur.execute(
+                    """
+                    SELECT node_code, region, role, status, last_heartbeat_at
+                    FROM detect_worker_nodes
+                    WHERE last_heartbeat_at < %s
+                    ORDER BY last_heartbeat_at ASC, node_code ASC
+                    """,
+                    (cutoff,),
+                )
+            rows = cur.fetchall()
+            for row in rows:
+                removed.append(
+                    {
+                        "node_code": row[0],
+                        "region": row[1],
+                        "role": row[2],
+                        "status": row[3],
+                        "last_heartbeat_at": row[4].isoformat(sep=" ", timespec="seconds") if row[4] else "",
+                    }
+                )
+            if removed and not args.dry_run:
+                cur.execute(
+                    "DELETE FROM detect_worker_nodes WHERE node_code = ANY(%s)",
+                    ([item["node_code"] for item in removed],),
+                )
+        if removed and not args.dry_run:
+            conn.commit()
+
+    if not removed:
+        print("no cluster nodes matched prune conditions")
+        return
+
+    print("matched cluster nodes:")
+    for item in removed:
+        print(
+            f"- {item['node_code']} | {item['region']} | {item['role']} | {item['status']} | {item['last_heartbeat_at']}"
+        )
+    if args.dry_run:
+        print("dry-run only, nothing deleted")
+    else:
+        print(f"deleted {len(removed)} cluster node record(s)")
+
+
+if __name__ == "__main__":
+    main()