Proven working Postgresql monitor

server.py

@@ -0,0 +1,255 @@
+#!/usr/bin/env python3
+
+import json
+import os
+import signal
+import socket
+import traceback
+
+import psycopg2
+import psycopg2.extras
+
+
+SOCKET_PATH = os.environ.get(
+    "MANAGED_CONTAINERD_SOCKET",
+    "/run/managed-containerd/managed-containerd.sock",
+)
+
+PGHOST = os.environ.get("PGHOST", "postgres")
+PGPORT = int(os.environ.get("PGPORT", "5432"))
+PGUSER = os.environ.get("PGUSER", "david")
+PGDATABASE = os.environ.get("PGDATABASE", "david")
+PGPASSWORD_FILE = os.environ.get("PGPASSWORD_FILE")
+
+server_socket = None
+
+
+def read_secret(path):
+    if not path:
+        return None
+
+    with open(path, "r", encoding="utf-8") as f:
+        return f.read().strip()
+
+
+def postgres_status():
+    password = read_secret(PGPASSWORD_FILE)
+
+    conn = psycopg2.connect(
+        host=PGHOST,
+        port=PGPORT,
+        user=PGUSER,
+        dbname=PGDATABASE,
+        password=password,
+        connect_timeout=5,
+        application_name="managed-containerd",
+    )
+
+    try:
+        conn.autocommit = True
+
+        with conn.cursor(
+            cursor_factory=psycopg2.extras.RealDictCursor
+        ) as cur:
+
+            #
+            # Establish role first.
+            #
+            cur.execute("""
+                SELECT pg_is_in_recovery() AS in_recovery
+            """)
+            in_recovery = bool(cur.fetchone()["in_recovery"])
+
+            #
+            # Primary-only state.
+            #
+            if not in_recovery:
+                cur.execute("""
+                    SELECT
+                        pg_current_wal_lsn()::text AS current_wal_lsn
+                """)
+                primary = cur.fetchone()
+            else:
+                primary = None
+
+            #
+            # Recovery state.
+            #
+            if in_recovery:
+                cur.execute("""
+                    SELECT
+                        pg_last_wal_receive_lsn()::text AS receive_lsn,
+                        pg_last_wal_replay_lsn()::text AS replay_lsn,
+                        pg_is_wal_replay_paused() AS replay_paused
+                """)
+                recovery = cur.fetchone()
+
+                cur.execute("""
+                    SELECT
+                        status,
+                        sender_host,
+                        sender_port,
+                        slot_name,
+                        latest_end_lsn::text AS latest_end_lsn
+                    FROM pg_stat_wal_receiver
+                    LIMIT 1
+                """)
+                receiver = cur.fetchone()
+
+            else:
+                recovery = {
+                    "receive_lsn": None,
+                    "replay_lsn": None,
+                    "replay_paused": False,
+                }
+                receiver = None
+
+            #
+            # General server state.
+            #
+            cur.execute("""
+                SELECT
+                    current_setting('server_version') AS server_version,
+                    current_setting('cluster_name', true) AS cluster_name,
+                    current_user AS current_user,
+                    current_database() AS current_database,
+                    inet_server_addr()::text AS server_address,
+                    inet_server_port() AS server_port
+            """)
+            server = cur.fetchone()
+
+        if in_recovery:
+            role = "standby"
+            streaming = bool(
+                receiver and receiver.get("status") == "streaming"
+            )
+        else:
+            role = "primary"
+            streaming = False
+
+        return {
+            "ok": True,
+            "application": "postgresql",
+            "role": role,
+            "in_recovery": in_recovery,
+            "streaming": streaming,
+            "wal_receiver": dict(receiver) if receiver else None,
+            "recovery": dict(recovery),
+            "primary": dict(primary) if primary else None,
+            "server": dict(server),
+        }
+
+    finally:
+        conn.close()
+
+
+def handle_request(req):
+    action = req.get("action")
+
+    if action == "ping":
+        return {
+            "ok": True,
+            "daemon": "managed-containerd",
+            "version": 1,
+        }
+
+    if action == "postgresql.status":
+        return postgres_status()
+
+    return {
+        "ok": False,
+        "error": f"unknown action: {action}",
+    }
+
+
+def cleanup_socket():
+    try:
+        os.unlink(SOCKET_PATH)
+    except FileNotFoundError:
+        pass
+
+
+def shutdown(signum, frame):
+    global server_socket
+
+    print(
+        f"managed-containerd received signal {signum}; shutting down",
+        flush=True,
+    )
+
+    if server_socket is not None:
+        try:
+            server_socket.close()
+        except OSError:
+            pass
+
+    cleanup_socket()
+    raise SystemExit(0)
+
+
+def serve():
+    global server_socket
+
+    os.makedirs(os.path.dirname(SOCKET_PATH), exist_ok=True)
+    cleanup_socket()
+
+    signal.signal(signal.SIGTERM, shutdown)
+    signal.signal(signal.SIGINT, shutdown)
+
+    server_socket = socket.socket(
+        socket.AF_UNIX,
+        socket.SOCK_STREAM,
+    )
+
+    server_socket.bind(SOCKET_PATH)
+    os.chmod(SOCKET_PATH, 0o660)
+    server_socket.listen(16)
+
+    print(
+        f"managed-containerd listening on {SOCKET_PATH}",
+        flush=True,
+    )
+
+    try:
+        while True:
+            conn, _ = server_socket.accept()
+
+            with conn:
+                try:
+                    raw = b""
+
+                    while b"\n" not in raw:
+                        chunk = conn.recv(65536)
+
+                        if not chunk:
+                            break
+
+                        raw += chunk
+
+                    req = json.loads(
+                        raw.decode("utf-8").strip()
+                    )
+
+                    response = handle_request(req)
+
+                except Exception as exc:
+                    traceback.print_exc()
+
+                    response = {
+                        "ok": False,
+                        "error": str(exc),
+                    }
+
+                conn.sendall(
+                    json.dumps(
+                        response,
+                        default=str,
+                    ).encode("utf-8") + b"\n"
+                )
+
+    finally:
+        cleanup_socket()
+
+
+if __name__ == "__main__":
+    serve()