Merge pull request #350 from Forest-Troll/master

Added prometheus integration

.env.example

@@ -34,3 +34,8 @@
 # Host and port to listen on
 SPOOLMAN_HOST=0.0.0.0
 SPOOLMAN_PORT=7912
+
+
+# Enable Collect Prometheus metrics at database
+# Default: FALSE
+#SPOOLMAN_METRICS_ENABLED=TRUE

README.md

@@ -115,7 +115,7 @@ These are either set in the .env file if you use the standalone installation, or
 If you want to connect with an external database instead, specify the `SPOOLMAN_DB_*` environment variables from the table below.
 
 | Variable                  | Description                                                                                                                  |
-| ------------------------- | ---------------------------------------------------------------------------------------------------------------------------- |
+|---------------------------|------------------------------------------------------------------------------------------------------------------------------|
 | SPOOLMAN_DB_TYPE          | Type of database, any of: `postgres`, `mysql`, `sqlite`, `cockroachdb`                                                       |
 | SPOOLMAN_DB_HOST          | Database hostname                                                                                                            |
 | SPOOLMAN_DB_PORT          | Database port                                                                                                                |
@@ -130,6 +130,7 @@ If you want to connect with an external database instead, specify the `SPOOLMAN_
 | PGID                      | (*docker only*) Set the GID of the user in the docker container. Default is 1000.                                            |
 | SPOOLMAN_PORT             | The port Spoolman should run on (default: 8000)                                                                              |
 | SPOOLMAN_HOST             | The hostname/ip Spoolman should bind to (default: 0.0.0.0)                                                                   |
+| SPOOLMAN_METRICS_ENABLED  | Enable collect Spoolman prometheus metrics at database. Default `False`                                                      |
 
 ## Frequently Asked Questions (FAQs)
 ### QR Code Does not work on HTTP / The page is not served over HTTPS
@@ -148,6 +149,7 @@ To setup yourself for Python development, do the following:
 1. Clone this repo
 2. CD into the repo
 3. Install PDM: `pip install --user pdm`
+> At pre-commit hook used pdm==2.7.4
 4. Install Spoolman dependencies: `pdm sync`
 
 And you should be all setup. Read the Style and Integration Testing sections below as well.

pyproject.toml

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

spoolman/database/database.py

@@ -1,5 +1,4 @@
 """SQLAlchemy database setup."""
-
 import datetime
 import logging
 import shutil
@@ -14,6 +13,7 @@
 from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, async_sessionmaker, create_async_engine
 
 from spoolman import env
+from spoolman.prometheus.metrics import filament_metrics, spool_metrics
 
 logger = logging.getLogger(__name__)
 
@@ -183,6 +183,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.debug("Start metrics collection")
+    async for session in get_db_session():
+        await filament_metrics(session)
+        await spool_metrics(session)
+    logger.debug("End metrics collection")
+
+
 def schedule_tasks(scheduler: Scheduler) -> None:
     """Schedule tasks to be executed by the provided scheduler.
 
@@ -191,6 +200,10 @@ def schedule_tasks(scheduler: Scheduler) -> None:
     """
     if __db is None:
         raise RuntimeError("DB is not setup.")
+    if env.is_metrics_enabled():
+        logger.info("Scheduling automatic metric collection.")
+        # Run every minute, may be needs specify timer
+        scheduler.minutely(datetime.time(second=0), _metrics)
     if not env.is_automatic_backup_enabled():
         return
     if "sqlite" in __db.connection_url.drivername:

spoolman/env.py

@@ -380,3 +380,21 @@ def is_data_dir_mounted() -> bool:
     mounts = subprocess.run("mount", check=True, stdout=subprocess.PIPE, text=True)  # noqa: S603, S607
     data_dir = str(get_data_dir().resolve())
     return any(data_dir in line for line in mounts.stdout.splitlines())
+
+
+def is_metrics_enabled() -> bool:
+    """Get whether collect prometheus metrics at database is enabled.
+
+        Returns False if no environment variable was set for collect metrics.
+
+    Returns:
+        bool: Whether collect metrics is enabled.
+    """
+    metrics_enabled = os.getenv("SPOOLMAN_METRICS_ENABLED", "FALSE").upper()
+    if metrics_enabled in {"FALSE", "0"}:
+        return False
+    if metrics_enabled in {"TRUE", "1"}:
+        return True
+    raise ValueError(
+        f"Failed to parse SPOOLMAN_METRICS_ENABLED variable: Unknown metrics enabled '{metrics_enabled}'.",
+    )

spoolman/main.py

@@ -1,5 +1,4 @@
 """Main entrypoint to the server."""
-
 import logging
 import subprocess
 from logging.handlers import TimedRotatingFileHandler
@@ -9,12 +8,15 @@
 from fastapi import FastAPI
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.middleware.gzip import GZipMiddleware
+from fastapi.responses import PlainTextResponse
+from prometheus_client import generate_latest
 from scheduler.asyncio.scheduler import Scheduler
 
 from spoolman import env
 from spoolman.api.v1.router import app as v1_app
 from spoolman.client import SinglePageApplication
 from spoolman.database import database
+from spoolman.prometheus.metrics import registry
 
 # Define a console logger
 console_handler = logging.StreamHandler()
@@ -37,7 +39,23 @@
 )
 app.add_middleware(GZipMiddleware)
 app.mount("/api/v1", v1_app)
-app.mount("/", app=SinglePageApplication(directory="client/dist"), name="client")
+
+
+# WA for prometheus /metrics bind with SinglePageApp at root
+@app.get(
+    "/metrics",
+    response_class=PlainTextResponse,
+    name="Get metrics for prometheus",
+    description=(
+        "Get app metrics for prometheusIf enabled SPOOLMAN_METRICS_ENABLED returned metrics by Spools and Filaments"
+    ),
+)
+def get_metrics() -> bytes:
+    """Return prometheus metrics."""
+    return generate_latest(registry)
+
+
+app.mount("", app=SinglePageApplication(directory="client/dist"))
 
 
 # Allow all origins if in debug mode

spoolman/prometheus/metrics.py

@@ -0,0 +1,86 @@
+"""Prometheus metrics collectors."""
+import logging
+from typing import Callable
+
+import sqlalchemy
+from prometheus_client import REGISTRY, Gauge, make_asgi_app
+from sqlalchemy.ext.asyncio import AsyncSession
+from sqlalchemy.orm import contains_eager
+
+from spoolman.database import models
+
+registry = REGISTRY
+
+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() -> Callable:
+    """Start ASGI prometheus app with global registry."""
+    logger.info("Start metrics app")
+    return make_asgi_app(registry=registry)
+
+
+metrics_app = make_asgi_app()
+
+
+async def spool_metrics(db: AsyncSession) -> None:
+    """Get metrics by Spools from DB and write to prometheus.
+
+    Args:
+        db: async db session
+    """
+    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:
+        if row.price is not None:
+            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:
+    """Get metrics and info by Filaments from DB and write to prometheus.
+
+    Args:
+        db: async db session
+    """
+    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:
+        vendor_name = "-"
+        if row.vendor is not None:
+            vendor_name = row.vendor.name
+        FILAMENT_INFO.labels(
+            str(row.id),
+            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)
+        if row.weight is not None:
+            FILAMENT_WEIGHT.labels(str(row.id)).set(row.weight)