Unable to add custom attribute to Celery trace

[schallis]

Description
I've been attempting to add a custom attribute to my Celery traces as per the docs.

When using the following snippet (as described in Issue 215), I do NOT see the expected attribute:

class CustomBaseTask(Task):
    """
    Annotates the new relic trace with args and kwargs.
    """

    def __call__(self, *args, **kwargs):
        if newrelic.agent.current_transaction():
            newrelic.agent.add_custom_parameter('custom_parameter': 'value')

        return super().__call__(*args, **kwargs)

However I am able to get the custom attribute to appear when I use the following snippet, however this appears as a second trace so I have a trace without the attribute and one with the attribute which is confusing:

class CustomBaseTask(Task):
    """
    Annotates the new relic trace with args and kwargs.
    """

    def __call__(self, *args, **kwargs):
        application = newrelic.agent.application()
        with newrelic.agent.BackgroundTask(application, name=self.name, group='Celery'):
            newrelic.agent.add_custom_parameter(
                    'custom_parameter',
                    'value',
                )

        return super().__call__(*args, **kwargs)

newrelic.agent.add_custom_parameter also works perfectly in my Web transactions.

Expected Behavior
Under Traces > Trace Groups > Trace Name > Attributes for my instrumented non-web transactions I expect to see my custom attribute appear but it does not with the first example, and for the second example it seems to create a new trace instead of modifying the existing one resulting in duplication.

Steps to Reproduce
See custom Celery Task class above

Your Environment
newrelic==8.5.0
python 3.9.6

Additional context
N/A

[workato-integration]

https://new-relic.atlassian.net/browse/NR-332298

[lrafeei]

Hello--sorry about the delay on this one. Basically, the gist is that we weren't supporting custom base tasks. Could you try using this branch and let me know if that works for you?

https://github.com/newrelic/newrelic-python-agent/tree/celery-base-class-fix

If it does, I can clean it up and include it in the next release

[schallis]

@lrafeei I tried out your branch and I unfortunately haven't seen any change in behavior.

I updated the Python dependency by setting the following in my pyproject.yoml and relocking Poetry etc. I verified the updated code with your changes is running in my test environment. Presumably this is all that was needed as it looks like newrelic-admin relies on this package.

newrelic = { git = "git@github.com:newrelic/newrelic-python-agent.git", rev = "celery-base-class-fix" }

Ultimately, my tasks are still reporting that current_transaction=None thus nothing to attach the custom attribute to. I tried with both my own application and the test application provided in the support ticket. Respective Celery logs for each:

[2024-11-20 23:53:23,222: INFO/MainProcess] Task pfm.tasks.worker_test_newrelic[614832c7-3420-4e54-a331-61536f7130e2] received
[2024-11-20 23:53:23,222: DEBUG/MainProcess] TaskPool: Apply <function fast_trace_task at 0x7f01f86facb0> (args:('pfm.tasks.worker_test_newrelic', '614832c7-3420-4e54-a331-61536f7130e2', {'lang': 'py', 'task': 'pfm.tasks.worker_test_newrelic', 'id': '614832c7-3420-4e54-a331-61536f7130e2', 'shadow': None, 'eta': None, 'expires': None, 'group': None, 'group_index': None, 'retries': 0, 'timelimit': [None, None], 'root_id': '614832c7-3420-4e54-a331-61536f7130e2', 'parent_id': None, 'argsrepr': '()', 'kwargsrepr': "{'albert_trace_id': 'ec6506a6-6e24-4e60-a5df-e293a2b9f999--23'}", 'origin': 'gen16@pfm-api-ff686dd8d-wxhfv', 'ignore_result': False, 'traceparent': '00-0a17e8455bfab89281e3e4e4a7c2f10e-1aa7576be7986d78-01', 'tracestate': '2252749@nr=0-0-2252749-516911477-1aa7576be7986d78-0a17e8455bfab892-1-1.840262-1732146803191', 'newrelic': 'eyJ2IjpbMCwxXSwiZCI6eyJ0eSI6IkFwcCIsImFjIjoiMjI1Mjc0OSIsImFwIjoiNTE2OTExNDc3IiwidHIiOiIwYTE3ZTg0NTViZmFiODkyODFlM2U0ZTRhN2MyZjEwZSIsInNhIjp0cnVlLCJwciI6MS44NDAyNjIsInR4IjoiMGExN2U4NDU1YmZhYjg5MiIsInRpIjoxNzMyMTQ2ODAzMTkxLCJpZCI6IjFhYTc1NzZiZTc5ODZkNzgifX0=', 'sentry-trace':... kwargs:{})
(target=pp trace=ec6506a6-6e24-4e60-a5df-e293a2b9f999--23 request=none profile=None level=INFO logger=albert_common.celery.albert.__call__:188) newrelic-test albert_trace_id set in __call__ method (before), albert_trace_id=ec6506a6-6e24-4e60-a5df-e293a2b9f999--23 current_transaction=None task_name=pfm.tasks.worker_test_newrelic
{"message": "newrelic-test albert_trace_id set in __call__ method (before), albert_trace_id=ec6506a6-6e24-4e60-a5df-e293a2b9f999--23 current_transaction=None task_name=pfm.tasks.worker_test_newrelic", "task_name": "pfm.tasks.worker_test_newrelic", "albert_trace_id": "ec6506a6-6e24-4e60-a5df-e293a2b9f999--23", "current_transaction": null, "request_id": "none", "profile_id": null, "target": "jsonlogger", "timestamp": "2024-11-20 23:53:23,224", "level": "INFO", "filename": "albert.py", "funcName": "albert.py/__call__", "pathname": "/opt/venv/lib/python3.10/site-packages/albert_common/celery/albert.py", "line": 188, "format": "json", "log.level": "INFO", "logger.name": "albert_common.celery.albert", "thread.id": 139646469498688, "thread.name": "MainThread", "process.id": 137, "process.name": "ForkPoolWorker-7", "file.name": "/opt/venv/lib/python3.10/site-packages/albert_common/celery/albert.py", "line.number": 188, "entity.type": "SERVICE", "entity.name": "staging-us-west-2-eks-3-pfm", "entity.guid": "MjI1Mjc0OXxBUE18QVBQTElDQVRJT058NTE2OTExNDc3", "hostname": "pfm-commonworker-667b49694-9w8gf", "message_orig": "newrelic-test albert_trace_id set in __call__ method (before)"}
(target=pp trace=ec6506a6-6e24-4e60-a5df-e293a2b9f999--23 request=none profile=None level=INFO logger=pfm.tasks.worker_test_newrelic:15) newrelic-test task running
{"message": "newrelic-test task running", "albert_trace_id": "ec6506a6-6e24-4e60-a5df-e293a2b9f999--23", "request_id": "none", "profile_id": null, "target": "jsonlogger", "timestamp": "2024-11-20 23:53:23,225", "level": "INFO", "filename": "tasks.py", "funcName": "tasks.py/worker_test_newrelic", "pathname": "/app/pfm/tasks.py", "line": 15, "format": "json", "log.level": "INFO", "logger.name": "pfm.tasks", "thread.id": 139646469498688, "thread.name": "MainThread", "process.id": 137, "process.name": "ForkPoolWorker-7", "file.name": "/app/pfm/tasks.py", "line.number": 15, "entity.type": "SERVICE", "entity.name": "staging-us-west-2-eks-3-pfm", "entity.guid": "MjI1Mjc0OXxBUE18QVBQTElDQVRJT058NTE2OTExNDc3", "hostname": "pfm-commonworker-667b49694-9w8gf", "span.id": "79d986e379ccf2f6", "trace.id": "fcef94e542e7413a52caa93ccb95ab0d", "message_orig": "newrelic-test task running"}
[2024-11-20 23:53:23,227: INFO/ForkPoolWorker-7] Task pfm.tasks.worker_test_newrelic[0a23fc32-c77b-4156-9f7c-b32c4aec5939] succeeded in 0.005267629996524192s: None
(target=pp trace=ec6506a6-6e24-4e60-a5df-e293a2b9f999--23 request=none profile=None level=INFO logger=albert_common.celery.albert.__call__:188) newrelic-test albert_trace_id set in __call__ method (before), albert_trace_id=ec6506a6-6e24-4e60-a5df-e293a2b9f999--23 current_transaction=None task_name=pfm.tasks.worker_test_newrelic
{"message": "newrelic-test albert_trace_id set in __call__ method (before), albert_trace_id=ec6506a6-6e24-4e60-a5df-e293a2b9f999--23 current_transaction=None task_name=pfm.tasks.worker_test_newrelic", "task_name": "pfm.tasks.worker_test_newrelic", "albert_trace_id": "ec6506a6-6e24-4e60-a5df-e293a2b9f999--23", "current_transaction": null, "request_id": "none", "profile_id": null, "target": "jsonlogger", "timestamp": "2024-11-20 23:53:25,435", "level": "INFO", "filename": "albert.py", "funcName": "albert.py/__call__", "pathname": "/opt/venv/lib/python3.10/site-packages/albert_common/celery/albert.py", "line": 188, "format": "json", "log.level": "INFO", "logger.name": "albert_common.celery.albert", "thread.id": 139646469498688, "thread.name": "MainThread", "process.id": 137, "process.name": "ForkPoolWorker-7", "file.name": "/opt/venv/lib/python3.10/site-packages/albert_common/celery/albert.py", "line.number": 188, "entity.type": "SERVICE", "entity.name": "staging-us-west-2-eks-3-pfm", "entity.guid": "MjI1Mjc0OXxBUE18QVBQTElDQVRJT058NTE2OTExNDc3", "hostname": "pfm-commonworker-667b49694-9w8gf", "message_orig": "newrelic-test albert_trace_id set in __call__ method (before)"}
(target=pp trace=ec6506a6-6e24-4e60-a5df-e293a2b9f999--23 request=none profile=None level=INFO logger=pfm.tasks.worker_test_newrelic:15) newrelic-test task running
{"message": "newrelic-test task running", "albert_trace_id": "ec6506a6-6e24-4e60-a5df-e293a2b9f999--23", "request_id": "none", "profile_id": null, "target": "jsonlogger", "timestamp": "2024-11-20 23:53:25,435", "level": "INFO", "filename": "tasks.py", "funcName": "tasks.py/worker_test_newrelic", "pathname": "/app/pfm/tasks.py", "line": 15, "format": "json", "log.level": "INFO", "logger.name": "pfm.tasks", "thread.id": 139646469498688, "thread.name": "MainThread", "process.id": 137, "process.name": "ForkPoolWorker-7", "file.name": "/app/pfm/tasks.py", "line.number": 15, "entity.type": "SERVICE", "entity.name": "staging-us-west-2-eks-3-pfm", "entity.guid": "MjI1Mjc0OXxBUE18QVBQTElDQVRJT058NTE2OTExNDc3", "hostname": "pfm-commonworker-667b49694-9w8gf", "span.id": "f14b03649d20cb3d", "trace.id": "c969829318334fd5d83f22c2e882082b", "message_orig": "newrelic-test task running"}
[2024-11-20 23:53:25,438: INFO/ForkPoolWorker-7] Task pfm.tasks.worker_test_newrelic[614832c7-3420-4e54-a331-61536f7130e2] succeeded in 0.006493738997960463s: None

[2024-11-20 23:55:29,064: INFO/MainProcess] celery@pfm-commonworker-667b49694-9w8gf ready.
[2024-11-20 23:55:29,068: INFO/MainProcess] Task custom_base_task[9087ee00-e540-4410-b2d0-6f563201bda5] received
[2024-11-20 23:55:29,068: DEBUG/MainProcess] TaskPool: Apply <function fast_trace_task at 0x7f7179406d40> (args:('custom_base_task', '9087ee00-e540-4410-b2d0-6f563201bda5', {'lang': 'py', 'task': 'custom_base_task', 'id': '9087ee00-e540-4410-b2d0-6f563201bda5', 'shadow': None, 'eta': None, 'expires': None, 'group': None, 'group_index': None, 'retries': 0, 'timelimit': [None, None], 'root_id': '9087ee00-e540-4410-b2d0-6f563201bda5', 'parent_id': None, 'argsrepr': '()', 'kwargsrepr': '{}', 'origin': 'gen183@pfm-api-ff686dd8d-wxhfv', 'ignore_result': False, 'properties': {'correlation_id': '9087ee00-e540-4410-b2d0-6f563201bda5', 'reply_to': 'bee0b1d8-8ae1-3aa2-86a0-13f0c8a54d28', 'delivery_mode': 2, 'delivery_info': {'exchange': '', 'routing_key': 'celery'}, 'priority': 0, 'body_encoding': 'base64', 'delivery_tag': '70ac2fcb-682b-4594-89d3-9eda670f0a46'}, 'reply_to': 'bee0b1d8-8ae1-3aa2-86a0-13f0c8a54d28', 'correlation_id': '9087ee00-e540-4410-b2d0-6f563201bda5', 'hostname': 'celery@pfm-commonworker-667b49694-9w8gf', 'delivery_info': {'exchange': '', 'routing_key': 'celery', 'priority': 0, 'redelivered': None}, 'args':... kwargs:{})
[2024-11-20 23:55:29,069: WARNING/ForkPoolWorker-8] No transaction
[2024-11-20 23:55:29,071: INFO/MainProcess] Task custom_base_task[70d5b07f-cbcc-402e-bfce-d3807ef63b53] received
[2024-11-20 23:55:29,071: DEBUG/MainProcess] TaskPool: Apply <function fast_trace_task at 0x7f7179406d40> (args:('custom_base_task', '70d5b07f-cbcc-402e-bfce-d3807ef63b53', {'lang': 'py', 'task': 'custom_base_task', 'id': '70d5b07f-cbcc-402e-bfce-d3807ef63b53', 'shadow': None, 'eta': None, 'expires': None, 'group': None, 'group_index': None, 'retries': 0, 'timelimit': [None, None], 'root_id': '70d5b07f-cbcc-402e-bfce-d3807ef63b53', 'parent_id': None, 'argsrepr': '()', 'kwargsrepr': '{}', 'origin': 'gen183@pfm-api-ff686dd8d-wxhfv', 'ignore_result': False, 'properties': {'correlation_id': '70d5b07f-cbcc-402e-bfce-d3807ef63b53', 'reply_to': 'bee0b1d8-8ae1-3aa2-86a0-13f0c8a54d28', 'delivery_mode': 2, 'delivery_info': {'exchange': '', 'routing_key': 'celery'}, 'priority': 0, 'body_encoding': 'base64', 'delivery_tag': '288a606d-f474-4416-ae24-90a45ec29f41'}, 'reply_to': 'bee0b1d8-8ae1-3aa2-86a0-13f0c8a54d28', 'correlation_id': '70d5b07f-cbcc-402e-bfce-d3807ef63b53', 'hostname': 'celery@pfm-commonworker-667b49694-9w8gf', 'delivery_info': {'exchange': '', 'routing_key': 'celery', 'priority': 0, 'redelivered': None}, 'args':... kwargs:{})
[2024-11-20 23:55:29,072: WARNING/ForkPoolWorker-1] No transaction
[2024-11-20 23:55:29,073: WARNING/ForkPoolWorker-8] running task
[2024-11-20 23:55:29,074: INFO/MainProcess] Task custom_base_task[a3bbab0c-e9ef-4b75-ada5-84c9ec434dca] received
[2024-11-20 23:55:29,074: DEBUG/MainProcess] TaskPool: Apply <function fast_trace_task at 0x7f7179406d40> (args:('custom_base_task', 'a3bbab0c-e9ef-4b75-ada5-84c9ec434dca', {'lang': 'py', 'task': 'custom_base_task', 'id': 'a3bbab0c-e9ef-4b75-ada5-84c9ec434dca', 'shadow': None, 'eta': None, 'expires': None, 'group': None, 'group_index': None, 'retries': 0, 'timelimit': [None, None], 'root_id': 'a3bbab0c-e9ef-4b75-ada5-84c9ec434dca', 'parent_id': None, 'argsrepr': '()', 'kwargsrepr': '{}', 'origin': 'gen183@pfm-api-ff686dd8d-wxhfv', 'ignore_result': False, 'properties': {'correlation_id': 'a3bbab0c-e9ef-4b75-ada5-84c9ec434dca', 'reply_to': 'bee0b1d8-8ae1-3aa2-86a0-13f0c8a54d28', 'delivery_mode': 2, 'delivery_info': {'exchange': '', 'routing_key': 'celery'}, 'priority': 0, 'body_encoding': 'base64', 'delivery_tag': 'f609e9f0-d605-4089-9986-aa2e6cacfe75'}, 'reply_to': 'bee0b1d8-8ae1-3aa2-86a0-13f0c8a54d28', 'correlation_id': 'a3bbab0c-e9ef-4b75-ada5-84c9ec434dca', 'hostname': 'celery@pfm-commonworker-667b49694-9w8gf', 'delivery_info': {'exchange': '', 'routing_key': 'celery', 'priority': 0, 'redelivered': None}, 'args':... kwargs:{})
[2024-11-20 23:55:29,075: WARNING/ForkPoolWorker-2] No transaction
[2024-11-20 23:55:29,076: WARNING/ForkPoolWorker-1] running task
[2024-11-20 23:55:29,079: WARNING/ForkPoolWorker-2] running task
(truncated...)

My custom task which generated this first logs currently looks like this:

class AlbertTask(Task):
    """Custom Celery Task class for Albert."""

    def __call__(self, *args, **kwargs):
        """Annotates the new relic trace with custom params."""
        albert_trace_id = getattr(local, 'albert_trace_id', None)
        current_transaction = newrelic.agent.current_transaction()

        if current_transaction:
            newrelic.agent.add_custom_parameter('albert_trace_id', albert_trace_id)

        logger.info(
            'newrelic-test albert_trace_id set in __call__ method (before)',
            task_name=self.name,
            albert_trace_id=albert_trace_id,
            current_transaction=current_transaction,
        )

        return super().__call__(*args, **kwargs)

[lrafeei]

Hello again--

Long story short, this fix will take a bit longer than we had hoped for. Until we get a proper fix released, here's a workaround you can use:

1. Add from newrelic.hooks.application_celery import CeleryTaskWrapper to the import list
2. Add something like CustomBaseTaskClass.__call__ = CeleryTaskWrapper(CustomBaseTaskClass.__call__)

My example app for reference:

import newrelic.agent
from newrelic.hooks.application_celery import CeleryTaskWrapper
from celery import Celery, Task

app = Celery('tasks', backend='redis://localhost', broker='redis://localhost')

# Suggested Template for Custom Base Task Classes
class DebugTask(Task):
    def __call__(self, *args, **kwargs):
        newrelic.agent.add_custom_attribute("custom_attribute", "suggested custom base")
        return self.run(*args, **kwargs)


# Alternative Method for Custom Base Task Classes
class ReproTask(Task):
    def __call__(self, *args, **kwargs):
        if newrelic.agent.current_transaction():
            newrelic.agent.add_custom_attribute("custom_attribute", "custom base with super")
    
        return super().__call__(*args, **kwargs)


# Manual wrapping
DebugTask.__call__ = CeleryTaskWrapper(DebugTask.__call__)
ReproTask.__call__ = CeleryTaskWrapper(ReproTask.__call__)


@app.task(base=DebugTask)
def add_custom_base(x,y):
    return x+y 


@app.task(base=ReproTask)   # Customer example
def add_custom_base_with_super(x,y):
    return x+y 


add_custom_base.delay(3,4)
add_custom_base_with_super.delay(5,6)