taskbadger 2.0.0__tar.gz → 2.1.0a2__tar.gz

{taskbadger-2.0.0 → taskbadger-2.1.0a2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: taskbadger
-Version: 2.0.0
+Version: 2.1.0a2
 Summary: The official Python SDK for Task Badger
 Project-URL: Changelog, https://github.com/taskbadger/taskbadger-python/releases
 Project-URL: homepage, https://taskbadger.net/
@@ -28,6 +28,8 @@ Requires-Dist: celery<6.0.0,>=4.0.0; extra == 'celery'
 Provides-Extra: cli
 Requires-Dist: rich>=13.0; extra == 'cli'
 Requires-Dist: typer>=0.12; extra == 'cli'
+Provides-Extra: procrastinate
+Requires-Dist: procrastinate>=3.0; extra == 'procrastinate'
 Description-Content-Type: text/markdown
 
 # Task Badger Python Client
@@ -73,3 +75,60 @@ taskbadger.init(
 $ export TASKBADGER_API_KEY=***
 $ taskbadger run "nightly-backup" -- ./backup.sh
 ```
+
+### Procrastinate Integration
+
+The SDK includes optional support for the [Procrastinate](https://procrastinate.readthedocs.io/) task queue.
+
+Install with the extra:
+
+```bash
+pip install 'taskbadger[procrastinate]'
+```
+
+Opt a single task into tracking with the `track` decorator:
+
+```python
+import procrastinate
+from taskbadger.procrastinate import track, current_task
+
+app = procrastinate.App(connector=...)
+
+@track
+@app.task(queue="default")
+async def add(a, b):
+    return a + b
+
+@track(name="report", value_max=100, tags={"env": "prod"})
+@app.task
+async def report(rows):
+    tb = current_task()
+    for i, row in enumerate(rows):
+        await process(row)
+        if i % 10 == 0:
+            tb.update(value=i)
+```
+
+To auto-track every task on an App, register the system integration:
+
+```python
+import taskbadger
+from taskbadger.systems.procrastinate import ProcrastinateSystemIntegration
+
+taskbadger.init(
+    token="***",
+    systems=[ProcrastinateSystemIntegration(
+        app=app,
+        auto_track_tasks=True,
+        includes=[r"myapp\..*"],
+        excludes=[r"myapp\.cleanup\..*"],
+        record_task_args=True,
+    )],
+)
+```
+
+#### Known limitations
+
+- **`task.configure(...).defer(...)` is not tracked.** Procrastinate's `configure()` returns a separate `JobDeferrer` whose methods bypass our wrapper. Use `task.defer(...)` directly for tracked deferrals. Tasks deferred via `configure().defer()` will run normally but will not appear in TaskBadger.
+- **`task.batch_defer*` is not tracked.** Same reason as `configure().defer()`.
+- **Tasks added via `app.add_tasks_from(blueprint)` after `ProcrastinateSystemIntegration` is constructed are not auto-instrumented.** Construct the integration after all blueprints are registered, or apply `@track` to those tasks explicitly.

taskbadger-2.1.0a2/README.md

@@ -0,0 +1,100 @@
+# Task Badger Python Client
+
+This is the official Python SDK for [Task Badger](https://taskbadger.net/).
+
+For full documentation go to https://docs.taskbadger.net/python/.
+
+![Integration Tests](https://github.com/taskbadger/taskbadger-python/actions/workflows/integration_tests.yml/badge.svg)
+
+---
+
+## Getting Started
+
+### Install
+
+```bash
+pip install taskbadger
+```
+
+To use the `taskbadger` command-line tool, install the `cli` extra:
+
+```bash
+pip install 'taskbadger[cli]'
+```
+
+### Client Usage
+
+```python
+import taskbadger
+from taskbadger.systems import CelerySystemIntegration
+
+taskbadger.init(
+    token="***",
+    systems=[CelerySystemIntegration()],
+    tags={"environment": "production"}
+)
+```
+
+### CLI Usage
+
+```shell
+$ export TASKBADGER_API_KEY=***
+$ taskbadger run "nightly-backup" -- ./backup.sh
+```
+
+### Procrastinate Integration
+
+The SDK includes optional support for the [Procrastinate](https://procrastinate.readthedocs.io/) task queue.
+
+Install with the extra:
+
+```bash
+pip install 'taskbadger[procrastinate]'
+```
+
+Opt a single task into tracking with the `track` decorator:
+
+```python
+import procrastinate
+from taskbadger.procrastinate import track, current_task
+
+app = procrastinate.App(connector=...)
+
+@track
+@app.task(queue="default")
+async def add(a, b):
+    return a + b
+
+@track(name="report", value_max=100, tags={"env": "prod"})
+@app.task
+async def report(rows):
+    tb = current_task()
+    for i, row in enumerate(rows):
+        await process(row)
+        if i % 10 == 0:
+            tb.update(value=i)
+```
+
+To auto-track every task on an App, register the system integration:
+
+```python
+import taskbadger
+from taskbadger.systems.procrastinate import ProcrastinateSystemIntegration
+
+taskbadger.init(
+    token="***",
+    systems=[ProcrastinateSystemIntegration(
+        app=app,
+        auto_track_tasks=True,
+        includes=[r"myapp\..*"],
+        excludes=[r"myapp\.cleanup\..*"],
+        record_task_args=True,
+    )],
+)
+```
+
+#### Known limitations
+
+- **`task.configure(...).defer(...)` is not tracked.** Procrastinate's `configure()` returns a separate `JobDeferrer` whose methods bypass our wrapper. Use `task.defer(...)` directly for tracked deferrals. Tasks deferred via `configure().defer()` will run normally but will not appear in TaskBadger.
+- **`task.batch_defer*` is not tracked.** Same reason as `configure().defer()`.
+- **Tasks added via `app.add_tasks_from(blueprint)` after `ProcrastinateSystemIntegration` is constructed are not auto-instrumented.** Construct the integration after all blueprints are registered, or apply `@track` to those tasks explicitly.

{taskbadger-2.0.0 → taskbadger-2.1.0a2}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "taskbadger"
-version = "2.0.0"
+version = "2.1.0a2"
 description = "The official Python SDK for Task Badger"
 requires-python = ">=3.10"
 authors = []
@@ -44,6 +44,9 @@ cli = [
     "typer >=0.12",
     "rich >=13.0",
 ]
+procrastinate = [
+    "procrastinate>=3.0",
+]
 
 [tool.uv]
 package = true
@@ -58,6 +61,7 @@ documentation = "https://docs.taskbadger.net/"
 [dependency-groups]
 dev = [
     "pre-commit",
+    "procrastinate",
     "pytest",
     "pytest-httpx",
     "invoke",

taskbadger-2.1.0a2/taskbadger/_integrations.py

@@ -0,0 +1,110 @@
+"""Shared internals for taskbadger's optional system integrations
+(Celery, Procrastinate). Not part of the public API.
+
+A single module-level ``TaskCache`` (``task_cache``) is shared across all
+integrations; task ids are UUIDs so cross-integration key collisions are not
+a concern. ``BaseSystemIntegration`` provides the common ctor/include-exclude
+shape; subclasses override ``track_task`` if they need to filter additional
+task names (e.g. Procrastinate built-ins).
+"""
+
+from __future__ import annotations
+
+import collections
+import logging
+import re
+
+from . import sdk
+from .internal.models import StatusEnum
+from .systems import System
+
+log = logging.getLogger("taskbadger")
+
+TERMINAL_STATES = {
+    StatusEnum.SUCCESS,
+    StatusEnum.ERROR,
+    StatusEnum.CANCELLED,
+    StatusEnum.STALE,
+}
+
+
+class TaskCache:
+    """Bounded LRU-ish cache for TaskBadger Task objects.
+
+    Keys are arbitrary hashable values chosen by the caller (typically the
+    task id). Auto-prunes on ``set`` when ``maxsize`` is exceeded.
+    """
+
+    def __init__(self, maxsize: int = 128):
+        self.cache: collections.OrderedDict = collections.OrderedDict()
+        self.maxsize = maxsize
+
+    def set(self, key, value) -> None:
+        self.cache[key] = value
+        if len(self.cache) > self.maxsize:
+            self.cache.popitem(last=False)
+
+    def get(self, key):
+        return self.cache.get(key)
+
+    def unset(self, key) -> None:
+        self.cache.pop(key, None)
+
+
+task_cache = TaskCache()
+
+
+def safe_get_task(task_id: str):
+    """Cache-aware ``get_task``: returns the cached entry if present, otherwise
+    fetches from the API and caches the result. Errors are logged and swallowed
+    (returns ``None``). ``None`` results are not cached.
+    """
+    cached = task_cache.get(task_id)
+    if cached is not None:
+        return cached
+    try:
+        task = sdk.get_task(task_id)
+    except Exception as e:
+        log.warning("Error fetching task '%s': %s", task_id, e)
+        return None
+    task_cache.set(task_id, task)
+    return task
+
+
+def match_task_name(task_name: str, includes, excludes) -> bool:
+    """Return True if ``task_name`` should be tracked under the given rules.
+
+    Excludes win over includes. Both lists contain regex strings matched with
+    ``re.fullmatch``. ``None`` means "no rule".
+    """
+    if excludes:
+        for exclude in excludes:
+            if re.fullmatch(exclude, task_name):
+                return False
+
+    if includes:
+        for include in includes:
+            if re.fullmatch(include, task_name):
+                return True
+        return False
+
+    return True
+
+
+class BaseSystemIntegration(System):
+    """Common ctor + ``track_task`` body for system integrations.
+
+    Subclasses set ``identifier`` and may override ``track_task`` to add
+    additional filtering (e.g. skipping built-in tasks).
+    """
+
+    def __init__(self, auto_track_tasks=True, includes=None, excludes=None, record_task_args=False):
+        self.auto_track_tasks = auto_track_tasks
+        self.includes = includes
+        self.excludes = excludes
+        self.record_task_args = record_task_args
+
+    def track_task(self, task_name: str) -> bool:
+        if not self.auto_track_tasks:
+            return False
+        return match_task_name(task_name, self.includes, self.excludes)

{taskbadger-2.0.0 → taskbadger-2.1.0a2}/taskbadger/celery.py

@@ -1,4 +1,3 @@
-import collections
 import functools
 import json
 import logging
@@ -13,66 +12,21 @@ from celery.signals import (
 )
 from kombu import serialization
 
+from . import sdk
+from ._integrations import TERMINAL_STATES, safe_get_task, task_cache
 from .internal.models import StatusEnum
 from .mug import Badger
 from .safe_sdk import create_task_safe, update_task_safe
-from .sdk import DefaultMergeStrategy, get_task
+from .sdk import DefaultMergeStrategy
 
 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"}
 TB_TASK_ID = f"{KWARG_PREFIX}task_id"
 
-TERMINAL_STATES = {
-    StatusEnum.SUCCESS,
-    StatusEnum.ERROR,
-    StatusEnum.CANCELLED,
-    StatusEnum.STALE,
-}
-
 log = logging.getLogger("taskbadger")
 
 
-class Cache:
-    def __init__(self, maxsize=128):
-        self.cache = collections.OrderedDict()
-        self.maxsize = maxsize
-
-    def set(self, key, value):
-        self.cache[key] = value
-
-    def unset(self, key):
-        self.cache.pop(key, None)
-
-    def get(self, key):
-        return self.cache.get(key)
-
-    def prune(self):
-        if len(self.cache) > self.maxsize:
-            self.cache.popitem(last=False)
-
-
-def cached(cache_none=True, maxsize=128):
-    cache = Cache(maxsize=maxsize)
-
-    def _wrapper(func):
-        @functools.wraps(func)
-        def _inner(*args, **kwargs):
-            key = args + tuple(sorted(kwargs.items()))
-            if key in cache.cache:
-                return cache.get(key)
-
-            result = func(*args, **kwargs)
-            if result is not None or cache_none:
-                cache.set(key, result)
-            return result
-
-        _inner.cache = cache
-        return _inner
-
-    return _wrapper
-
-
 class Task(celery.Task):
     """A Celery Task that tracks itself with TaskBadger.
 
@@ -135,7 +89,7 @@ class Task(celery.Task):
         tb_task_id = info.get(TB_TASK_ID) if isinstance(info, dict) else None
         setattr(result, TB_TASK_ID, tb_task_id)
 
-        _get_task = functools.partial(get_task, tb_task_id) if tb_task_id else lambda: None
+        _get_task = functools.partial(sdk.get_task, tb_task_id) if tb_task_id else lambda: None
         setattr(result, "get_taskbadger_task", _get_task)
 
         return result
@@ -292,7 +246,7 @@ def _maybe_create_task(signal_sender):
     if task:
         # Store the task ID in the request so _update_task can find it
         signal_sender.request.update({TB_TASK_ID: task.id})
-        safe_get_task.cache.set((task.id,), task)
+        task_cache.set(task.id, task)
 
 
 @task_prerun.connect
@@ -344,7 +298,7 @@ def _update_task(signal_sender, status, einfo=None):
         data = DefaultMergeStrategy().merge(task.data, {"exception": str(einfo)})
     task = update_task_safe(task.id, status=status, data=data)
     if task:
-        safe_get_task.cache.set((task_id,), task)
+        task_cache.set(task_id, task)
 
 
 def enter_session():
@@ -364,22 +318,13 @@ def exit_session(signal_sender):
     if not task_id or not Badger.is_configured():
         return
 
-    safe_get_task.cache.unset((task_id,))
-    safe_get_task.cache.prune()
+    task_cache.unset(task_id)
 
     session = Badger.current.session()
     if session.client:
         session.__exit__()
 
 
-@cached(cache_none=False)
-def safe_get_task(task_id: str):
-    try:
-        return get_task(task_id)
-    except Exception as e:
-        log.warning("Error fetching task '%s': %s", task_id, e)
-
-
 def _get_taskbadger_task_id(request):
     if not request:
         return

taskbadger-2.1.0a2/taskbadger/procrastinate.py

@@ -0,0 +1,335 @@
+"""TaskBadger integration for the Procrastinate task queue.
+
+This module is opt-in. Users install Procrastinate themselves (or via the
+``taskbadger[procrastinate]`` extra) and import from here.
+
+Public API:
+    - ``track``: decorator to opt a single task into TaskBadger tracking.
+    - ``current_task()``: accessor for the TaskBadger task associated with the
+      currently-running Procrastinate job (returns ``None`` if not tracked).
+"""
+
+from __future__ import annotations
+
+import functools
+import inspect
+import json
+import logging
+from contextvars import ContextVar
+
+from ._integrations import TERMINAL_STATES, safe_get_task, task_cache
+from .internal.models import StatusEnum
+from .mug import Badger
+from .safe_sdk import create_task_safe, update_task_safe
+from .sdk import DefaultMergeStrategy
+
+log = logging.getLogger("taskbadger")
+
+# Reserved key used to smuggle the TaskBadger task id through Procrastinate's
+# task_kwargs from the deferring process to the worker. Stripped before the
+# user function is called.
+TB_TASK_ID_KWARG = "__taskbadger_task_id__"
+
+# Sentinel attribute names set on a Procrastinate Task object once it has been
+# instrumented. Used to make instrumentation idempotent.
+_INSTRUMENTED_ATTR = "_taskbadger_instrumented"
+_MANUAL_ATTR = "_taskbadger_manual"
+_OPTS_ATTR = "_taskbadger_opts"
+
+_current_tb_task_id: ContextVar[str | None] = ContextVar("_current_tb_task_id", default=None)
+
+
+def _instrument_task(task, system=None, manual=False, opts=None):
+    """Wrap a Procrastinate Task's ``func`` so the worker side updates TaskBadger.
+
+    Idempotent: a second call on the same task is a no-op (but ``manual`` and
+    ``opts`` will be merged onto the existing attributes if provided).
+    """
+    if opts is not None:
+        existing_opts = getattr(task, _OPTS_ATTR, {}) or {}
+        merged = {**existing_opts, **opts}
+        setattr(task, _OPTS_ATTR, merged)
+    elif not hasattr(task, _OPTS_ATTR):
+        setattr(task, _OPTS_ATTR, {})
+
+    if manual:
+        setattr(task, _MANUAL_ATTR, True)
+
+    if getattr(task, _INSTRUMENTED_ATTR, False):
+        return
+
+    original_func = task.func
+    is_async = inspect.iscoroutinefunction(original_func)
+
+    # pass_context=True works transparently: Procrastinate passes the context
+    # object as the first positional arg; our *args/**kwargs wrapper forwards it.
+    if is_async:
+
+        @functools.wraps(original_func)
+        async def wrapped(*args, **kwargs):
+            tb_id = kwargs.pop(TB_TASK_ID_KWARG, None)
+            if tb_id is None:
+                return await original_func(*args, **kwargs)
+            token = _current_tb_task_id.set(tb_id)
+            try:
+                _update_status(tb_id, StatusEnum.PROCESSING)
+                try:
+                    result = await original_func(*args, **kwargs)
+                except Exception as exc:
+                    _update_status(tb_id, StatusEnum.ERROR, exception=exc)
+                    raise
+                _update_status(tb_id, StatusEnum.SUCCESS)
+                return result
+            finally:
+                _current_tb_task_id.reset(token)
+    else:
+
+        @functools.wraps(original_func)
+        def wrapped(*args, **kwargs):
+            tb_id = kwargs.pop(TB_TASK_ID_KWARG, None)
+            if tb_id is None:
+                return original_func(*args, **kwargs)
+            token = _current_tb_task_id.set(tb_id)
+            try:
+                _update_status(tb_id, StatusEnum.PROCESSING)
+                try:
+                    result = original_func(*args, **kwargs)
+                except Exception as exc:
+                    _update_status(tb_id, StatusEnum.ERROR, exception=exc)
+                    raise
+                _update_status(tb_id, StatusEnum.SUCCESS)
+                return result
+            finally:
+                _current_tb_task_id.reset(token)
+
+    _wrap_defer(task)
+    task.func = wrapped
+    setattr(task, _INSTRUMENTED_ATTR, True)
+    setattr(task, "_taskbadger_system", system)
+
+
+def _update_status(tb_id, status, exception=None):
+    """Update the TaskBadger task to ``status``. Skips if already terminal."""
+    if not Badger.is_configured():
+        return
+
+    if exception is not None or status in TERMINAL_STATES:
+        # Bypass the cache for the terminal-state check: the user may have
+        # updated the task to a terminal state via the regular SDK during
+        # the body, which wouldn't be reflected in our local cache.
+        task_cache.unset(tb_id)
+        current = safe_get_task(tb_id)
+        if current is not None and current.status in TERMINAL_STATES:
+            return
+        data = None
+        if exception is not None and current is not None:
+            base = dict(current.data) if current.data else None
+            data = DefaultMergeStrategy().merge(base, {"exception": str(exception)})
+        if data is not None:
+            updated = update_task_safe(tb_id, status=status, data=data)
+        else:
+            updated = update_task_safe(tb_id, status=status)
+    else:
+        updated = update_task_safe(tb_id, status=status)
+
+    if updated is not None:
+        task_cache.set(tb_id, updated)
+
+
+def _wrap_defer(task):
+    """Wrap ``task.defer`` and ``task.defer_async`` so they create a TaskBadger
+    task in PENDING state and inject its id into the job's task_kwargs.
+
+    Not idempotent on its own — the caller (``_instrument_task``) gates this
+    via ``_INSTRUMENTED_ATTR`` so each task is wrapped at most once.
+    """
+    original_defer = task.defer
+    original_defer_async = task.defer_async
+
+    @functools.wraps(original_defer)
+    def defer(**kwargs):
+        kwargs = _maybe_create_pending(task, kwargs)
+        return original_defer(**kwargs)
+
+    @functools.wraps(original_defer_async)
+    async def defer_async(**kwargs):
+        kwargs = _maybe_create_pending(task, kwargs)
+        return await original_defer_async(**kwargs)
+
+    task.defer = defer
+    task.defer_async = defer_async
+
+
+def _create_pending_task(task, task_kwargs):
+    """Create a PENDING TaskBadger task for ``task`` if it should be tracked.
+
+    Returns the created TaskBadger task, or ``None`` if Badger isn't
+    configured, the task isn't tracked (neither manual nor auto), or the
+    create call failed. ``task_kwargs`` is used only for the
+    ``record_task_args`` data capture.
+    """
+    if not Badger.is_configured():
+        return None
+
+    system = getattr(task, "_taskbadger_system", None)
+    manual = getattr(task, _MANUAL_ATTR, False)
+    auto = bool(system) and system.track_task(task.name)
+    if not manual and not auto:
+        return None
+
+    opts = dict(getattr(task, _OPTS_ATTR, {}) or {})
+    name = opts.pop("name", None) or task.name
+    create_kwargs = {"status": StatusEnum.PENDING}
+    for key in ("value_max", "tags"):
+        if key in opts and opts[key] is not None:
+            create_kwargs[key] = opts[key]
+
+    data = dict(opts.get("data") or {})
+
+    record_args = opts.get("record_task_args")
+    if record_args is None:
+        record_args = bool(system) and system.record_task_args
+    if record_args:
+        data["procrastinate_task_kwargs"] = _serialize_kwargs(task_kwargs)
+
+    if data:
+        create_kwargs["data"] = data
+
+    return create_task_safe(name, **create_kwargs)
+
+
+def _maybe_create_pending(task, kwargs):
+    """Decide whether to track this defer, and if so create the TaskBadger
+    task and inject its id into ``kwargs``. Always returns the kwargs dict."""
+    tb_task = _create_pending_task(task, kwargs)
+    if tb_task is None:
+        return kwargs
+
+    new_kwargs = dict(kwargs)
+    new_kwargs[TB_TASK_ID_KWARG] = tb_task.id
+    return new_kwargs
+
+
+def _serialize_kwargs(kwargs):
+    """Return a JSON-roundtrippable copy of the defer kwargs.
+
+    Procrastinate already requires kwargs be JSON-serializable, so a json
+    dumps/loads roundtrip is safe. Non-serializable values are dropped with
+    a warning."""
+    try:
+        return json.loads(json.dumps(kwargs))
+    except (TypeError, ValueError) as e:
+        log.warning("Error serializing task arguments: %s", e)
+        return {}
+
+
+_TRACK_OPT_KEYS = ("name", "value_max", "tags", "data", "record_task_args")
+
+
+def track(original_task=None, **opts):
+    """Opt a Procrastinate task into TaskBadger tracking.
+
+    Usage:
+
+        @track
+        @app.task(...)
+        def my_task(...): ...
+
+        @track(name="custom", value_max=100, tags={"env": "prod"})
+        @app.task(...)
+        async def big_job(...): ...
+
+    Accepted keyword options (all optional):
+        name: TaskBadger task name (defaults to the Procrastinate task's name).
+        value_max: Maximum value for the TaskBadger task.
+        tags: Dict of tags applied to the TaskBadger task.
+        data: Dict of initial data merged into the TaskBadger task.
+        record_task_args: If True, serialize the Procrastinate job kwargs and
+            store them under ``data["procrastinate_task_kwargs"]``. Defaults to
+            ``None`` meaning "inherit from system integration if any, else False".
+    """
+    unknown = set(opts) - set(_TRACK_OPT_KEYS)
+    if unknown:
+        raise TypeError(f"track() got unexpected keyword arguments: {sorted(unknown)}")
+
+    def wrap(task):
+        _instrument_task(task, system=None, manual=True, opts=opts)
+        return task
+
+    if original_task is None:
+        return wrap
+    return wrap(original_task)
+
+
+def current_task():
+    """Return the TaskBadger Task for the currently-running Procrastinate job.
+
+    Returns ``None`` outside of a tracked task or if the task can't be fetched.
+    Result is cached for the lifetime of the worker process via an LRU.
+    """
+    tb_id = _current_tb_task_id.get()
+    if tb_id is None:
+        return None
+    return safe_get_task(tb_id)
+
+
+def _patch_job_manager(app, system):
+    """Patch ``app.job_manager.defer_periodic_job`` so periodic tasks are tracked.
+
+    Procrastinate's ``PeriodicDeferrer`` enqueues jobs by calling
+    ``job_manager.defer_periodic_job(job=..., ...)`` directly, bypassing
+    ``task.defer``/``defer_async``. Without this hook, ``@app.periodic`` tasks
+    would never get a PENDING TaskBadger task created at enqueue time.
+
+    Idempotent: a second call updates the system reference but doesn't
+    re-wrap.
+    """
+    jm = app.job_manager
+    if not getattr(jm, "_taskbadger_original_defer_periodic_job", None):
+        original = jm.defer_periodic_job
+        jm._taskbadger_original_defer_periodic_job = original
+
+        @functools.wraps(original)
+        async def patched(*, job, periodic_id, defer_timestamp):
+            task = app.tasks.get(job.task_name)
+            if task is not None:
+                tb_task = _create_pending_task(task, job.task_kwargs)
+                if tb_task is not None:
+                    new_kwargs = {**job.task_kwargs, TB_TASK_ID_KWARG: tb_task.id}
+                    job = job.evolve(task_kwargs=new_kwargs)
+            return await jm._taskbadger_original_defer_periodic_job(
+                job=job, periodic_id=periodic_id, defer_timestamp=defer_timestamp
+            )
+
+        jm.defer_periodic_job = patched
+
+
+def _patch_app_task(app, system):
+    """Replace ``app.task`` with a wrapper that instruments newly-registered
+    tasks under the supplied ``system``. Idempotent — a second call replaces
+    the wrapper but keeps the same original task method."""
+    original = getattr(app, "_taskbadger_original_task", None) or app.task
+    if not getattr(app, "_taskbadger_original_task", None):
+        app._taskbadger_original_task = original
+
+    @functools.wraps(original)
+    def patched(*args, **kwargs):
+        task = original(*args, **kwargs)
+        # ``original`` may return the Task directly or a decorator depending on
+        # call form. Procrastinate's ``app.task`` always returns a decorator
+        # when called with arguments and the Task when called bare.
+        if callable(task) and not hasattr(task, "name"):
+            # decorator form: wrap the returned decorator
+            inner_decorator = task
+
+            @functools.wraps(inner_decorator)
+            def outer(func):
+                t = inner_decorator(func)
+                _instrument_task(t, system=system)
+                return t
+
+            return outer
+        _instrument_task(task, system=system)
+        return task
+
+    app.task = patched

{taskbadger-2.0.0 → taskbadger-2.1.0a2}/taskbadger/systems/celery.py

@@ -1,9 +1,7 @@
-import re
+from taskbadger._integrations import BaseSystemIntegration
 
-from taskbadger.systems import System
 
-
-class CelerySystemIntegration(System):
+class CelerySystemIntegration(BaseSystemIntegration):
     identifier = "celery"
 
     def __init__(self, auto_track_tasks=True, includes=None, excludes=None, record_task_args=False):
@@ -18,29 +16,13 @@ class CelerySystemIntegration(System):
                 the full task name or a regular expression. Exclusions take precedence over inclusions.
             record_task_args: Record the arguments passed to each task.
         """
-        self.auto_track_tasks = auto_track_tasks
-        self.includes = includes
-        self.excludes = excludes
-        self.record_task_args = record_task_args
+        super().__init__(
+            auto_track_tasks=auto_track_tasks,
+            includes=includes,
+            excludes=excludes,
+            record_task_args=record_task_args,
+        )
 
         if auto_track_tasks:
             # Importing this here ensures that the Celery signal handlers are registered
-            import taskbadger.celery  # noqa
-
-    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
+            import taskbadger.celery  # noqa: F401

taskbadger-2.1.0a2/taskbadger/systems/procrastinate.py

@@ -0,0 +1,51 @@
+"""ProcrastinateSystemIntegration — auto-track tasks on a Procrastinate App."""
+
+from __future__ import annotations
+
+from taskbadger._integrations import BaseSystemIntegration
+from taskbadger.procrastinate import _instrument_task, _patch_app_task, _patch_job_manager
+
+
+class ProcrastinateSystemIntegration(BaseSystemIntegration):
+    identifier = "procrastinate"
+
+    def __init__(
+        self,
+        app,
+        auto_track_tasks=True,
+        includes=None,
+        excludes=None,
+        record_task_args=False,
+    ):
+        """
+        Args:
+            app: The ``procrastinate.App`` instance to instrument.
+            auto_track_tasks: Track all tasks regardless of whether they use
+                the ``@taskbadger.procrastinate.track`` decorator.
+            includes: List of task names to include in auto-tracking. Each
+                entry can be a full name or a regex (matched with
+                ``re.fullmatch``).
+            excludes: List of task names to exclude. Same semantics as
+                ``includes``. Exclusions take precedence.
+            record_task_args: Record the task's defer kwargs into the
+                TaskBadger task's ``data`` under ``procrastinate_task_kwargs``.
+        """
+        super().__init__(
+            auto_track_tasks=auto_track_tasks,
+            includes=includes,
+            excludes=excludes,
+            record_task_args=record_task_args,
+        )
+        self.app = app
+
+        for task in list(app.tasks.values()):
+            _instrument_task(task, system=self)
+        _patch_app_task(app, system=self)
+        _patch_job_manager(app, system=self)
+
+    def track_task(self, task_name):
+        # Never auto-track Procrastinate's built-in housekeeping tasks
+        # (e.g. ``builtin:procrastinate.builtin_tasks.remove_old_jobs``).
+        if task_name.startswith("builtin:") or task_name.startswith("procrastinate."):
+            return False
+        return super().track_task(task_name)

taskbadger-2.0.0/README.md

@@ -1,43 +0,0 @@
-# Task Badger Python Client
-
-This is the official Python SDK for [Task Badger](https://taskbadger.net/).
-
-For full documentation go to https://docs.taskbadger.net/python/.
-
-![Integration Tests](https://github.com/taskbadger/taskbadger-python/actions/workflows/integration_tests.yml/badge.svg)
-
----
-
-## Getting Started
-
-### Install
-
-```bash
-pip install taskbadger
-```
-
-To use the `taskbadger` command-line tool, install the `cli` extra:
-
-```bash
-pip install 'taskbadger[cli]'
-```
-
-### Client Usage
-
-```python
-import taskbadger
-from taskbadger.systems import CelerySystemIntegration
-
-taskbadger.init(
-    token="***",
-    systems=[CelerySystemIntegration()],
-    tags={"environment": "production"}
-)
-```
-
-### CLI Usage
-
-```shell
-$ export TASKBADGER_API_KEY=***
-$ taskbadger run "nightly-backup" -- ./backup.sh
-```