diff --git a/README.md b/README.md index 327cec3..828da7b 100644 --- a/README.md +++ b/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}) +``` diff --git a/taskbadger/celery.py b/taskbadger/celery.py index 14f0802..26c4e67 100644 --- a/taskbadger/celery.py +++ b/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 diff --git a/tests/test_celery.py b/tests/test_celery.py index e86d565..56ca7bc 100644 --- a/tests/test_celery.py +++ b/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."""