feat(tracing): Improve @trace decorator. (#4648)

Update the `@trace` decorator and make it more powerful. It accepts now
the following parameters: `op`, `name`, `attributes`.

Example usage:
```python
import sentry_sdk
from sentry_sdk.consts import OP

# Simple usage (like before)
@sentry_sdk.trace
def process_data():
    # Function implementation
    pass

# With custom parameters
@sentry_sdk.trace(
    op=OP.DB_QUERY,
    name="Get user data",
    attributes={"postgres": True}
)
def make_db_query(sql):
    # Function implementation
    pass
```

This creates better DX for our users.

---------

Co-authored-by: Daniel Szoke <[email protected]>

Author: antonpirker

docs/api.rst

@@ -37,6 +37,7 @@ Enriching Events
 Performance Monitoring
 ======================
 
+.. autofunction:: sentry_sdk.api.trace
 .. autofunction:: sentry_sdk.api.continue_trace
 .. autofunction:: sentry_sdk.api.get_current_span
 .. autofunction:: sentry_sdk.api.start_span

sentry_sdk/integrations/starlite.py

@@ -17,7 +17,7 @@
     from starlite.plugins.base import get_plugin_for_value  # type: ignore
     from starlite.routes.http import HTTPRoute  # type: ignore
     from starlite.utils import ConnectionDataExtractor, is_async_callable, Ref  # type: ignore
-    from pydantic import BaseModel  # type: ignore
+    from pydantic import BaseModel
 except ImportError:
     raise DidNotEnable("Starlite is not installed")
 

sentry_sdk/tracing.py

@@ -1340,43 +1340,85 @@ def _set_initial_sampling_decision(self, sampling_context):
 if TYPE_CHECKING:
 
     @overload
-    def trace(func=None):
-        # type: (None) -> Callable[[Callable[P, R]], Callable[P, R]]
+    def trace(func=None, *, op=None, name=None, attributes=None):
+        # type: (None, Optional[str], Optional[str], Optional[dict[str, Any]]) -> Callable[[Callable[P, R]], Callable[P, R]]
+        # Handles: @trace() and @trace(op="custom")
         pass
 
     @overload
     def trace(func):
         # type: (Callable[P, R]) -> Callable[P, R]
+        # Handles: @trace
         pass
 
 
-def trace(func=None):
-    # type: (Optional[Callable[P, R]]) -> Union[Callable[P, R], Callable[[Callable[P, R]], Callable[P, R]]]
+def trace(func=None, *, op=None, name=None, attributes=None):
+    # type: (Optional[Callable[P, R]], Optional[str], Optional[str], Optional[dict[str, Any]]) -> Union[Callable[P, R], Callable[[Callable[P, R]], Callable[P, R]]]
     """
-    Decorator to start a child span under the existing current transaction.
-    If there is no current transaction, then nothing will be traced.
+    Decorator to start a child span around a function call.
 
-    .. code-block::
-        :caption: Usage
+    This decorator automatically creates a new span when the decorated function
+    is called, and finishes the span when the function returns or raises an exception.
+
+    :param func: The function to trace. When used as a decorator without parentheses,
+        this is the function being decorated. When used with parameters (e.g.,
+        ``@trace(op="custom")``, this should be None.
+    :type func: Callable or None
+
+    :param op: The operation name for the span. This is a high-level description
+        of what the span represents (e.g., "http.client", "db.query").
+        You can use predefined constants from :py:class:`sentry_sdk.consts.OP`
+        or provide your own string. If not provided, a default operation will
+        be assigned based on the template.
+    :type op: str or None
+
+    :param name: The human-readable name/description for the span. If not provided,
+        defaults to the function name. This provides more specific details about
+        what the span represents (e.g., "GET /api/users", "process_user_data").
+    :type name: str or None
+
+    :param attributes: A dictionary of key-value pairs to add as attributes to the span.
+        Attribute values must be strings, integers, floats, or booleans. These
+        attributes provide additional context about the span's execution.
+    :type attributes: dict[str, Any] or None
+
+    :returns: When used as ``@trace``, returns the decorated function. When used as
+        ``@trace(...)`` with parameters, returns a decorator function.
+    :rtype: Callable or decorator function
+
+    Example::
 
         import sentry_sdk
+        from sentry_sdk.consts import OP
 
+        # Simple usage with default values
         @sentry_sdk.trace
-        def my_function():
-            ...
+        def process_data():
+            # Function implementation
+            pass
 
-        @sentry_sdk.trace
-        async def my_async_function():
-            ...
+        # With custom parameters
+        @sentry_sdk.trace(
+            op=OP.DB_QUERY,
+            name="Get user data",
+            attributes={"postgres": True}
+        )
+        def make_db_query(sql):
+            # Function implementation
+            pass
     """
-    from sentry_sdk.tracing_utils import start_child_span_decorator
+    from sentry_sdk.tracing_utils import create_span_decorator
+
+    decorator = create_span_decorator(
+        op=op,
+        name=name,
+        attributes=attributes,
+    )
 
-    # 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)
+        return decorator(func)
     else:
-        return start_child_span_decorator
+        return decorator
 
 
 # Circular imports

sentry_sdk/tracing_utils.py

@@ -1,12 +1,12 @@
 import contextlib
+import functools
 import inspect
 import os
 import re
 import sys
 from collections.abc import Mapping
 from datetime import timedelta
 from decimal import ROUND_DOWN, Decimal, DefaultContext, localcontext
-from functools import wraps
 from random import Random
 from urllib.parse import quote, unquote
 import uuid
@@ -770,70 +770,86 @@ def normalize_incoming_data(incoming_data):
     return data
 
 
-def start_child_span_decorator(func):
-    # type: (Any) -> Any
+def create_span_decorator(op=None, name=None, attributes=None):
+    # type: (Optional[str], Optional[str], Optional[dict[str, Any]]) -> Any
     """
-    Decorator to add child spans for functions.
+    Create a span decorator that can wrap both sync and async functions.
 
-    See also ``sentry_sdk.tracing.trace()``.
+    :param op: The operation type for the span.
+    :param name: The name of the span.
+    :param attributes: Additional attributes to set on the span.
     """
-    # Asynchronous case
-    if inspect.iscoroutinefunction(func):
 
-        @wraps(func)
-        async def func_with_tracing(*args, **kwargs):
-            # type: (*Any, **Any) -> Any
+    def span_decorator(f):
+        # type: (Any) -> Any
+        """
+        Decorator to create a span for the given function.
+        """
 
-            span = get_current_span()
+        @functools.wraps(f)
+        async def async_wrapper(*args, **kwargs):
+            # type: (*Any, **Any) -> Any
+            current_span = get_current_span()
 
-            if span is None:
+            if current_span is None:
                 logger.debug(
                     "Cannot create a child span for %s. "
                     "Please start a Sentry transaction before calling this function.",
-                    qualname_from_function(func),
+                    qualname_from_function(f),
                 )
-                return await func(*args, **kwargs)
+                return await f(*args, **kwargs)
+
+            span_op = op or OP.FUNCTION
+            span_name = name or qualname_from_function(f) or ""
 
-            with span.start_child(
-                op=OP.FUNCTION,
-                name=qualname_from_function(func),
-            ):
-                return await func(*args, **kwargs)
+            with current_span.start_child(
+                op=span_op,
+                name=span_name,
+            ) as span:
+                span.update_data(attributes or {})
+                result = await f(*args, **kwargs)
+                return result
 
         try:
-            func_with_tracing.__signature__ = inspect.signature(func)  # type: ignore[attr-defined]
+            async_wrapper.__signature__ = inspect.signature(f)  # type: ignore[attr-defined]
         except Exception:
             pass
 
-    # Synchronous case
-    else:
-
-        @wraps(func)
-        def func_with_tracing(*args, **kwargs):
+        @functools.wraps(f)
+        def sync_wrapper(*args, **kwargs):
             # type: (*Any, **Any) -> Any
+            current_span = get_current_span()
 
-            span = get_current_span()
-
-            if span is None:
+            if current_span is None:
                 logger.debug(
                     "Cannot create a child span for %s. "
                     "Please start a Sentry transaction before calling this function.",
-                    qualname_from_function(func),
+                    qualname_from_function(f),
                 )
-                return func(*args, **kwargs)
+                return f(*args, **kwargs)
+
+            span_op = op or OP.FUNCTION
+            span_name = name or qualname_from_function(f) or ""
 
-            with span.start_child(
-                op=OP.FUNCTION,
-                name=qualname_from_function(func),
-            ):
-                return func(*args, **kwargs)
+            with current_span.start_child(
+                op=span_op,
+                name=span_name,
+            ) as span:
+                span.update_data(attributes or {})
+                result = f(*args, **kwargs)
+                return result
 
         try:
-            func_with_tracing.__signature__ = inspect.signature(func)  # type: ignore[attr-defined]
+            sync_wrapper.__signature__ = inspect.signature(f)  # type: ignore[attr-defined]
         except Exception:
             pass
 
-    return func_with_tracing
+        if inspect.iscoroutinefunction(f):
+            return async_wrapper
+        else:
+            return sync_wrapper
+
+    return span_decorator
 
 
 def get_current_span(scope=None):

tests/tracing/test_decorator.py

@@ -4,7 +4,7 @@
 import pytest
 
 from sentry_sdk.tracing import trace
-from sentry_sdk.tracing_utils import start_child_span_decorator
+from sentry_sdk.tracing_utils import create_span_decorator
 from sentry_sdk.utils import logger
 from tests.conftest import patch_start_tracing_child
 
@@ -24,6 +24,7 @@ def test_trace_decorator():
         fake_start_child.assert_not_called()
         assert result == "return_of_sync_function"
 
+        start_child_span_decorator = create_span_decorator()
         result2 = start_child_span_decorator(my_example_function)()
         fake_start_child.assert_called_once_with(
             op="function", name="test_decorator.my_example_function"
@@ -38,6 +39,7 @@ def test_trace_decorator_no_trx():
             fake_debug.assert_not_called()
             assert result == "return_of_sync_function"
 
+            start_child_span_decorator = create_span_decorator()
             result2 = start_child_span_decorator(my_example_function)()
             fake_debug.assert_called_once_with(
                 "Cannot create a child span for %s. "
@@ -55,6 +57,7 @@ async def test_trace_decorator_async():
         fake_start_child.assert_not_called()
         assert result == "return_of_async_function"
 
+        start_child_span_decorator = create_span_decorator()
         result2 = await start_child_span_decorator(my_async_example_function)()
         fake_start_child.assert_called_once_with(
             op="function",
@@ -71,6 +74,7 @@ async def test_trace_decorator_async_no_trx():
             fake_debug.assert_not_called()
             assert result == "return_of_async_function"
 
+            start_child_span_decorator = create_span_decorator()
             result2 = await start_child_span_decorator(my_async_example_function)()
             fake_debug.assert_called_once_with(
                 "Cannot create a child span for %s. "