From f309eb75e7cbe4effa81e578163e4df09f4721d3 Mon Sep 17 00:00:00 2001
From: sankit <sankit.a@logically.co.uk>
Date: Thu, 3 Jun 2021 14:35:29 +0530
Subject: [PATCH 1/2] feat: Custom route

get_conf function only checks routes that have queue as a key
Added --route-options to added custom route if defined in celery
---
 celery_exporter/__main__.py | 20 +++++++++++++++++++-
 celery_exporter/core.py     |  3 +++
 celery_exporter/utils.py    |  7 +++++--
 3 files changed, 27 insertions(+), 3 deletions(-)

diff --git a/celery_exporter/__main__.py b/celery_exporter/__main__.py
index 2737dcd..6803773 100644
--- a/celery_exporter/__main__.py
+++ b/celery_exporter/__main__.py
@@ -12,7 +12,6 @@
 
 LOG_FORMAT = "[%(asctime)s] %(name)s:%(levelname)s: %(message)s"
 
-
 @click.command(context_settings={"auto_envvar_prefix": "CELERY_EXPORTER"})
 @click.option(
     "--broker-url",
@@ -56,6 +55,12 @@
     allow_from_autoenv=False,
     help="JSON object with additional options passed to the underlying transport.",
 )
+@click.option(
+    "--route-options",
+    type=str,
+    allow_from_autoenv=False,
+    help="JSON object with additional options passed to the underlying route.",
+)
 @click.option(
     "--enable-events",
     is_flag=True,
@@ -108,6 +113,7 @@ def main(
     max_tasks,
     namespace,
     transport_options,
+    route_options,
     enable_events,
     use_ssl,
     ssl_verify,
@@ -137,6 +143,17 @@ def main(
                 )
             )
             sys.exit(1)
+    
+    if route_options:
+        try:
+            route_option = json.loads(route_options)
+        except ValueError:
+            logging.error(
+                "Error parsing broker transport options from JSON '{}'".format(
+                    route_options
+                )
+            )
+            sys.exit(1)
 
     broker_use_ssl = generate_broker_use_ssl(
         use_ssl,
@@ -153,6 +170,7 @@ def main(
         max_tasks,
         namespace,
         transport_options,
+        route_options,
         enable_events,
         broker_use_ssl,
     )
diff --git a/celery_exporter/core.py b/celery_exporter/core.py
index f0612ee..64efdd9 100644
--- a/celery_exporter/core.py
+++ b/celery_exporter/core.py
@@ -23,6 +23,7 @@ def __init__(
         transport_options=None,
         enable_events=False,
         broker_use_ssl=None,
+        route_options=None
     ):
         self._listen_address = listen_address
         self._max_tasks = max_tasks
@@ -31,6 +32,8 @@ def __init__(
 
         self._app = celery.Celery(broker=broker_url, broker_use_ssl=broker_use_ssl)
         self._app.conf.broker_transport_options = transport_options or {}
+        if route_options:
+            self._app.conf.task_routes = route_options
 
     def start(self):
 
diff --git a/celery_exporter/utils.py b/celery_exporter/utils.py
index 7c2317c..825d603 100644
--- a/celery_exporter/utils.py
+++ b/celery_exporter/utils.py
@@ -36,8 +36,11 @@ def get_config(app):
                 routes = conf["task_routes"]
                 res[task_name] = default
                 for i in task_wildcard_names:
-                    if i in routes and "queue" in routes[i]:
-                        res[task_name] = routes[i]["queue"]
+                    if i in routes:
+                        if 'queue' in routes[i]:
+                            res[task_name] = routes[i]['queue']
+                        else:
+                            res[task_name] = routes[i]
                         break
             else:
                 res[task_name] = default

From 9736fe0d233bfc27995d1ea514b15b9ce67cdc01 Mon Sep 17 00:00:00 2001
From: sankit <sankit.a@logically.co.uk>
Date: Sat, 12 Jun 2021 13:12:06 +0530
Subject: [PATCH 2/2] fix key mapping in utils

---
 celery_exporter/utils.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/celery_exporter/utils.py b/celery_exporter/utils.py
index 825d603..71182d2 100644
--- a/celery_exporter/utils.py
+++ b/celery_exporter/utils.py
@@ -37,7 +37,7 @@ def get_config(app):
                 res[task_name] = default
                 for i in task_wildcard_names:
                     if i in routes:
-                        if 'queue' in routes[i]:
+                        if 'queue' in routes:
                             res[task_name] = routes[i]['queue']
                         else:
                             res[task_name] = routes[i]