From 2707d7dd0a1b0178f49296a0fca5774cd61dbd7f Mon Sep 17 00:00:00 2001
From: Simon Kelly <skelly@dimagi.com>
Date: Fri, 3 Nov 2023 14:15:35 +0200
Subject: [PATCH] allow specifying task names to include / exclude from
 tracking

---
 taskbadger/celery.py                    |  2 +-
 taskbadger/systems/celery.py            | 32 +++++++++++++++++++++----
 tests/test_celery_system_integration.py | 17 +++++++++++++
 3 files changed, 46 insertions(+), 5 deletions(-)

diff --git a/taskbadger/celery.py b/taskbadger/celery.py
index 565b1b7..6fe467f 100644
--- a/taskbadger/celery.py
+++ b/taskbadger/celery.py
@@ -144,7 +144,7 @@ def task_publish_handler(sender=None, headers=None, body=None, **kwargs):
         return
 
     celery_system = Badger.current.settings.get_system_by_id("celery")
-    auto_track = celery_system and celery_system.auto_track_tasks
+    auto_track = celery_system and celery_system.track_task(sender)
     manual_track = headers.get("taskbadger_track")
     if not manual_track and not auto_track:
         return
diff --git a/taskbadger/systems/celery.py b/taskbadger/systems/celery.py
index 966fdec..70c2b42 100644
--- a/taskbadger/systems/celery.py
+++ b/taskbadger/systems/celery.py
@@ -1,16 +1,40 @@
+import re
+
 from taskbadger.systems import System
 
 
 class CelerySystemIntegration(System):
     identifier = "celery"
 
-    def __init__(self, auto_track_tasks=True):
+    def __init__(self, auto_track_tasks=True, includes=None, excludes=None):
         """
         Args:
             auto_track_tasks: Automatically track all Celery tasks regardless of whether they are using the
                 `taskbadger.celery.Task` base class.
+            includes: A list of task names to include in tracking. These can be either the full task name
+                (e.g. `myapp.tasks.export_data`) or a regular expression (e.g. `export_.*`). If a task name
+                matches both an include and an exclude, it will be excluded.
+            excludes: A list of task names to exclude from tracking. As with `includes`, these can be either
+                the full task name or a regular expression. Exclusions take precedence over inclusions.
         """
         self.auto_track_tasks = auto_track_tasks
-        if auto_track_tasks:
-            # Importing this here ensures that the Celery signal handlers are registered
-            import taskbadger.celery  # noqa
+        self.includes = includes
+        self.excludes = excludes
+
+    def track_task(self, task_name):
+        if not self.auto_track_tasks:
+            return False
+
+        if self.excludes:
+            for exclude in self.excludes:
+                if re.fullmatch(exclude, task_name):
+                    return False
+
+        if self.includes:
+            for include in self.includes:
+                if re.fullmatch(include, task_name):
+                    break
+            else:
+                return False
+
+        return True
diff --git a/tests/test_celery_system_integration.py b/tests/test_celery_system_integration.py
index 8e43df5..5c834be 100644
--- a/tests/test_celery_system_integration.py
+++ b/tests/test_celery_system_integration.py
@@ -60,3 +60,20 @@ def add_normal(self, a, b):
     assert get_task.call_count == 1
     assert update.call_count == 2
     assert Badger.current.session().client is None
+
+
+@pytest.mark.parametrize(
+    "include,exclude,expected",
+    [
+        (None, None, True),
+        (["myapp.tasks.export_data"], None, True),
+        ([".*export_data"], [], True),
+        ([".*export_da"], [], False),
+        (["myapp.tasks.export_data"], ["myapp.tasks.export_data"], False),
+        ([".*"], ["myapp.tasks.export_data"], False),
+        ([".*"], [".*tasks.*"], False),
+    ],
+)
+def test_task_name_matching(include, exclude, expected: bool):
+    integration = CelerySystemIntegration(includes=include, excludes=exclude)
+    assert integration.track_task("myapp.tasks.export_data") is expected