Merge pull request #62 from taskbadger/sk/celery-track-opt-out

Honour taskbadger_track=False in Celery

Author: snopoke

README.md

@@ -125,3 +125,18 @@ def slow_job():
 
 Unless `stale_timeout` is given explicitly it is set to twice the interval. All running tasks are
 pinged from a single background thread, started the first time a task with a heartbeat runs.
+
+### Skipping tracking for a single Celery call
+
+Pass `taskbadger_track=False` to leave one execution untracked. This overrides auto-tracking as
+well as the `taskbadger.Task` base class:
+
+```python
+noisy_job.apply_async(args, taskbadger_track=False)
+```
+
+Canvas primitives don't go through `apply_async` on the task itself, so pass it in the headers:
+
+```python
+noisy_job.map(items).apply_async(headers={"taskbadger_track": False})
+```

taskbadger/celery.py

@@ -25,7 +25,16 @@
 
 KWARG_PREFIX = "taskbadger_"
 TB_KWARGS_ARG = f"{KWARG_PREFIX}kwargs"
-IGNORE_ARGS = {TB_KWARGS_ARG, f"{KWARG_PREFIX}task", f"{KWARG_PREFIX}task_id", f"{KWARG_PREFIX}record_task_args"}
+# Per-execution tracking switch carried on the message headers. Not a task field, so it
+# must never end up in the kwargs passed to `create_task`.
+TB_TRACK = f"{KWARG_PREFIX}track"
+IGNORE_ARGS = {
+    TB_KWARGS_ARG,
+    TB_TRACK,
+    f"{KWARG_PREFIX}task",
+    f"{KWARG_PREFIX}task_id",
+    f"{KWARG_PREFIX}record_task_args",
+}
 TB_TASK_ID = f"{KWARG_PREFIX}task_id"
 TB_HEARTBEAT_INTERVAL = f"{KWARG_PREFIX}heartbeat_interval"
 TB_STALE_TIMEOUT = f"{KWARG_PREFIX}stale_timeout"
@@ -61,6 +70,12 @@ class Task(celery.Task):
     keeps tasks with a `stale_timeout` from going stale. Unless `taskbadger_stale_timeout`
     is also given it is set to twice the interval.
 
+    A single execution can opt out of tracking with `taskbadger_track=False`, either as an
+    argument to `apply_async` or in its `headers`. This also overrides auto-tracking, and
+    works for canvas tasks (`.map()` / `.starmap()`), which only accept it via `headers`.
+    It is per-execution only — to exclude a task permanently use
+    `CelerySystemIntegration(excludes=[...])` rather than setting it on the task.
+
     Access to the task is provided via the `taskbadger_task` property of the Celery task.
     The task ID may also be accessed via the `taskbadger_task_id` property. These may
     be `None` if the task is not being tracked (e.g. Task Badger is not configured or
@@ -99,7 +114,15 @@ def apply_async(self, *args, **kwargs):
             tb_kwargs.update(self._get_tb_kwargs(args[1]))
 
         if Badger.is_configured():
-            headers["taskbadger_track"] = True
+            # An explicit `taskbadger_track=False` opts this execution out of tracking and
+            # must survive to the signal handlers. It arrives either as a `taskbadger_`
+            # prefixed argument (already extracted into `tb_kwargs`) or straight in the
+            # headers, hence the `setdefault` for the latter.
+            track = tb_kwargs.pop("track", None)
+            if track is None:
+                headers.setdefault(TB_TRACK, True)
+            else:
+                headers[TB_TRACK] = track
             headers[TB_KWARGS_ARG] = tb_kwargs
             if "record_task_args" in tb_kwargs:
                 headers["taskbadger_record_task_args"] = tb_kwargs.pop("record_task_args")
@@ -186,7 +209,10 @@ def task_publish_handler(sender=None, headers=None, body=None, **kwargs):
 
     celery_system = Badger.current.settings.get_system_by_id("celery")
     auto_track = celery_system and celery_system.track_task(sender)
-    manual_track = headers.get("taskbadger_track")
+    manual_track = headers.get(TB_TRACK)
+    if manual_track is False:
+        # explicit opt-out for this execution, which also overrides auto-tracking
+        return
     if not manual_track and not auto_track:
         return
 
@@ -271,7 +297,12 @@ def _maybe_create_task(signal_sender):
     # Badger wasn't configured at publish time but has stale config in worker.
     headers = signal_sender.request.headers or {}
     is_canvas_task = task_name in ("celery.map", "celery.starmap")
-    if not is_canvas_task and not headers.get("taskbadger_track"):
+    track_header = headers.get(TB_TRACK)
+    if track_header is False:
+        # explicit opt-out, which canvas tasks honour too: they are only ever created
+        # here, so this is the one place their header is checked
+        return
+    if not is_canvas_task and not track_header:
         return
 
     # NOW it's safe to check Badger configuration

tests/test_celery.py

@@ -18,7 +18,8 @@
 
 from taskbadger import Action, EmailIntegration, StatusEnum
 from taskbadger.celery import Task, task_publish_handler
-from taskbadger.mug import Badger
+from taskbadger.mug import Badger, Settings
+from taskbadger.systems.celery import CelerySystemIntegration
 from tests.utils import task_for_test
 
 
@@ -297,6 +298,85 @@ def test_celery_publish_handler_task_not_registered_locally():
     assert headers["taskbadger_task_id"] == create.return_value.id
 
 
+@pytest.mark.parametrize(("track_header", "expect_created"), [({}, True), ({"taskbadger_track": False}, False)])
+def test_celery_publish_handler_opt_out_beats_auto_track(track_header, expect_created):
+    """An explicit `taskbadger_track=False` header opts a single execution out, even
+    when auto-tracking would otherwise pick the task up."""
+    settings = Settings("https://taskbadger.net", "token", "org", "proj", systems={"celery": CelerySystemIntegration()})
+    Badger.current.bind(settings)
+    try:
+        with mock.patch("taskbadger.celery.create_task_safe") as create:
+            create.return_value = task_for_test()
+            headers = {"id": "abc123", "task": "auto.tracked.task", **track_header}
+            task_publish_handler(sender="auto.tracked.task", headers=headers, body=[[], {}, {}])
+    finally:
+        Badger.current.bind(None)
+
+    assert create.called is expect_created
+
+
+@pytest.mark.usefixtures("_bind_settings")
+def test_celery_track_attr_not_passed_to_create(celery_session_app):
+    """`taskbadger_track` is a tracking switch, not a task field, so it must never reach
+    `create_task` — which would raise on the unexpected kwarg and lose the task."""
+
+    @celery_session_app.task(base=Task, name="track.attr.task", taskbadger_track=False)
+    def track_attr_task():
+        return 1
+
+    with mock.patch("taskbadger.celery.create_task_safe") as create:
+        create.return_value = task_for_test()
+        headers = {"id": "abc123", "task": "track.attr.task", "taskbadger_track": True}
+        task_publish_handler(sender="track.attr.task", headers=headers, body=[[], {}, {}])
+
+    assert "track" not in create.call_args.kwargs
+
+
+@pytest.mark.usefixtures("_bind_settings")
+def test_celery_task_opt_out(celery_session_app, celery_session_worker):
+    """`headers={"taskbadger_track": False}` prevents tracking of a single execution."""
+
+    @celery_session_app.task(bind=True, base=Task)
+    def add_opt_out(self, a, b):
+        assert self.taskbadger_task_id is None, "task should not be tracked"
+        return a + b
+
+    celery_session_worker.reload()
+
+    with (
+        mock.patch("taskbadger.celery.create_task_safe") as create,
+        mock.patch("taskbadger.celery.update_task_safe") as update,
+    ):
+        result = add_opt_out.apply_async((2, 2), headers={"taskbadger_track": False})
+        assert result.get(timeout=10, propagate=True) == 4
+
+    create.assert_not_called()
+    update.assert_not_called()
+
+
+@pytest.mark.usefixtures("_bind_settings")
+def test_celery_task_opt_out_kwarg(celery_session_app, celery_session_worker):
+    """`taskbadger_track=False` also works as a `taskbadger_`-prefixed option, the way the
+    other per-call options are passed, and never leaks into the create_task kwargs."""
+
+    @celery_session_app.task(bind=True, base=Task)
+    def add_opt_out_kwarg(self, a, b):
+        assert self.taskbadger_task_id is None, "task should not be tracked"
+        return a + b
+
+    celery_session_worker.reload()
+
+    with (
+        mock.patch("taskbadger.celery.create_task_safe") as create,
+        mock.patch("taskbadger.celery.update_task_safe") as update,
+    ):
+        result = add_opt_out_kwarg.apply_async((2, 2), taskbadger_track=False)
+        assert result.get(timeout=10, propagate=True) == 4
+
+    create.assert_not_called()
+    update.assert_not_called()
+
+
 @pytest.mark.usefixtures("_bind_settings")
 def test_celery_task_custom_queue(celery_session_app, celery_session_worker):
     @celery_session_app.task(bind=True, base=Task)
@@ -542,6 +622,30 @@ def task_map_fn(self, a):
     assert Badger.current.session().client is None
 
 
+@pytest.mark.usefixtures("_bind_settings")
+def test_task_map_opt_out(celery_session_worker):
+    """Canvas tasks honour the opt-out too. They are created in the worker rather than
+    at publish time, so the header is checked there."""
+
+    @celery.shared_task(bind=True, base=Task)
+    def task_map_opt_out_fn(self, a):
+        return a * 2
+
+    celery_session_worker.reload()
+
+    map_canvas = task_map_opt_out_fn.map(list(range(3)))
+
+    with (
+        mock.patch("taskbadger.celery.create_task_safe") as create,
+        mock.patch("taskbadger.celery.update_task_safe") as update,
+    ):
+        result = map_canvas.apply_async(headers={"taskbadger_track": False})
+        assert result.get(timeout=10, propagate=True) == [0, 2, 4]
+
+    create.assert_not_called()
+    update.assert_not_called()
+
+
 @pytest.mark.usefixtures("_bind_settings")
 def test_task_starmap(celery_session_worker):
     """Tasks executed via starmap canvas primitive should be tracked."""