healthcheck.py

"""Container HEALTHCHECK: probe the Flight server with a trivial query."""
from __future__ import annotations

import os
import sys

import pyarrow.flight as flight


def main() -> int:
    host = os.environ.get("FLIGHT_HOST", "127.0.0.1")
    if host == "0.0.0.0":
        host = "127.0.0.1"
    port = os.environ.get("FLIGHT_PORT", "50051")
    try:
        client = flight.FlightClient(f"grpc://{host}:{port}")
        descriptor = flight.FlightDescriptor.for_command(b"RETURN 1 AS ok")
        info = client.get_flight_info(descriptor, options=flight.FlightCallOptions(timeout=4))
        reader = client.do_get(
            info.endpoints[0].ticket, options=flight.FlightCallOptions(timeout=4)
        )
        table = reader.read_all()
        return 0 if table.num_rows >= 1 else 1
    except Exception as e:  # noqa: BLE001 -- healthcheck must not crash, just fail
        print(f"unhealthy: {e}", file=sys.stderr)
        return 1


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