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():
+            ...
+
+        @senrty_sdk.trace
+        async def my_async_function():
+            ...
+    """
+    if PY2:
+        from sentry_sdk.tracing_utils_python2 import start_child_span_decorator
+    else:
+        from sentry_sdk.tracing_utils_python3 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_python2.py

@@ -0,0 +1,53 @@
+from functools import wraps
+
+import sentry_sdk
+from sentry_sdk._types import TYPE_CHECKING
+from sentry_sdk.consts import OP
+from sentry_sdk.utils import logger, qualname_from_function
+
+
+if TYPE_CHECKING:
+    from typing import Any, Optional
+
+    from sentry_sdk.tracing import Transaction
+
+
+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()``.
+    """
+
+    def _get_transaction():
+        # type: () -> Optional[Transaction]
+        transaction = sentry_sdk.Hub.current.scope.transaction
+        return transaction
+
+    @wraps(func)
+    def func_with_tracing(*args, **kwargs):
+        # type: (*Any, **Any) -> Any
+
+        transaction = _get_transaction()
+
+        # If no transaction, do nothing
+        if transaction 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)
+
+        # If we have a transaction, we decorate the function!
+        with transaction.start_child(
+            op=OP.FUNCTION,
+            description=qualname_from_function(func),
+        ):
+            return func(*args, **kwargs)
+
+    return func_with_tracing

sentry_sdk/tracing_utils_python3.py

@@ -0,0 +1,82 @@
+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.utils import logger, qualname_from_function
+
+
+if TYPE_CHECKING:
+    from typing import Any, Optional
+
+    from sentry_sdk.tracing import Transaction
+
+
+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()``.
+    """
+
+    def _get_transaction():
+        # type: () -> Optional[Transaction]
+        transaction = sentry_sdk.Hub.current.scope.transaction
+        return transaction
+
+    # Asynchronous case
+    if inspect.iscoroutinefunction(func):
+
+        @wraps(func)
+        async def func_with_tracing(*args, **kwargs):
+            # type: (*Any, **Any) -> Any
+
+            transaction = _get_transaction()
+
+            # If no transaction, do nothing
+            if transaction 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)
+
+            # If we have a transaction, we wrap the function.
+            with transaction.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
+
+            transaction = _get_transaction()
+
+            # If no transaction, do nothing
+            if transaction 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)
+
+            # If we have a transaction, we decorate the function!
+            with transaction.start_child(
+                op=OP.FUNCTION,
+                description=qualname_from_function(func),
+            ):
+                return func(*args, **kwargs)
+
+    return func_with_tracing