Add Prometheus metrics

pyproject.toml

@@ -18,6 +18,7 @@ dependencies = [
     "structlog>=23.1.0",
     "asgi-correlation-id>=4.2.0",
     "sentry-sdk[fastapi]>=1.25.1",
+    "prometheus-client>=0.20.0",
     "structlog-sentry>=2.0.3",
     "requests>=2.30.0",
     "httpx>=0.24.1",

src/mainframe/endpoints/package.py

@@ -14,6 +14,12 @@
 from mainframe.database import get_db
 from mainframe.dependencies import get_pypi_client, validate_token
 from mainframe.json_web_token import AuthenticationData
+from mainframe.metrics import (
+    package_ingest_counter,
+    package_scan_fail_counter,
+    package_scan_success_counter,
+    package_scanned_counter,
+)
 from mainframe.models.orm import DownloadURL, Rule, Scan, Status
 from mainframe.models.schemas import (
     Error,
@@ -42,6 +48,7 @@ def submit_results(
     name = result.name
     version = result.version
 
+    package_scanned_counter.inc()
     scan = session.scalar(
         select(Scan).where(Scan.name == name).where(Scan.version == version).options(joinedload(Scan.rules))
     )
@@ -65,6 +72,7 @@ def submit_results(
     if isinstance(result, PackageScanResultFail):
         scan.status = Status.FAILED
         scan.fail_reason = result.reason
+        package_scan_fail_counter.inc()
 
         session.commit()
         return
@@ -101,6 +109,7 @@ def submit_results(
         tag="scan_submitted",
     )
 
+    package_scan_success_counter.inc()
     session.commit()
 
 
@@ -221,6 +230,7 @@ def batch_queue_package(
             ],
         )
 
+        package_ingest_counter.inc()
         session.add(scan)
 
     session.commit()
@@ -281,6 +291,7 @@ def queue_package(
 
     try:
         session.commit()
+        package_ingest_counter.inc()
     except IntegrityError:
         log.warn(f"Package {name}@{version} already queued for scanning.", tag="already_queued")
         raise HTTPException(409, f"Package {name}@{version} is already queued for scanning")

src/mainframe/endpoints/report.py

@@ -13,6 +13,7 @@
 from mainframe.database import get_db
 from mainframe.dependencies import get_pypi_client, validate_token
 from mainframe.json_web_token import AuthenticationData
+from mainframe.metrics import package_scan_report_counter
 from mainframe.models.orm import Scan
 from mainframe.models.schemas import (
     EmailReport,
@@ -248,4 +249,5 @@ def report_package(
 
     scan.reported_by = auth.subject
     scan.reported_at = dt.datetime.now(dt.timezone.utc)
+    package_scan_report_counter.inc()
     session.commit()

src/mainframe/metrics.py

@@ -0,0 +1,12 @@
+from prometheus_client import Counter, Histogram, make_asgi_app  # type: ignore
+
+package_ingest_counter = Counter("packages_ingested", "Total amount of packages ingested so far")
+package_scanned_counter = Counter("packages_scanned", "Total amount of packages scanned")
+package_scan_success_counter = Counter("package_scan_success", "Total amount of packages successfully scanned")
+package_scan_fail_counter = Counter("package_scan_fail", "Total amount of packages that failed while scanning")
+package_scan_report_counter = Counter("package_scan_report", "Total amount of packages that have been reported")
+
+request_time = Histogram("request_time", "Time taken to serve a request")
+request_counter = Counter("requests", "Total amount of requests served")
+
+metrics_app = make_asgi_app()  # type: ignore

src/mainframe/server.py

@@ -16,6 +16,7 @@
 from mainframe.constants import GIT_SHA, Sentry
 from mainframe.dependencies import validate_token, validate_token_override
 from mainframe.endpoints import routers
+from mainframe.metrics import metrics_app, request_counter, request_time  # type: ignore
 from mainframe.models.schemas import ServerMetadata
 from mainframe.rules import Rules, fetch_rules
 
@@ -119,6 +120,8 @@ async def lifespan(app_: FastAPI):
     version=__version__,
 )
 
+app.mount("/metrics", metrics_app)  # type: ignore
+
 if GIT_SHA in ("development", "testing"):
     app.dependency_overrides[validate_token] = validate_token_override
 
@@ -148,6 +151,8 @@ async def logging_middleware(request: Request, call_next: Callable[[Request], Aw
         raise
     finally:
         process_time = time.perf_counter_ns() - start_time
+        request_time.observe(process_time / 1_000_000_000)
+        request_counter.inc()
         status_code = response.status_code
         http_method = request.method
         http_version = request.scope["http_version"]