From ad15f7bc8271dad44774cfb6e510c49d34c4100c Mon Sep 17 00:00:00 2001
From: natashz <nata7che@gmail.com>
Date: Thu, 9 Jun 2022 11:42:54 +0300
Subject: [PATCH] Logging instr hook (#1117)

---
 CHANGELOG.md                                  |  5 ++--
 .../instrumentation/logging/__init__.py       | 15 +++++++++++
 .../tests/test_logging.py                     | 25 +++++++++++++++++++
 3 files changed, 43 insertions(+), 2 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index eb4f59e119..5ee2b41196 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -15,11 +15,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - cleanup type hints for textmap `Getter` and `Setter` classes
   ([1106](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/1106))
   
+
 ### Added
+- `opentelemetry-instrumentation-logging` add log hook support
+  ([#1117](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/1117))
 - `opentelemetry-instrumentation-remoulade` Initial release
   ([#1082](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/1082))
-
-### Added
 - Added `opentelemetry-instrumention-confluent-kafka`
   ([#1111](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/1111))
 
diff --git a/instrumentation/opentelemetry-instrumentation-logging/src/opentelemetry/instrumentation/logging/__init__.py b/instrumentation/opentelemetry-instrumentation-logging/src/opentelemetry/instrumentation/logging/__init__.py
index ae6b992dd1..b532a08883 100644
--- a/instrumentation/opentelemetry-instrumentation-logging/src/opentelemetry/instrumentation/logging/__init__.py
+++ b/instrumentation/opentelemetry-instrumentation-logging/src/opentelemetry/instrumentation/logging/__init__.py
@@ -55,6 +55,10 @@ class LoggingInstrumentor(BaseInstrumentor):  # pylint: disable=empty-docstring
 
         {DEFAULT_LOGGING_FORMAT}
 
+        def log_hook(span: Span, record: LogRecord):
+                if span and span.is_recording():
+                    record.custom_user_attribute_from_log_hook = "some-value"
+
     Args:
         tracer_provider: Tracer provider instance that can be used to fetch a tracer.
         set_logging_format: When set to True, it calls logging.basicConfig() and sets a logging format.
@@ -66,11 +70,13 @@ class LoggingInstrumentor(BaseInstrumentor):  # pylint: disable=empty-docstring
             logging.WARN
             logging.ERROR
             logging.FATAL
+        log_hook: execute custom logic when record is created
 
     See `BaseInstrumentor`
     """
 
     _old_factory = None
+    _log_hook = None
 
     def instrumentation_dependencies(self) -> Collection[str]:
         return _instruments
@@ -80,6 +86,7 @@ def _instrument(self, **kwargs):
         provider = kwargs.get("tracer_provider", None) or get_tracer_provider()
         old_factory = logging.getLogRecordFactory()
         LoggingInstrumentor._old_factory = old_factory
+        LoggingInstrumentor._log_hook = kwargs.get("log_hook", None)
 
         service_name = None
 
@@ -107,6 +114,14 @@ def record_factory(*args, **kwargs):
                 if ctx != INVALID_SPAN_CONTEXT:
                     record.otelSpanID = format(ctx.span_id, "016x")
                     record.otelTraceID = format(ctx.trace_id, "032x")
+                    if callable(LoggingInstrumentor._log_hook):
+                        try:
+                            LoggingInstrumentor._log_hook(  # pylint: disable=E1102
+                                span, record
+                            )
+                        except Exception:  # pylint: disable=W0703
+                            pass
+
             return record
 
         logging.setLogRecordFactory(record_factory)
diff --git a/instrumentation/opentelemetry-instrumentation-logging/tests/test_logging.py b/instrumentation/opentelemetry-instrumentation-logging/tests/test_logging.py
index 5a0f1511b6..bddee9bf84 100644
--- a/instrumentation/opentelemetry-instrumentation-logging/tests/test_logging.py
+++ b/instrumentation/opentelemetry-instrumentation-logging/tests/test_logging.py
@@ -64,6 +64,10 @@ def test_trace_context_injection(self):
             self.assertEqual(record.otelServiceName, "")
 
 
+def log_hook(span, record):
+    record.custom_user_attribute_from_log_hook = "some-value"
+
+
 class TestLoggingInstrumentor(TestBase):
     @pytest.fixture(autouse=True)
     def inject_fixtures(self, caplog):
@@ -150,6 +154,27 @@ def test_custom_format_and_level_api(
             format="%(message)s span_id=%(otelSpanID)s", level=logging.WARNING
         )
 
+    def test_log_hook(self):
+        LoggingInstrumentor().uninstrument()
+        LoggingInstrumentor().instrument(
+            set_logging_format=True,
+            log_hook=log_hook,
+        )
+        with self.tracer.start_as_current_span("s1") as span:
+            span_id = format(span.get_span_context().span_id, "016x")
+            trace_id = format(span.get_span_context().trace_id, "032x")
+            with self.caplog.at_level(level=logging.INFO):
+                logger = logging.getLogger("test logger")
+                logger.info("hello")
+                self.assertEqual(len(self.caplog.records), 1)
+                record = self.caplog.records[0]
+                self.assertEqual(record.otelSpanID, span_id)
+                self.assertEqual(record.otelTraceID, trace_id)
+                self.assertEqual(record.otelServiceName, "unknown_service")
+                self.assertEqual(
+                    record.custom_user_attribute_from_log_hook, "some-value"
+                )
+
     def test_uninstrumented(self):
         with self.tracer.start_as_current_span("s1") as span:
             span_id = format(span.get_span_context().span_id, "016x")