From 60811032dd3fd8fb8505606b52ce5df40fde1c60 Mon Sep 17 00:00:00 2001 From: Ivana Kellyer Date: Fri, 11 Sep 2026 15:52:30 +0200 Subject: [PATCH] chore: Drop dual tracing helpers --- sentry_sdk/ai/monitoring.py | 24 +++----- sentry_sdk/ai/utils.py | 20 ++----- sentry_sdk/integrations/aiohttp.py | 3 +- sentry_sdk/integrations/aiomysql.py | 23 ++----- sentry_sdk/integrations/google_genai/utils.py | 28 ++++----- sentry_sdk/integrations/huggingface_hub.py | 19 +++--- sentry_sdk/integrations/langgraph.py | 20 +++---- sentry_sdk/integrations/mcp.py | 60 ++++++++----------- .../integrations/pydantic_ai/spans/utils.py | 22 ++++--- .../integrations/redis/modules/caches.py | 26 ++++---- sentry_sdk/integrations/rust_tracing.py | 29 ++++----- 11 files changed, 104 insertions(+), 170 deletions(-) diff --git a/sentry_sdk/ai/monitoring.py b/sentry_sdk/ai/monitoring.py index 5f096cc8cc..dfadc290db 100644 --- a/sentry_sdk/ai/monitoring.py +++ b/sentry_sdk/ai/monitoring.py @@ -1,18 +1,17 @@ from typing import TYPE_CHECKING -from sentry_sdk.ai.utils import _set_span_data_attribute from sentry_sdk.consts import SPANDATA -from sentry_sdk.traces import StreamedSpan -from sentry_sdk.tracing import Span if TYPE_CHECKING: from typing import Any, Awaitable, Callable, Optional, TypeVar, Union + from sentry_sdk.traces import StreamedSpan + F = TypeVar("F", bound=Union[Callable[..., Any], Callable[..., Awaitable[Any]]]) def record_token_usage( - span: "Union[Span, StreamedSpan]", + span: "StreamedSpan", input_tokens: "Optional[int]" = None, input_tokens_cached: "Optional[int]" = None, input_tokens_cache_write: "Optional[int]" = None, @@ -21,30 +20,25 @@ def record_token_usage( total_tokens: "Optional[int]" = None, ) -> None: if input_tokens is not None: - _set_span_data_attribute(span, SPANDATA.GEN_AI_USAGE_INPUT_TOKENS, input_tokens) + span.set_attribute(SPANDATA.GEN_AI_USAGE_INPUT_TOKENS, input_tokens) if input_tokens_cached is not None: - _set_span_data_attribute( - span, + span.set_attribute( SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHED, input_tokens_cached, ) if input_tokens_cache_write is not None: - _set_span_data_attribute( - span, + span.set_attribute( SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHE_WRITE, input_tokens_cache_write, ) if output_tokens is not None: - _set_span_data_attribute( - span, SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS, output_tokens - ) + span.set_attribute(SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS, output_tokens) if output_tokens_reasoning is not None: - _set_span_data_attribute( - span, + span.set_attribute( SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS_REASONING, output_tokens_reasoning, ) @@ -53,4 +47,4 @@ def record_token_usage( total_tokens = input_tokens + output_tokens if total_tokens is not None: - _set_span_data_attribute(span, SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS, total_tokens) + span.set_attribute(SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS, total_tokens) diff --git a/sentry_sdk/ai/utils.py b/sentry_sdk/ai/utils.py index a158d1e9a8..006c05bec3 100644 --- a/sentry_sdk/ai/utils.py +++ b/sentry_sdk/ai/utils.py @@ -3,12 +3,11 @@ from typing import TYPE_CHECKING if TYPE_CHECKING: - from typing import Any, Dict, Optional, Tuple, Union + from typing import Any, Dict, Optional, Tuple - from sentry_sdk.tracing import Span + from sentry_sdk.traces import StreamedSpan import sentry_sdk -from sentry_sdk.traces import StreamedSpan from sentry_sdk.utils import logger @@ -449,25 +448,16 @@ def _normalize_data(data: "Any", unpack: bool = True) -> "Any": def set_data_normalized( - span: "Union[Span, StreamedSpan]", + span: "StreamedSpan", key: str, value: "Any", unpack: bool = True, ) -> None: normalized = _normalize_data(value, unpack=unpack) if isinstance(normalized, (int, float, bool, str)): - _set_span_data_attribute(span, key, normalized) + span.set_attribute(key, normalized) else: - _set_span_data_attribute(span, key, json.dumps(normalized)) - - -def _set_span_data_attribute( - span: "Union[Span, StreamedSpan]", key: str, value: "Any" -) -> None: - if isinstance(span, StreamedSpan): - span.set_attribute(key, value) - else: - span.set_data(key, value) + span.set_attribute(key, json.dumps(normalized)) def normalize_message_role(role: str) -> str: diff --git a/sentry_sdk/integrations/aiohttp.py b/sentry_sdk/integrations/aiohttp.py index 778f5c1e28..9970dda6c7 100644 --- a/sentry_sdk/integrations/aiohttp.py +++ b/sentry_sdk/integrations/aiohttp.py @@ -73,7 +73,6 @@ from aiohttp.web_urldispatcher import UrlMappingMatchInfo from sentry_sdk._types import Attributes, Event, EventProcessor - from sentry_sdk.tracing import Span from sentry_sdk.utils import ExcInfo @@ -343,7 +342,7 @@ async def on_request_start( parsed_url.url if parsed_url else SENSITIVE_DATA_SUBSTITUTE, ) - span: "Union[Span, StreamedSpan, None]" = None + span: "Optional[StreamedSpan]" = None attributes: "Attributes" = { "sentry.op": OP.HTTP_CLIENT, "sentry.origin": AioHttpIntegration.origin, diff --git a/sentry_sdk/integrations/aiomysql.py b/sentry_sdk/integrations/aiomysql.py index 5dccc4144d..d8d637c0c9 100644 --- a/sentry_sdk/integrations/aiomysql.py +++ b/sentry_sdk/integrations/aiomysql.py @@ -241,32 +241,21 @@ def _get_connect_data(conn: Any, *, use_streaming_keys: bool = False) -> dict[st def _set_db_data(span: Any, conn: Any) -> None: """Set database-related span data from connection object.""" - if isinstance(span, StreamedSpan): - set_value = span.set_attribute - db_system = SPANDATA.DB_SYSTEM_NAME - db_name = SPANDATA.DB_NAMESPACE - else: - # Remove this else block once we've completely migrated to streamed spans - # The use of deprecated attributes here is to ensure backwards compatibility - set_value = span.set_data - db_system = SPANDATA.DB_SYSTEM - db_name = SPANDATA.DB_NAME - - set_value(db_system, "mysql") - set_value(SPANDATA.DB_DRIVER_NAME, "aiomysql") + span.set_attribute(SPANDATA.DB_SYSTEM_NAME, "mysql") + span.set_attribute(SPANDATA.DB_DRIVER_NAME, "aiomysql") host = getattr(conn, "host", None) if host is not None: - set_value(SPANDATA.SERVER_ADDRESS, host) + span.set_attribute(SPANDATA.SERVER_ADDRESS, host) port = getattr(conn, "port", None) if port is not None: - set_value(SPANDATA.SERVER_PORT, port) + span.set_attribute(SPANDATA.SERVER_PORT, port) database = getattr(conn, "db", None) if database is not None: - set_value(db_name, database) + span.set_attribute(SPANDATA.DB_NAMESPACE, database) user = getattr(conn, "user", None) if user is not None: - set_value(SPANDATA.DB_USER, user) + span.set_attribute(SPANDATA.DB_USER, user) diff --git a/sentry_sdk/integrations/google_genai/utils.py b/sentry_sdk/integrations/google_genai/utils.py index 27ba665d2e..6d075edb8f 100644 --- a/sentry_sdk/integrations/google_genai/utils.py +++ b/sentry_sdk/integrations/google_genai/utils.py @@ -28,7 +28,6 @@ ) from sentry_sdk.consts import OP, SPANDATA from sentry_sdk.scope import should_send_default_pii -from sentry_sdk.traces import StreamedSpan from sentry_sdk.utils import ( capture_internal_exceptions, event_from_exception, @@ -52,6 +51,7 @@ ) from sentry_sdk._types import TextPart + from sentry_sdk.traces import StreamedSpan _is_PIL_available = False try: @@ -708,22 +708,21 @@ def wrapped_tool(tool: "Tool | Callable[..., Any]") -> "Tool | Callable[..., Any @wraps(tool) async def async_wrapped(*args: "Any", **kwargs: "Any") -> "Any": with _create_tool_span(tool_name, tool_doc) as span: - set_on_span = ( - span.set_attribute - if isinstance(span, StreamedSpan) - else span.set_data - ) # Capture tool input tool_input = _capture_tool_input(args, kwargs, tool) with capture_internal_exceptions(): - set_on_span(SPANDATA.GEN_AI_TOOL_INPUT, safe_serialize(tool_input)) + span.set_attribute( + SPANDATA.GEN_AI_TOOL_INPUT, safe_serialize(tool_input) + ) try: result = await tool(*args, **kwargs) # Capture tool output with capture_internal_exceptions(): - set_on_span(SPANDATA.GEN_AI_TOOL_OUTPUT, safe_serialize(result)) + span.set_attribute( + SPANDATA.GEN_AI_TOOL_OUTPUT, safe_serialize(result) + ) return result except Exception as exc: @@ -736,22 +735,21 @@ async def async_wrapped(*args: "Any", **kwargs: "Any") -> "Any": @wraps(tool) def sync_wrapped(*args: "Any", **kwargs: "Any") -> "Any": with _create_tool_span(tool_name, tool_doc) as span: - set_on_span = ( - span.set_attribute - if isinstance(span, StreamedSpan) - else span.set_data - ) # Capture tool input tool_input = _capture_tool_input(args, kwargs, tool) with capture_internal_exceptions(): - set_on_span(SPANDATA.GEN_AI_TOOL_INPUT, safe_serialize(tool_input)) + span.set_attribute( + SPANDATA.GEN_AI_TOOL_INPUT, safe_serialize(tool_input) + ) try: result = tool(*args, **kwargs) # Capture tool output with capture_internal_exceptions(): - set_on_span(SPANDATA.GEN_AI_TOOL_OUTPUT, safe_serialize(result)) + span.set_attribute( + SPANDATA.GEN_AI_TOOL_OUTPUT, safe_serialize(result) + ) return result except Exception as exc: diff --git a/sentry_sdk/integrations/huggingface_hub.py b/sentry_sdk/integrations/huggingface_hub.py index 93fd813467..1fb951c5dc 100644 --- a/sentry_sdk/integrations/huggingface_hub.py +++ b/sentry_sdk/integrations/huggingface_hub.py @@ -5,10 +5,7 @@ import sentry_sdk from sentry_sdk.ai.monitoring import record_token_usage -from sentry_sdk.ai.utils import ( - _set_span_data_attribute, - set_data_normalized, -) +from sentry_sdk.ai.utils import set_data_normalized from sentry_sdk.consts import OP, SPANDATA from sentry_sdk.integrations import DidNotEnable, Integration, _check_minimum_version from sentry_sdk.scope import should_send_default_pii @@ -106,10 +103,10 @@ def new_huggingface_task(*args: "Any", **kwargs: "Any") -> "Any": }, ) - _set_span_data_attribute(span, SPANDATA.GEN_AI_OPERATION_NAME, operation_name) + span.set_attribute(SPANDATA.GEN_AI_OPERATION_NAME, operation_name) if model: - _set_span_data_attribute(span, SPANDATA.GEN_AI_REQUEST_MODEL, model) + span.set_attribute(SPANDATA.GEN_AI_REQUEST_MODEL, model) attribute_mapping = { "frequency_penalty": SPANDATA.GEN_AI_REQUEST_FREQUENCY_PENALTY, @@ -143,7 +140,7 @@ def new_huggingface_task(*args: "Any", **kwargs: "Any") -> "Any": value = kwargs.get(attribute, None) if value is not None: if isinstance(value, (int, float, bool, str)): - _set_span_data_attribute(span, span_attribute, value) + span.set_attribute(span_attribute, value) else: set_data_normalized(span, span_attribute, value, unpack=False) @@ -204,9 +201,7 @@ def new_huggingface_task(*args: "Any", **kwargs: "Any") -> "Any": response_text_buffer.append(choice.message.content) if response_model is not None: - _set_span_data_attribute( - span, SPANDATA.GEN_AI_RESPONSE_MODEL, response_model - ) + span.set_attribute(SPANDATA.GEN_AI_RESPONSE_MODEL, response_model) if finish_reason is not None: set_data_normalized( @@ -380,8 +375,8 @@ def new_iterator() -> "Iterable[ChatCompletionStreamOutput]": yield chunk if response_model is not None: - _set_span_data_attribute( - span, SPANDATA.GEN_AI_RESPONSE_MODEL, response_model + span.set_attribute( + SPANDATA.GEN_AI_RESPONSE_MODEL, response_model ) if finish_reason is not None: diff --git a/sentry_sdk/integrations/langgraph.py b/sentry_sdk/integrations/langgraph.py index 65db4d4f9c..f6a85f4dc3 100644 --- a/sentry_sdk/integrations/langgraph.py +++ b/sentry_sdk/integrations/langgraph.py @@ -1,5 +1,5 @@ from functools import wraps -from typing import Any, Callable, List, Optional +from typing import TYPE_CHECKING, Any, Callable, List, Optional import sentry_sdk from sentry_sdk.ai.utils import ( @@ -12,13 +12,15 @@ # This is fine because langgraph depends on langchain-base, and LangchainIntegration only imports from langchain-base. from sentry_sdk.integrations.langchain import LangchainIntegration from sentry_sdk.scope import should_send_default_pii -from sentry_sdk.traces import StreamedSpan from sentry_sdk.utils import ( has_data_collection_enabled, package_version, safe_serialize, ) +if TYPE_CHECKING: + from sentry_sdk.traces import StreamedSpan + try: from langgraph.errors import GraphBubbleUp from langgraph.pregel import Pregel @@ -266,7 +268,7 @@ def _extract_tool_calls(messages: "Optional[List[Any]]") -> "Optional[List[Any]] return tool_calls if tool_calls else None -def _set_usage_data(span: "sentry_sdk.tracing.Span", messages: "Any") -> None: +def _set_usage_data(span: "StreamedSpan", messages: "Any") -> None: input_tokens = 0 output_tokens = 0 total_tokens = 0 @@ -284,24 +286,20 @@ def _set_usage_data(span: "sentry_sdk.tracing.Span", messages: "Any") -> None: output_tokens += int(token_usage.get("completion_tokens", 0)) total_tokens += int(token_usage.get("total_tokens", 0)) - set_on_span = ( - span.set_attribute if isinstance(span, StreamedSpan) else span.set_data - ) - if input_tokens > 0: - set_on_span(SPANDATA.GEN_AI_USAGE_INPUT_TOKENS, input_tokens) + span.set_attribute(SPANDATA.GEN_AI_USAGE_INPUT_TOKENS, input_tokens) if output_tokens > 0: - set_on_span(SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS, output_tokens) + span.set_attribute(SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS, output_tokens) if total_tokens > 0: - set_on_span( + span.set_attribute( SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS, total_tokens, ) -def _set_response_model_name(span: "sentry_sdk.tracing.Span", messages: "Any") -> None: +def _set_response_model_name(span: "StreamedSpan", messages: "Any") -> None: if len(messages) == 0: return diff --git a/sentry_sdk/integrations/mcp.py b/sentry_sdk/integrations/mcp.py index 1d0bb6927f..a6fff93198 100644 --- a/sentry_sdk/integrations/mcp.py +++ b/sentry_sdk/integrations/mcp.py @@ -14,11 +14,9 @@ from typing import TYPE_CHECKING import sentry_sdk -from sentry_sdk.ai.utils import _set_span_data_attribute from sentry_sdk.consts import OP, SPANDATA from sentry_sdk.integrations import DidNotEnable, Integration, _check_minimum_version from sentry_sdk.scope import should_send_default_pii -from sentry_sdk.traces import StreamedSpan from sentry_sdk.utils import ( capture_internal_exceptions, event_from_exception, @@ -71,7 +69,6 @@ from starlette.types import Receive, Scope, Send from sentry_sdk.traces import StreamedSpan - from sentry_sdk.tracing import Span class MCPIntegration(Integration): @@ -182,7 +179,7 @@ def _get_request_context_data( def _set_span_input_data( - span: "Union[StreamedSpan, Span]", + span: "StreamedSpan", handler_name: str, span_data_key: str, mcp_method_name: str, @@ -194,28 +191,27 @@ def _set_span_input_data( """Set input span data for MCP handlers.""" # Set handler identifier - _set_span_data_attribute(span, span_data_key, handler_name) - _set_span_data_attribute(span, SPANDATA.MCP_METHOD_NAME, mcp_method_name) + span.set_attribute(span_data_key, handler_name) + span.set_attribute(SPANDATA.MCP_METHOD_NAME, mcp_method_name) # Set transport/MCP transport type - _set_span_data_attribute( - span, + span.set_attribute( SPANDATA.NETWORK_TRANSPORT, "pipe" if mcp_transport == "stdio" else "tcp", ) - _set_span_data_attribute(span, SPANDATA.MCP_TRANSPORT, mcp_transport) + span.set_attribute(SPANDATA.MCP_TRANSPORT, mcp_transport) # Set request_id if provided if request_id: - _set_span_data_attribute(span, SPANDATA.MCP_REQUEST_ID, request_id) + span.set_attribute(SPANDATA.MCP_REQUEST_ID, request_id) # Set session_id if provided if session_id: - _set_span_data_attribute(span, SPANDATA.MCP_SESSION_ID, session_id) + span.set_attribute(SPANDATA.MCP_SESSION_ID, session_id) # Set request arguments (excluding common request context objects) for k, v in arguments.items(): - _set_span_data_attribute(span, f"mcp.request.argument.{k}", safe_serialize(v)) + span.set_attribute(f"mcp.request.argument.{k}", safe_serialize(v)) def _extract_tool_result_content(result: "Any") -> "Any": @@ -369,13 +365,13 @@ async def _tool_handler_wrapper( extracted = _extract_tool_result_content(result) if extracted is not None and should_include_data: - _set_span_data_attribute( - span, SPANDATA.MCP_TOOL_RESULT_CONTENT, safe_serialize(extracted) + span.set_attribute( + SPANDATA.MCP_TOOL_RESULT_CONTENT, safe_serialize(extracted) ) # Set content count if result is a dict if isinstance(extracted, dict): - _set_span_data_attribute( - span, SPANDATA.MCP_TOOL_RESULT_CONTENT_COUNT, len(extracted) + span.set_attribute( + SPANDATA.MCP_TOOL_RESULT_CONTENT_COUNT, len(extracted) ) return result @@ -457,15 +453,13 @@ async def _instrument_v2_tool_call( result_content = _extract_text_from_content_blocks(result["content"]) if result_content is not None and should_include_result_data: - _set_span_data_attribute( - span, + span.set_attribute( SPANDATA.MCP_TOOL_RESULT_CONTENT, safe_serialize(result_content), ) # Set content count if result is a dict if isinstance(result_content, dict): - _set_span_data_attribute( - span, + span.set_attribute( SPANDATA.MCP_TOOL_RESULT_CONTENT_COUNT, len(result_content), ) @@ -586,8 +580,8 @@ async def _prompt_handler_wrapper( # Always set message count if we found messages if message_count > 0: - _set_span_data_attribute( - span, SPANDATA.MCP_PROMPT_RESULT_MESSAGE_COUNT, message_count + span.set_attribute( + SPANDATA.MCP_PROMPT_RESULT_MESSAGE_COUNT, message_count ) # Only set role and content for single-message prompts if PII is allowed @@ -601,9 +595,7 @@ async def _prompt_handler_wrapper( role = first_message["role"] if role: - _set_span_data_attribute( - span, SPANDATA.MCP_PROMPT_RESULT_MESSAGE_ROLE, role - ) + span.set_attribute(SPANDATA.MCP_PROMPT_RESULT_MESSAGE_ROLE, role) # Extract content text content_text = None @@ -624,8 +616,7 @@ async def _prompt_handler_wrapper( content_text = msg_content if content_text: - _set_span_data_attribute( - span, + span.set_attribute( SPANDATA.MCP_PROMPT_RESULT_MESSAGE_CONTENT, content_text, ) @@ -718,8 +709,8 @@ async def _instrument_v2_prompt_get( # Always set message count if we found messages if message_count > 0: - _set_span_data_attribute( - span, SPANDATA.MCP_PROMPT_RESULT_MESSAGE_COUNT, message_count + span.set_attribute( + SPANDATA.MCP_PROMPT_RESULT_MESSAGE_COUNT, message_count ) # Only set role and content for single-message prompts if PII is allowed @@ -731,9 +722,7 @@ async def _instrument_v2_prompt_get( role = first_message["role"] if role: - _set_span_data_attribute( - span, SPANDATA.MCP_PROMPT_RESULT_MESSAGE_ROLE, role - ) + span.set_attribute(SPANDATA.MCP_PROMPT_RESULT_MESSAGE_ROLE, role) content_text = None if "content" in first_message: @@ -742,8 +731,7 @@ async def _instrument_v2_prompt_get( content_text = msg_content["text"] if content_text: - _set_span_data_attribute( - span, + span.set_attribute( SPANDATA.MCP_PROMPT_RESULT_MESSAGE_CONTENT, content_text, ) @@ -821,7 +809,7 @@ async def _resource_handler_wrapper( elif handler_name and "://" in handler_name: protocol = handler_name.split("://")[0] if protocol: - _set_span_data_attribute(span, SPANDATA.MCP_RESOURCE_PROTOCOL, protocol) + span.set_attribute(SPANDATA.MCP_RESOURCE_PROTOCOL, protocol) try: # Execute the async handler @@ -880,7 +868,7 @@ async def _instrument_v2_resource_read( if handler_name and "://" in handler_name: protocol = handler_name.split("://")[0] if protocol: - _set_span_data_attribute(span, SPANDATA.MCP_RESOURCE_PROTOCOL, protocol) + span.set_attribute(SPANDATA.MCP_RESOURCE_PROTOCOL, protocol) try: result = await call_next(ctx) diff --git a/sentry_sdk/integrations/pydantic_ai/spans/utils.py b/sentry_sdk/integrations/pydantic_ai/spans/utils.py index cbce0c0c77..05e34aaa4d 100644 --- a/sentry_sdk/integrations/pydantic_ai/spans/utils.py +++ b/sentry_sdk/integrations/pydantic_ai/spans/utils.py @@ -2,18 +2,18 @@ from typing import TYPE_CHECKING -import sentry_sdk from sentry_sdk._types import BLOB_DATA_SUBSTITUTE from sentry_sdk.ai.consts import DATA_URL_BASE64_REGEX from sentry_sdk.ai.utils import get_modality_from_mime_type from sentry_sdk.consts import SPANDATA -from sentry_sdk.traces import StreamedSpan if TYPE_CHECKING: from typing import Any, Dict, Union from pydantic_ai.usage import RequestUsage, RunUsage + from sentry_sdk.traces import StreamedSpan + def _serialize_image_url_item(item: "Any") -> "Dict[str, Any]": """Serialize an ImageUrl content item for span data. @@ -47,7 +47,7 @@ def _serialize_binary_content_item(item: "Any") -> "Dict[str, Any]": def _set_usage_data( - span: "Union[sentry_sdk.tracing.Span, StreamedSpan]", + span: "StreamedSpan", usage: "Union[RequestUsage, RunUsage]", ) -> None: """Set token usage data on a span. @@ -62,26 +62,24 @@ def _set_usage_data( if usage is None: return - set_on_span = ( - span.set_attribute if isinstance(span, StreamedSpan) else span.set_data - ) - if hasattr(usage, "input_tokens") and usage.input_tokens is not None: - set_on_span(SPANDATA.GEN_AI_USAGE_INPUT_TOKENS, usage.input_tokens) + span.set_attribute(SPANDATA.GEN_AI_USAGE_INPUT_TOKENS, usage.input_tokens) # Pydantic AI uses cache_read_tokens (not input_tokens_cached) if hasattr(usage, "cache_read_tokens") and usage.cache_read_tokens is not None: - set_on_span(SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHED, usage.cache_read_tokens) + span.set_attribute( + SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHED, usage.cache_read_tokens + ) # Pydantic AI uses cache_write_tokens (not input_tokens_cache_write) if hasattr(usage, "cache_write_tokens") and usage.cache_write_tokens is not None: - set_on_span( + span.set_attribute( SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHE_WRITE, usage.cache_write_tokens, ) if hasattr(usage, "output_tokens") and usage.output_tokens is not None: - set_on_span(SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS, usage.output_tokens) + span.set_attribute(SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS, usage.output_tokens) if hasattr(usage, "total_tokens") and usage.total_tokens is not None: - set_on_span(SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS, usage.total_tokens) + span.set_attribute(SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS, usage.total_tokens) diff --git a/sentry_sdk/integrations/redis/modules/caches.py b/sentry_sdk/integrations/redis/modules/caches.py index e7946dd8ce..9f618ef26d 100644 --- a/sentry_sdk/integrations/redis/modules/caches.py +++ b/sentry_sdk/integrations/redis/modules/caches.py @@ -4,7 +4,6 @@ from sentry_sdk.consts import OP, SPANDATA from sentry_sdk.integrations.redis.utils import _get_safe_key, _key_as_string -from sentry_sdk.traces import StreamedSpan from sentry_sdk.utils import capture_internal_exceptions GET_COMMANDS = ("get", "mget") @@ -13,10 +12,10 @@ from typing import TYPE_CHECKING if TYPE_CHECKING: - from typing import Any, Optional, Union + from typing import Any, Optional from sentry_sdk.integrations.redis import RedisIntegration - from sentry_sdk.tracing import Span + from sentry_sdk.traces import StreamedSpan def _get_op(name: str) -> "Optional[str]": @@ -78,30 +77,25 @@ def _get_cache_span_description( def _set_cache_data( - span: "Union[Span, StreamedSpan]", + span: "StreamedSpan", redis_client: "Any", properties: "dict[str, Any]", return_value: "Optional[Any]", ) -> None: - if isinstance(span, StreamedSpan): - set_on_span = span.set_attribute - else: - set_on_span = span.set_data - with capture_internal_exceptions(): - set_on_span(SPANDATA.CACHE_KEY, properties["key"]) + span.set_attribute(SPANDATA.CACHE_KEY, properties["key"]) if properties["redis_command"] in GET_COMMANDS: if return_value is not None: - set_on_span(SPANDATA.CACHE_HIT, True) + span.set_attribute(SPANDATA.CACHE_HIT, True) size = ( len(str(return_value).encode("utf-8")) if not isinstance(return_value, bytes) else len(return_value) ) - set_on_span(SPANDATA.CACHE_ITEM_SIZE, size) + span.set_attribute(SPANDATA.CACHE_ITEM_SIZE, size) else: - set_on_span(SPANDATA.CACHE_HIT, False) + span.set_attribute(SPANDATA.CACHE_HIT, False) elif properties["redis_command"] in SET_COMMANDS: if properties["value"] is not None: @@ -110,7 +104,7 @@ def _set_cache_data( if not isinstance(properties["value"], bytes) else len(properties["value"]) ) - set_on_span(SPANDATA.CACHE_ITEM_SIZE, size) + span.set_attribute(SPANDATA.CACHE_ITEM_SIZE, size) try: connection_params = redis_client.connection_pool.connection_kwargs @@ -125,8 +119,8 @@ def _set_cache_data( host = connection_params.get("host") if host is not None: - set_on_span(SPANDATA.NETWORK_PEER_ADDRESS, host) + span.set_attribute(SPANDATA.NETWORK_PEER_ADDRESS, host) port = connection_params.get("port") if port is not None: - set_on_span(SPANDATA.NETWORK_PEER_PORT, port) + span.set_attribute(SPANDATA.NETWORK_PEER_PORT, port) diff --git a/sentry_sdk/integrations/rust_tracing.py b/sentry_sdk/integrations/rust_tracing.py index 4c06954521..cde4b4716f 100644 --- a/sentry_sdk/integrations/rust_tracing.py +++ b/sentry_sdk/integrations/rust_tracing.py @@ -32,15 +32,16 @@ import json from enum import Enum, auto -from typing import Any, Callable, Dict, Optional, Union +from typing import TYPE_CHECKING, Any, Callable, Dict, Optional import sentry_sdk from sentry_sdk.integrations import Integration from sentry_sdk.scope import should_send_default_pii -from sentry_sdk.traces import StreamedSpan -from sentry_sdk.tracing import Span as SentrySpan from sentry_sdk.utils import SENSITIVE_DATA_SUBSTITUTE +if TYPE_CHECKING: + from sentry_sdk.traces import StreamedSpan + class RustTracingLevel(Enum): Trace = "TRACE" @@ -169,7 +170,7 @@ def _include_tracing_fields(self) -> bool: else self.include_tracing_fields ) - def on_event(self, event: str, sentry_span: "SentrySpan") -> None: + def on_event(self, event: str, sentry_span: "StreamedSpan") -> None: deserialized_event = json.loads(event) metadata = deserialized_event.get("metadata", {}) @@ -183,9 +184,7 @@ def on_event(self, event: str, sentry_span: "SentrySpan") -> None: elif event_type == EventTypeMapping.Event: process_event(deserialized_event) - def on_new_span( - self, attrs: str, span_id: str - ) -> "Optional[Union[SentrySpan, StreamedSpan]]": + def on_new_span(self, attrs: str, span_id: str) -> "Optional[StreamedSpan]": attrs = json.loads(attrs) metadata = attrs.get("metadata", {}) @@ -224,30 +223,22 @@ def on_new_span( return sentry_span - def on_close(self, span_id: str, sentry_span: "SentrySpan") -> None: + def on_close(self, span_id: str, sentry_span: "StreamedSpan") -> None: if sentry_span is None: return sentry_span.__exit__(None, None, None) - def on_record( - self, span_id: str, values: str, sentry_span: "Union[SentrySpan, StreamedSpan]" - ) -> None: + def on_record(self, span_id: str, values: str, sentry_span: "StreamedSpan") -> None: if sentry_span is None: return - set_on_span = ( - sentry_span.set_attribute - if isinstance(sentry_span, StreamedSpan) - else sentry_span.set_data - ) - deserialized_values = json.loads(values) for key, value in deserialized_values.items(): if self._include_tracing_fields(): - set_on_span(key, value) + sentry_span.set_attribute(key, value) else: - set_on_span(key, SENSITIVE_DATA_SUBSTITUTE) + sentry_span.set_attribute(key, SENSITIVE_DATA_SUBSTITUTE) class RustTracingIntegration(Integration):