diff --git a/.sampo/changesets/celery-exception-autocapture.md b/.sampo/changesets/celery-exception-autocapture.md new file mode 100644 index 000000000..40525c59d --- /dev/null +++ b/.sampo/changesets/celery-exception-autocapture.md @@ -0,0 +1,5 @@ +--- +pypi/posthog: minor +--- + +Exception autocapture now captures Celery task failures and retry exceptions. `PosthogCeleryIntegration` also captures the exception of a retried task. diff --git a/posthog/__init__.py b/posthog/__init__.py index 66fb9ea19..fbcec64da 100644 --- a/posthog/__init__.py +++ b/posthog/__init__.py @@ -356,7 +356,9 @@ def get_tags() -> Dict[str, Any]: serverless handlers should call ``flush()`` before returning. Applied when ``setup()`` builds the global client, or on a later ``setup()`` call if no span has been started yet. - enable_exception_autocapture: Automatically capture uncaught exceptions. + enable_exception_autocapture: Automatically capture uncaught exceptions, + and Celery task exceptions when Celery is imported before the client + is created. log_captured_exceptions: Also log exceptions captured by error tracking. project_root: Root path used to determine in-app exception stack frames. privacy_mode: Capture AI usage metadata without prompt inputs or outputs. diff --git a/posthog/client.py b/posthog/client.py index 2f33a29ae..b313cc037 100644 --- a/posthog/client.py +++ b/posthog/client.py @@ -774,7 +774,8 @@ def __init__( to 1. Set to 0 to disable retries. super_properties: Properties merged into every captured event. enable_exception_autocapture: Automatically capture uncaught - exceptions. + exceptions, and Celery task exceptions when Celery is + imported before the client is created. log_captured_exceptions: Also log exceptions captured by error tracking. project_root: Root path used to determine in-app stack frames for diff --git a/posthog/exception_capture.py b/posthog/exception_capture.py index 6dac4057f..95d1419a7 100644 --- a/posthog/exception_capture.py +++ b/posthog/exception_capture.py @@ -44,6 +44,13 @@ def __init__( self._threading_excepthook = self.thread_exception_handler sys.excepthook = self._sys_excepthook threading.excepthook = self._threading_excepthook + self._disconnect_celery = None + try: + from .integrations.celery import _instrument_exception_autocapture + + self._disconnect_celery = _instrument_exception_autocapture(self) + except Exception: + self.log.exception("Failed to set up Celery exception autocapture") # opt-in client-side rate limiting: per exception type, allow a burst # of captures, then refill over time self._rate_limiter = None @@ -79,6 +86,9 @@ def close(self): if threading.excepthook is self._threading_excepthook: threading.excepthook = original_threading_excepthook + if self._disconnect_celery is not None: + self._disconnect_celery() + if self._rate_limiter is not None: self._rate_limiter.stop() diff --git a/posthog/integrations/celery.py b/posthog/integrations/celery.py index 59a513be4..4844ceb2a 100644 --- a/posthog/integrations/celery.py +++ b/posthog/integrations/celery.py @@ -60,14 +60,23 @@ Additional properties on specific states: - **failure**: ``error_type``, ``error_message`` - - **retry**: ``celery_reason`` + - **retry**: ``celery_reason``, and ``error_type``, ``error_message`` when + the retry has an exception + +Exception autocapture: + When ``enable_exception_autocapture`` is on and Celery is imported before + the PostHog client is created, the client captures task failures and retry + exceptions without this integration. An instrumented integration with + ``capture_exceptions=True`` replaces that capture and applies its + ``task_filter``. """ import atexit import json import logging +import sys import time -from typing import Any, Callable, Optional, cast +from typing import TYPE_CHECKING, Any, Callable, Optional, cast from .. import contexts from ..client import Client @@ -76,6 +85,9 @@ _capture_exception_with_metadata, ) +if TYPE_CHECKING: + from ..exception_capture import ExceptionCapture + CONTEXT_DISTINCT_ID_HEADER = "X-POSTHOG-DISTINCT-ID" CONTEXT_SESSION_ID_HEADER = "X-POSTHOG-SESSION-ID" @@ -83,6 +95,10 @@ logger = logging.getLogger("posthog") +# Instrumented integrations. Exception autocapture defers to them, so a task +# exception is captured once and respects the integration's task filter. +_instrumented_integrations: set["PosthogCeleryIntegration"] = set() + class PosthogCeleryIntegration: """Celery integration that captures task lifecycle events and exceptions. @@ -158,6 +174,7 @@ def instrument(self) -> None: atexit.register(self.shutdown) self._instrumented = True + _instrumented_integrations.add(self) def _disconnect_signals(self) -> None: if not self._instrumented or not self._signals: @@ -176,6 +193,7 @@ def _disconnect_signals(self) -> None: self._signals = None self._instrumented = False + _instrumented_integrations.discard(self) def uninstrument(self) -> None: """Disconnect Celery signal handlers and unregister exit cleanup. @@ -338,6 +356,7 @@ def _on_task_retry(self, *args, **kwargs): extra_properties={ "celery_reason": str(kwargs.get("reason")), }, + exception=_retry_exception(kwargs.get("reason")), **kwargs, ) @@ -378,7 +397,7 @@ def _handle_task_end( should_track = self._should_track(task_name, task_properties) if exception and self.capture_exceptions and should_track: - self._capture_exception(exception) + self._capture_exception(exception, _capture_metadata(state)) if self.capture_task_lifecycle_events and should_track: self._capture_event(f"celery task {state}", properties=task_properties) @@ -478,12 +497,11 @@ def _capture_event(self, event: str, properties: dict[str, Any]) -> None: capture(event, properties=properties) - def _capture_exception(self, exception: Exception) -> None: - capture_metadata: _ExceptionCaptureMetadata = { - "level": "error", - "source": "celery.task_failure", - "mechanism": {"type": "task", "handled": False}, - } + def _capture_exception( + self, + exception: BaseException, + capture_metadata: _ExceptionCaptureMetadata, + ) -> None: if self.client: _capture_exception_with_metadata(self.client, exception, capture_metadata) else: @@ -492,6 +510,67 @@ def _capture_exception(self, exception: Exception) -> None: cast(Any, capture_exception)(exception, _capture_metadata=capture_metadata) +def _retry_exception(reason: Any) -> Optional[BaseException]: + # Celery sends a ``Retry`` as the reason. Its ``exc`` is the exception + # that caused the retry, or ``None`` when the task retried without one. + exc = getattr(reason, "exc", None) + return exc if isinstance(exc, BaseException) else None + + +def _capture_metadata(state: str) -> _ExceptionCaptureMetadata: + if state == "retry": + # Celery handles the exception and runs the task again. + return { + "level": "warning", + "source": "celery.task_retry", + "mechanism": {"type": "task", "handled": True}, + } + return { + "level": "error", + "source": "celery.task_failure", + "mechanism": {"type": "task", "handled": False}, + } + + +def _instrument_exception_autocapture( + exception_capture: "ExceptionCapture", +) -> Optional[Callable[[], None]]: + """Capture Celery task exceptions with exception autocapture. + + Celery catches task exceptions, so ``sys.excepthook`` never sees them. + Returns a callable that disconnects the handlers, or ``None`` when the + process does not use Celery. + """ + if "celery" not in sys.modules: + return None + + from celery import signals + + def capture(exception: Optional[BaseException], state: str) -> None: + if exception is None or exception_capture._closed: + return + if any(i.capture_exceptions for i in _instrumented_integrations): + return + exception_capture._capture_exception( + exception, capture_metadata=_capture_metadata(state) + ) + + def on_task_failure(*args, **kwargs): + capture(kwargs.get("exception"), "failure") + + def on_task_retry(*args, **kwargs): + capture(_retry_exception(kwargs.get("reason")), "retry") + + signals.task_failure.connect(on_task_failure, weak=False) + signals.task_retry.connect(on_task_retry, weak=False) + + def disconnect() -> None: + signals.task_failure.disconnect(on_task_failure) + signals.task_retry.disconnect(on_task_retry) + + return disconnect + + __all__ = [ "PosthogCeleryIntegration", ] diff --git a/posthog/test/integrations/test_celery_integration.py b/posthog/test/integrations/test_celery_integration.py index 0e06314eb..7128ad67e 100644 --- a/posthog/test/integrations/test_celery_integration.py +++ b/posthog/test/integrations/test_celery_integration.py @@ -1,14 +1,17 @@ +import sys import unittest from types import ModuleType from types import SimpleNamespace from unittest.mock import Mock, patch from posthog import contexts +from posthog.exception_capture import ExceptionCapture from posthog.integrations.celery import ( CONTEXT_DISTINCT_ID_HEADER, CONTEXT_SESSION_ID_HEADER, CONTEXT_TAGS_HEADER, PosthogCeleryIntegration, + _capture_metadata, ) @@ -24,6 +27,19 @@ def disconnect(self, handler): self.disconnected.append(handler) +class FakeRetry(Exception): + def __init__(self, message, exc=None): + super().__init__(message) + self.exc = exc + + +RETRY_CAPTURE_METADATA = { + "level": "warning", + "source": "celery.task_retry", + "mechanism": {"type": "task", "handled": True}, +} + + class TestPosthogCeleryIntegration(unittest.TestCase): SIGNAL_NAMES = ( "task_prerun", @@ -506,6 +522,45 @@ def test_task_retry_captures_event_with_reason(self): props = retry_call.kwargs["properties"] self.assertEqual(props["celery_reason"], "broker down") + def test_task_retry_captures_retry_exception(self): + mock_client = Mock() + integration = PosthogCeleryIntegration(client=mock_client) + task = SimpleNamespace( + name="app.tasks.retrying", request=SimpleNamespace(delivery_info={}) + ) + exception = KeyError("missing") + + integration._on_task_retry( + sender=task, + task_id="task-retry", + reason=FakeRetry("Retry in 5s", exc=exception), + ) + + mock_client.capture_exception.assert_called_once_with( + exception, _capture_metadata=RETRY_CAPTURE_METADATA + ) + retry_call = [ + c + for c in mock_client.capture.call_args_list + if c.args[0] == "celery task retry" + ][0] + self.assertEqual(retry_call.kwargs["properties"]["error_type"], "KeyError") + + def test_task_retry_without_exception_skips_exception_capture(self): + mock_client = Mock() + integration = PosthogCeleryIntegration(client=mock_client) + task = SimpleNamespace( + name="app.tasks.retrying", request=SimpleNamespace(delivery_info={}) + ) + + integration._on_task_retry( + sender=task, + task_id="task-retry", + reason=FakeRetry("Retry in 5s"), + ) + + mock_client.capture_exception.assert_not_called() + def test_task_filter_applies_to_worker_lifecycle_events(self): mock_client = Mock() integration = PosthogCeleryIntegration( @@ -668,7 +723,7 @@ def test_capture_exception_falls_back_to_global_capture_exception(self): exception = ValueError("boom") with patch("posthog.capture_exception") as mock_capture_exception: - integration._capture_exception(exception) + integration._capture_exception(exception, _capture_metadata("failure")) mock_capture_exception.assert_called_once_with( exception, @@ -716,3 +771,98 @@ def test_extract_propagated_tags_invalid_json_returns_empty_dict(self): tags = integration._extract_propagated_tags(request) self.assertEqual(tags, {}) + + +class TestCeleryExceptionAutocapture(unittest.TestCase): + def setUp(self): + self.fake_signals = SimpleNamespace( + **{ + signal_name: FakeSignal() + for signal_name in TestPosthogCeleryIntegration.SIGNAL_NAMES + } + ) + fake_celery = ModuleType("celery") + fake_celery.signals = self.fake_signals + fake_celery.__version__ = "5.0.0" + modules_patch = patch.dict("sys.modules", {"celery": fake_celery}) + modules_patch.start() + self.addCleanup(modules_patch.stop) + registry_patch = patch( + "posthog.integrations.celery._instrumented_integrations", set() + ) + registry_patch.start() + self.addCleanup(registry_patch.stop) + + self.client = Mock() + self.capture = ExceptionCapture(self.client) + self.addCleanup(self.capture.close) + + def _failure_handler(self): + return self.fake_signals.task_failure.connected[0][0] + + def _retry_handler(self): + return self.fake_signals.task_retry.connected[0][0] + + def test_captures_task_failure(self): + exception = KeyError("missing") + + self._failure_handler()(sender=Mock(), task_id="t1", exception=exception) + + self.client.capture_exception.assert_called_once_with( + exception, + _capture_metadata=_capture_metadata("failure"), + distinct_id=None, + ) + + def test_captures_task_retry_exception(self): + exception = ConnectionError("broker down") + + self._retry_handler()( + sender=Mock(), reason=FakeRetry("Retry in 5s", exc=exception) + ) + + self.client.capture_exception.assert_called_once_with( + exception, + _capture_metadata=RETRY_CAPTURE_METADATA, + distinct_id=None, + ) + + def test_defers_to_instrumented_integration(self): + integration = PosthogCeleryIntegration(client=Mock()) + with patch("posthog.integrations.celery.atexit"): + integration.instrument() + self.addCleanup(integration.uninstrument) + + self._failure_handler()(sender=Mock(), exception=KeyError("missing")) + self.client.capture_exception.assert_not_called() + + integration.uninstrument() + self._failure_handler()(sender=Mock(), exception=KeyError("missing")) + self.client.capture_exception.assert_called_once() + + def test_captures_when_integration_does_not_capture_exceptions(self): + integration = PosthogCeleryIntegration(client=Mock(), capture_exceptions=False) + with patch("posthog.integrations.celery.atexit"): + integration.instrument() + self.addCleanup(integration.uninstrument) + + self._failure_handler()(sender=Mock(), exception=KeyError("missing")) + + self.client.capture_exception.assert_called_once() + + def test_close_disconnects_signals(self): + failure_handler = self._failure_handler() + retry_handler = self._retry_handler() + + self.capture.close() + + self.assertEqual(self.fake_signals.task_failure.disconnected, [failure_handler]) + self.assertEqual(self.fake_signals.task_retry.disconnected, [retry_handler]) + + def test_skips_celery_when_not_imported(self): + with patch.dict("sys.modules"): + sys.modules.pop("celery") + capture = ExceptionCapture(Mock()) + self.addCleanup(capture.close) + + self.assertIsNone(capture._disconnect_celery)