fix(aiomysql): Handle add_query_source for StreamedSpan and fix test mocks

- Call add_query_source inside context for StreamedSpan, outside for
  regular Span (following asyncpg pattern)
- Fix tests to mock record_sql_queries_supporting_streaming instead of
  record_sql_queries

Addresses cursor bot review comments.

Co-Authored-By: Claude Opus 4.7 <[email protected]>

Author: tonal

sentry_sdk/integrations/aiomysql.py

@@ -96,9 +96,13 @@ async def _inner(*args: Any, **kwargs: Any) -> T:
             if conn:
                 _set_db_data(span, conn)
             res = await f(*args, **kwargs)
+            if isinstance(span, StreamedSpan):
+                with capture_internal_exceptions():
+                    add_query_source(span)
 
-        with capture_internal_exceptions():
-            add_query_source(span)
+        if not isinstance(span, StreamedSpan):
+            with capture_internal_exceptions():
+                add_query_source(span)
 
         return res
 
@@ -139,9 +143,13 @@ async def _inner(*args: Any, **kwargs: Any) -> T:
                 if conn:
                     _set_db_data(span, conn)
                 res = await f(*args, **kwargs)
+                if isinstance(span, StreamedSpan):
+                    with capture_internal_exceptions():
+                        add_query_source(span)
 
-            with capture_internal_exceptions():
-                add_query_source(span)
+            if not isinstance(span, StreamedSpan):
+                with capture_internal_exceptions():
+                    add_query_source(span)
 
             return res
         finally:

tests/integrations/aiomysql/test_aiomysql.py

@@ -22,7 +22,7 @@
 from sentry_sdk import capture_message, start_transaction
 from sentry_sdk.consts import SPANDATA
 from sentry_sdk.integrations.aiomysql import AioMySQLIntegration
-from sentry_sdk.tracing_utils import record_sql_queries
+from sentry_sdk.tracing_utils import record_sql_queries_supporting_streaming
 from tests.conftest import ApproxDict
 
 MYSQL_HOST = os.getenv("SENTRY_PYTHON_TEST_MYSQL_HOST", "localhost")
@@ -569,15 +569,15 @@ async def test_no_query_source_if_duration_too_short(sentry_init, capture_events
 
         @contextmanager
         def fake_record_sql_queries(*args, **kwargs):
-            with record_sql_queries(*args, **kwargs) as span:
+            with record_sql_queries_supporting_streaming(*args, **kwargs) as span:
                 pass
             span.start_timestamp = datetime.datetime(2024, 1, 1, microsecond=0)
             span.timestamp = datetime.datetime(2024, 1, 1, microsecond=99999)
             yield span
 
         async with conn.cursor() as cur:
             with mock.patch(
-                "sentry_sdk.integrations.aiomysql.record_sql_queries",
+                "sentry_sdk.integrations.aiomysql.record_sql_queries_supporting_streaming",
                 fake_record_sql_queries,
             ):
                 await cur.execute(
@@ -615,15 +615,15 @@ async def test_query_source_if_duration_over_threshold(sentry_init, capture_even
 
         @contextmanager
         def fake_record_sql_queries(*args, **kwargs):
-            with record_sql_queries(*args, **kwargs) as span:
+            with record_sql_queries_supporting_streaming(*args, **kwargs) as span:
                 pass
             span.start_timestamp = datetime.datetime(2024, 1, 1, microsecond=0)
             span.timestamp = datetime.datetime(2024, 1, 1, microsecond=100001)
             yield span
 
         async with conn.cursor() as cur:
             with mock.patch(
-                "sentry_sdk.integrations.aiomysql.record_sql_queries",
+                "sentry_sdk.integrations.aiomysql.record_sql_queries_supporting_streaming",
                 fake_record_sql_queries,
             ):
                 await cur.execute(