diff --git a/scripts/populate_tox/config.py b/scripts/populate_tox/config.py index 71b6cfa14a..b2db1058bb 100644 --- a/scripts/populate_tox/config.py +++ b/scripts/populate_tox/config.py @@ -509,6 +509,11 @@ }, "sqlalchemy": { "package": "sqlalchemy", + "deps": { + # Needed to test the asyncio extension. SQLAlchemy 1.4 installs + # greenlet only on some platforms, 2.x only with an extra. + ">=1.4": ["greenlet"], + }, "python": { "<1.4": "<3.10", }, diff --git a/sentry_sdk/tracing_utils.py b/sentry_sdk/tracing_utils.py index 4d4c35ce39..3a1b665be4 100644 --- a/sentry_sdk/tracing_utils.py +++ b/sentry_sdk/tracing_utils.py @@ -292,6 +292,50 @@ def _should_be_included( ) +def _is_sqlalchemy_async_bridge(current: "Any") -> bool: + # SQLAlchemy marks its bridge greenlet with this attribute and checks for + # it itself, from 1.4.53 and 2.0.31. Earlier releases check the class. + if getattr(current, "__sqlalchemy_greenlet_provider__", False): + return True + cls = type(current) + return cls.__name__ == "_AsyncIoGreenlet" and cls.__module__.startswith( + "sqlalchemy." + ) + + +def _caller_frames(frame: "Optional[FrameType]") -> "Iterator[FrameType]": + """ + Yields ``frame`` and the frames that called it. + + Async SQLAlchemy runs the sync driver, and so the cursor events that + create query spans, in a greenlet it spawns from the coroutine awaiting + the query. That greenlet's own frames end at the bridge, so when the + current greenlet is SQLAlchemy's bridge, the walk continues into the + frames of its parent, which is suspended in the awaiting caller. Any other + greenlet's parent need not be its caller, so its frames are not followed. + """ + while frame is not None: + yield frame + frame = frame.f_back + + greenlet = sys.modules.get("greenlet") + if greenlet is None: + return + + try: + current = greenlet.getcurrent() + if not _is_sqlalchemy_async_bridge(current): + return + parent = current.parent + frame = parent.gr_frame if parent is not None else None + except Exception: + return + + while frame is not None: + yield frame + frame = frame.f_back + + def add_source( span: "Union[sentry_sdk.tracing.Span, sentry_sdk.traces.StreamedSpan]", project_root: "Optional[str]", @@ -302,8 +346,8 @@ def add_source( Adds OTel compatible source code information to the span """ # Find the correct frame - frame: "Union[FrameType, None]" = sys._getframe() - while frame is not None: + frame: "Union[FrameType, None]" = None + for frame in _caller_frames(sys._getframe()): abs_path = _get_frame_module_abs_path(frame) try: @@ -325,8 +369,6 @@ def add_source( ) if should_be_included: break - - frame = frame.f_back else: frame = None diff --git a/tests/integrations/sqlalchemy/test_sqlalchemy.py b/tests/integrations/sqlalchemy/test_sqlalchemy.py index e08bb0a423..66edb10ec4 100644 --- a/tests/integrations/sqlalchemy/test_sqlalchemy.py +++ b/tests/integrations/sqlalchemy/test_sqlalchemy.py @@ -1,8 +1,10 @@ +import asyncio import os from datetime import datetime from unittest import mock import pytest +import sqlalchemy.util from sqlalchemy import Column, ForeignKey, Integer, String, create_engine, text from sqlalchemy.exc import IntegrityError from sqlalchemy.ext.declarative import declarative_base @@ -760,6 +762,66 @@ class Person(Base): raise AssertionError("No db span found") +@pytest.mark.parametrize("span_streaming", [True, False]) +def test_query_source_through_async_bridge( + sentry_init, + capture_events, + capture_items, + span_streaming, +): + # Not skipped when greenlet is missing: from 1.4 on, the test environments + # install it, and greenlet_spawn fails loudly without it. + greenlet_spawn = getattr(sqlalchemy.util, "greenlet_spawn", None) + if greenlet_spawn is None: + pytest.skip("SQLAlchemy has no asyncio support before 1.4") + + sentry_init( + integrations=[SqlalchemyIntegration()], + traces_sample_rate=1.0, + enable_db_query_source=True, + db_query_source_threshold_ms=0, + trace_lifecycle="stream" if span_streaming else "static", + ) + engine = create_engine("sqlite:///:memory:") + + async def query_from_coroutine(): + with engine.connect() as connection: + # As SQLAlchemy's asyncio extension does, run the sync call in + # SQLAlchemy's bridge greenlet, whose own frames hold no + # application code. + await greenlet_spawn(connection.execute, text("SELECT 1")) + + def run_query(): + loop = asyncio.new_event_loop() + try: + loop.run_until_complete(query_from_coroutine()) + finally: + loop.close() + + if span_streaming: + items = capture_items("span") + with sentry_sdk.traces.start_span(name="custom parent"): + run_query() + sentry_sdk.flush() + sources = [ + item.payload["attributes"].get(SPANDATA.CODE_FUNCTION) + for item in items + if item.payload["name"] == "SELECT 1" + ] + else: + events = capture_events() + with start_transaction(name="test_transaction", sampled=True): + run_query() + (event,) = events + sources = [ + span["data"].get(SPANDATA.CODE_FUNCTION) + for span in event["spans"] + if span.get("description") == "SELECT 1" + ] + + assert sources == ["query_from_coroutine"] + + @pytest.mark.parametrize("span_streaming", [True, False]) def test_query_source_with_module_in_search_path( sentry_init, diff --git a/tox.ini b/tox.ini index c8af226e19..ec72269d10 100644 --- a/tox.ini +++ b/tox.ini @@ -12872,26 +12872,43 @@ deps = sqlalchemy-v1.4.54: sqlalchemy==1.4.54 + py3.11-sqlalchemy-v1.4.54: greenlet==3.5.6 + + py3.12-sqlalchemy-v1.4.54: greenlet==3.5.6 + sqlalchemy-v2.1.6: sqlalchemy==2.1.6 + py3.11-sqlalchemy-v2.1.6: greenlet==3.5.6 py3.11-sqlalchemy-v2.1.6: typing_extensions==4.16.0 + py3.14-sqlalchemy-v2.1.6: greenlet==3.5.6 py3.14-sqlalchemy-v2.1.6: typing_extensions==4.16.0 + py3.15-sqlalchemy-v2.1.6: greenlet==3.5.6 py3.15-sqlalchemy-v2.1.6: typing_extensions==4.16.0 + py3.15t-sqlalchemy-v2.1.6: greenlet==3.5.6 py3.15t-sqlalchemy-v2.1.6: typing_extensions==4.16.0 sqlalchemy-latest: sqlalchemy==2.1.6 + sqlalchemy-latest: greenlet + + py3.11-sqlalchemy-latest: greenlet==3.5.6 py3.11-sqlalchemy-latest: typing_extensions==4.16.0 + py3.14-sqlalchemy-latest: greenlet==3.5.6 py3.14-sqlalchemy-latest: typing_extensions==4.16.0 + py3.15-sqlalchemy-latest: greenlet==3.5.6 py3.15-sqlalchemy-latest: typing_extensions==4.16.0 + py3.15t-sqlalchemy-latest: greenlet==3.5.6 py3.15t-sqlalchemy-latest: typing_extensions==4.16.0 + sqlalchemy-v1.4.54: greenlet + sqlalchemy-v2.1.6: greenlet + sqlalchemy-latest: greenlet # ~~~ Flags ~~~