NewRelicContextFormatter update to support adding a stack trace (#1168)

* NewRelicContextFormatter update to support adding a stack trace

* Adding test coverage

* Remove typing for Python 2 compatibility

* Refactor instance methods to class methods and export

* Add truthy safeguard to log tests

* Fix Python 2 issues in tests

* lint:fix

---------

Co-authored-by: Tim Pansino <timpansino@gmail.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

newrelic/api/log.py

@@ -16,6 +16,7 @@
 import logging
 import re
 import warnings
+from traceback import format_tb
 
 from newrelic.api.application import application_instance
 from newrelic.api.time_trace import get_linking_metadata
@@ -28,22 +29,6 @@
 from newrelic.packages import six
 
 
-def format_exc_info(exc_info):
-    _, _, fullnames, message = parse_exc_info(exc_info)
-    fullname = fullnames[0]
-
-    formatted = {
-        "error.class": fullname,
-        "error.message": message,
-    }
-
-    expected = is_expected_error(exc_info)
-    if expected is not None:
-        formatted["error.expected"] = expected
-
-    return formatted
-
-
 def safe_json_encode(obj, ignore_string_types=False, **kwargs):
     # Performs the same operation as json_encode but replaces unserializable objects with a string containing their class name.
     # If ignore_string_types is True, do not encode string types further.
@@ -70,10 +55,48 @@ 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__()
+        """
+        :param Optional[int] stack_trace_limit:
+            Specifies the maximum number of frames to include for stack traces.
+            Defaults to `0` to suppress stack traces.
+            Setting this to `None` will make it so all available frames are included.
+        """
+        stack_trace_limit = kwargs.pop("stack_trace_limit", 0)
+
+        if stack_trace_limit is not None:
+            if not isinstance(stack_trace_limit, int):
+                raise TypeError("stack_trace_limit must be None or a non-negative integer")
+            if stack_trace_limit < 0:
+                raise ValueError("stack_trace_limit must be None or a non-negative integer")
+        self._stack_trace_limit = stack_trace_limit
+
+        super(NewRelicContextFormatter, self).__init__(*args, **kwargs)
 
     @classmethod
-    def log_record_to_dict(cls, record):
+    def format_exc_info(cls, exc_info, stack_trace_limit=0):
+        _, _, fullnames, message = parse_exc_info(exc_info)
+        fullname = fullnames[0]
+
+        formatted = {
+            "error.class": fullname,
+            "error.message": message,
+        }
+
+        expected = is_expected_error(exc_info)
+        if expected is not None:
+            formatted["error.expected"] = expected
+
+        if stack_trace_limit is None or stack_trace_limit > 0:
+            if exc_info[2] is not None:
+                stack_trace = "".join(format_tb(exc_info[2], limit=stack_trace_limit)) or None
+            else:
+                stack_trace = None
+            formatted["error.stack_trace"] = stack_trace
+
+        return formatted
+
+    @classmethod
+    def log_record_to_dict(cls, record, stack_trace_limit=0):
         output = {
             "timestamp": int(record.created * 1000),
             "message": record.getMessage(),
@@ -96,12 +119,17 @@ def log_record_to_dict(cls, record):
             output["extra." + key] = getattr(record, key)
 
         if record.exc_info:
-            output.update(format_exc_info(record.exc_info))
+            output.update(cls.format_exc_info(record.exc_info, stack_trace_limit))
 
         return output
 
     def format(self, record):
-        return json.dumps(self.log_record_to_dict(record), default=safe_json_encode, separators=(",", ":"))
+        return json.dumps(self.log_record_to_dict(record, self._stack_trace_limit), default=safe_json_encode, separators=(",", ":"))
+
+
+# Export class methods as top level functions for compatibility
+log_record_to_dict = NewRelicContextFormatter.log_record_to_dict
+format_exc_info = NewRelicContextFormatter.format_exc_info
 
 
 class NewRelicLogForwardingHandler(logging.Handler):

tests/agent_features/test_logs_in_context.py

@@ -14,6 +14,8 @@
 
 import json
 import logging
+import sys
+from traceback import format_tb
 
 import pytest
 
@@ -47,6 +49,22 @@ def log_buffer(caplog):
     _logger.removeHandler(_handler)
 
 
+@pytest.fixture
+def log_buffer_with_stack_trace(caplog):
+    buf = Buffer()
+
+    _formatter = NewRelicContextFormatter("", datefmt="ISO8601", stack_trace_limit=None)
+    _handler = logging.StreamHandler(buf)
+    _handler.setFormatter(_formatter)
+
+    _logger.addHandler(_handler)
+    caplog.set_level(logging.INFO, logger=__name__)
+
+    yield buf
+
+    _logger.removeHandler(_handler)
+
+
 class NonPrintableObject(object):
     def __str__(self):
         raise RuntimeError("Unable to print object.")
@@ -174,7 +192,7 @@ class ExceptionForTest(ValueError):
 
 
 @background_task()
-def test_newrelic_logger_error_inside_transaction(log_buffer):
+def test_newrelic_logger_error_inside_transaction_no_stack_trace(log_buffer):
     try:
         raise ExceptionForTest
     except ExceptionForTest:
@@ -220,7 +238,59 @@ def test_newrelic_logger_error_inside_transaction(log_buffer):
     assert set(message.keys()) == set(expected_extra_txn_keys)
 
 
-def test_newrelic_logger_error_outside_transaction(log_buffer):
+@background_task()
+def test_newrelic_logger_error_inside_transaction_with_stack_trace(log_buffer_with_stack_trace):
+    expected_stack_trace = ""
+    try:
+        raise ExceptionForTest
+    except ExceptionForTest:
+        _logger.exception("oops")
+        expected_stack_trace = "".join(format_tb(sys.exc_info()[2]))
+
+    log_buffer_with_stack_trace.seek(0)
+    message = json.load(log_buffer_with_stack_trace)
+
+    timestamp = message.pop("timestamp")
+    thread_id = message.pop("thread.id")
+    process_id = message.pop("process.id")
+    filename = message.pop("file.name")
+    line_number = message.pop("line.number")
+    stack_trace = message.pop("error.stack_trace")
+
+    assert isinstance(timestamp, int)
+    assert isinstance(thread_id, int)
+    assert isinstance(process_id, int)
+    assert filename.endswith("/test_logs_in_context.py")
+    assert isinstance(line_number, int)
+    assert isinstance(stack_trace, six.string_types)
+    assert stack_trace and stack_trace == expected_stack_trace
+
+    expected = {
+        "entity.name": "Python Agent Test (agent_features)",
+        "entity.type": "SERVICE",
+        "message": "oops",
+        "log.level": "ERROR",
+        "logger.name": "test_logs_in_context",
+        "thread.name": "MainThread",
+        "process.name": "MainProcess",
+        "error.class": "test_logs_in_context:ExceptionForTest",
+        "error.message": "",
+        "error.expected": False
+    }
+    expected_extra_txn_keys = (
+        "trace.id",
+        "span.id",
+        "entity.guid",
+        "hostname"
+    )
+
+    for k, v in expected.items():
+        assert message.pop(k) == v
+
+    assert set(message.keys()) == set(expected_extra_txn_keys)
+
+
+def test_newrelic_logger_error_outside_transaction_no_stack_trace(log_buffer):
     try:
         raise ExceptionForTest
     except ExceptionForTest:
@@ -263,6 +333,54 @@ def test_newrelic_logger_error_outside_transaction(log_buffer):
     assert set(message.keys()) == set(expected_extra_txn_keys)
 
 
+def test_newrelic_logger_error_outside_transaction_with_stack_trace(log_buffer_with_stack_trace):
+    expected_stack_trace = ""
+    try:
+        raise ExceptionForTest
+    except ExceptionForTest:
+        _logger.exception("oops")
+        expected_stack_trace = "".join(format_tb(sys.exc_info()[2]))
+
+    log_buffer_with_stack_trace.seek(0)
+    message = json.load(log_buffer_with_stack_trace)
+
+    timestamp = message.pop("timestamp")
+    thread_id = message.pop("thread.id")
+    process_id = message.pop("process.id")
+    filename = message.pop("file.name")
+    line_number = message.pop("line.number")
+    stack_trace = message.pop("error.stack_trace")
+
+    assert isinstance(timestamp, int)
+    assert isinstance(thread_id, int)
+    assert isinstance(process_id, int)
+    assert filename.endswith("/test_logs_in_context.py")
+    assert isinstance(line_number, int)
+    assert isinstance(stack_trace, six.string_types)
+    assert stack_trace and stack_trace == expected_stack_trace
+
+    expected = {
+        "entity.name": "Python Agent Test (agent_features)",
+        "entity.type": "SERVICE",
+        "message": "oops",
+        "log.level": "ERROR",
+        "logger.name": "test_logs_in_context",
+        "thread.name": "MainThread",
+        "process.name": "MainProcess",
+        "error.class": "test_logs_in_context:ExceptionForTest",
+        "error.message": "",
+    }
+    expected_extra_txn_keys = (
+        "entity.guid",
+        "hostname",
+    )
+
+    for k, v in expected.items():
+        assert message.pop(k) == v
+
+    assert set(message.keys()) == set(expected_extra_txn_keys)
+
+
 EXPECTED_KEYS_TXN = (
     "trace.id",
     "span.id",