Add Prometheus metrics

pyproject.toml

@@ -22,6 +22,7 @@ dependencies = [
     "structlog>=23.1.0",
     "asgi-correlation-id>=4.2.0",
     "sentry-sdk[fastapi]>=1.25.1",
+    "prometheus-client>=0.17.1",
 ]
 dynamic = ["version"]
 

src/mainframe/endpoints/package.py

@@ -13,6 +13,12 @@
 from mainframe.database import get_db
 from mainframe.dependencies import 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,
@@ -41,6 +47,8 @@ async def submit_results(
     name = result.name
     version = result.version
 
+    package_scanned_counter.inc()
+
     scan = await session.scalar(
         select(Scan).where(Scan.name == name).where(Scan.version == version).options(selectinload(Scan.rules))
     )
@@ -64,6 +72,7 @@ async def submit_results(
     if isinstance(result, PackageScanResultFail):
         scan.status = Status.FAILED
         scan.fail_reason = result.reason
+        package_scan_fail_counter.inc()
 
         await session.commit()
         return
@@ -100,6 +109,7 @@ async def submit_results(
         tag="scan_submitted",
     )
 
+    package_scan_success_counter.inc()
     await session.commit()
 
 
@@ -210,6 +220,9 @@ async def batch_queue_package(
                     session.add(row)
             except IntegrityError:
                 pass
+            else:
+                print("Incrementing!")
+                package_ingest_counter.inc()
 
 
 @router.post(
@@ -268,6 +281,7 @@ async def queue_package(
 
     try:
         await session.commit()
+        package_ingest_counter.inc()
     except IntegrityError:
         await 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

@@ -15,6 +15,7 @@
 from mainframe.database import get_db
 from mainframe.dependencies import get_ms_graph_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 Error, PackageSpecifier
 from mainframe.utils.mailer import send_email
@@ -228,4 +229,5 @@ async def report_package(
 
     row.reported_by = auth.subject
     row.reported_at = dt.datetime.now(dt.timezone.utc)
+    package_scan_report_counter.inc()
     await session.commit()

src/mainframe/metrics.py

@@ -0,0 +1,9 @@
+from prometheus_client import Counter, 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")
+
+metrics_app = make_asgi_app()  # type: ignore

src/mainframe/server.py

@@ -17,6 +17,7 @@
 from mainframe.database import async_session
 from mainframe.dependencies import validate_token, validate_token_override
 from mainframe.endpoints import routers
+from mainframe.metrics import metrics_app  # type: ignore
 from mainframe.models.schemas import ServerMetadata
 from mainframe.rules import Rules, fetch_rules
 
@@ -126,6 +127,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