add query source needs a live span

Author: sentrivana

sentry_sdk/integrations/asyncpg.py

@@ -79,8 +79,9 @@ 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():
+                print("ADDING QUERY SOURCE")
+                add_query_source(span)
 
         return res
 

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
@@ -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 span is just ending now. We can't use
+    # the actual end timestamp of the span because 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
 

tests/integrations/asyncpg/test_asyncpg.py

@@ -264,7 +264,7 @@ async def test_cursor(sentry_init, capture_events) -> None:
         async for record in conn.cursor(
             "SELECT * FROM users WHERE dob > $1", datetime.date(1970, 1, 1)
         ):
-            print(record)
+            pass
 
     await conn.close()