getsentry · sentrivana · Jan 22, 2024 · Jan 18, 2024 · Jan 18, 2024 · Jan 19, 2024
diff --git a/MIGRATION_GUIDE.md b/MIGRATION_GUIDE.md
@@ -8,6 +8,8 @@
 
 - The `BackgroundWorker` thread used to process events was renamed from `raven-sentry.BackgroundWorker` to `sentry-sdk.BackgroundWorker`.
 - The `reraise` function was moved from `sentry_sdk._compat` to `sentry_sdk.utils`.
+- Moved the contents of `tracing_utils_py3.py` to `tracing_utils.py`. The `start_child_span_decorator` is now in `sentry_sdk.tracing_utils`.
+- The actual implementation of `get_current_span` was moved to `sentry_sdk.tracing_utils`. `sentry_sdk.get_current_span` is still accessible as part of the top-level API.
 
 ## Removed
 
@@ -17,5 +19,6 @@
 - Removed support for Flask 0.\*.
 - `sentry_sdk._functools` was removed.
 - A number of compatibility utilities were removed from `sentry_sdk._compat`: the constants `PY2` and `PY33`; the functions `datetime_utcnow`, `utc_from_timestamp`, `implements_str`, `contextmanager`; and the aliases `text_type`, `string_types`, `number_types`, `int_types`, `iteritems`, `binary_sequence_types`.
+- Removed `tracing_utils_py2.py`. The `start_child_span_decorator` is now in `sentry_sdk.tracing_utils`.
 
 ## Deprecated
diff --git a/sentry_sdk/api.py b/sentry_sdk/api.py
@@ -1,5 +1,6 @@
 import inspect
 
+from sentry_sdk import tracing_utils
 from sentry_sdk._types import TYPE_CHECKING
 from sentry_sdk.hub import Hub
 from sentry_sdk.scope import Scope
@@ -238,11 +239,7 @@ def get_current_span(hub=None):
     """
     Returns the currently active span if there is one running, otherwise `None`
     """
-    if hub is None:
-        hub = Hub.current
-
-    current_span = hub.scope.span
-    return current_span
+    return tracing_utils.get_current_span(hub)
 
 
 def get_traceparent():

diff --git a/sentry_sdk/tracing.py b/sentry_sdk/tracing.py
@@ -999,7 +999,7 @@ def my_function():
         async def my_async_function():
             ...
     """
-    from sentry_sdk.tracing_utils_py3 import start_child_span_decorator
+    from sentry_sdk.tracing_utils import start_child_span_decorator
 
     # This patterns allows usage of both @sentry_traced and @sentry_traced(...)
     # See https://stackoverflow.com/questions/52126071/decorator-with-arguments-avoid-parenthesis-when-no-arguments/52126278

diff --git a/sentry_sdk/tracing_utils.py b/sentry_sdk/tracing_utils.py
@@ -1,15 +1,19 @@
 import contextlib
+import inspect
 import re
 import sys
 from collections.abc import Mapping
+from functools import wraps
 from urllib.parse import quote, unquote
 
 import sentry_sdk
 from sentry_sdk.consts import OP, SPANDATA
 from sentry_sdk.utils import (
     capture_internal_exceptions,
     Dsn,
+    logger,
     match_regex_list,
+    qualname_from_function,
     to_string,
     is_sentry_url,
     _is_external_source,
@@ -501,5 +505,76 @@ def normalize_incoming_data(incoming_data):
     return data
 
 
+def start_child_span_decorator(func):
+    # type: (Any) -> Any
+    """
+    Decorator to add child spans for functions.
+
+    See also ``sentry_sdk.tracing.trace()``.
+    """
+    # Asynchronous case
+    if inspect.iscoroutinefunction(func):
+
+        @wraps(func)
+        async def func_with_tracing(*args, **kwargs):
+            # type: (*Any, **Any) -> Any
+
+            span = get_current_span(sentry_sdk.Hub.current)
+
+            if span is None:
+                logger.warning(
+                    "Can not create a child span for %s. "
+                    "Please start a Sentry transaction before calling this function.",
+                    qualname_from_function(func),
+                )
+                return await func(*args, **kwargs)
+
+            with span.start_child(
+                op=OP.FUNCTION,
+                description=qualname_from_function(func),
+            ):
+                return await func(*args, **kwargs)
+
+    # Synchronous case
+    else:
+
+        @wraps(func)
+        def func_with_tracing(*args, **kwargs):
+            # type: (*Any, **Any) -> Any
+
+            span = get_current_span(sentry_sdk.Hub.current)
+
+            if span is None:
+                logger.warning(
+                    "Can not create a child span for %s. "
+                    "Please start a Sentry transaction before calling this function.",
+                    qualname_from_function(func),
+                )
+                return func(*args, **kwargs)
+
+            with span.start_child(
+                op=OP.FUNCTION,
+                description=qualname_from_function(func),
+            ):
+                return func(*args, **kwargs)
+
+    return func_with_tracing
+
+
+def get_current_span(hub=None):
+    # type: (Optional[sentry_sdk.Hub]) -> Optional[Span]
+    """
+    Returns the currently active span if there is one running, otherwise `None`
+    """
+    if hub is None:
+        hub = sentry_sdk.Hub.current
+
+    current_span = hub.scope.span
+    return current_span
+
+
 # Circular imports
 from sentry_sdk.tracing import LOW_QUALITY_TRANSACTION_SOURCES
+
+if TYPE_CHECKING:
+    from sentry_sdk.tracing import Span
diff --git a/sentry_sdk/tracing_utils_py2.py b/sentry_sdk/tracing_utils_py2.py
diff --git a/sentry_sdk/tracing_utils_py3.py b/sentry_sdk/tracing_utils_py3.py
diff --git a/tests/conftest.py b/tests/conftest.py
@@ -639,7 +639,6 @@ def patch_start_tracing_child(fake_transaction_is_none=False):
         fake_start_child = None
 
     with mock.patch(
-        "sentry_sdk.tracing_utils_py3.get_current_span",
-        return_value=fake_transaction,
+        "sentry_sdk.tracing_utils.get_current_span", return_value=fake_transaction
     ):
         yield fake_start_child
diff --git a/tests/tracing/test_decorator.py b/tests/tracing/test_decorator.py
@@ -0,0 +1,76 @@
+from unittest import mock
+
+import pytest
+
+from sentry_sdk.tracing_utils import start_child_span_decorator
+from sentry_sdk.utils import logger
+from tests.conftest import patch_start_tracing_child
+
+
+def my_example_function():
+    return "return_of_sync_function"
+
+
+async def my_async_example_function():
+    return "return_of_async_function"
+
+
+def test_trace_decorator():
+    with patch_start_tracing_child() as fake_start_child:
+        result = my_example_function()
+        fake_start_child.assert_not_called()
+        assert result == "return_of_sync_function"
+
+        result2 = start_child_span_decorator(my_example_function)()
+        fake_start_child.assert_called_once_with(
+            op="function", description="test_decorator.my_example_function"
+        )
+        assert result2 == "return_of_sync_function"
+
+
+def test_trace_decorator_no_trx():
+    with patch_start_tracing_child(fake_transaction_is_none=True):
+        with mock.patch.object(logger, "warning", mock.Mock()) as fake_warning:
+            result = my_example_function()
+            fake_warning.assert_not_called()
+            assert result == "return_of_sync_function"
+
+            result2 = start_child_span_decorator(my_example_function)()
+            fake_warning.assert_called_once_with(
+                "Can not create a child span for %s. "
+                "Please start a Sentry transaction before calling this function.",
+                "test_decorator.my_example_function",
+            )
+            assert result2 == "return_of_sync_function"
+
+
+@pytest.mark.asyncio
+async def test_trace_decorator_async():
+    with patch_start_tracing_child() as fake_start_child:
+        result = await my_async_example_function()
+        fake_start_child.assert_not_called()
+        assert result == "return_of_async_function"
+
+        result2 = await start_child_span_decorator(my_async_example_function)()
+        fake_start_child.assert_called_once_with(
+            op="function",
+            description="test_decorator.my_async_example_function",
+        )
+        assert result2 == "return_of_async_function"
+
+
+@pytest.mark.asyncio
+async def test_trace_decorator_async_no_trx():
+    with patch_start_tracing_child(fake_transaction_is_none=True):
+        with mock.patch.object(logger, "warning", mock.Mock()) as fake_warning:
+            result = await my_async_example_function()
+            fake_warning.assert_not_called()
+            assert result == "return_of_async_function"
+
+            result2 = await start_child_span_decorator(my_async_example_function)()
+            fake_warning.assert_called_once_with(
+                "Can not create a child span for %s. "
+                "Please start a Sentry transaction before calling this function.",
+                "test_decorator.my_async_example_function",
+            )
+            assert result2 == "return_of_async_function"
diff --git a/tests/tracing/test_decorator_async_py3.py b/tests/tracing/test_decorator_async_py3.py