Add middleware for healthcheck

runner_manager/auth.py

@@ -0,0 +1,16 @@
+
+from fastapi.middleware.trustedhost import TrustedHostMiddleware
+from starlette.types import Receive, Scope, Send
+
+
+class TrustedHostHealthRoutes(TrustedHostMiddleware):
+    """A healthcheck endpoint that answers to GET requests on /_health"""
+
+
+    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
+        """If the request is made on the path _health then execute the check on hosts"""
+        if scope["path"] == "/_health/":
+            print(scope)
+            await super().__call__(scope, receive, send)
+        else:
+            await self.app(scope, receive, send)

runner_manager/main.py

@@ -1,25 +1,28 @@
 import logging
 
-from fastapi import Depends, FastAPI, HTTPException, Response, Security, status
+from fastapi import Depends, FastAPI, HTTPException, Security, status
 from fastapi.security import APIKeyHeader, APIKeyQuery
+from runner_manager.auth import TrustedHostHealthRoutes
 
 from runner_manager.dependencies import get_queue, get_settings
 from runner_manager.jobs.startup import startup
 from runner_manager.models.settings import Settings
-from runner_manager.routers import webhook
+from runner_manager.routers import webhook, _health
 
 log = logging.getLogger(__name__)
 
 app = FastAPI()
 api_key_query = APIKeyQuery(name="api-key", auto_error=False)
 api_key_header = APIKeyHeader(name="x-api-key", auto_error=False)
+settings = get_settings()
 
 
 def get_api_key(
     api_key_query: str = Security(api_key_query),
     api_key_header: str = Security(api_key_header),
     settings: Settings = Depends(get_settings),
 ) -> str:
+    """Get the API key from either the query parameter or the header"""
     if not settings.api_key:
         return ""
     if api_key_query in [settings.api_key.get_secret_value()]:
@@ -33,6 +36,11 @@ def get_api_key(
 
 
 app.include_router(webhook.router)
+app.include_router(_health.router)
+app.add_middleware(
+    TrustedHostHealthRoutes,
+    allowed_hosts=settings.allowed_hosts
+)
 
 
 @app.on_event("startup")
@@ -43,9 +51,6 @@ def startup_event():
     log.info(f"Startup job is {status}")
 
 
-@app.get("/_health")
-def health():
-    return Response(status_code=200)
 
 
 @app.get("/public")

runner_manager/models/settings.py

@@ -1,5 +1,5 @@
 from pathlib import Path
-from typing import Any, Dict, Optional
+from typing import Any, Dict, Optional, Sequence
 
 import yaml
 from pydantic import AnyHttpUrl, BaseSettings, RedisDsn, SecretStr
@@ -26,6 +26,10 @@ class Settings(BaseSettings):
     redis_om_url: Optional[RedisDsn] = None
     github_base_url: Optional[AnyHttpUrl] = None
     api_key: Optional[SecretStr] = None
+    allowed_hosts: Optional[Sequence[str]] = [
+        "localhost",
+        "testserver",
+    ]
 
     class Config:
         smart_union = True

runner_manager/routers/_health.py

@@ -0,0 +1,10 @@
+
+from fastapi import APIRouter, Response
+
+router = APIRouter(prefix="/_health")
+
+
+@router.get("/", status_code=200)
+def healthcheck():
+    """Healthcheck endpoint that answers to GET requests on /_healthz"""
+    return Response(status_code=200)

tests/api/tests_auth.py

@@ -27,5 +27,4 @@ def test_private_endpoint_with_valid_api_key(fastapp, client):
     fastapp.dependency_overrides[get_settings] = settings_api_key
     headers = {"x-api-key": "secret"}
     response = client.get("/private", headers=headers)
-    print(response.text)
     assert response.status_code == 200