Tracing of Celery tasks does not work with Signatures

[alexpersin]

How do you use Sentry?

Sentry Saas (sentry.io)

Version

1.5.10

Steps to Reproduce

The python sentry_sdk celery integration uses _wrap_apply_async to wrap the celery.Task.apply_async method. However it does not wrap celery Signatures as these have their own apply_async method defined in the celery package in celery/canvas.py. This means that tasks defined using signatures are not traced.

To reproduce:

import sentry_sdk
import os
from sentry_sdk.integrations.celery import CeleryIntegration

sentry_sdk.init(os.getenv('SENTRY_DSN'), integrations=[CeleryIntegration()],  traces_sample_rate=1.0, debug=True)

group_result = celery.group([
    celery.signature('foo', kwargs=kwargs),
    celery.signature('bar', kwargs=kwargs)
]).apply_async()

This does not produce traces. Alternatively,

import sentry_sdk
import os
from sentry_sdk.integrations.celery import CeleryIntegration
from celery.app.task import Task 

sentry_sdk.init(os.getenv('SENTRY_DSN'), integrations=[CeleryIntegration()],  traces_sample_rate=1.0, debug=True)

# Not wrapped by sentry sdk
print(celery.group.apply_async.__code__)

# Wrapped by sentry sdk
print(Task.apply_async.__code__)

Running the following fixes the issue:

import celery
from sentry_sdk.integrations.celery import _wrap_apply_async

celery.group.apply_async = _wrap_apply_async(celery.group.apply_async)
celery.chunks.apply_async = _wrap_apply_async(celery.chunks.apply_async)
celery.chain.apply_async = _wrap_apply_async(celery.chain.apply_async)
celery.chord.apply_async = _wrap_apply_async(celery.chord.apply_async)
Signature.apply_async = _wrap_apply_async(Signature.apply_async)

Expected Result

Transactions visible in the UI.

Actual Result

No transaction visible. Stepping through the code with a debugger, the function created with _wrap_apply_async is not called as the Signature.apply_async method is not wrapped.

[sl0thentr0py]

@alexpersin thx for the report, since you already identified the problem, PR's are always welcome!

[alexpersin]

My initial fix was slightly wrong, it should be

from celery.canvas import Signature
import celery
from sentry_sdk.integrations.celery import _wrap_apply_async

celery.group.apply_async = _wrap_apply_async(celery.group.apply_async)
celery.chunks.apply_async = _wrap_apply_async(celery.chunks.apply_async)
celery.canvas._chain.apply_async = _wrap_apply_async(celery.canvas._chain.apply_async)
celery.canvas._chord.apply_async = _wrap_apply_async(celery.canvas._chord.apply_async)
Signature.apply_async = _wrap_apply_async(Signature.apply_async)

Confirmed that this works for groups within chains within groups. Chained tasks show up as child transactions, and it works for retried tasks too.