diff --git a/sentry_sdk/traces.py b/sentry_sdk/traces.py index 1739f8b25c..bede851e0c 100644 --- a/sentry_sdk/traces.py +++ b/sentry_sdk/traces.py @@ -14,9 +14,12 @@ from sentry_sdk.utils import format_attribute, logger if TYPE_CHECKING: - from typing import Any, Optional, Union + from typing import Any, Callable, Optional, ParamSpec, TypeVar, Union from sentry_sdk._types import Attributes, AttributeValue + P = ParamSpec("P") + R = TypeVar("R") + class SpanStatus(str, Enum): OK = "ok" @@ -228,6 +231,14 @@ def __repr__(self) -> str: f"active={self._active})>" ) + def __enter__(self) -> "StreamedSpan": + return self + + def __exit__( + self, ty: "Optional[Any]", value: "Optional[Any]", tb: "Optional[Any]" + ) -> None: + pass + def get_attributes(self) -> "Attributes": return self._attributes @@ -299,6 +310,14 @@ def __init__(self) -> None: def __repr__(self) -> str: return f"<{self.__class__.__name__}(sampled={self.sampled})>" + def __enter__(self) -> "NoOpStreamedSpan": + return self + + def __exit__( + self, ty: "Optional[Any]", value: "Optional[Any]", tb: "Optional[Any]" + ) -> None: + pass + def get_attributes(self) -> "Attributes": return {} @@ -342,3 +361,74 @@ def trace_id(self) -> str: @property def sampled(self) -> "Optional[bool]": return False + + +def trace( + func: "Optional[Callable[P, R]]" = None, + *, + name: "Optional[str]" = None, + attributes: "Optional[dict[str, Any]]" = None, + active: bool = True, +) -> "Union[Callable[P, R], Callable[[Callable[P, R]], Callable[P, R]]]": + """ + Decorator to start a span around a function call. + + 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 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 + + :param active: Controls whether spans started while this span is running + will automatically become its children. That's the default behavior. If + you want to create a span that shouldn't have any children (unless + provided explicitly via the `parent_span` argument), set this to False. + :type active: bool + + :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 + + # Simple usage with default values + @sentry_sdk.trace + def process_data(): + # Function implementation + pass + + # With custom parameters + @sentry_sdk.trace( + name="Get user data", + attributes={"postgres": True} + ) + def make_db_query(sql): + # Function implementation + pass + """ + from sentry_sdk.tracing_utils import create_streaming_span_decorator + + decorator = create_streaming_span_decorator( + name=name, + attributes=attributes, + active=active, + ) + + if func: + return decorator(func) + else: + return decorator diff --git a/sentry_sdk/tracing_utils.py b/sentry_sdk/tracing_utils.py index c1d6c44535..80b4628153 100644 --- a/sentry_sdk/tracing_utils.py +++ b/sentry_sdk/tracing_utils.py @@ -942,6 +942,57 @@ def sync_wrapper(*args: "Any", **kwargs: "Any") -> "Any": return span_decorator +def create_streaming_span_decorator( + name: "Optional[str]" = None, + attributes: "Optional[dict[str, Any]]" = None, + active: bool = True, +) -> "Any": + """ + Create a span creating decorator that can wrap both sync and async functions. + """ + + def span_decorator(f: "Any") -> "Any": + """ + Decorator to create a span for the given function. + """ + + @functools.wraps(f) + async def async_wrapper(*args: "Any", **kwargs: "Any") -> "Any": + span_name = name or qualname_from_function(f) or "" + + with start_streaming_span( + name=span_name, attributes=attributes, active=active + ): + result = await f(*args, **kwargs) + return result + + try: + async_wrapper.__signature__ = inspect.signature(f) # type: ignore[attr-defined] + except Exception: + pass + + @functools.wraps(f) + def sync_wrapper(*args: "Any", **kwargs: "Any") -> "Any": + span_name = name or qualname_from_function(f) or "" + + with start_streaming_span( + name=span_name, attributes=attributes, active=active + ): + return f(*args, **kwargs) + + try: + sync_wrapper.__signature__ = inspect.signature(f) # type: ignore[attr-defined] + except Exception: + pass + + if inspect.iscoroutinefunction(f): + return async_wrapper + else: + return sync_wrapper + + return span_decorator + + def get_current_span(scope: "Optional[sentry_sdk.Scope]" = None) -> "Optional[Span]": """ Returns the currently active span if there is one running, otherwise `None` @@ -1317,6 +1368,11 @@ def add_sentry_baggage_to_headers( LOW_QUALITY_TRANSACTION_SOURCES, SENTRY_TRACE_HEADER_NAME, ) +from sentry_sdk.traces import ( + LOW_QUALITY_SEGMENT_SOURCES, + start_span as start_streaming_span, + StreamedSpan, +) if TYPE_CHECKING: from sentry_sdk.tracing import Span