Added prometheus integration

pyproject.toml

@@ -20,6 +20,7 @@ dependencies = [
     "psycopg2-binary~=2.9",
     "setuptools~=68.0",
     "WebSockets>=11.0.3",
+    "prometheus-client"
 ]
 requires-python = ">=3.9"
 

spoolman/database/database.py

@@ -1,10 +1,11 @@
 """SQLAlchemy database setup."""
-
+import asyncio
 import datetime
 import logging
 import shutil
 import sqlite3
 from collections.abc import AsyncGenerator
+from contextlib import asynccontextmanager
 from os import PathLike
 from pathlib import Path
 from typing import Optional, Union
@@ -15,6 +16,8 @@
 
 from spoolman import env
 
+from spoolman.prometheus.metrics import filament_metrics, spool_metrics
+
 logger = logging.getLogger(__name__)
 
 
@@ -183,6 +186,15 @@ async def _backup_task() -> Optional[Path]:
     return __db.backup_and_rotate(env.get_backups_dir(), num_backups=5)
 
 
+async def _metrics() -> None:
+    """Create some useful prometheus metrics"""
+    logger.info("Start metrics collection")
+    async with get_session() as session:
+        await filament_metrics(session)
+        await spool_metrics(session)
+    logger.info("End metrics collection")
+
+
 def schedule_tasks(scheduler: Scheduler) -> None:
     """Schedule tasks to be executed by the provided scheduler.
 
@@ -197,6 +209,9 @@ def schedule_tasks(scheduler: Scheduler) -> None:
         logger.info("Scheduling automatic database backup for midnight.")
         # Schedule for midnight
         scheduler.daily(datetime.time(hour=0, minute=0, second=0), _backup_task)  # type: ignore[arg-type]
+        logger.info("Scheduling automatic metric collection.")
+        logger.info("%s", datetime.time(minute=1))
+        scheduler.minutely(datetime.time(second=0), _metrics)
 
 
 async def get_db_session() -> AsyncGenerator[AsyncSession, None]:
@@ -216,3 +231,18 @@ async def get_db_session() -> AsyncGenerator[AsyncSession, None]:
             raise exc
         finally:
             await session.close()
+
+
+@asynccontextmanager
+async def get_session() -> AsyncSession:
+    if __db is None or __db.session_maker is None:
+        raise RuntimeError("DB is not setup.")
+    async with __db.session_maker() as session:
+        try:
+            yield session
+            await session.commit()
+        except Exception as exc:
+            await session.rollback()
+            raise exc
+        finally:
+            await session.close()

spoolman/main.py

@@ -16,6 +16,8 @@
 from spoolman.client import SinglePageApplication
 from spoolman.database import database
 
+from spoolman.prometheus.metrics import metrics_app
+
 # Define a console logger
 console_handler = logging.StreamHandler()
 console_handler.setFormatter(logging.Formatter("%(name)-26s %(levelname)-8s %(message)s"))
@@ -37,9 +39,9 @@
 )
 app.add_middleware(GZipMiddleware)
 app.mount("/api/v1", v1_app)
+app.mount("/metrics", metrics_app)
 app.mount("/", app=SinglePageApplication(directory="client/dist"), name="client")
 
-
 # Allow all origins if in debug mode
 if env.is_debug_mode():
     logger.warning("Running in debug mode, allowing all origins.")

spoolman/prometheus/metrics.py

@@ -0,0 +1,71 @@
+import sqlalchemy
+import logging
+
+from sqlalchemy.ext.asyncio import AsyncSession
+from sqlalchemy.orm import contains_eager
+
+from spoolman.database import models
+
+from prometheus_client import REGISTRY, make_asgi_app, Gauge
+
+PREFIX = "spoolman"
+
+SPOOL_PRICE = Gauge("%s_spool_price" % PREFIX, "Total Spool price", ['spool_id', 'filament_id'])
+SPOOL_USED_WEIGHT = Gauge("%s_spool_weight_used" % PREFIX, "Spool Used Weight", ['spool_id', 'filament_id'])
+FILAMENT_INFO = Gauge("%s_filament_info" % PREFIX, "Filament information", [
+            'filament_id',
+            'vendor',
+            'name',
+            'material',
+            'color'])
+FILAMENT_DENSITY = Gauge("%s_filament_density" % PREFIX, "Density of filament", ["filament_id"])
+FILAMENT_DIAMETER = Gauge("%s_filament_diameter" % PREFIX, "Diameter of filament", ["filament_id"])
+FILAMENT_WEIGHT = Gauge("%s_filament_weight" % PREFIX, "Net weight of filament", ["filament_id"])
+
+logger = logging.getLogger(__name__)
+
+
+def make_metrics_app():
+    registry = REGISTRY
+    logger.info("Start metrics app")
+    return make_asgi_app(registry=registry)
+
+
+metrics_app = make_metrics_app()
+
+
+async def spool_metrics(db: AsyncSession) -> None:
+    stmt = (
+        sqlalchemy.select(models.Spool)
+            .where(
+                sqlalchemy.or_(
+                    models.Spool.archived.is_(False),
+                    models.Spool.archived.is_(None),
+                )
+            )
+    )
+    rows = await db.execute(stmt)
+    result = list(rows.unique().scalars().all())
+    for row in result:
+        SPOOL_PRICE.labels(str(row.id), str(row.filament_id)).set(row.price)
+        SPOOL_USED_WEIGHT.labels(str(row.id), str(row.filament_id)).set(row.used_weight)
+
+
+async def filament_metrics(db: AsyncSession) -> None:
+    stmt = (
+        sqlalchemy.select(models.Filament)
+            .options(contains_eager(models.Filament.vendor))
+            .join(models.Filament.vendor, isouter=True)
+    )
+    rows = await db.execute(stmt)
+    result = list(rows.unique().scalars().all())
+    for row in result:
+        FILAMENT_INFO.labels(str(row.id),
+                             row.vendor.name,
+                             row.name,
+                             row.material,
+                             row.color_hex).set(1)
+        FILAMENT_DENSITY.labels(str(row.id)).set(row.density)
+        FILAMENT_DIAMETER.labels(str(row.id)).set(row.diameter)
+        FILAMENT_WEIGHT.labels(str(row.id)).set(row.weight)
+