Add integration for FastAPI

setup.py

@@ -31,6 +31,7 @@ def read(*parts):
         "Framework :: Django :: 4.0",
         "Framework :: Django :: 4.1",
         "Framework :: Flask",
+        "Framework :: FastAPI",
         "Intended Audience :: Developers",
         "License :: OSI Approved :: Mozilla Public License 2.0 (MPL 2.0)",
         "Programming Language :: Python",
@@ -47,6 +48,7 @@ def read(*parts):
         "django": ["django"],
         "flask": ["flask", "blinker"],
         "sanic": ["sanic"],
+        "fastapi": ["fastapi", "asgiref"],
     },
     zip_safe=False,
     python_requires=">=3.7,<4",

src/dockerflow/fastapi/__init__.py

@@ -0,0 +1,14 @@
+from fastapi import APIRouter
+from fastapi.routing import APIRoute
+
+from .checks import register_heartbeat_check  # noqa
+from .views import heartbeat, lbheartbeat, version
+
+router = APIRouter(
+    tags=["Dockerflow"],
+    routes=[
+        APIRoute("/__lbheartbeat__", endpoint=lbheartbeat, methods=["GET", "HEAD"]),
+        APIRoute("/__heartbeat__", endpoint=heartbeat, methods=["GET", "HEAD"]),
+        APIRoute("/__version__", endpoint=version, methods=["GET"]),
+    ],
+)

src/dockerflow/fastapi/checks.py

@@ -0,0 +1,52 @@
+import functools
+import logging
+from dataclasses import dataclass
+from typing import Dict, List, Tuple
+
+from ..checks import level_to_text
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class CheckDetail:
+    status: str
+    level: int
+    messages: Dict[int, str]
+
+
+registered_checks = dict()
+
+
+def register_heartbeat_check(func=None, *, name=None):
+    if func is None:
+        return functools.partial(register_heartbeat_check, name=name)
+
+    if name is None:
+        name = func.__name__
+
+    logger.debug("Registered Dockerflow check %s", name)
+
+    @functools.wraps(func)
+    def decorated_function(*args, **kwargs):
+        logger.debug("Called Dockerflow check %s", name)
+        return func(*args, **kwargs)
+
+    registered_checks[name] = decorated_function
+    return decorated_function
+
+
+def _heartbeat_check_detail(check):
+    errors = check()
+    level = max([0] + [e.level for e in errors])
+    return CheckDetail(
+        status=level_to_text(level), level=level, messages={e.id: e.msg for e in errors}
+    )
+
+
+def run_heartbeat_checks():
+    check_details: List[Tuple[str, CheckDetail]] = []
+    for name, check in registered_checks.items():
+        detail = _heartbeat_check_detail(check)
+        check_details.append((name, detail))
+    return check_details

src/dockerflow/fastapi/middleware.py

@@ -0,0 +1,77 @@
+from __future__ import annotations
+
+import logging
+import sys
+import time
+from typing import Any, Dict
+
+from asgiref.typing import (
+    ASGI3Application,
+    ASGIReceiveCallable,
+    ASGISendCallable,
+    ASGISendEvent,
+    HTTPScope,
+)
+
+from ..logging import JsonLogFormatter
+
+
+class MozlogRequestSummaryLogger:
+    def __init__(
+        self,
+        app: ASGI3Application,
+        logger: logging.Logger | None = None,
+    ) -> None:
+        self.app = app
+        if logger is None:
+            self.logger = logging.getLogger("request.summary")
+            self.logger.setLevel(logging.INFO)
+            handler = logging.StreamHandler(sys.stdout)
+            handler.setLevel(logging.INFO)
+            handler.setFormatter(JsonLogFormatter)
+            self.logger.addHandler(handler)
+        else:
+            self.logger = logger
+
+    async def __call__(
+        self, scope: HTTPScope, receive: ASGIReceiveCallable, send: ASGISendCallable
+    ) -> None:
+        if scope["type"] != "http":
+            return await self.app(scope, receive, send)
+
+        info = dict(request_headers={}, response={})
+
+        async def inner_send(message: ASGISendEvent) -> None:
+            if message["type"] == "http.response.start":
+                info["response"] = message
+
+            await send(message)
+
+        try:
+            info["start_time"] = time.time()
+            await self.app(scope, receive, inner_send)
+        except Exception as exc:
+            info["response"]["status"] = 500
+            raise exc
+        finally:
+            info["end_time"] = time.time()
+            self._log(scope, info)
+
+    def _log(self, scope: HTTPScope, info) -> None:
+        self.logger.info("", extra=self._format(scope, info))
+
+    def _format(self, scope: HTTPScope, info) -> Dict[str, Any]:
+        for name, value in scope["headers"]:
+            header_key = name.decode("latin1").lower()
+            header_val = value.decode("latin1")
+            info["request_headers"][header_key] = header_val
+
+        request_duration_ms = (info["end_time"] - info["start_time"]) * 1000.0
+        return {
+            "agent": info["request_headers"].get("user-agent", ""),
+            "path": scope["path"],
+            "method": scope["method"],
+            "code": info["response"]["status"],
+            "lang": info["request_headers"].get("accept-language"),
+            "t": int(request_duration_ms),
+        }

src/dockerflow/fastapi/views.py

@@ -0,0 +1,46 @@
+import os
+
+from fastapi import Request, Response
+
+from dockerflow import checks
+
+from ..version import get_version
+from .checks import run_heartbeat_checks
+
+
+def lbheartbeat():
+    return {"status": "ok"}
+
+
+def heartbeat(response: Response):
+    check_results = run_heartbeat_checks()
+    details = {}
+    statuses = {}
+    level = 0
+
+    for name, detail in check_results:
+        statuses[name] = detail.status
+        level = max(level, detail.level)
+        if detail.level > 0:
+            details[name] = detail
+
+    if level < checks.ERROR:
+        response.status_code = 200
+    else:
+        response.status_code = 500
+
+    return {
+        "status": checks.level_to_text(level),
+        "checks": statuses,
+        "details": details,
+    }
+
+
+def version(request: Request):
+    if getattr(request.app.state, "APP_DIR", None):
+        root = request.app.state.APP_DIR
+    elif os.getenv("APP_DIR"):
+        root = os.getenv("APP_DIR")
+    else:
+        root = "/app"
+    return get_version(root)

tests/constraints/fastapi-0.100.txt

@@ -0,0 +1 @@
+fastapi>=0.100,<0.101

tests/fastapi/test_fastapi.py

@@ -0,0 +1,133 @@
+# This Source Code Form is subject to the terms of the Mozilla Public
+# License, v. 2.0. If a copy of the MPL was not distributed with this
+# file, you can obtain one at http://mozilla.org/MPL/2.0/.
+import json
+import logging
+
+import pytest
+from fastapi import FastAPI
+from fastapi.testclient import TestClient
+
+from dockerflow.checks import Error
+from dockerflow.fastapi import register_heartbeat_check
+from dockerflow.fastapi import router as dockerflow_router
+from dockerflow.fastapi.middleware import MozlogRequestSummaryLogger
+
+
+def create_app():
+    app = FastAPI()
+    app.include_router(dockerflow_router)
+    app.add_middleware(MozlogRequestSummaryLogger)
+    return app
+
+
+@pytest.fixture
+def app():
+    return create_app()
+
+
+@pytest.fixture
+def client(app):
+    return TestClient(app)
+
+
+def test_lbheartbeat_get(client):
+    response = client.get("/__lbheartbeat__")
+    assert response.status_code == 200
+    assert response.json() == {"status": "ok"}
+
+
+def test_lbheartbeat_head(client):
+    response = client.head("/__lbheartbeat__")
+    assert response.status_code == 200
+    assert response.content == b""
+
+
+def test_mozlog(client, caplog):
+    client.get(
+        "/__lbheartbeat__",
+        headers={
+            "User-Agent": "dockerflow/tests",
+            "Accept-Language": "en-US",
+        },
+    )
+    record = caplog.records[0]
+
+    assert record.levelno == logging.INFO
+    assert record.agent == "dockerflow/tests"
+    assert record.lang == "en-US"
+    assert record.method == "GET"
+    assert record.path == "/__lbheartbeat__"
+    assert isinstance(record.t, int)
+
+
+VERSION_CONTENT = {"foo": "bar"}
+
+
+def test_version_app_state(client, tmp_path, app):
+    version_path = tmp_path / "version.json"
+    version_path.write_text(json.dumps(VERSION_CONTENT))
+
+    app.state.APP_DIR = tmp_path.resolve()
+    response = client.get("/__version__")
+    assert response.status_code == 200
+    assert response.json() == VERSION_CONTENT
+
+
+def test_version_env_var(client, tmp_path, monkeypatch):
+    version_path = tmp_path / "version.json"
+    version_path.write_text(json.dumps(VERSION_CONTENT))
+
+    monkeypatch.setenv("APP_DIR", tmp_path.resolve())
+
+    response = client.get("/__version__")
+    assert response.status_code == 200
+    assert response.json() == VERSION_CONTENT
+
+
+def test_version_default(client, mocker):
+    mock_get_version = mocker.MagicMock(return_value=VERSION_CONTENT)
+    mocker.patch("dockerflow.fastapi.views.get_version", mock_get_version)
+
+    response = client.get("/__version__")
+    assert response.status_code == 200
+    assert response.json() == VERSION_CONTENT
+    mock_get_version.assert_called_with("/app")
+
+
+def test_heartbeat_get(client):
+    @register_heartbeat_check
+    def return_error():
+        return [Error("BOOM", id="foo")]
+
+    response = client.get("/__heartbeat__")
+    assert response.status_code == 500
+    assert response.json() == {
+        "status": "error",
+        "checks": {"return_error": "error"},
+        "details": {
+            "return_error": {
+                "level": 40,
+                "messages": {"foo": "BOOM"},
+                "status": "error",
+            }
+        },
+    }
+
+
+def test_heartbeat_head(client):
+    @register_heartbeat_check
+    def return_error():
+        return [Error("BOOM", id="foo")]
+
+    response = client.head("/__heartbeat__")
+    assert response.content == b""
+
+
+def test_heartbeat_custom_name(client):
+    @register_heartbeat_check(name="my_check_name")
+    def return_error():
+        return [Error("BOOM", id="foo")]
+
+    response = client.get("/__heartbeat__")
+    assert response.json()["checks"]["my_check_name"]

tests/requirements/fastapi.txt

@@ -0,0 +1,3 @@
+fastapi
+asgiref
+httpx

tox.ini

@@ -6,6 +6,7 @@ envlist =
     py38-docs
     py{37,38,39,310,311}-dj32
     py{38,39,310,311}-dj{40,41,42}
+    py{37,38,39,310,311}-fa100
     py{37,38,39,310,311}-fl{20,21,22}
     py{37,38,39,310,311}-s{21,22}
 
@@ -26,12 +27,14 @@ setenv =
 deps =
     -rtests/requirements/default.txt
     dj{32,40,41,42}: -rtests/requirements/django.txt
+    fa100: -rtests/requirements/fastapi.txt
     fl{20,21,22}: -rtests/requirements/flask.txt
     s{21,22}: -rtests/requirements/sanic.txt
     dj32: -ctests/constraints/django-3.2.txt
     dj40: -ctests/constraints/django-4.0.txt
     dj41: -ctests/constraints/django-4.1.txt
     dj42: -ctests/constraints/django-4.2.txt
+    fa100: -ctests/constraints/fastapi-0.100.txt
     fl20: -ctests/constraints/flask-2.0.txt
     fl21: -ctests/constraints/flask-2.1.txt
     fl22: -ctests/constraints/flask-2.2.txt
@@ -40,7 +43,8 @@ deps =
 commands =
     python --version
     dj{32,40,41,42}: pytest tests/core/ tests/django --nomigrations {posargs:}
-    fl{012,10,11,20,21,22}: pytest tests/core/ tests/flask/ {posargs:}
+    fa{100}: pytest tests/core/ tests/fastapi/ {posargs:}
+    fl{20,21,22}: pytest tests/core/ tests/flask/ {posargs:}
     s{20,21,22}: pytest tests/core/ tests/sanic/ {posargs:}
 
 [testenv:py38-docs]