Add decorator for Sentry tracing

sentry_sdk/__init__.py

@@ -7,6 +7,8 @@
 
 from sentry_sdk.consts import VERSION  # noqa
 
+from sentry_sdk.tracing import trace  # noqa
+
 __all__ = [  # noqa
     "Hub",
     "Scope",

sentry_sdk/tracing.py

@@ -6,21 +6,23 @@
 import sentry_sdk
 from sentry_sdk.consts import INSTRUMENTER
 from sentry_sdk.utils import logger, nanosecond_time
+from sentry_sdk._compat import PY2
 from sentry_sdk._types import TYPE_CHECKING
 
 
 if TYPE_CHECKING:
     import typing
 
-    from typing import Optional
     from typing import Any
     from typing import Dict
+    from typing import Iterator
     from typing import List
+    from typing import Optional
     from typing import Tuple
-    from typing import Iterator
 
     import sentry_sdk.profiler
-    from sentry_sdk._types import Event, SamplingContext, MeasurementUnit
+    from sentry_sdk._types import Event, MeasurementUnit, SamplingContext
+
 
 BAGGAGE_HEADER_NAME = "baggage"
 SENTRY_TRACE_HEADER_NAME = "sentry-trace"
@@ -803,6 +805,36 @@ def finish(self, hub=None, end_timestamp=None):
         pass
 
 
+def trace(func=None):
+    # type: (Any) -> Any
+    """
+    Decorator to start a child span under the existing current transaction.
+    If there is no current transaction, than nothing will be traced.
+
+    Usage:
+        import sentry_sdk
+
+        @sentry_sdk.trace
+        def my_function():
+            ...
+
+        @sentry_sdk.trace
+        async def my_async_function():
+            ...
+    """
+    if PY2:
+        from sentry_sdk.tracing_utils_py2 import start_child_span_decorator
+    else:
+        from sentry_sdk.tracing_utils_py3 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
+    if func:
+        return start_child_span_decorator(func)
+    else:
+        return start_child_span_decorator
+
+
 # Circular imports
 
 from sentry_sdk.tracing_utils import (

sentry_sdk/tracing_utils.py

@@ -395,6 +395,16 @@ def should_propagate_trace(hub, url):
     return False
 
 
+def get_running_span_or_transaction(hub):
+    # type: (sentry_sdk.Hub) -> Optional[Union[Span, Transaction]]
+    current_span = hub.scope.span
+    if current_span is not None:
+        return current_span
+
+    transaction = hub.scope.transaction
+    return transaction
+
+
 # Circular imports
 from sentry_sdk.tracing import LOW_QUALITY_TRANSACTION_SOURCES
 

sentry_sdk/tracing_utils_py2.py

@@ -0,0 +1,45 @@
+from functools import wraps
+
+import sentry_sdk
+from sentry_sdk._types import TYPE_CHECKING
+from sentry_sdk.tracing_utils import get_running_span_or_transaction
+from sentry_sdk.consts import OP
+from sentry_sdk.utils import logger, qualname_from_function
+
+
+if TYPE_CHECKING:
+    from typing import Any
+
+
+def start_child_span_decorator(func):
+    # type: (Any) -> Any
+    """
+    Decorator to add child spans for functions.
+
+    This is the Python 2 compatible version of the decorator.
+    Duplicated code from ``sentry_sdk.tracing_utils_python3.start_child_span_decorator``.
+
+    See also ``sentry_sdk.tracing.trace()``.
+    """
+
+    @wraps(func)
+    def func_with_tracing(*args, **kwargs):
+        # type: (*Any, **Any) -> Any
+
+        span_or_trx = get_running_span_or_transaction(sentry_sdk.Hub.current)
+
+        if span_or_trx is None:
+            logger.warning(
+                "No transaction found. Not creating a child span for %s. "
+                "Please start a Sentry transaction before calling this function.",
+                qualname_from_function(func),
+            )
+            return func(*args, **kwargs)
+
+        with span_or_trx.start_child(
+            op=OP.FUNCTION,
+            description=qualname_from_function(func),
+        ):
+            return func(*args, **kwargs)
+
+    return func_with_tracing

sentry_sdk/tracing_utils_py3.py

@@ -0,0 +1,72 @@
+import inspect
+from functools import wraps
+
+import sentry_sdk
+from sentry_sdk._types import TYPE_CHECKING
+from sentry_sdk.consts import OP
+from sentry_sdk.tracing_utils import get_running_span_or_transaction
+from sentry_sdk.utils import logger, qualname_from_function
+
+
+if TYPE_CHECKING:
+    from typing import Any
+
+
+def start_child_span_decorator(func):
+    # type: (Any) -> Any
+    """
+    Decorator to add child spans for functions.
+
+    This is the Python 3 compatible version of the decorator.
+    For Python 2 there is duplicated code here: ``sentry_sdk.tracing_utils_python2.start_child_span_decorator()``.
+
+    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_or_trx = get_running_span_or_transaction(sentry_sdk.Hub.current)
+
+            if span_or_trx is None:
+                logger.warning(
+                    "No transaction found. Not creating 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_or_trx.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_or_trx = get_running_span_or_transaction(sentry_sdk.Hub.current)
+
+            if span_or_trx is None:
+                logger.warning(
+                    "No transaction found. Not creating a child span for %s. "
+                    "Please start a Sentry transaction before calling this function.",
+                    qualname_from_function(func),
+                )
+                return func(*args, **kwargs)
+
+            with span_or_trx.start_child(
+                op=OP.FUNCTION,
+                description=qualname_from_function(func),
+            ):
+                return func(*args, **kwargs)
+
+    return func_with_tracing

tests/tracing/test_decorator_py2.py

@@ -0,0 +1,50 @@
+import mock
+
+from sentry_sdk.tracing_utils_py2 import (
+    start_child_span_decorator as start_child_span_decorator_py2,
+)
+from sentry_sdk.utils import logger
+
+
+def my_example_function():
+    return "return_of_sync_function"
+
+
+def test_trace_decorator_py2():
+    fake_start_child = mock.MagicMock()
+    fake_transaction = mock.MagicMock()
+    fake_transaction.start_child = fake_start_child
+
+    with mock.patch(
+        "sentry_sdk.tracing_utils_py2.get_running_span_or_transaction",
+        return_value=fake_transaction,
+    ):
+        result = my_example_function()
+        fake_start_child.assert_not_called()
+        assert result == "return_of_sync_function"
+
+        result2 = start_child_span_decorator_py2(my_example_function)()
+        fake_start_child.assert_called_once_with(
+            op="function", description="test_decorator_py2.my_example_function"
+        )
+        assert result2 == "return_of_sync_function"
+
+
+def test_trace_decorator_py2_no_trx():
+    fake_transaction = None
+
+    with mock.patch(
+        "sentry_sdk.tracing_utils_py2.get_running_span_or_transaction",
+        return_value=fake_transaction,
+    ):
+        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_py2(my_example_function)()
+            fake_warning.assert_called_once_with(
+                "No transaction found. Not creating a child span for %s. Please start a Sentry transaction before calling this function.",
+                "test_decorator_py2.my_example_function",
+            )
+            assert result2 == "return_of_sync_function"

tests/tracing/test_decorator_py3.py

@@ -0,0 +1,101 @@
+import mock
+import pytest
+import sys
+
+from sentry_sdk.tracing_utils_py3 import (
+    start_child_span_decorator as start_child_span_decorator_py3,
+)
+from sentry_sdk.utils import logger
+
+if sys.version_info < (3, 6):
+    pytest.skip("Async decorator only works on Python 3.6+", allow_module_level=True)
+
+
+def my_example_function():
+    return "return_of_sync_function"
+
+
+async def my_async_example_function():
+    return "return_of_async_function"
+
+
+def test_trace_decorator_sync_py3():
+    fake_start_child = mock.MagicMock()
+    fake_transaction = mock.MagicMock()
+    fake_transaction.start_child = fake_start_child
+
+    with mock.patch(
+        "sentry_sdk.tracing_utils_py3.get_running_span_or_transaction",
+        return_value=fake_transaction,
+    ):
+        result = my_example_function()
+        fake_start_child.assert_not_called()
+        assert result == "return_of_sync_function"
+
+        result2 = start_child_span_decorator_py3(my_example_function)()
+        fake_start_child.assert_called_once_with(
+            op="function", description="test_decorator_py3.my_example_function"
+        )
+        assert result2 == "return_of_sync_function"
+
+
+def test_trace_decorator_sync_py3_no_trx():
+    fake_transaction = None
+
+    with mock.patch(
+        "sentry_sdk.tracing_utils_py3.get_running_span_or_transaction",
+        return_value=fake_transaction,
+    ):
+        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_py3(my_example_function)()
+            fake_warning.assert_called_once_with(
+                "No transaction found. Not creating a child span for %s. Please start a Sentry transaction before calling this function.",
+                "test_decorator_py3.my_example_function",
+            )
+            assert result2 == "return_of_sync_function"
+
+
+@pytest.mark.asyncio
+async def test_trace_decorator_async_py3():
+    fake_start_child = mock.MagicMock()
+    fake_transaction = mock.MagicMock()
+    fake_transaction.start_child = fake_start_child
+
+    with mock.patch(
+        "sentry_sdk.tracing_utils_py3.get_running_span_or_transaction",
+        return_value=fake_transaction,
+    ):
+        result = await my_async_example_function()
+        fake_start_child.assert_not_called()
+        assert result == "return_of_async_function"
+
+        result2 = await start_child_span_decorator_py3(my_async_example_function)()
+        fake_start_child.assert_called_once_with(
+            op="function", description="test_decorator_py3.my_async_example_function"
+        )
+        assert result2 == "return_of_async_function"
+
+
+@pytest.mark.asyncio
+async def test_trace_decorator_async_py3_no_trx():
+    fake_transaction = None
+
+    with mock.patch(
+        "sentry_sdk.tracing_utils_py3.get_running_span_or_transaction",
+        return_value=fake_transaction,
+    ):
+        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_py3(my_async_example_function)()
+            fake_warning.assert_called_once_with(
+                "No transaction found. Not creating a child span for %s. Please start a Sentry transaction before calling this function.",
+                "test_decorator_py3.my_async_example_function",
+            )
+            assert result2 == "return_of_async_function"

tests/tracing/test_misc.py

@@ -8,7 +8,10 @@
 from sentry_sdk import Hub, start_span, start_transaction, set_measurement
 from sentry_sdk.consts import MATCH_ALL
 from sentry_sdk.tracing import Span, Transaction
-from sentry_sdk.tracing_utils import should_propagate_trace
+from sentry_sdk.tracing_utils import (
+    should_propagate_trace,
+    get_running_span_or_transaction,
+)
 
 try:
     from unittest import mock  # python 3.3 and above
@@ -306,3 +309,20 @@ def test_should_propagate_trace(
     hub.client.options = {"trace_propagation_targets": trace_propagation_targets}
 
     assert should_propagate_trace(hub, url) == expected_propagation_decision
+
+
+def test_get_running_span_or_transaction():
+    fake_hub = mock.MagicMock()
+    fake_hub.scope = mock.MagicMock()
+
+    fake_hub.scope.span = mock.MagicMock()
+    fake_hub.scope.transaction = None
+    assert get_running_span_or_transaction(fake_hub) == fake_hub.scope.span
+
+    fake_hub.scope.span = None
+    fake_hub.scope.transaction = mock.MagicMock()
+    assert get_running_span_or_transaction(fake_hub) == fake_hub.scope.transaction
+
+    fake_hub.scope.span = None
+    fake_hub.scope.transaction = None
+    assert get_running_span_or_transaction(fake_hub) is None