Hybrid agent set status

newrelic/api/opentelemetry.py

@@ -21,6 +21,7 @@
 from opentelemetry.propagate import set_global_textmap
 from opentelemetry.propagators.composite import CompositePropagator
 from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
+from opentelemetry.trace.status import Status, StatusCode
 
 from newrelic.api.application import application_instance
 from newrelic.api.background_task import BackgroundTask
@@ -111,6 +112,8 @@ def __init__(
         nr_transaction=None,
         nr_trace_type=FunctionTrace,
         instrumenting_module=None,
+        record_exception=True,
+        set_status_on_exception=True,
         *args,
         **kwargs,
     ):
@@ -123,6 +126,9 @@ def __init__(
         )  # This attribute is purely to prevent garbage collection
         self.nr_trace = None
         self.instrumenting_module = instrumenting_module
+        self.status = Status(StatusCode.UNSET)
+        self._record_exception = record_exception
+        self.set_status_on_exception = set_status_on_exception
 
         self.nr_parent = None
         current_nr_trace = current_trace()
@@ -238,19 +244,57 @@ def is_recording(self):
         if getattr(self.nr_trace, "end_time", None):
             return False
 
-        return getattr(self.nr_transaction, "priority", 1) > 0
+        # If priority is either not set at this point
+        # or greater than 0, we are recording.
+        priority = self.nr_transaction.priority
+        return (priority is None) or (priority > 0)
 
     def set_status(self, status, description=None):
-        # TODO: not implemented yet
-        raise NotImplementedError("Not implemented yet")
+        """
+        This code is modeled after the OpenTelemetry SDK's
+        status implementation:
+        https://github.com/open-telemetry/opentelemetry-python/blob/main/opentelemetry-sdk/src/opentelemetry/sdk/trace/__init__.py#L979
+
+        Additional Notes:
+        1. Ignore future calls if status is already set to OK
+            since span should be completed if status is OK.
+        2. Similarly, ignore calls to set to StatusCode.UNSET
+            since this will be either invalid or unnecessary.
+        """
+        if isinstance(status, Status):
+            if (self.status.status_code is StatusCode.OK) or status.is_unset:
+                return
+            if description is not None:
+                # `description` should only exist if status is StatusCode.ERROR
+                _logger.warning(
+                    "Description %s ignored. Use either `Status` or `(StatusCode, Description)`", description
+                )
+            self.status = status
+        elif isinstance(status, StatusCode):
+            if (self.status.status_code is StatusCode.OK) or (status is StatusCode.UNSET):
+                return
+            self.status = Status(status, description)
+        else:
+            _logger.warning("Invalid status type %s. Expected Status or StatusCode.", type(status))
+            return
+
+        # Add status as attribute
+        self.set_attribute("status_code", self.status.status_code.name)
+        self.set_attribute("status_description", self.status.description)
 
     def record_exception(self, exception, attributes=None, timestamp=None, escaped=False):
         error_args = sys.exc_info() if not exception else (type(exception), exception, exception.__traceback__)
 
-        if not hasattr(self, "nr_trace"):
-            notice_error(error_args, attributes=attributes)
+        # `escaped` indicates whether the exception has not
+        # been unhandled by the time the span has ended.
+        if attributes:
+            attributes["exception.escaped"] = escaped
         else:
-            self.nr_trace.notice_error(error_args, attributes=attributes)
+            attributes = {"exception.escaped": escaped}
+
+        self.set_attributes(attributes)
+
+        notice_error(error_args, attributes=attributes)
 
     def end(self, end_time=None, *args, **kwargs):
         # We will ignore the end_time parameter and use NR's end_time
@@ -283,6 +327,20 @@ def end(self, end_time=None, *args, **kwargs):
 
         self.nr_trace.__exit__(*sys.exc_info())
 
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """
+        Ends context manager and calls `end` on the `Span`.
+        This is used when span is called as a context manager
+        i.e. `with tracer.start_span() as span:`
+        """
+        if exc_val and self.is_recording():
+            if self._record_exception:
+                self.record_exception(exception=exc_val, escaped=True)
+            if self.set_status_on_exception:
+                self.set_status(Status(status_code=StatusCode.ERROR, description=f"{exc_type.__name__}: {exc_val}"))
+
+        super().__exit__(exc_type, exc_val, exc_tb)
+
 
 class Tracer(otel_api_trace.Tracer):
     def __init__(self, resource=None, instrumentation_library=None, *args, **kwargs):
@@ -311,6 +369,9 @@ def start_span(
             # Force application registration if not already active
             self.nr_application.activate()
 
+        self._record_exception = record_exception
+        self.set_status_on_exception = set_status_on_exception
+
         if not self.nr_application.settings.otel_bridge.enabled:
             return otel_api_trace.INVALID_SPAN
 
@@ -439,6 +500,8 @@ def start_span(
             nr_transaction=transaction,
             nr_trace_type=nr_trace_type,
             instrumenting_module=self.instrumentation_library,
+            record_exception=self._record_exception,
+            set_status_on_exception=self.set_status_on_exception,
         )
 
         return span

newrelic/hooks/hybridagent_opentelemetry.py

@@ -30,7 +30,6 @@
 os.environ["OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_REQUEST"] = ".*"
 os.environ["OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_RESPONSE"] = ".*"
 
-
 ###########################################
 #   Context Instrumentation
 ###########################################

tests/hybridagent_opentelemetry/test_attributes.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 from opentelemetry import trace
+from testing_support.fixtures import dt_enabled
 from testing_support.validators.validate_span_events import validate_span_events
 
 from newrelic.api.background_task import background_task
@@ -22,6 +23,7 @@
 
 
 def test_trace_with_span_attributes(tracer):
+    @dt_enabled
     @validate_span_events(
         count=1,
         exact_intrinsics={
@@ -69,6 +71,7 @@ def newrelic_function_trace():
         otel_span = trace.get_current_span()
         otel_span.set_attribute("otel_span_attribute_FT", "OTel span attribute from FT")
 
+    @dt_enabled
     @validate_span_events(
         count=1,
         exact_intrinsics={

tests/hybridagent_opentelemetry/test_context_propagation.py

@@ -26,7 +26,6 @@
 _override_settings = {"trusted_account_key": "1", "distributed_tracing.enabled": True, "span_events.enabled": True}
 
 
-# @dt_enabled
 @pytest.mark.parametrize("telemetry", ["newrelic", "otel"])
 @pytest.mark.parametrize("propagation", [accept_distributed_trace_headers, PROPAGATOR.extract])
 def test_distributed_trace_header_compatibility_full_granularity(telemetry, propagation):

tests/hybridagent_opentelemetry/test_status.py

@@ -0,0 +1,220 @@
+# Copyright 2010 New Relic, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pytest
+from opentelemetry.trace.status import Status, StatusCode
+from testing_support.fixtures import dt_enabled
+from testing_support.util import conditional_decorator
+from testing_support.validators.validate_error_event_attributes import validate_error_event_attributes
+
+from newrelic.api.background_task import background_task
+
+
+# `set_status` takes in a status argument that can be either
+# a Status or StatusCode type and a description argument.
+# Status has a StatusCode attribute and an optional description attribute.
+# If current StatusCode is StatusCode.OK, calls to set_status on this span are no-ops.
+# If status_to_set is StatusCode.UNSET, this is also a no-op.
+@pytest.mark.parametrize(
+    "current_status,status_to_set,expected_status_code",
+    [
+        (Status(StatusCode.UNSET), Status(StatusCode.OK), StatusCode.OK),  # current_status==UNSET -> status_to_set
+        (
+            Status(StatusCode.UNSET),
+            Status(StatusCode.ERROR),
+            StatusCode.ERROR,
+        ),  # current_status==UNSET -> status_to_set
+        (
+            Status(StatusCode.OK),
+            Status(StatusCode.UNSET),
+            StatusCode.OK,
+        ),  # current_status==OK -> No-Op / status_to_set==UNSET -> No-Op
+        (Status(StatusCode.OK), Status(StatusCode.ERROR), StatusCode.OK),  # current_status==OK -> No-Op
+        (Status(StatusCode.ERROR), Status(StatusCode.UNSET), StatusCode.ERROR),  # status_to_set==UNSET -> No-Op
+        (Status(StatusCode.ERROR), Status(StatusCode.OK), StatusCode.OK),  # current_status==ERROR -> status_to_set
+        (Status(StatusCode.UNSET), StatusCode.OK, StatusCode.OK),  # current_status==UNSET -> status_to_set
+        (Status(StatusCode.UNSET), StatusCode.ERROR, StatusCode.ERROR),  # current_status==UNSET -> status_to_set
+        (
+            Status(StatusCode.OK),
+            StatusCode.UNSET,
+            StatusCode.OK,
+        ),  # current_status==OK -> No-Op / status_to_set==UNSET -> No-Op
+        (Status(StatusCode.OK), StatusCode.ERROR, StatusCode.OK),  # current_status==OK -> No-Op
+        (Status(StatusCode.ERROR), StatusCode.UNSET, StatusCode.ERROR),  # status_to_set==UNSET -> No-Op
+        (Status(StatusCode.ERROR), StatusCode.OK, StatusCode.OK),  # current_status==ERROR -> status_to_set
+    ],
+    ids=(
+        "status_unset_to_ok",
+        "status_unset_to_error",
+        "status_ok_to_unset",
+        "status_ok_to_error",
+        "status_error_to_unset",
+        "status_error_to_ok",
+        "status_code_unset_to_ok",
+        "status_code_unset_to_error",
+        "status_code_ok_to_unset",
+        "status_code_ok_to_error",
+        "status_code_error_to_unset",
+        "status_code_error_to_ok",
+    ),
+)
+def test_status_setting(tracer, current_status, status_to_set, expected_status_code):
+    @background_task()
+    def _test():
+        with tracer.start_as_current_span(name="TestSpan") as span:
+            # First, set to the current status to simulate the initial state
+            span.set_status(current_status)
+
+            # Then, attempt to set the new status
+            span.set_status(status_to_set)
+            assert span.status.status_code == expected_status_code
+
+    _test()
+
+
+@pytest.mark.parametrize("_record_exception", [True, False])
+@pytest.mark.parametrize("_set_status_on_exception", [True, False])
+def test_set_status_with_start_as_current_span(tracer, _record_exception, _set_status_on_exception):
+    @dt_enabled
+    @conditional_decorator(
+        condition=_record_exception,
+        decorator=validate_error_event_attributes(
+            exact_attrs={
+                "agent": {},
+                "intrinsic": {"error.message": "Test exception message", "error.class": "builtins:ValueError"},
+                "user": {"exception.escaped": False},
+            }
+        ),
+    )
+    @background_task()
+    def _test():
+        with pytest.raises(ValueError):
+            with tracer.start_as_current_span(
+                name="TestSpan", record_exception=_record_exception, set_status_on_exception=_set_status_on_exception
+            ) as span:
+                raise ValueError("Test exception message")
+
+        assert span.status.status_code == StatusCode.ERROR if _set_status_on_exception else StatusCode.UNSET
+
+    _test()
+
+
+@pytest.mark.parametrize("_record_exception", [True, False])
+@pytest.mark.parametrize("_set_status_on_exception", [True, False])
+def test_set_status_with_start_span(tracer, _record_exception, _set_status_on_exception):
+    @dt_enabled
+    @conditional_decorator(
+        condition=_record_exception,
+        decorator=validate_error_event_attributes(
+            exact_attrs={
+                "agent": {},
+                "intrinsic": {"error.message": "Test exception message", "error.class": "builtins:ValueError"},
+                "user": {"exception.escaped": True},
+            }
+        ),
+    )
+    @background_task()
+    def _test():
+        with pytest.raises(ValueError):
+            with tracer.start_span(
+                name="TestSpan", record_exception=_record_exception, set_status_on_exception=_set_status_on_exception
+            ) as span:
+                raise ValueError("Test exception message")
+
+        assert span.status.status_code == StatusCode.ERROR if _set_status_on_exception else StatusCode.UNSET
+
+    _test()
+
+
+# `set_status` takes in a status argument that can be either
+# a Status or StatusCode type and a description argument.
+# Status has a StatusCode attribute and an optional description attribute.
+# If Status type is passed in, the description argument is ignored.
+# If StatusCode type is passed in, the description argument is used to
+# create a Status object.
+# Only StatusCode.ERROR should have a description, so if StatusCode.OK is passed
+# in with a description, the description will be ignored within the Status object.
+@pytest.mark.parametrize(
+    "status_to_set,description,expected_status_code",
+    [
+        (Status(StatusCode.OK), None, StatusCode.OK),  # OK_Status_no_description, no description
+        (
+            Status(StatusCode.OK),
+            "I will be ignored in set_status",
+            StatusCode.OK,
+        ),  # OK_Status_no_description, description
+        (
+            Status(StatusCode.OK, "I will be ignored in Status"),
+            None,
+            StatusCode.OK,
+        ),  # OK_Status_with_description, no description
+        (
+            Status(StatusCode.OK, "I will be ignored in Status"),
+            "I will be ignored in set_status",
+            StatusCode.OK,
+        ),  # OK_Status_with_description, description
+        (Status(StatusCode.ERROR), None, StatusCode.ERROR),  # Error_Status_no_description, no description
+        (
+            Status(StatusCode.ERROR),
+            "I will be ignored in set_status",
+            StatusCode.ERROR,
+        ),  # Error_Status_no_description, description
+        (
+            Status(StatusCode.ERROR, "This is where I belong"),
+            None,
+            StatusCode.ERROR,
+        ),  # Error_Status_with_description, no description
+        (
+            Status(StatusCode.ERROR, "This is where I belong"),
+            "I will be ignored in set_status",
+            StatusCode.ERROR,
+        ),  # Error_Status_with_description, description
+        (StatusCode.OK, None, StatusCode.OK),  # OK_StatusCode, no description
+        (StatusCode.OK, "I will be ignored in Status", StatusCode.OK),  # OK_StatusCode, description
+        (StatusCode.ERROR, None, StatusCode.ERROR),  # Error_StatusCode, no description
+        (StatusCode.ERROR, "This is where I belong", StatusCode.ERROR),  # Error_StatusCode, description
+    ],
+    ids=(
+        "ok_status_no_description-no_description",
+        "ok_status_no_description-description",
+        "ok_status_description-no_description",
+        "ok_status_description-description",
+        "error_status_no_description-no_description",
+        "error_status_no_description-description",
+        "error_status_description-no_description",
+        "error_status_description-description",
+        "ok_status_code-no_description",
+        "ok_status_code-description",
+        "error_status_code-no_description",
+        "error_status_code-description",
+    ),
+)
+def test_status_setting(tracer, status_to_set, description, expected_status_code):
+    @background_task()
+    def _test():
+        with tracer.start_as_current_span(name="TestSpan") as span:
+            # Set the new status
+            span.set_status(status_to_set, description)
+
+            # If status code is OK, do not have description
+            # if expected_status_code == StatusCode.OK:
+            if span.status.status_code == StatusCode.OK:
+                assert span.status.description is None
+
+            # If status code is ERROR, make sure description
+            # is set correctly (if provided):
+            if span.status.status_code == StatusCode.ERROR:
+                assert span.status.description in [None, "This is where I belong"]
+
+    _test()