Feat: Add prometheus metrics

src/backend/core/apps.py

@@ -14,6 +14,9 @@ class CoreConfig(AppConfig):
     def ready(self):
         """Register signal handlers when the app is ready."""
         # Import signal handlers to register them
+        from prometheus_client.core import REGISTRY
+        from .metrics import CustomDBMetricsCollector
+        REGISTRY.register(CustomDBMetricsCollector())
 
         # pylint: disable=unused-import, import-outside-toplevel
         import core.signals  # noqa

src/backend/core/metrics.py

@@ -0,0 +1,25 @@
+from django.apps import apps
+from prometheus_client.core import GaugeMetricFamily
+from django.db import ProgrammingError, connection
+from .models import MessageRecipient
+from .enums import MessageDeliveryStatusChoices
+from django.db import models
+
+class CustomDBMetricsCollector:
+
+    def get_messages_with_status(self):
+        messages_statuses_count = MessageRecipient.objects.values('delivery_status').annotate(count=models.Count('id'))
+        gauges = [GaugeMetricFamily(
+            f"message_{MessageDeliveryStatusChoices(message_status_count['delivery_status'])._label_}_count",
+            f"Number of messages with status {MessageDeliveryStatusChoices(message_status_count['delivery_status'])._label_}",
+            value=message_status_count['count']
+        ) for message_status_count in messages_statuses_count]
+
+        yield from gauges
+
+    def collect(self):
+        # Only run if apps are ready and model is migrated
+        if not apps.ready or not apps.is_installed("core"):
+            return
+
+        yield from self.get_messages_with_status()

src/backend/messages/settings.py

@@ -549,6 +549,8 @@ class Base(Configuration):
 
     ENABLE_DOCKERFLOW = values.BooleanValue(default=True)
 
+    ENABLE_PROMETHEUS = values.BooleanValue(default=True)
+
     # AI
     AI_API_KEY = values.Value(None, environ_name="AI_API_KEY", environ_prefix=None)
     AI_BASE_URL = values.Value(None, environ_name="AI_BASE_URL", environ_prefix=None)
@@ -622,6 +624,14 @@ def __init__(self, *args, **kwargs):
             self.INSTALLED_APPS += ["dockerflow.django"]
             self.MIDDLEWARE += ["dockerflow.django.middleware.DockerflowMiddleware"]
 
+        if self.ENABLE_PROMETHEUS:
+            self.INSTALLED_APPS += ["django_prometheus"]
+            self.MIDDLEWARE = [
+                "django_prometheus.middleware.PrometheusBeforeMiddleware",
+                *self.MIDDLEWARE,
+                "django_prometheus.middleware.PrometheusAfterMiddleware",
+            ]
+
     # pylint: disable=invalid-name
     @property
     def ENVIRONMENT(self):

src/backend/messages/urls.py

@@ -46,3 +46,11 @@
             name="redoc-schema",
         ),
     ]
+
+if settings.ENABLE_PROMETHEUS:
+    urlpatterns += [
+        path(
+            "",
+            include("django_prometheus.urls"),
+        ),
+    ]

src/backend/pyproject.toml

@@ -63,6 +63,7 @@ dependencies = [
     "sentry-sdk[django]==2.27.0",
     "url-normalize==1.4.3",
     "whitenoise==6.8.2",
+    "django-prometheus (>=2.4.1,<3.0.0)",
 ]
 
 [project.urls]