.

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

sentry_sdk/integrations/django/__init__.py

@@ -9,6 +9,7 @@
 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,
@@ -720,14 +721,21 @@ def _rollback(self: "BaseDatabaseWrapper") -> None:
 
 
 def _set_db_data(
-    span: "Span", cursor_or_db: "Any", db_operation: "Optional[str]" = None
+    span: "Union[Span, StreamedSpan]",
+    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
-    span.set_data(SPANDATA.DB_SYSTEM, vendor)
+    set_on_span(SPANDATA.DB_SYSTEM, vendor)
 
     if db_operation is not None:
-        span.set_data(SPANDATA.DB_OPERATION, db_operation)
+        set_on_span(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`
@@ -760,19 +768,19 @@ def _set_db_data(
 
     db_name = connection_params.get("dbname") or connection_params.get("database")
     if db_name is not None:
-        span.set_data(SPANDATA.DB_NAME, db_name)
+        set_on_span(SPANDATA.DB_NAME, db_name)
 
     server_address = connection_params.get("host")
     if server_address is not None:
-        span.set_data(SPANDATA.SERVER_ADDRESS, server_address)
+        set_on_span(SPANDATA.SERVER_ADDRESS, server_address)
 
     server_port = connection_params.get("port")
     if server_port is not None:
-        span.set_data(SPANDATA.SERVER_PORT, str(server_port))
+        set_on_span(SPANDATA.SERVER_PORT, str(server_port))
 
     server_socket_address = connection_params.get("unix_socket")
     if server_socket_address is not None:
-        span.set_data(SPANDATA.SERVER_SOCKET_ADDRESS, server_socket_address)
+        set_on_span(SPANDATA.SERVER_SOCKET_ADDRESS, server_socket_address)
 
 
 def add_template_context_repr_sequence() -> None:

sentry_sdk/integrations/sqlalchemy.py

@@ -138,52 +138,33 @@ def _get_db_system(name: str) -> "Optional[str]":
 
 
 def _set_db_data(span: "Union[Span, StreamedSpan]", conn: "Any") -> None:
-    db_system = _get_db_system(conn.engine.name)
     if isinstance(span, StreamedSpan):
-        if db_system is not None:
-            span.set_attribute(SPANDATA.DB_SYSTEM, db_system)
-
-        try:
-            driver = conn.dialect.driver
-            if driver:
-                span.set_attribute(SPANDATA.DB_DRIVER_NAME, driver)
-        except Exception:
-            pass
+        set_on_span = span.set_attribute
     else:
-        if db_system is not None:
-            span.set_data(SPANDATA.DB_SYSTEM, db_system)
+        set_on_span = span.set_data
+
+    db_system = _get_db_system(conn.engine.name)
+    if db_system is not None:
+        set_on_span(SPANDATA.DB_SYSTEM, db_system)
 
-        try:
-            driver = conn.dialect.driver
-            if driver:
-                span.set_data(SPANDATA.DB_DRIVER_NAME, driver)
-        except Exception:
-            pass
+    try:
+        driver = conn.dialect.driver
+        if driver:
+            set_on_span(SPANDATA.DB_DRIVER_NAME, driver)
+    except Exception:
+        pass
 
     if conn.engine.url is None:
         return
 
-    if isinstance(span, StreamedSpan):
-        db_name = conn.engine.url.database
-        if db_name is not None:
-            span.set_attribute(SPANDATA.DB_NAME, db_name)
-
-        server_address = conn.engine.url.host
-        if server_address is not None:
-            span.set_attribute(SPANDATA.SERVER_ADDRESS, server_address)
-
-        server_port = conn.engine.url.port
-        if server_port is not None:
-            span.set_attribute(SPANDATA.SERVER_PORT, server_port)
-    else:
-        db_name = conn.engine.url.database
-        if db_name is not None:
-            span.set_data(SPANDATA.DB_NAME, db_name)
+    db_name = conn.engine.url.database
+    if db_name is not None:
+        set_on_span(SPANDATA.DB_NAME, db_name)
 
-        server_address = conn.engine.url.host
-        if server_address is not None:
-            span.set_data(SPANDATA.SERVER_ADDRESS, server_address)
+    server_address = conn.engine.url.host
+    if server_address is not None:
+        set_on_span(SPANDATA.SERVER_ADDRESS, server_address)
 
-        server_port = conn.engine.url.port
-        if server_port is not None:
-            span.set_data(SPANDATA.SERVER_PORT, server_port)
+    server_port = conn.engine.url.port
+    if server_port is not None:
+        set_on_span(SPANDATA.SERVER_PORT, server_port)