Update NewRelicLogForwardingHandler (#1038)

TimPansino · mergify[bot] · hmstepanek · web-flow · commit d6c1a91aa3f4 · 2024-01-24T11:53:01.000-08:00
* Update log forwarding handler

* Parametrize logger in conftest

* Update tests for logging

* Change NewRelicLogForwardingHandler to reproduce instrumentation

* Replace testing for forwarding handler

* Clean out imports

* Add test for custom formatter which returns a dict

* Refactor to reduce code complexity

---------

Co-authored-by: mergify[bot] &lt;37929162+mergify[bot]@users.noreply.github.com&gt;
Co-authored-by: Hannah Stepanek &lt;hstepanek@newrelic.com&gt;
diff --git a/newrelic/api/log.py b/newrelic/api/log.py
@@ -16,8 +16,8 @@
 import logging
 import re
 import warnings
-from logging import Formatter, LogRecord
 
+from newrelic.api.application import application_instance
 from newrelic.api.time_trace import get_linking_metadata
 from newrelic.api.transaction import current_transaction, record_log_event
 from newrelic.common import agent_http
@@ -66,8 +66,8 @@ def safe_json_encode(obj, ignore_string_types=False, **kwargs):
         return "<unprintable %s object>" % type(obj).__name__
 
 
-class NewRelicContextFormatter(Formatter):
-    DEFAULT_LOG_RECORD_KEYS = frozenset(set(vars(LogRecord("", 0, "", 0, "", (), None))) | {"message"})
+class NewRelicContextFormatter(logging.Formatter):
+    DEFAULT_LOG_RECORD_KEYS = frozenset(set(vars(logging.LogRecord("", 0, "", 0, "", (), None))) | {"message"})
 
     def __init__(self, *args, **kwargs):
         super(NewRelicContextFormatter, self).__init__()
@@ -104,29 +104,61 @@ def format(self, record):
 
 
 class NewRelicLogForwardingHandler(logging.Handler):
-    DEFAULT_LOG_RECORD_KEYS = frozenset(set(vars(LogRecord("", 0, "", 0, "", (), None))) | {"message"})
+    IGNORED_LOG_RECORD_KEYS = set(["message", "msg"])
 
     def emit(self, record):
         try:
-            # Avoid getting local log decorated message
-            if hasattr(record, "_nr_original_message"):
-                message = record._nr_original_message()
+            nr = None
+            transaction = current_transaction()
+            # Retrieve settings
+            if transaction:
+                settings = transaction.settings
+                nr = transaction
             else:
-                message = record.getMessage()
-
-            attrs = self.filter_record_attributes(record)
-            record_log_event(message, record.levelname, int(record.created * 1000), attributes=attrs)
+                application = application_instance(activate=False)
+                if application and application.enabled:
+                    nr = application
+                    settings = application.settings
+                else:
+                    # If no settings have been found, fallback to global settings
+                    settings = global_settings()
+
+            # If logging is enabled and the application or transaction is not None.
+            if settings and settings.application_logging.enabled and nr:
+                level_name = str(getattr(record, "levelname", "UNKNOWN"))
+                if settings.application_logging.metrics.enabled:
+                    nr.record_custom_metric("Logging/lines", {"count": 1})
+                    nr.record_custom_metric("Logging/lines/%s" % level_name, {"count": 1})
+
+                if settings.application_logging.forwarding.enabled:
+                    if self.formatter:
+                        # Formatter supplied, allow log records to be formatted into a string
+                        message = self.format(record)
+                    else:
+                        # No formatter supplied, attempt to handle dict log records
+                        message = record.msg
+                        if not isinstance(message, dict):
+                            # Allow python to convert the message to a string and template it with args.
+                            message = record.getMessage()
+
+                    # Grab and filter context attributes from log record
+                    context_attrs = self.filter_record_attributes(record)
+
+                    record_log_event(
+                        message=message,
+                        level=level_name,
+                        timestamp=int(record.created * 1000),
+                        attributes=context_attrs,
+                    )
+        except RecursionError:  # Emulates behavior of CPython.
+            raise
         except Exception:
             self.handleError(record)
 
     @classmethod
     def filter_record_attributes(cls, record):
         record_attrs = vars(record)
-        DEFAULT_LOG_RECORD_KEYS = cls.DEFAULT_LOG_RECORD_KEYS
-        if len(record_attrs) > len(DEFAULT_LOG_RECORD_KEYS):
-            return {k: v for k, v in six.iteritems(vars(record)) if k not in DEFAULT_LOG_RECORD_KEYS}
-        else:
-            return None
+        return {k: record_attrs[k] for k in record_attrs if k not in cls.IGNORED_LOG_RECORD_KEYS}
 
 
 class NewRelicLogHandler(logging.Handler):
diff --git a/tests/logger_logging/conftest.py b/tests/logger_logging/conftest.py
@@ -20,6 +20,8 @@
     collector_available_fixture,
 )
 
+from newrelic.api.log import NewRelicLogForwardingHandler
+
 _default_settings = {
     "transaction_tracer.explain_threshold": 0.0,
     "transaction_tracer.transaction_threshold": 0.0,
@@ -54,13 +56,45 @@ def emit(self, record):
         self.records.append(self.format(record))
 
 
+@pytest.fixture(scope="function", params=["instrumented_logger", "forwarding_handler"])
+def logger(request):
+    _logger = logging.getLogger("my_app")
+    caplog = CaplogHandler()
+    _logger.addHandler(caplog)
+    _logger.caplog = caplog
+    _logger.setLevel(logging.WARNING)
+
+    # Save instrumentation so we can disable it
+    instrumented = logging.Logger.callHandlers  
+
+    forwarding_handler = None
+    if request.param == "forwarding_handler":
+        forwarding_handler = NewRelicLogForwardingHandler()
+        _logger.addHandler(forwarding_handler)
+
+        # Uninstrument Logging
+        logging.Logger.callHandlers = logging.Logger.callHandlers.__wrapped__  # noqa, pylint: disable=E1101
+
+    yield _logger
+    del caplog.records[:]
+    
+    _logger.removeHandler(caplog)
+    if forwarding_handler:
+        _logger.removeHandler(forwarding_handler)
+    
+    # Reinstrument logging in case it was uninstrumented
+    logging.Logger.callHandlers = instrumented
+
+
 @pytest.fixture(scope="function")
-def logger():
+def instrumented_logger():
     _logger = logging.getLogger("my_app")
     caplog = CaplogHandler()
     _logger.addHandler(caplog)
     _logger.caplog = caplog
     _logger.setLevel(logging.WARNING)
+
     yield _logger
     del caplog.records[:]
+    
     _logger.removeHandler(caplog)
diff --git a/tests/logger_logging/test_local_decorating.py b/tests/logger_logging/test_local_decorating.py
@@ -14,6 +14,8 @@
 
 import platform
 
+import pytest
+
 from testing_support.fixtures import reset_core_stats_engine
 from testing_support.validators.validate_log_event_count import validate_log_event_count
 from testing_support.validators.validate_log_event_count_outside_transaction import (
@@ -64,32 +66,32 @@ def get_metadata_string(log_message, is_txn):
 
 
 @reset_core_stats_engine()
-def test_local_log_decoration_inside_transaction(logger):
+def test_local_log_decoration_inside_transaction(instrumented_logger):
     @validate_log_event_count(1)
     @background_task()
     def test():
-        exercise_logging(logger)
-        assert logger.caplog.records[0] == get_metadata_string("C", True)
+        exercise_logging(instrumented_logger)
+        assert instrumented_logger.caplog.records[0] == get_metadata_string("C", True)
 
     test()
 
 
 @reset_core_stats_engine()
-def test_local_log_decoration_outside_transaction(logger):
+def test_local_log_decoration_outside_transaction(instrumented_logger):
     @validate_log_event_count_outside_transaction(1)
     def test():
-        exercise_logging(logger)
-        assert logger.caplog.records[0] == get_metadata_string("C", False)
+        exercise_logging(instrumented_logger)
+        assert instrumented_logger.caplog.records[0] == get_metadata_string("C", False)
 
     test()
 
 
 @reset_core_stats_engine()
-def test_local_log_decoration_dict_message(logger):
+def test_local_log_decoration_dict_message(instrumented_logger):
     @validate_log_event_count(1)
     @background_task()
     def test():
-        exercise_logging(logger, {"message": "dict_message"})
-        assert logger.caplog.records[0] == get_metadata_string("{'message': 'dict_message'}", True)
+        exercise_logging(instrumented_logger, {"message": "dict_message"})
+        assert instrumented_logger.caplog.records[0] == get_metadata_string("{'message': 'dict_message'}", True)
 
     test()
diff --git a/tests/logger_logging/test_logging_handler.py b/tests/logger_logging/test_logging_handler.py
@@ -15,7 +15,7 @@
 import logging
 
 import pytest
-from conftest import logger as conf_logger  # noqa, pylint: disable=W0611
+from conftest import instrumented_logger as conf_logger  # noqa, pylint: disable=W0611
 from testing_support.fixtures import (
     override_application_settings,
     reset_core_stats_engine,
@@ -46,8 +46,9 @@ def uninstrument_logging():
 
 
 @pytest.fixture(scope="function")
-def logger(conf_logger, uninstrument_logging):
+def formatting_logger(conf_logger, uninstrument_logging):
     handler = NewRelicLogForwardingHandler()
+    handler.setFormatter(logging.Formatter(r"%(levelname)s - %(message)s"))
     conf_logger.addHandler(handler)
     yield conf_logger
     conf_logger.removeHandler(handler)
@@ -62,50 +63,44 @@ def set_trace_ids():
         trace.guid = "abcdefgh"
 
 
-def exercise_logging(logger):
-    set_trace_ids()
-    logger.warning("C", extra={"key": "value"})
-    assert len(logger.caplog.records) == 1
+class DictMessageFormatter(logging.Formatter):
+    def format(self, record):
+        message = record.msg
+        if isinstance(message, dict):
+            message["formatter_attr"] = 1
+        return message
 
 
-@override_application_settings(
-    {
-        "application_logging.forwarding.context_data.enabled": True,
-    }
-)
-def test_handler_inside_transaction(logger):
+def test_handler_with_formatter(formatting_logger):
     @validate_log_events(
         [
             {
-                "message": "C",
+                "message": "WARNING - C",
                 "level": "WARNING",
                 "timestamp": None,
                 "hostname": None,
                 "entity.name": "Python Agent Test (logger_logging)",
                 "entity.guid": None,
                 "span.id": "abcdefgh",
                 "trace.id": "abcdefgh12345678",
-                "context.key": "value",
+                "context.key": "value",  # Extra attr
+                "context.module": "test_logging_handler",  # Default attr
             }
         ]
     )
     @validate_log_event_count(1)
     @validate_function_called("newrelic.api.log", "NewRelicLogForwardingHandler.emit")
     @background_task()
     def test():
-        exercise_logging(logger)
+        set_trace_ids()
+        formatting_logger.warning("C", extra={"key": "value"})
+        assert len(formatting_logger.caplog.records) == 1
 
     test()
 
 
-@reset_core_stats_engine()
-@override_application_settings(
-    {
-        "application_logging.forwarding.context_data.enabled": True,
-    }
-)
-def test_handler_outside_transaction(logger):
-    @validate_log_events_outside_transaction(
+def test_handler_dict_message_no_formatter(formatting_logger):
+    @validate_log_events(
         [
             {
                 "message": "C",
@@ -114,13 +109,75 @@ def test_handler_outside_transaction(logger):
                 "hostname": None,
                 "entity.name": "Python Agent Test (logger_logging)",
                 "entity.guid": None,
-                "context.key": "value",
+                "span.id": "abcdefgh",
+                "trace.id": "abcdefgh12345678",
+                "message.attr": 3,  # Message attr
             }
-        ]
+        ],
+    )
+    @validate_log_event_count(1)
+    @validate_function_called("newrelic.api.log", "NewRelicLogForwardingHandler.emit")
+    @background_task()
+    def test():
+        set_trace_ids()
+        formatting_logger.handlers[1].setFormatter(None)  # Turn formatter off to enable dict message support
+        formatting_logger.warning({"message": "C", "attr": 3})
+        assert len(formatting_logger.caplog.records) == 1
+
+    test()
+
+
+def test_handler_dict_message_with_formatter(formatting_logger):
+    @validate_log_events(
+        [
+            {
+                "message": "WARNING - {'message': 'C', 'attr': 3}",
+                "level": "WARNING",
+                "timestamp": None,
+                "hostname": None,
+                "entity.name": "Python Agent Test (logger_logging)",
+                "entity.guid": None,
+                "span.id": "abcdefgh",
+                "trace.id": "abcdefgh12345678",
+            }
+        ],
+        forgone_attrs=["message.attr"]  # Explicit formatters take precedence over dict message support
     )
-    @validate_log_event_count_outside_transaction(1)
+    @validate_log_event_count(1)
     @validate_function_called("newrelic.api.log", "NewRelicLogForwardingHandler.emit")
+    @background_task()
+    def test():
+        set_trace_ids()
+        formatting_logger.warning({"message": "C", "attr": 3})
+        assert len(formatting_logger.caplog.records) == 1
+
+    test()
+
+
+def test_handler_formatter_returns_dict_message(formatting_logger):
+    @validate_log_events(
+        [
+            {
+                "message": "C",
+                "level": "WARNING",
+                "timestamp": None,
+                "hostname": None,
+                "entity.name": "Python Agent Test (logger_logging)",
+                "entity.guid": None,
+                "span.id": "abcdefgh",
+                "trace.id": "abcdefgh12345678",
+                "message.attr": 3,  # Message attr
+                "message.formatter_attr": 1,  # Formatter message attr
+            }
+        ],
+    )
+    @validate_log_event_count(1)
+    @validate_function_called("newrelic.api.log", "NewRelicLogForwardingHandler.emit")
+    @background_task()
     def test():
-        exercise_logging(logger)
+        set_trace_ids()
+        formatting_logger.handlers[1].setFormatter(DictMessageFormatter())  # Set formatter to return a dict
+        formatting_logger.warning({"message": "C", "attr": 3})
+        assert len(formatting_logger.caplog.records) == 1
 
     test()
diff --git a/tests/logger_logging/test_settings.py b/tests/logger_logging/test_settings.py
@@ -51,16 +51,16 @@ def test():
 
 @pytest.mark.parametrize("feature_setting,subfeature_setting,expected", _settings_matrix)
 @reset_core_stats_engine()
-def test_local_decorating_settings(logger, feature_setting, subfeature_setting, expected):
+def test_local_decorating_settings(instrumented_logger, feature_setting, subfeature_setting, expected):
     @override_application_settings({
         "application_logging.enabled": feature_setting,
         "application_logging.local_decorating.enabled": subfeature_setting,
     })
     @background_task()
     def test():
-        basic_logging(logger)
-        assert len(logger.caplog.records) == 1
-        message = logger.caplog.records.pop()
+        basic_logging(instrumented_logger)
+        assert len(instrumented_logger.caplog.records) == 1
+        message = instrumented_logger.caplog.records.pop()
         if expected:
             assert len(message) > 1
         else:
