diff --git a/aurweb/asgi.py b/aurweb/asgi.py index 2ba2afd0..16de771e 100644 --- a/aurweb/asgi.py +++ b/aurweb/asgi.py @@ -9,6 +9,7 @@ from urllib.parse import quote_plus from fastapi import FastAPI, HTTPException, Request from fastapi.responses import HTMLResponse, RedirectResponse from fastapi.staticfiles import StaticFiles +from prometheus_client import multiprocess from sqlalchemy import and_, or_ from starlette.middleware.authentication import AuthenticationMiddleware from starlette.middleware.sessions import SessionMiddleware @@ -29,7 +30,7 @@ app = FastAPI(exception_handlers=errors.exceptions) # library with custom collectors and expose /metrics. instrumentator().add(http_api_requests_total()) instrumentator().add(http_requests_total()) -instrumentator().instrument(app).expose(app) +instrumentator().instrument(app) @app.on_event("startup") @@ -79,6 +80,12 @@ async def app_startup(): get_engine() +def child_exit(server, worker): # pragma: no cover + """ This function is required for gunicorn customization + of prometheus multiprocessing. """ + multiprocess.mark_process_dead(worker.pid) + + @app.exception_handler(HTTPException) async def http_exception_handler(request, exc): """ diff --git a/aurweb/routers/html.py b/aurweb/routers/html.py index c749ca67..4cee5f99 100644 --- a/aurweb/routers/html.py +++ b/aurweb/routers/html.py @@ -1,11 +1,14 @@ """ AURWeb's primary routing module. Define all routes via @app.app.{get,post} decorators in some way; more complex routes should be defined in their own modules and imported here. """ +import os + from datetime import datetime from http import HTTPStatus -from fastapi import APIRouter, Form, HTTPException, Request +from fastapi import APIRouter, Form, HTTPException, Request, Response from fastapi.responses import HTMLResponse, RedirectResponse +from prometheus_client import CONTENT_TYPE_LATEST, CollectorRegistry, generate_latest, multiprocess from sqlalchemy import and_, case, or_ import aurweb.config @@ -203,7 +206,21 @@ async def index(request: Request): return render_template(request, "index.html", context) -# A route that returns a error 503. For testing purposes. +@router.get("/metrics") +async def metrics(request: Request): + registry = CollectorRegistry() + if os.environ.get("FASTAPI_BACKEND", "") == "gunicorn": # pragma: no cover + # This case only ever happens in production, when we are running + # gunicorn. We don't test with gunicorn, so we don't cover this path. + multiprocess.MultiProcessCollector(registry) + data = generate_latest(registry) + headers = { + "Content-Type": CONTENT_TYPE_LATEST, + "Content-Length": str(len(data)) + } + return Response(data, headers=headers) + + @router.get("/raisefivethree", response_class=HTMLResponse) async def raise_service_unavailable(request: Request): raise HTTPException(status_code=503) diff --git a/test/test_html.py b/test/test_html.py index 2018840b..8e7cb2d1 100644 --- a/test/test_html.py +++ b/test/test_html.py @@ -117,3 +117,10 @@ def test_get_successes(): """ successes = get_successes(html) assert successes[0].text.strip() == "Test" + + +def test_metrics(client: TestClient): + with client as request: + resp = request.get("/metrics") + assert resp.status_code == int(HTTPStatus.OK) + assert resp.headers.get("Content-Type").startswith("text/plain")