From 477477e4f7e31aabf1c3b625cf03504bcf9e09a3 Mon Sep 17 00:00:00 2001
From: Michael Derynck <michael.derynck@grafana.com>
Date: Mon, 6 May 2024 15:55:09 -0600
Subject: [PATCH] Add logging of arguments for celery tasks (#4314)

# What this PR does
Enables logging of arguments for celery tasks. Currently it can be
difficult to troubleshoot issues as many tasks give no context to what
they are operating on in the logs.

## Which issue(s) this PR closes

<!--
*Note*: if you have more than one GitHub issue that this PR closes, be
sure to preface
each issue link with a [closing
keyword](https://docs.github.com/en/get-started/writing-on-github/working-with-advanced-formatting/using-keywords-in-issues-and-pull-requests#linking-a-pull-request-to-an-issue).
This ensures that the issue(s) are auto-closed once the PR has been
merged.
-->

## Checklist

- [x] Unit, integration, and e2e (if applicable) tests updated
- [x] Documentation added (or `pr:no public docs` PR label added if not
required)
- [x] Added the relevant release notes label (see labels prefixed w/
`release:`). These labels dictate how your PR will
    show up in the autogenerated release notes.
---
 engine/engine/celery.py | 14 ++++++++++++++
 engine/settings/base.py |  2 +-
 2 files changed, 15 insertions(+), 1 deletion(-)

diff --git a/engine/engine/celery.py b/engine/engine/celery.py
index 480a5570..d6f17fa4 100644
--- a/engine/engine/celery.py
+++ b/engine/engine/celery.py
@@ -95,3 +95,17 @@ if settings.PYROSCOPE_PROFILER_ENABLED:
             detect_subprocesses=True,  # detect subprocesses started by the main process; default is False
             tags={"type": "celery", "celery_worker": os.environ.get("CELERY_WORKER_QUEUE", "no_queue_specified")},
         )
+
+
+if settings.LOG_CELERY_TASK_ARGUMENTS:
+    """
+    Note: Task ID and name are already provided in TaskFormatter prefix, arguments get listed in message
+    """
+
+    @celery.signals.task_prerun.connect
+    def log_started_task_arguments(sender=None, task_id=None, task=None, args=None, kwargs=None, **extras):
+        logger.info(f"task started args={args} kwargs={kwargs}")
+
+    @celery.signals.task_postrun.connect
+    def log_finished_task_arguments(sender=None, task_id=None, task=None, args=None, kwargs=None, **extras):
+        logger.info(f"task finished args={args} kwargs={kwargs}")
diff --git a/engine/settings/base.py b/engine/settings/base.py
index 36de8578..ecdac1a1 100644
--- a/engine/settings/base.py
+++ b/engine/settings/base.py
@@ -352,7 +352,7 @@ if OTEL_TRACING_ENABLED:
     MIDDLEWARE.insert(0, "engine.middlewares.LogRequestHeadersMiddleware")
 
 LOG_REQUEST_ID_HEADER = "HTTP_X_CLOUD_TRACE_CONTEXT"
-
+LOG_CELERY_TASK_ARGUMENTS = getenv_boolean("LOG_CELERY_TASK_ARGUMENTS", default=True)
 
 log_fmt = "source=engine:app google_trace_id=%(request_id)s logger=%(name)s %(message)s"