feat: Add flagsmith healthcheck command (#60)

Author: khvn26

poetry.lock

@@ -16,6 +16,7 @@ dependencies = [
     "gunicorn (>=19.1)",
     "prometheus-client (>=0.0.16)",
     "psycopg2-binary (>=2.9,<3)",
+    "requests",
     "simplejson (>=3,<4)",
 ]
 optional-dependencies = { test-tools = ["pyfakefs (>=5,<6)"] }
@@ -71,8 +72,8 @@ pytest-asyncio = "^0.25.3"
 pytest-cov = "^6.0.0"
 pytest-django = "^4.10.0"
 pytest-freezegun = "^0.4.2"
+pytest-httpserver = "^1.1.3"
 pytest-mock = "^3.14.0"
-requests = "^2.32.3"
 ruff = "*"
 setuptools = "^77.0.3"
 types-simplejson = "^3.20.0.20250326"

pyproject.toml

@@ -0,0 +1,120 @@
+import argparse
+import socket
+import urllib.parse
+
+import requests
+
+DEFAULT_PORT = 8000
+DEFAULT_TIMEOUT_SECONDS = 1
+
+
+def get_args(
+    argv: list[str],
+    *,
+    prog: str,
+) -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description=(
+            "Perform health checks. "
+            f"If ran without subcommand, defaults to a TCP check of port {DEFAULT_PORT}."
+        ),
+        prog=prog,
+    )
+    subcommands = parser.add_subparsers(dest="subcommand")
+    tcp_parser = subcommands.add_parser(
+        "tcp",
+        help="Check if the API is able to accept local TCP connections",
+    )
+    tcp_parser.add_argument(
+        "--port",
+        "-p",
+        type=int,
+        default=DEFAULT_PORT,
+        help=f"Port to check the API on (default: {DEFAULT_PORT})",
+    )
+    tcp_parser.add_argument(
+        "--timeout",
+        "-t",
+        type=int,
+        default=DEFAULT_TIMEOUT_SECONDS,
+        help=f"Socket timeout for the connection attempt in seconds (default: {DEFAULT_TIMEOUT_SECONDS})",
+    )
+    http_parser = subcommands.add_parser(
+        "http", help="Check if the API is able to serve HTTP requests"
+    )
+    http_parser.add_argument(
+        "--port",
+        "-p",
+        type=int,
+        default=DEFAULT_PORT,
+        help=f"Port to check the API on (default: {DEFAULT_PORT})",
+    )
+    http_parser.add_argument(
+        "--timeout",
+        "-t",
+        type=int,
+        default=DEFAULT_TIMEOUT_SECONDS,
+        help=f"Request timeout in seconds (default: {DEFAULT_TIMEOUT_SECONDS})",
+    )
+    http_parser.add_argument(
+        "path",
+        nargs="?",
+        type=str,
+        default="/health/liveness",
+        help="Request path (default: /health/liveness)",
+    )
+    return parser.parse_args(argv)
+
+
+def check_tcp_connection(
+    port: int,
+    timeout_seconds: int,
+) -> None:
+    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+    sock.settimeout(timeout_seconds)
+    try:
+        sock.connect(("127.0.0.1", port))
+    except socket.error as e:
+        print(f"Failed: {e} {port=}")
+        exit(1)
+    else:
+        exit(0)
+    finally:
+        sock.close()
+
+
+def check_http_response(
+    port: int,
+    timeout_seconds: int,
+    path: str,
+) -> None:
+    url = urllib.parse.urljoin(f"http://127.0.0.1:{port}", path)
+    requests.get(
+        url,
+        timeout=timeout_seconds,
+    ).raise_for_status()
+
+
+def main(
+    argv: list[str],
+    *,
+    prog: str,
+) -> None:
+    args = get_args(argv, prog=prog)
+    match args.subcommand:
+        case None:
+            check_tcp_connection(
+                port=DEFAULT_PORT,
+                timeout_seconds=DEFAULT_TIMEOUT_SECONDS,
+            )
+        case "tcp":
+            check_tcp_connection(
+                port=args.port,
+                timeout_seconds=args.timeout,
+            )
+        case "http":
+            check_http_response(
+                port=args.port,
+                timeout_seconds=args.timeout,
+                path=args.path,
+            )

src/common/core/cli/__init__.py

@@ -5,7 +5,11 @@
 import tempfile
 import typing
 
-from django.core.management import execute_from_command_line
+from django.core.management import (
+    execute_from_command_line as django_execute_from_command_line,
+)
+
+from common.core.cli import healthcheck
 
 logger = logging.getLogger(__name__)
 
@@ -56,7 +60,25 @@ def ensure_cli_env() -> typing.Generator[None, None, None]:
         yield
 
 
-def main() -> None:
+def execute_from_command_line(argv: list[str]) -> None:
+    try:
+        subcommand = argv[1]
+        subcommand_main = {
+            "healthcheck": healthcheck.main,
+            # Backwards compatibility for task-processor health checks
+            # See https://github.com/Flagsmith/flagsmith-task-processor/issues/24
+            "checktaskprocessorthreadhealth": healthcheck.main,
+        }[subcommand]
+    except (IndexError, KeyError):
+        django_execute_from_command_line(argv)
+    else:
+        subcommand_main(
+            argv[2:],
+            prog=f"{os.path.basename(argv[0])} {subcommand}",
+        )
+
+
+def main(argv: list[str] = sys.argv) -> None:
     """
     The main entry point to the Flagsmith application.
 
@@ -72,5 +94,5 @@ def main() -> None:
     `flagsmith <command> [options]`
     """
     with ensure_cli_env():
-        # Run Django
-        execute_from_command_line(sys.argv)
+        # Run own commands and Django
+        execute_from_command_line(argv)

src/common/core/cli/healthcheck.py

@@ -0,0 +1,10 @@
+from typing import Generator
+
+import pytest
+from pytest_httpserver import HTTPServer
+
+
+@pytest.fixture()
+def http_server(unused_tcp_port: int) -> Generator[HTTPServer, None, None]:
+    with HTTPServer(port=unused_tcp_port) as server:
+        yield server