Improved handling of span status (#3261)

antonpirker · web-flow · commit b7fd54aaea4e · 2024-07-10T07:17:03.000Z
diff --git a/sentry_sdk/consts.py b/sentry_sdk/consts.py
@@ -386,6 +386,32 @@ class SPANDATA:
     """
 
 
+class SPANSTATUS:
+    """
+    The status of a Sentry span.
+
+    See: https://develop.sentry.dev/sdk/event-payloads/contexts/#trace-context
+    """
+
+    ABORTED = "aborted"
+    ALREADY_EXISTS = "already_exists"
+    CANCELLED = "cancelled"
+    DATA_LOSS = "data_loss"
+    DEADLINE_EXCEEDED = "deadline_exceeded"
+    FAILED_PRECONDITION = "failed_precondition"
+    INTERNAL_ERROR = "internal_error"
+    INVALID_ARGUMENT = "invalid_argument"
+    NOT_FOUND = "not_found"
+    OK = "ok"
+    OUT_OF_RANGE = "out_of_range"
+    PERMISSION_DENIED = "permission_denied"
+    RESOURCE_EXHAUSTED = "resource_exhausted"
+    UNAUTHENTICATED = "unauthenticated"
+    UNAVAILABLE = "unavailable"
+    UNIMPLEMENTED = "unimplemented"
+    UNKNOWN_ERROR = "unknown_error"
+
+
 class OP:
     ANTHROPIC_MESSAGES_CREATE = "ai.messages.create.anthropic"
     CACHE_GET = "cache.get"
diff --git a/sentry_sdk/integrations/aiohttp.py b/sentry_sdk/integrations/aiohttp.py
@@ -3,7 +3,7 @@
 
 import sentry_sdk
 from sentry_sdk.api import continue_trace
-from sentry_sdk.consts import OP, SPANDATA
+from sentry_sdk.consts import OP, SPANSTATUS, SPANDATA
 from sentry_sdk.integrations import Integration, DidNotEnable
 from sentry_sdk.integrations.logging import ignore_logger
 from sentry_sdk.scope import Scope
@@ -133,7 +133,7 @@ async def sentry_app_handle(self, request, *args, **kwargs):
                             transaction.set_http_status(e.status_code)
                             raise
                         except (asyncio.CancelledError, ConnectionResetError):
-                            transaction.set_status("cancelled")
+                            transaction.set_status(SPANSTATUS.CANCELLED)
                             raise
                         except Exception:
                             # This will probably map to a 500 but seems like we
diff --git a/sentry_sdk/integrations/arq.py b/sentry_sdk/integrations/arq.py
@@ -2,7 +2,7 @@
 
 import sentry_sdk
 from sentry_sdk._types import TYPE_CHECKING
-from sentry_sdk.consts import OP
+from sentry_sdk.consts import OP, SPANSTATUS
 from sentry_sdk.integrations import DidNotEnable, Integration
 from sentry_sdk.integrations.logging import ignore_logger
 from sentry_sdk.scope import Scope, should_send_default_pii
@@ -119,10 +119,10 @@ def _capture_exception(exc_info):
 
     if scope.transaction is not None:
         if exc_info[0] in ARQ_CONTROL_FLOW_EXCEPTIONS:
-            scope.transaction.set_status("aborted")
+            scope.transaction.set_status(SPANSTATUS.ABORTED)
             return
 
-        scope.transaction.set_status("internal_error")
+        scope.transaction.set_status(SPANSTATUS.INTERNAL_ERROR)
 
     event, hint = event_from_exception(
         exc_info,
diff --git a/sentry_sdk/integrations/celery/__init__.py b/sentry_sdk/integrations/celery/__init__.py
@@ -5,7 +5,7 @@
 import sentry_sdk
 from sentry_sdk import isolation_scope
 from sentry_sdk.api import continue_trace
-from sentry_sdk.consts import OP, SPANDATA
+from sentry_sdk.consts import OP, SPANSTATUS, SPANDATA
 from sentry_sdk.integrations import Integration, DidNotEnable
 from sentry_sdk.integrations.celery.beat import (
     _patch_beat_apply_entry,
@@ -317,7 +317,7 @@ def _inner(*args, **kwargs):
                     origin=CeleryIntegration.origin,
                 )
                 transaction.name = task.name
-                transaction.set_status("ok")
+                transaction.set_status(SPANSTATUS.OK)
 
             if transaction is None:
                 return f(*args, **kwargs)
diff --git a/sentry_sdk/integrations/huey.py b/sentry_sdk/integrations/huey.py
@@ -4,7 +4,7 @@
 import sentry_sdk
 from sentry_sdk._types import TYPE_CHECKING
 from sentry_sdk.api import continue_trace, get_baggage, get_traceparent
-from sentry_sdk.consts import OP
+from sentry_sdk.consts import OP, SPANSTATUS
 from sentry_sdk.integrations import DidNotEnable, Integration
 from sentry_sdk.scope import Scope, should_send_default_pii
 from sentry_sdk.tracing import (
@@ -109,10 +109,10 @@ def _capture_exception(exc_info):
     scope = Scope.get_current_scope()
 
     if exc_info[0] in HUEY_CONTROL_FLOW_EXCEPTIONS:
-        scope.transaction.set_status("aborted")
+        scope.transaction.set_status(SPANSTATUS.ABORTED)
         return
 
-    scope.transaction.set_status("internal_error")
+    scope.transaction.set_status(SPANSTATUS.INTERNAL_ERROR)
     event, hint = event_from_exception(
         exc_info,
         client_options=Scope.get_client().options,
@@ -161,7 +161,7 @@ def _sentry_execute(self, task, timestamp=None):
                 source=TRANSACTION_SOURCE_TASK,
                 origin=HueyIntegration.origin,
             )
-            transaction.set_status("ok")
+            transaction.set_status(SPANSTATUS.OK)
 
             if not getattr(task, "_sentry_is_patched", False):
                 task.execute = _wrap_task_execute(task.execute)
diff --git a/sentry_sdk/integrations/opentelemetry/span_processor.py b/sentry_sdk/integrations/opentelemetry/span_processor.py
@@ -16,7 +16,7 @@
     INVALID_TRACE_ID,
 )
 from sentry_sdk import get_client, start_transaction
-from sentry_sdk.consts import INSTRUMENTER
+from sentry_sdk.consts import INSTRUMENTER, SPANSTATUS
 from sentry_sdk.integrations.opentelemetry.consts import (
     SENTRY_BAGGAGE_KEY,
     SENTRY_TRACE_KEY,
@@ -299,10 +299,10 @@ def _update_span_with_otel_status(self, sentry_span, otel_span):
             return
 
         if otel_span.status.is_ok:
-            sentry_span.set_status("ok")
+            sentry_span.set_status(SPANSTATUS.OK)
             return
 
-        sentry_span.set_status("internal_error")
+        sentry_span.set_status(SPANSTATUS.INTERNAL_ERROR)
 
     def _update_span_with_otel_data(self, sentry_span, otel_span):
         # type: (SentrySpan, OTelSpan) -> None
diff --git a/sentry_sdk/integrations/pymongo.py b/sentry_sdk/integrations/pymongo.py
@@ -1,7 +1,7 @@
 import copy
 
 import sentry_sdk
-from sentry_sdk.consts import SPANDATA, OP
+from sentry_sdk.consts import SPANSTATUS, SPANDATA, OP
 from sentry_sdk.integrations import DidNotEnable, Integration
 from sentry_sdk.scope import should_send_default_pii
 from sentry_sdk.tracing import Span
@@ -181,7 +181,7 @@ def failed(self, event):
 
         try:
             span = self._ongoing_operations.pop(self._operation_key(event))
-            span.set_status("internal_error")
+            span.set_status(SPANSTATUS.INTERNAL_ERROR)
             span.__exit__(None, None, None)
         except KeyError:
             return
@@ -193,7 +193,7 @@ def succeeded(self, event):
 
         try:
             span = self._ongoing_operations.pop(self._operation_key(event))
-            span.set_status("ok")
+            span.set_status(SPANSTATUS.OK)
             span.__exit__(None, None, None)
         except KeyError:
             pass
diff --git a/sentry_sdk/integrations/sqlalchemy.py b/sentry_sdk/integrations/sqlalchemy.py
@@ -1,6 +1,6 @@
 import sentry_sdk
 from sentry_sdk._types import TYPE_CHECKING
-from sentry_sdk.consts import SPANDATA
+from sentry_sdk.consts import SPANSTATUS, SPANDATA
 from sentry_sdk.db.explain_plan.sqlalchemy import attach_explain_plan_to_span
 from sentry_sdk.integrations import Integration, DidNotEnable
 from sentry_sdk.tracing_utils import add_query_source, record_sql_queries
@@ -107,7 +107,7 @@ def _handle_error(context, *args):
     span = getattr(execution_context, "_sentry_sql_span", None)  # type: Optional[Span]
 
     if span is not None:
-        span.set_status("internal_error")
+        span.set_status(SPANSTATUS.INTERNAL_ERROR)
 
     # _after_cursor_execute does not get called for crashing SQL stmts. Judging
     # from SQLAlchemy codebase it does seem like any error coming into this
diff --git a/sentry_sdk/tracing.py b/sentry_sdk/tracing.py
@@ -3,7 +3,7 @@
 from datetime import datetime, timedelta, timezone
 
 import sentry_sdk
-from sentry_sdk.consts import INSTRUMENTER, SPANDATA
+from sentry_sdk.consts import INSTRUMENTER, SPANSTATUS, SPANDATA
 from sentry_sdk.profiler.continuous_profiler import get_profiler_id
 from sentry_sdk.utils import (
     get_current_thread_meta,
@@ -149,6 +149,45 @@ class TransactionKwargs(SpanKwargs, total=False):
 }
 
 
+def get_span_status_from_http_code(http_status_code):
+    # type: (int) -> str
+    """
+    Returns the Sentry status corresponding to the given HTTP status code.
+
+    See: https://develop.sentry.dev/sdk/event-payloads/contexts/#trace-context
+    """
+    if http_status_code < 400:
+        return SPANSTATUS.OK
+
+    elif 400 <= http_status_code < 500:
+        if http_status_code == 403:
+            return SPANSTATUS.PERMISSION_DENIED
+        elif http_status_code == 404:
+            return SPANSTATUS.NOT_FOUND
+        elif http_status_code == 429:
+            return SPANSTATUS.RESOURCE_EXHAUSTED
+        elif http_status_code == 413:
+            return SPANSTATUS.FAILED_PRECONDITION
+        elif http_status_code == 401:
+            return SPANSTATUS.UNAUTHENTICATED
+        elif http_status_code == 409:
+            return SPANSTATUS.ALREADY_EXISTS
+        else:
+            return SPANSTATUS.INVALID_ARGUMENT
+
+    elif 500 <= http_status_code < 600:
+        if http_status_code == 504:
+            return SPANSTATUS.DEADLINE_EXCEEDED
+        elif http_status_code == 501:
+            return SPANSTATUS.UNIMPLEMENTED
+        elif http_status_code == 503:
+            return SPANSTATUS.UNAVAILABLE
+        else:
+            return SPANSTATUS.INTERNAL_ERROR
+
+    return SPANSTATUS.UNKNOWN_ERROR
+
+
 class _SpanRecorder:
     """Limits the number of spans recorded in a transaction."""
 
@@ -317,7 +356,7 @@ def __enter__(self):
     def __exit__(self, ty, value, tb):
         # type: (Optional[Any], Optional[Any], Optional[Any]) -> None
         if value is not None:
-            self.set_status("internal_error")
+            self.set_status(SPANSTATUS.INTERNAL_ERROR)
 
         scope, old_span = self._context_manager_state
         del self._context_manager_state
@@ -540,37 +579,9 @@ def set_http_status(self, http_status):
         # type: (int) -> None
         self.set_tag(
             "http.status_code", str(http_status)
-        )  # we keep this for backwards compatability
+        )  # we keep this for backwards compatibility
         self.set_data(SPANDATA.HTTP_STATUS_CODE, http_status)
-
-        if http_status < 400:
-            self.set_status("ok")
-        elif 400 <= http_status < 500:
-            if http_status == 403:
-                self.set_status("permission_denied")
-            elif http_status == 404:
-                self.set_status("not_found")
-            elif http_status == 429:
-                self.set_status("resource_exhausted")
-            elif http_status == 413:
-                self.set_status("failed_precondition")
-            elif http_status == 401:
-                self.set_status("unauthenticated")
-            elif http_status == 409:
-                self.set_status("already_exists")
-            else:
-                self.set_status("invalid_argument")
-        elif 500 <= http_status < 600:
-            if http_status == 504:
-                self.set_status("deadline_exceeded")
-            elif http_status == 501:
-                self.set_status("unimplemented")
-            elif http_status == 503:
-                self.set_status("unavailable")
-            else:
-                self.set_status("internal_error")
-        else:
-            self.set_status("unknown_error")
+        self.set_status(get_span_status_from_http_code(http_status))
 
     def is_success(self):
         # type: () -> bool
diff --git a/tests/tracing/test_integration_tests.py b/tests/tracing/test_integration_tests.py
@@ -10,6 +10,7 @@
     start_span,
     start_transaction,
 )
+from sentry_sdk.consts import SPANSTATUS
 from sentry_sdk.transport import Transport
 from sentry_sdk.tracing import Transaction
 
@@ -20,7 +21,7 @@ def test_basic(sentry_init, capture_events, sample_rate):
     events = capture_events()
 
     with start_transaction(name="hi") as transaction:
-        transaction.set_status("ok")
+        transaction.set_status(SPANSTATUS.OK)
         with pytest.raises(ZeroDivisionError):
             with start_span(op="foo", description="foodesc"):
                 1 / 0
