Fix asyncpg spans (#3639)

- Change `span.set_data` to `span.set_attribute`
- Change `start_transaction`s in the tests to `start_span`s
- Stringify `db.params`, `db.paramstyle`, `db.cursor`
- ❗ Change how adding query source works.
  Previously, we first let the DB span finish and then looked at its start and end timestamps to figure out if the query was slow enough to attach query data.
  
  With the switch to OTel, we can no longer put stuff on a finished span, meaning the slow query span still has to be alive when adding data to it. So instead of looking at the end timestamp of a finished span, we keep the span alive and instead look at current time at the point when the query is finished.
  
  This means the DB span will be longer than the actual query (since it now includes adding query data).

Author: sentrivana

sentry_sdk/integrations/asyncpg.py

@@ -40,7 +40,6 @@ def setup_once() -> None:
         asyncpg.Connection.execute = _wrap_execute(
             asyncpg.Connection.execute,
         )
-
         asyncpg.Connection._execute = _wrap_connection_method(
             asyncpg.Connection._execute
         )
@@ -80,8 +79,8 @@ async def _inner(*args: Any, **kwargs: Any) -> T:
         ) as span:
             res = await f(*args, **kwargs)
 
-        with capture_internal_exceptions():
-            add_query_source(span)
+            with capture_internal_exceptions():
+                add_query_source(span)
 
         return res
 
@@ -148,7 +147,7 @@ def _inner(*args: Any, **kwargs: Any) -> T:  # noqa: N807
         ) as span:
             _set_db_data(span, args[0])
             res = f(*args, **kwargs)
-            span.set_data("db.cursor", res)
+            span.set_attribute("db.cursor", str(res))
 
         return res
 
@@ -168,21 +167,37 @@ async def _inner(*args: Any, **kwargs: Any) -> T:
             name="connect",
             origin=AsyncPGIntegration.origin,
         ) as span:
-            span.set_data(SPANDATA.DB_SYSTEM, "postgresql")
+            span.set_attribute(SPANDATA.DB_SYSTEM, "postgresql")
             addr = kwargs.get("addr")
             if addr:
                 try:
-                    span.set_data(SPANDATA.SERVER_ADDRESS, addr[0])
-                    span.set_data(SPANDATA.SERVER_PORT, addr[1])
+                    span.set_attribute(SPANDATA.SERVER_ADDRESS, addr[0])
+                    span.set_attribute(SPANDATA.SERVER_PORT, addr[1])
                 except IndexError:
                     pass
-            span.set_data(SPANDATA.DB_NAME, database)
-            span.set_data(SPANDATA.DB_USER, user)
+
+            span.set_attribute(SPANDATA.DB_NAME, database)
+            span.set_attribute(SPANDATA.DB_USER, user)
 
             with capture_internal_exceptions():
+                data = {}
+                for attr in (
+                    "db.cursor",
+                    "db.params",
+                    "db.paramstyle",
+                    SPANDATA.DB_NAME,
+                    SPANDATA.DB_SYSTEM,
+                    SPANDATA.DB_USER,
+                    SPANDATA.SERVER_ADDRESS,
+                    SPANDATA.SERVER_PORT,
+                ):
+                    if span.get_attribute(attr):
+                        data[attr] = span.get_attribute(attr)
+
                 sentry_sdk.add_breadcrumb(
-                    message="connect", category="query", data=span._data
+                    message="connect", category="query", data=data
                 )
+
             res = await f(*args, **kwargs)
 
         return res
@@ -191,20 +206,20 @@ async def _inner(*args: Any, **kwargs: Any) -> T:
 
 
 def _set_db_data(span: Span, conn: Any) -> None:
-    span.set_data(SPANDATA.DB_SYSTEM, "postgresql")
+    span.set_attribute(SPANDATA.DB_SYSTEM, "postgresql")
 
     addr = conn._addr
     if addr:
         try:
-            span.set_data(SPANDATA.SERVER_ADDRESS, addr[0])
-            span.set_data(SPANDATA.SERVER_PORT, addr[1])
+            span.set_attribute(SPANDATA.SERVER_ADDRESS, addr[0])
+            span.set_attribute(SPANDATA.SERVER_PORT, addr[1])
         except IndexError:
             pass
 
     database = conn._params.database
     if database:
-        span.set_data(SPANDATA.DB_NAME, database)
+        span.set_attribute(SPANDATA.DB_NAME, database)
 
     user = conn._params.user
     if user:
-        span.set_data(SPANDATA.DB_USER, user)
+        span.set_attribute(SPANDATA.DB_USER, user)

sentry_sdk/tracing_utils.py

@@ -4,7 +4,7 @@
 import re
 import sys
 from collections.abc import Mapping
-from datetime import timedelta
+from datetime import datetime, timedelta, timezone
 from functools import wraps
 from urllib.parse import quote, unquote
 import uuid
@@ -133,13 +133,13 @@ def record_sql_queries(
 
     data = {}
     if params_list is not None:
-        data["db.params"] = params_list
+        data["db.params"] = str(params_list)
     if paramstyle is not None:
-        data["db.paramstyle"] = paramstyle
+        data["db.paramstyle"] = str(paramstyle)
     if executemany:
         data["db.executemany"] = True
     if record_cursor_repr and cursor is not None:
-        data["db.cursor"] = cursor
+        data["db.cursor"] = str(cursor)
 
     with capture_internal_exceptions():
         sentry_sdk.add_breadcrumb(message=query, category="query", data=data)
@@ -209,14 +209,17 @@ def add_query_source(span):
     if not client.is_active():
         return
 
-    if span.timestamp is None or span.start_timestamp is None:
+    if span.start_timestamp is None:
         return
 
     should_add_query_source = client.options.get("enable_db_query_source", True)
     if not should_add_query_source:
         return
 
-    duration = span.timestamp - span.start_timestamp
+    # We assume here that the query is just ending now. We can't use
+    # the actual end timestamp of the span because in OTel the span
+    # can't be finished in order to set any attributes on it.
+    duration = datetime.now(tz=timezone.utc) - span.start_timestamp
     threshold = client.options.get("db_query_source_threshold_ms", 0)
     slow_query = duration / timedelta(milliseconds=1) > threshold