fix(loguru): ensure loguru events are automatically patched (#7995)

Resolves: #7988

## Checklist

- [x] Change(s) are motivated and described in the PR description.
- [x] Testing strategy is described if automated tests are not included
in the PR.
- [x] Risk is outlined (performance impact, potential for breakage,
maintainability, etc).
- [x] Change is maintainable (easy to change, telemetry, documentation).
- [x] [Library release note
guidelines](https://ddtrace.readthedocs.io/en/stable/releasenotes.html)
are followed. If no release note is required, add label
`changelog/no-changelog`.
- [x] Documentation is included (in-code, generated user docs, [public
corp docs](https://github.com/DataDog/documentation/)).
- [x] Backport labels are set (if
[applicable](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting))

## Reviewer Checklist

- [ ] Title is accurate.
- [ ] No unnecessary changes are introduced.
- [ ] Description motivates each change.
- [ ] Avoids breaking
[API](https://ddtrace.readthedocs.io/en/stable/versioning.html#interfaces)
changes unless absolutely necessary.
- [ ] Testing strategy adequately addresses listed risk(s).
- [ ] Change is maintainable (easy to change, telemetry, documentation).
- [ ] Release note makes sense to a user of the library.
- [ ] Reviewer has explicitly acknowledged and discussed the performance
implications of this PR as reported in the benchmarks PR comment.
- [ ] Backport labels are set in a manner that is consistent with the
[release branch maintenance
policy](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting)
- [ ] If this PR touches code that signs or publishes builds or
packages, or handles credentials of any kind, I've requested a review
from `@DataDog/security-design-and-guidance`.
- [ ] This PR doesn't touch any of that.

Author: mabdinur

ddtrace/contrib/loguru/patch.py

@@ -47,9 +47,18 @@ def _tracer_injection(event_dict):
     return event_dict
 
 
-def _w_add(func, instance, args, kwargs):
-    # patch logger to include datadog info before logging
-    instance.configure(patcher=lambda record: record.update(_tracer_injection(record["extra"])))
+def _w_configure(func, instance, args, kwargs):
+    original_patcher = kwargs.get("patcher", None)
+    instance._dd_original_patcher = original_patcher
+    if not original_patcher:
+        # no patcher, we do not need to worry about ddtrace fields being overridden
+        return func(*args, **kwargs)
+
+    def _wrapped_patcher(record):
+        original_patcher(record)
+        record.update(_tracer_injection(record["extra"]))
+
+    kwargs["patcher"] = _wrapped_patcher
     return func(*args, **kwargs)
 
 
@@ -61,12 +70,16 @@ def patch():
     if getattr(loguru, "_datadog_patch", False):
         return
     loguru._datadog_patch = True
-
-    _w(loguru.logger, "add", _w_add)
+    # Adds ddtrace fields to loguru logger
+    loguru.logger.configure(patcher=lambda record: record.update(_tracer_injection(record["extra"])))
+    # Ensures that calling loguru.logger.configure(..) does not overwrite ddtrace fields
+    _w(loguru.logger, "configure", _w_configure)
 
 
 def unpatch():
     if getattr(loguru, "_datadog_patch", False):
         loguru._datadog_patch = False
 
-        _u(loguru.logger, "add")
+        _u(loguru.logger, "configure")
+        if hasattr(loguru.logger, "_dd_original_patcher"):
+            loguru.logger.configure(patcher=loguru.logger._dd_original_patcher)

releasenotes/notes/fix-loguru-auto-patching-ec9b6a5dd3863ec7.yaml

@@ -0,0 +1,4 @@
+---
+fixes:
+  - |
+    loguru: Ensures log correlation is enabled when the root logger is initialized. Previously, log correlation was only enabled when a new sink was added.

tests/contrib/loguru/test_loguru_logging.py

@@ -1,4 +1,5 @@
 import json
+import os
 
 from loguru import logger
 import pytest
@@ -13,9 +14,6 @@
 from tests.utils import override_global_config
 
 
-captured_logs = []
-
-
 def _test_logging(output, span, env, service, version):
     dd_trace_id, dd_span_id = (span.trace_id, span.span_id) if span else (0, 0)
 
@@ -30,12 +28,13 @@ def _test_logging(output, span, env, service, version):
     assert json.loads(output[0])["record"]["extra"]["dd.version"] == version or ""
 
 
-@pytest.fixture(autouse=True)
-def patch_loguru():
+@pytest.fixture()
+def captured_logs():
+    captured_logs = []
     try:
         patch()
         logger.add(captured_logs.append, serialize=True)
-        yield
+        yield captured_logs
     finally:
         unpatch()
 
@@ -44,10 +43,9 @@ def patch_loguru():
 def global_config():
     with override_global_config({"service": "logging", "env": "global.env", "version": "global.version"}):
         yield
-    captured_logs.clear()
 
 
-def test_log_trace_global_values():
+def test_log_trace_global_values(captured_logs):
     """
     Check trace info includes global values over local span values
     """
@@ -61,13 +59,68 @@ def test_log_trace_global_values():
     _test_logging(captured_logs, span, config.env, config.service, config.version)
 
 
-def test_log_no_trace():
+def test_log_no_trace(captured_logs):
     logger.info("Hello!")
 
     _test_logging(captured_logs, None, config.env, config.service, config.version)
 
 
-@pytest.mark.subprocess(env=dict(DD_TRACE_128_BIT_TRACEID_GENERATION_ENABLED="False"))
+def test_log_with_default_sink(ddtrace_run_python_code_in_subprocess):
+    code = """
+from loguru import logger
+from ddtrace import tracer
+
+with tracer.trace("test.logging") as span:
+    logger.info("Hello!")
+    """
+
+    env = os.environ.copy()
+    env.update(dict(LOGURU_SERIALIZE="1", DD_LOGS_INJECTION="1", DD_SERVICE="dds", DD_ENV="ddenv", DD_VERSION="vv"))
+    out, err, status, _ = ddtrace_run_python_code_in_subprocess(code, env=env)
+
+    assert status == 0, err + out
+
+    # default sink is stderr
+    result = json.loads(err)
+    assert "Hello!" in result["text"]
+    assert result["record"]["extra"]["dd.trace_id"] != ""
+    assert result["record"]["extra"]["dd.span_id"] != ""
+    assert result["record"]["extra"]["dd.env"] == "ddenv"
+    assert result["record"]["extra"]["dd.service"] == "dds"
+    assert result["record"]["extra"]["dd.version"] == "vv"
+
+
+def test_log_with_default_sink_and_configure(ddtrace_run_python_code_in_subprocess):
+    code = """
+from loguru import logger
+from ddtrace import tracer
+
+logger.configure(patcher=lambda r: r.update({"extra": {"dd.new": "cc"}}))
+
+with tracer.trace("test.logging") as span:
+    logger.info("Hello!")
+    """
+
+    env = os.environ.copy()
+    env.update(dict(LOGURU_SERIALIZE="1", DD_LOGS_INJECTION="1", DD_SERVICE="dds", DD_ENV="ddenv", DD_VERSION="vv"))
+    out, err, status, _ = ddtrace_run_python_code_in_subprocess(code, env=env)
+
+    assert status == 0, err + out
+
+    # default sink is stderr
+    result = json.loads(err)
+    assert "Hello!" in result["text"]
+    assert result["record"]["extra"]["dd.trace_id"] != ""
+    assert result["record"]["extra"]["dd.span_id"] != ""
+    assert result["record"]["extra"]["dd.env"] == "ddenv"
+    assert result["record"]["extra"]["dd.service"] == "dds"
+    assert result["record"]["extra"]["dd.version"] == "vv"
+    assert result["record"]["extra"]["dd.new"] == "cc"
+
+
+@pytest.mark.subprocess(
+    ddtrace_run=True, env=dict(DD_TRACE_128_BIT_TRACEID_GENERATION_ENABLED="False", DD_LOGS_INJECTION="1")
+)
 def test_log_trace():
     """
     Check logging patched and formatter including trace info when 64bit trace ids are generated.
@@ -79,15 +132,11 @@ def test_log_trace():
 
     from ddtrace import config
     from ddtrace import tracer
-    from ddtrace.contrib.loguru import patch
-    from ddtrace.contrib.loguru import unpatch
 
     config.service = "logging"
     config.env = "global.env"
     config.version = "global.version"
 
-    patch()
-
     captured_logs = []
     logger.remove()
     logger.add(captured_logs.append, serialize=True)
@@ -103,12 +152,14 @@ def test_log_trace():
     assert json.loads(captured_logs[0])["record"]["extra"]["dd.service"] == "logging"
     assert json.loads(captured_logs[0])["record"]["extra"]["dd.version"] == "global.version"
 
-    captured_logs.clear()
-    unpatch()
-
 
 @pytest.mark.subprocess(
-    env=dict(DD_TRACE_128_BIT_TRACEID_GENERATION_ENABLED="True", DD_TRACE_128_BIT_TRACEID_LOGGING_ENABLED="True")
+    ddtrace_run=True,
+    env=dict(
+        DD_TRACE_128_BIT_TRACEID_GENERATION_ENABLED="True",
+        DD_TRACE_128_BIT_TRACEID_LOGGING_ENABLED="True",
+        DD_LOGS_INJECTION="1",
+    ),
 )
 def test_log_trace_128bit_trace_ids():
     """
@@ -121,16 +172,12 @@ def test_log_trace_128bit_trace_ids():
 
     from ddtrace import config
     from ddtrace import tracer
-    from ddtrace.contrib.loguru import patch
-    from ddtrace.contrib.loguru import unpatch
     from ddtrace.internal.constants import MAX_UINT_64BITS
 
     config.service = "logging"
     config.env = "global.env"
     config.version = "global.version"
 
-    patch()
-
     captured_logs = []
     logger.remove()
     logger.add(captured_logs.append, serialize=True)
@@ -147,12 +194,14 @@ def test_log_trace_128bit_trace_ids():
     assert json.loads(captured_logs[0])["record"]["extra"]["dd.service"] == "logging"
     assert json.loads(captured_logs[0])["record"]["extra"]["dd.version"] == "global.version"
 
-    captured_logs.clear()
-    unpatch()
-
 
 @pytest.mark.subprocess(
-    env=dict(DD_TRACE_128_BIT_TRACEID_GENERATION_ENABLED="True", DD_TRACE_128_BIT_TRACEID_LOGGING_ENABLED="False")
+    ddtrace_run=True,
+    env=dict(
+        DD_TRACE_128_BIT_TRACEID_GENERATION_ENABLED="True",
+        DD_TRACE_128_BIT_TRACEID_LOGGING_ENABLED="False",
+        DD_LOGS_INJECTION="1",
+    ),
 )
 def test_log_trace_128bit_trace_ids_log_64bits():
     """
@@ -165,16 +214,12 @@ def test_log_trace_128bit_trace_ids_log_64bits():
 
     from ddtrace import config
     from ddtrace import tracer
-    from ddtrace.contrib.loguru import patch
-    from ddtrace.contrib.loguru import unpatch
     from ddtrace.internal.constants import MAX_UINT_64BITS
 
     config.service = "logging"
     config.env = "global.env"
     config.version = "global.version"
 
-    patch()
-
     captured_logs = []
     logger.remove()
     logger.add(captured_logs.append, serialize=True)
@@ -191,21 +236,16 @@ def test_log_trace_128bit_trace_ids_log_64bits():
     assert json.loads(captured_logs[0])["record"]["extra"]["dd.service"] == "logging"
     assert json.loads(captured_logs[0])["record"]["extra"]["dd.version"] == "global.version"
 
-    captured_logs.clear()
-    unpatch()
 
-
-@pytest.mark.subprocess(env=dict(DD_TAGS="service:ddtagservice,env:ddenv,version:ddversion"))
+@pytest.mark.subprocess(
+    ddtrace_run=True, env=dict(DD_TAGS="service:ddtagservice,env:ddenv,version:ddversion", DD_LOGS_INJECTION="1")
+)
 def test_log_DD_TAGS():
     import json
 
     from loguru import logger
 
     from ddtrace import tracer
-    from ddtrace.contrib.loguru import patch
-    from ddtrace.contrib.loguru import unpatch
-
-    patch()
 
     captured_logs = []
     logger.remove()
@@ -222,11 +262,8 @@ def test_log_DD_TAGS():
     assert json.loads(captured_logs[0])["record"]["extra"]["dd.service"] == "ddtagservice"
     assert json.loads(captured_logs[0])["record"]["extra"]["dd.version"] == "ddversion"
 
-    captured_logs.clear()
-    unpatch()
-
 
-@pytest.mark.subprocess
+@pytest.mark.subprocess(ddtrace_run=True, env=dict(DD_LOGS_INJECTION="1"))
 def test_configured_format():
     """
     Ensure injection works when user configured format is used
@@ -254,15 +291,11 @@ def serialize(record):
 
     from ddtrace import config
     from ddtrace import tracer
-    from ddtrace.contrib.loguru import patch
-    from ddtrace.contrib.loguru import unpatch
 
     config.service = "logging"
     config.env = "global.env"
     config.version = "global.version"
 
-    patch()
-
     captured_logs = []
     logger.remove()
     logger.add(captured_logs.append, format=log_format)
@@ -277,6 +310,3 @@ def serialize(record):
     assert json.loads(captured_logs[0])["dd.env"] == "global.env"
     assert json.loads(captured_logs[0])["dd.service"] == "logging"
     assert json.loads(captured_logs[0])["dd.version"] == "global.version"
-
-    captured_logs.clear()
-    unpatch()

tests/contrib/loguru/test_loguru_patch.py

@@ -22,10 +22,10 @@ class TestLoguruPatch(PatchTestCase.Base):
     __get_version__ = get_version
 
     def assert_module_patched(self, loguru):
-        self.assert_wrapped(loguru.logger.add)
+        self.assert_wrapped(loguru.logger.configure)
 
     def assert_not_module_patched(self, loguru):
-        self.assert_not_wrapped(loguru.logger.add)
+        self.assert_not_wrapped(loguru.logger.configure)
 
     def assert_not_module_double_patched(self, loguru):
-        self.assert_not_double_wrapped(loguru.logger.add)
+        self.assert_not_double_wrapped(loguru.logger.configure)