From 4c6d5e998aa862a3551d24e5bec03cbd7f6d39a4 Mon Sep 17 00:00:00 2001
From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com>
Date: Tue, 13 Jul 2021 23:07:56 -0500
Subject: [PATCH] Add support for custom attributes in OTLPHandler

---
 .../src/opentelemetry/sdk/logs/__init__.py    | 21 ++++++++++++---
 opentelemetry-sdk/tests/logs/test_handler.py  | 26 +++++++++++++++++++
 2 files changed, 44 insertions(+), 3 deletions(-)

diff --git a/opentelemetry-sdk/src/opentelemetry/sdk/logs/__init__.py b/opentelemetry-sdk/src/opentelemetry/sdk/logs/__init__.py
index 8b0da9e22a..b32637b38a 100644
--- a/opentelemetry-sdk/src/opentelemetry/sdk/logs/__init__.py
+++ b/opentelemetry-sdk/src/opentelemetry/sdk/logs/__init__.py
@@ -244,20 +244,35 @@ def force_flush(self, timeout_millis: int = 30000) -> bool:
         return True
 
 
+# skip natural LogRecord attributes
+# http://docs.python.org/library/logging.html#logrecord-attributes
+_RESERVED_ATTRS = frozenset((
+    'args', 'asctime', 'created', 'exc_info', 'exc_text', 'filename',
+    'funcName', 'levelname', 'levelno', 'lineno', 'module',
+    'msecs', 'message', 'msg', 'name', 'pathname', 'process',
+    'processName', 'relativeCreated', 'stack_info', 'thread', 'threadName')
+)
+
+
 class OTLPHandler(logging.Handler):
     """A handler class which writes logging records, in OTLP format, to
     a network destination or file.
     """
 
-    def __init__(self, level=logging.NOTSET, log_emitter=None) -> None:
+    def __init__(self, level=logging.NOTSET, log_emitter=None, *, attributes_key: Optional[str] = None) -> None:
         super().__init__(level=level)
         self._log_emitter = log_emitter or get_log_emitter(__name__)
+        self.attributes_key = attributes_key
+
+    def _get_attributes(self, record: logging.LogRecord) -> Attributes:
+        if self.attributes_key is not None:
+            return vars(record)[self.attributes_key]
+        return {k: v for k, v in vars(record).items() if k not in _RESERVED_ATTRS}
 
     def _translate(self, record: logging.LogRecord) -> LogRecord:
         timestamp = int(record.created * 1e9)
         span_context = get_current_span().get_span_context()
-        # TODO: attributes (or resource attributes?) from record metadata
-        attributes: Attributes = {}
+        attributes = self._get_attributes(record)
         severity_number = std_to_otlp(record.levelno)
         return LogRecord(
             timestamp=timestamp,
diff --git a/opentelemetry-sdk/tests/logs/test_handler.py b/opentelemetry-sdk/tests/logs/test_handler.py
index 1d1b84f0fd..3b3daf48f8 100644
--- a/opentelemetry-sdk/tests/logs/test_handler.py
+++ b/opentelemetry-sdk/tests/logs/test_handler.py
@@ -65,6 +65,32 @@ def test_log_record_no_span_context(self):
             log_record.trace_flags, INVALID_SPAN_CONTEXT.trace_flags
         )
 
+    def test_log_record_user_attributes(self):
+        """Attributes can be injected into logs by adding them as attributes to the LogRecord"""
+        emitter_mock = Mock(spec=LogEmitter)
+        logger = get_logger(log_emitter=emitter_mock)
+        # Assert emit gets called for warning message
+        logger.warning("Wanrning message", extra={"http.status_code": 200})
+        args, _ = emitter_mock.emit.call_args_list[0]
+        log_record = args[0]
+
+        self.assertIsNotNone(log_record)
+        self.assertEqual(log_record.attributes, {"http.status_code": 200})
+
+    def test_log_record_user_attributes_key(self):
+        """Users can specify a key to extract attributes from"""
+        emitter_mock = Mock(spec=LogEmitter)
+        logger = logging.getLogger(__name__)
+        handler = OTLPHandler(level=logging.NOTSET, log_emitter=emitter_mock, attributes_key="opentelemetry")
+        logger.addHandler(handler)
+        # Assert emit gets called for warning message
+        logger.warning("Wanrning message", extra={"opentelemetry": {"http.status_code": 200}, "ignored": True})
+        args, _ = emitter_mock.emit.call_args_list[0]
+        log_record = args[0]
+
+        self.assertIsNotNone(log_record)
+        self.assertEqual(log_record.attributes, {"http.status_code": 200})
+
     def test_log_record_trace_correlation(self):
         emitter_mock = Mock(spec=LogEmitter)
         logger = get_logger(log_emitter=emitter_mock)