managed-containerd/server.py

#!/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()