Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions scripts/populate_tox/config.py
Original file line number Diff line number Diff line change
Expand Up @@ -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",
},
Expand Down
50 changes: 46 additions & 4 deletions sentry_sdk/tracing_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Comment thread
cursor[bot] marked this conversation as resolved.
Resolved
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]",
Expand All @@ -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:
Expand All @@ -325,8 +369,6 @@ def add_source(
)
if should_be_included:
break

frame = frame.f_back
else:
frame = None

Expand Down
62 changes: 62 additions & 0 deletions tests/integrations/sqlalchemy/test_sqlalchemy.py
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -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,
Expand Down
17 changes: 17 additions & 0 deletions tox.ini

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.