.

Author: alexander-alderman-webb

sentry_sdk/integrations/asyncpg.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 import contextlib
 import re
-from typing import Any, TypeVar, Callable, Awaitable, Iterator, Union
+from typing import Any, TypeVar, Callable, Awaitable, Iterator
 
 import sentry_sdk
 from sentry_sdk.consts import OP, SPANDATA
@@ -13,7 +13,6 @@
     parse_version,
     capture_internal_exceptions,
 )
-from sentry_sdk.traces import StreamedSpan
 
 try:
     import asyncpg  # type: ignore[import-not-found]
@@ -102,7 +101,7 @@ def _record(
     params_list: "tuple[Any, ...] | None",
     *,
     executemany: bool = False,
-) -> "Iterator[Union[Span, StreamedSpan]]":
+) -> "Iterator[Span]":
     integration = sentry_sdk.get_client().get_integration(AsyncPGIntegration)
     if integration is not None and not integration._record_params:
         params_list = None
@@ -198,27 +197,22 @@ async def _inner(*args: "Any", **kwargs: "Any") -> "T":
     return _inner
 
 
-def _set_db_data(span: "Union[Span, StreamedSpan]", conn: "Any") -> None:
-    if isinstance(span, StreamedSpan):
-        set_on_span = span.set_attribute
-    else:
-        set_on_span = span.set_data
-
-    set_on_span(SPANDATA.DB_SYSTEM, "postgresql")
-    set_on_span(SPANDATA.DB_DRIVER_NAME, "asyncpg")
+def _set_db_data(span: "Span", conn: "Any") -> None:
+    span.set_data(SPANDATA.DB_SYSTEM, "postgresql")
+    span.set_data(SPANDATA.DB_DRIVER_NAME, "asyncpg")
 
     addr = conn._addr
     if addr:
         try:
-            set_on_span(SPANDATA.SERVER_ADDRESS, addr[0])
-            set_on_span(SPANDATA.SERVER_PORT, addr[1])
+            span.set_data(SPANDATA.SERVER_ADDRESS, addr[0])
+            span.set_data(SPANDATA.SERVER_PORT, addr[1])
         except IndexError:
             pass
 
     database = conn._params.database
     if database:
-        set_on_span(SPANDATA.DB_NAME, database)
+        span.set_data(SPANDATA.DB_NAME, database)
 
     user = conn._params.user
     if user:
-        set_on_span(SPANDATA.DB_USER, user)
+        span.set_data(SPANDATA.DB_USER, user)

sentry_sdk/integrations/django/__init__.py

@@ -9,7 +9,6 @@
 from sentry_sdk.scope import add_global_event_processor, should_send_default_pii
 from sentry_sdk.serializer import add_global_repr_processor, add_repr_sequence_type
 from sentry_sdk.tracing import SOURCE_FOR_STYLE, TransactionSource
-from sentry_sdk.traces import StreamedSpan
 from sentry_sdk.tracing_utils import add_query_source, record_sql_queries
 from sentry_sdk.utils import (
     AnnotatedValue,
@@ -721,21 +720,14 @@ def _rollback(self: "BaseDatabaseWrapper") -> None:
 
 
 def _set_db_data(
-    span: "Union[Span, StreamedSpan]",
-    cursor_or_db: "Any",
-    db_operation: "Optional[str]" = None,
+    span: "Span", cursor_or_db: "Any", db_operation: "Optional[str]" = None
 ) -> None:
-    if isinstance(span, StreamedSpan):
-        set_on_span = span.set_attribute
-    else:
-        set_on_span = span.set_data
-
     db = cursor_or_db.db if hasattr(cursor_or_db, "db") else cursor_or_db
     vendor = db.vendor
-    set_on_span(SPANDATA.DB_SYSTEM, vendor)
+    span.set_data(SPANDATA.DB_SYSTEM, vendor)
 
     if db_operation is not None:
-        set_on_span(SPANDATA.DB_OPERATION, db_operation)
+        span.set_data(SPANDATA.DB_OPERATION, db_operation)
 
     # Some custom backends override `__getattr__`, making it look like `cursor_or_db`
     # actually has a `connection` and the `connection` has a `get_dsn_parameters`
@@ -768,19 +760,19 @@ def _set_db_data(
 
     db_name = connection_params.get("dbname") or connection_params.get("database")
     if db_name is not None:
-        set_on_span(SPANDATA.DB_NAME, db_name)
+        span.set_data(SPANDATA.DB_NAME, db_name)
 
     server_address = connection_params.get("host")
     if server_address is not None:
-        set_on_span(SPANDATA.SERVER_ADDRESS, server_address)
+        span.set_data(SPANDATA.SERVER_ADDRESS, server_address)
 
     server_port = connection_params.get("port")
     if server_port is not None:
-        set_on_span(SPANDATA.SERVER_PORT, str(server_port))
+        span.set_data(SPANDATA.SERVER_PORT, str(server_port))
 
     server_socket_address = connection_params.get("unix_socket")
     if server_socket_address is not None:
-        set_on_span(SPANDATA.SERVER_SOCKET_ADDRESS, server_socket_address)
+        span.set_data(SPANDATA.SERVER_SOCKET_ADDRESS, server_socket_address)
 
 
 def add_template_context_repr_sequence() -> None: