fix: enable tracing upon import / improve logging setup (#7859)

* fix: fix auto-tracing

* feat: add context var logging to structlog

* docs: add release notes

haystack/__init__.py

@@ -3,6 +3,7 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import haystack.logging
+import haystack.tracing
 from haystack.core.component import component
 from haystack.core.errors import ComponentError, DeserializationError
 from haystack.core.pipeline import Pipeline, PredefinedPipeline
@@ -13,6 +14,9 @@ from haystack.dataclasses import Answer, Document, ExtractedAnswer, GeneratedAns
 # This is a no-op unless `structlog` is installed
 haystack.logging.configure_logging()
 
+# Same for tracing (no op if `opentelemetry` or `ddtrace` is not installed)
+haystack.tracing.auto_enable_tracing()
+
 __all__ = [
     "component",
     "default_from_dict",

haystack/logging.py

@@ -298,6 +298,7 @@ def configure_logging(use_json: Optional[bool] = None) -> None:
 
     try:
         import structlog
+        import structlog.contextvars
         from structlog.processors import ExceptionRenderer
         from structlog.tracebacks import ExceptionDictTransformer
 
@@ -330,6 +331,7 @@ def configure_logging(use_json: Optional[bool] = None) -> None:
         structlog.stdlib.add_log_level,
         # Adds the current timestamp in ISO format to logs
         structlog.processors.TimeStamper(fmt="iso"),
+        structlog.contextvars.merge_contextvars,
         add_line_and_file,
     ]
 

releasenotes/notes/fix-auto-tracing-51ed3a590000d6c8.yaml

@@ -0,0 +1,4 @@
+---
+fixes:
+  - |
+    Auto enable tracing upon import if `ddtrace` or `opentelemetry` is installed.

releasenotes/notes/structlog-contextvars-c13bb3c59a6a92c7.yaml

@@ -0,0 +1,4 @@
+---
+enhancements:
+  - |
+    Add support for [structlog context variables](https://www.structlog.org/en/24.2.0/contextvars.html) to structured logging.

test/test_logging.py

@@ -8,6 +8,10 @@ import os
 import sys
 from datetime import datetime, timezone
 from pathlib import Path
+from typing import Generator
+
+import structlog
+
 from test.tracing.utils import SpyingTracer
 from unittest.mock import ANY
 
@@ -28,6 +32,15 @@ def reset_logging_config() -> None:
     logging.root.handlers = old_handlers
 
 
+@pytest.fixture()
+def set_context_var_key() -> Generator[str, None, None]:
+    structlog.contextvars.bind_contextvars(context_var="value")
+
+    yield "context_var"
+
+    structlog.contextvars.unbind_contextvars("context_var")
+
+
 class TestSkipLoggingConfiguration:
     def test_skip_logging_configuration(
         self, monkeypatch: MonkeyPatch, capfd: CaptureFixture, caplog: LogCaptureFixture
@@ -173,6 +186,17 @@ class TestStructuredLoggingConsoleRendering:
 
         assert "An error happened" in output
 
+    def test_logging_of_contextvars(self, capfd: CaptureFixture, set_context_var_key: str) -> None:
+        haystack_logging.configure_logging()
+
+        logger = logging.getLogger(__name__)
+        logger.warning("Hello, structured logging!", extra={"key1": "value1", "key2": "value2"})
+
+        # Use `capfd` to capture the output of the final structlog rendering result
+        output = capfd.readouterr().err
+
+        assert set_context_var_key in output
+
 
 class TestStructuredLoggingJSONRendering:
     def test_logging_as_json_if_not_atty(self, capfd: CaptureFixture, monkeypatch: MonkeyPatch) -> None:
@@ -237,6 +261,30 @@ class TestStructuredLoggingJSONRendering:
             "module": "test.test_logging",
         }
 
+    def test_logging_of_contextvars(
+        self, capfd: CaptureFixture, monkeypatch: MonkeyPatch, set_context_var_key: str
+    ) -> None:
+        monkeypatch.setenv("HAYSTACK_LOGGING_USE_JSON", "true")
+        haystack_logging.configure_logging()
+
+        logger = logging.getLogger(__name__)
+        logger.warning("Hello, structured logging!", extra={"key1": "value1", "key2": "value2"})
+
+        # Use `capfd` to capture the output of the final structlog rendering result
+        output = capfd.readouterr().err
+        parsed_output = json.loads(output)  # should not raise an error
+
+        assert parsed_output == {
+            "event": "Hello, structured logging!",
+            "key1": "value1",
+            "key2": "value2",
+            set_context_var_key: "value",
+            "level": "warning",
+            "timestamp": ANY,
+            "lineno": ANY,
+            "module": "test.test_logging",
+        }
+
     def test_logging_exceptions_json(self, capfd: CaptureFixture) -> None:
         haystack_logging.configure_logging(use_json=True)