[CHANGE] added celery exporter to accept different queue name

celery_exporter/__main__.py

@@ -11,7 +11,7 @@
 from .utils import generate_broker_use_ssl, get_transport_scheme
 
 LOG_FORMAT = "[%(asctime)s] %(name)s:%(levelname)s: %(message)s"
-
+CELERY_DEFAULT_QUEUE = "celery"
 
 @click.command(context_settings={"auto_envvar_prefix": "CELERY_EXPORTER"})
 @click.option(
@@ -102,6 +102,15 @@
 @click.option(
     "--verbose", is_flag=True, allow_from_autoenv=False, help="Enable verbose logging."
 )
+@click.option(
+    "--queue",
+    "-q",
+    type=str,
+    show_default=True,
+    show_envvar=True,
+    default=CELERY_DEFAULT_QUEUE,
+    help="Queue name for metrics.",
+)
 def main(
     broker_url,
     listen_address,
@@ -116,6 +125,7 @@ def main(
     ssl_keyfile,
     tz,
     verbose,
+    queue,
 ):  # pragma: no cover
 
     if verbose:
@@ -155,6 +165,7 @@ def main(
         transport_options,
         enable_events,
         broker_use_ssl,
+        queue,
     )
 
     celery_exporter.start()

celery_exporter/core.py

@@ -23,6 +23,7 @@ def __init__(
         transport_options=None,
         enable_events=False,
         broker_use_ssl=None,
+        queue='celery',
     ):
         self._listen_address = listen_address
         self._max_tasks = max_tasks
@@ -31,17 +32,19 @@ def __init__(
 
         self._app = celery.Celery(broker=broker_url, broker_use_ssl=broker_use_ssl)
         self._app.conf.broker_transport_options = transport_options or {}
+        self._queue = queue
 
     def start(self):
 
-        setup_metrics(self._app, self._namespace)
+        setup_metrics(self._app, self._namespace, self._queue)
 
         self._start_httpd()
 
         t = TaskThread(
             app=self._app,
             namespace=self._namespace,
             max_tasks_in_memory=self._max_tasks,
+            queue=self._queue
         )
         t.daemon = True
         t.start()

celery_exporter/monitor.py

@@ -18,9 +18,10 @@ class TaskThread(threading.Thread):
     exposed from Celery using its eventing system.
     """
 
-    def __init__(self, app, namespace, max_tasks_in_memory, *args, **kwargs):
+    def __init__(self, app, namespace, max_tasks_in_memory, queue, *args, **kwargs):
         self._app = app
         self._namespace = namespace
+        self._queue = queue
         self.log = logging.getLogger("task-thread")
         self._state = CeleryState(max_tasks_in_memory=max_tasks_in_memory)
         self._known_states = set()
@@ -56,12 +57,12 @@ def _monitor(self):  # pragma: no cover
                     recv = self._app.events.Receiver(
                         conn, handlers={"*": self._process_event}
                     )
-                    setup_metrics(self._app, self._namespace)
+                    setup_metrics(self._app, self._namespace, self._queue)
                     self.log.info("Start capturing events...")
                     recv.capture(limit=None, timeout=None, wakeup=True)
             except Exception:
                 self.log.exception("Connection failed")
-                setup_metrics(self._app, self._namespace)
+                setup_metrics(self._app, self._namespace, self._queue)
                 time.sleep(5)
 
 
@@ -109,13 +110,13 @@ def enable_events(self):
         self._app.control.enable_events()
 
 
-def setup_metrics(app, namespace):
+def setup_metrics(app, namespace, queue='celery'):
     """
     This initializes the available metrics with default values so that
     even before the first event is received, data can be exposed.
     """
     WORKERS.labels(namespace=namespace)
-    config = get_config(app)
+    config = get_config(app, queue=queue)
 
     if not config:  # pragma: no cover
         for metric in TASKS.collect():

celery_exporter/utils.py

@@ -2,7 +2,6 @@
 from itertools import chain
 from urllib.parse import urlparse
 
-CELERY_DEFAULT_QUEUE = "celery"
 CELERY_MISSING_DATA = "undefined"
 
 
@@ -15,7 +14,7 @@ def _gen_wildcards(name):
     return res
 
 
-def get_config(app):
+def get_config(app, queue):
     res = dict()
     try:
         registered_tasks = app.control.inspect().registered_tasks().values()
@@ -26,7 +25,7 @@ def get_config(app):
     default_queues = []
     for task_name in set(chain.from_iterable(registered_tasks)):
         for conf in confs.values():
-            default = conf.get("task_default_queue", CELERY_DEFAULT_QUEUE)
+            default = conf.get("task_default_queue", queue)
             default_queues.append(default)
             if task_name in res and res[task_name] not in default_queues:
                 break

test/test_func.py

@@ -28,6 +28,7 @@ def setUp(self):
             max_tasks=TestCeleryExporter.max_tasks,
             namespace=TestCeleryExporter.namespace,
             enable_events=True,
+            queue='celery',
         )
 
     def test_setup_metrics(self):

test/test_unit.py

@@ -36,7 +36,7 @@ def setUp(self):
                     "celery@12311847jsa2": {},
                 }
                 registered.return_value = {"celery@d6f95e9e24fc": [self.task, "trial"]}
-                setup_metrics(self.app, self.namespace)  # reset metrics
+                setup_metrics(self.app, self.namespace, queue=self.queue)  # reset metrics
 
     def test_initial_metric_values(self):
         self._assert_task_states(celery.states.ALL_STATES, 0)
@@ -70,7 +70,7 @@ def test_workers_count(self):
         )
 
         with patch.object(self.app.control, "ping") as mock_ping:
-            w = WorkerMonitoringThread(app=self.app, namespace=self.namespace)
+            w = WorkerMonitoringThread(app=self.app, namespace=self.namespace, queue=self.queue)
 
             mock_ping.return_value = []
             w.update_workers_count()
@@ -116,7 +116,7 @@ def test_tasks_events(self):
         runtime = 234.5
 
         m = TaskThread(
-            app=self.app, namespace=self.namespace, max_tasks_in_memory=self.max_tasks
+            app=self.app, namespace=self.namespace, max_tasks_in_memory=self.max_tasks, queue=self.queue
         )
 
         self._assert_task_states(celery.states.ALL_STATES, 0)