utils: add command.run_uid/name to command logs

to easily regroup them

itou/utils/command.py

@@ -1,10 +1,25 @@
+import collections
 import contextlib
 import logging
 import time
+import uuid
 
+from asgiref.local import Local  # NOQA
 from django.core.management import base
 
 
+local = Local()
+
+CommandInfo = collections.namedtuple("CommandInfo", ["run_uid", "name"])
+
+
+def get_current_command_info():
+    try:
+        return local.command_info
+    except AttributeError:
+        return None
+
+
 def _log_command_result(command, duration_in_ns, result):
     command.logger.info(
         f"Management command %s {result} in %0.2f seconds",
@@ -29,13 +44,26 @@ def _command_duration_logger(command):
     _log_command_result(command, time.perf_counter_ns() - before, "succeeded")
 
 
+@contextlib.contextmanager
+def _command_info_manager(command):
+    command_info_before = get_current_command_info()
+    if command_info_before is None:
+        local.command_info = CommandInfo(str(command.run_uid), command.__class__.__module__)
+    try:
+        yield
+    finally:
+        if command_info_before is None:
+            del local.command_info
+
+
 class LoggedCommandMixin:
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
+        self.run_uid = uuid.uuid4()
         self.logger = logging.getLogger(self.__class__.__module__)
 
     def execute(self, *args, **kwargs):
-        with _command_duration_logger(self):
+        with _command_info_manager(self), _command_duration_logger(self):
             try:
                 return super().execute(*args, **kwargs)
             except Exception:

itou/utils/logging.py

@@ -2,6 +2,8 @@
 
 from django_datadog_logger.formatters import datadog
 
+from itou.utils.command import get_current_command_info
+
 
 logger = logging.getLogger(__name__)
 
@@ -36,4 +38,7 @@ def json_record(self, message, extra, record):
                         logger.warning(
                             "Request using token (%r) without datadog_info() method: please define one", token
                         )
+        if (command_info := get_current_command_info()) is not None:
+            log_entry_dict["command.run_uid"] = command_info.run_uid
+            log_entry_dict["command.name"] = command_info.name
         return log_entry_dict

requirements/base.in

@@ -87,6 +87,9 @@ xlrd==2.0.*  # https://pypi.org/project/xlrd/
 # Third-party applications
 # ------------------------------------------------------------------------------
 
+# Django dependency that we directly use for a future-proof threading.local equivalent
+asgiref  # https://github.com/django/asgiref/
+
 # Huey for async processing
 # https://huey.readthedocs.io/en/latest/
 huey  # https://github.com/coleifer/huey

requirements/base.txt

@@ -8,6 +8,7 @@ asgiref==3.8.1 \
     --hash=sha256:3e1e3ecc849832fe52ccf2cb6686b7a55f82bb1d6aee72a58826471390335e47 \
     --hash=sha256:c343bd80a0bec947a9860adb4c432ffa7db769836c64238fc34bdc3fec84d590
     # via
+    #   -r requirements/base.in
     #   django
     #   django-allauth
     #   django-htmx

tests/utils/tests.py

@@ -20,6 +20,7 @@
 from django.contrib.contenttypes.models import ContentType
 from django.contrib.messages.middleware import MessageMiddleware
 from django.contrib.sessions.middleware import SessionMiddleware
+from django.core import management
 from django.core.exceptions import ValidationError
 from django.http import HttpResponse
 from django.template import Context, Template
@@ -1680,6 +1681,26 @@ def test_log_current_organization(client):
     assert f'"usr.organization_id": {membership.company_id}' in captured.getvalue()
 
 
+def test_log_current_command(client):
+    root_logger = logging.getLogger()
+    stream_handler = root_logger.handlers[0]
+    captured = io.StringIO()
+    assert isinstance(stream_handler, logging.StreamHandler)
+    # caplog cannot be used since the command infos are written by the log formatter
+    # capsys/capfd did not want to work because https://github.com/pytest-dev/pytest/issues/5997
+    with patch.object(stream_handler, "stream", captured):
+        management.call_command(
+            # This could have been any other command inheriting from LoggedCommandMixin
+            "delete_old_emails"
+        )
+    # Check that the organization_id is properly logged to stdout
+    lines = captured.getvalue().splitlines()
+    log = json.loads(lines[0])
+    assert log["command.name"] == "itou.emails.management.commands.delete_old_emails"
+    assert "command.run_uid" in log
+    assert uuid.UUID(log["command.run_uid"])
+
+
 def test_create_fake_postcode():
     with mock.patch.dict(DEPARTMENTS, {"2A": "2A - Corse-du-Sud"}):
         postcode = create_fake_postcode()