From 1801c18d611c1c0399cea31476fb9dc8e5af9918 Mon Sep 17 00:00:00 2001
From: "Terence D. Honles" <terence@honles.com>
Date: Fri, 16 Aug 2024 14:32:46 +0200
Subject: [PATCH] export RQ status as prometheus metrics

fixes: https://github.com/rq/django-rq/issues/503
---
 django_rq/apps.py  | 27 +++++++++++++++++++++++++++
 django_rq/urls.py  | 10 ++++++++++
 django_rq/views.py | 21 ++++++++++++++++++++-
 setup.py           |  1 +
 4 files changed, 58 insertions(+), 1 deletion(-)

diff --git a/django_rq/apps.py b/django_rq/apps.py
index f3f2e3b1..a11a5465 100644
--- a/django_rq/apps.py
+++ b/django_rq/apps.py
@@ -1,6 +1,33 @@
 from django.apps import AppConfig
 
+from .queues import filter_connection_params, get_connection, get_queue_class, get_unique_connection_configs
+from .workers import get_worker_class
+
+try:
+    import prometheus_client
+
+    from rq_exporter.collector import RQCollector
+except ImportError:
+    RQCollector = None
+
 
 class DjangoRqAdminConfig(AppConfig):
     default_auto_field = "django.db.models.AutoField"
     name = "django_rq"
+
+    def ready(self):
+        if RQCollector is None:
+            return
+
+        from .settings import QUEUES
+
+        worker_class = get_worker_class()
+        registered = {}
+        unique_configs = get_unique_connection_configs()
+        for name, config in QUEUES.items():
+            index = unique_configs.index(filter_connection_params(config))
+            queue_class = get_queue_class(config)
+            key = index, queue_class
+            if key not in registered:
+                registered[key] = True
+                prometheus_client.REGISTRY.register(RQCollector(get_connection(name), worker_class, queue_class))
diff --git a/django_rq/urls.py b/django_rq/urls.py
index 1aff8d9c..15cce015 100644
--- a/django_rq/urls.py
+++ b/django_rq/urls.py
@@ -2,9 +2,19 @@
 
 from . import views
 
+try:
+    import rq_exporter
+
+    metrics_view = [
+        re_path(r'^metrics/?$', views.prometheus_metrics, name='rq_metrics'),
+    ]
+except ImportError:
+    metrics_view = []
+
 urlpatterns = [
     re_path(r'^$', views.stats, name='rq_home'),
     re_path(r'^stats.json/(?P<token>[\w]+)?/?$', views.stats_json, name='rq_home_json'),
+    *metrics_view,
     re_path(r'^queues/(?P<queue_index>[\d]+)/$', views.jobs, name='rq_jobs'),
     re_path(r'^workers/(?P<queue_index>[\d]+)/$', views.workers, name='rq_workers'),
     re_path(r'^workers/(?P<queue_index>[\d]+)/(?P<key>[-\w\.\:\$]+)/$', views.worker_details, name='rq_worker_details'),
diff --git a/django_rq/views.py b/django_rq/views.py
index 95f5a43f..5c5a1add 100644
--- a/django_rq/views.py
+++ b/django_rq/views.py
@@ -4,7 +4,7 @@
 
 from django.contrib import admin, messages
 from django.contrib.admin.views.decorators import staff_member_required
-from django.http import Http404, JsonResponse
+from django.http import Http404, HttpResponse, JsonResponse
 from django.shortcuts import redirect, render
 from django.urls import reverse
 from django.views.decorators.cache import never_cache
@@ -27,6 +27,11 @@
 from .settings import API_TOKEN, QUEUES_MAP
 from .utils import get_jobs, get_scheduler_statistics, get_statistics, stop_jobs
 
+try:
+    import prometheus_client
+except ImportError:
+    prometheus_client = None
+
 
 @never_cache
 @staff_member_required
@@ -48,6 +53,20 @@ def stats_json(request, token=None):
     )
 
 
+@never_cache
+@staff_member_required
+def prometheus_metrics(request):
+    if not prometheus_client:
+        raise Http404
+
+    registry = prometheus_client.REGISTRY
+    encoder, content_type = prometheus_client.exposition.choose_encoder(request.META.get('HTTP_ACCEPT', ''))
+    if 'name[]' in request.GET:
+        registry = registry.restricted_registry(request.GET.getlist('name[]'))
+
+    return HttpResponse(encoder(registry), headers={'Content-Type': content_type})
+
+
 @never_cache
 @staff_member_required
 def jobs(request, queue_index):
diff --git a/setup.py b/setup.py
index cf02edb0..e464be5d 100644
--- a/setup.py
+++ b/setup.py
@@ -16,6 +16,7 @@
     package_data={'': ['README.rst']},
     install_requires=['django>=3.2', 'rq>=1.14', 'redis>=3'],
     extras_require={
+        'prometheus-metrics': ['prometheus_client>=0.4.0', 'rq-exporter'],
         'Sentry': ['sentry-sdk>=1.0.0'],
         'testing': [],
     },