Fix potel tests in potel

sentry_sdk/integrations/sqlalchemy.py

@@ -76,15 +76,15 @@ def _after_cursor_execute(conn, cursor, statement, parameters, context, *args):
         context, "_sentry_sql_span_manager", None
     )  # type: Optional[ContextManager[Any]]
 
-    if ctx_mgr is not None:
-        context._sentry_sql_span_manager = None
-        ctx_mgr.__exit__(None, None, None)
-
     span = getattr(context, "_sentry_sql_span", None)  # type: Optional[Span]
     if span is not None:
         with capture_internal_exceptions():
             add_query_source(span)
 
+    if ctx_mgr is not None:
+        context._sentry_sql_span_manager = None
+        ctx_mgr.__exit__(None, None, None)
+
 
 def _handle_error(context, *args):
     # type: (Any, *Any) -> None

sentry_sdk/tracing_utils.py

@@ -147,6 +147,7 @@ def record_sql_queries(
         op=OP.DB,
         name=query,
         origin=span_origin,
+        only_if_parent=True,
     ) as span:
         for k, v in data.items():
             span.set_data(k, v)

tests/integrations/sqlalchemy/test_sqlalchemy.py

@@ -1,3 +1,4 @@
+import contextlib
 import os
 from datetime import datetime
 from unittest import mock
@@ -11,7 +12,6 @@
 from sqlalchemy import text
 
 import sentry_sdk
-from sentry_sdk import capture_message, start_transaction
 from sentry_sdk.consts import DEFAULT_MAX_VALUE_LENGTH, SPANDATA
 from sentry_sdk.integrations.sqlalchemy import SqlalchemyIntegration
 from sentry_sdk.serializer import MAX_EVENT_BYTES
@@ -54,7 +54,7 @@ class Address(Base):
 
     assert session.query(Person).first() == bob
 
-    capture_message("hi")
+    sentry_sdk.capture_message("hi")
 
     (event,) = events
 
@@ -111,7 +111,7 @@ class Address(Base):
     Session = sessionmaker(bind=engine)  # noqa: N806
     session = Session()
 
-    with start_transaction(name="test_transaction", sampled=True):
+    with sentry_sdk.start_span(name="test_transaction", sampled=True):
         with session.begin_nested():
             session.query(Person).first()
 
@@ -135,7 +135,7 @@ class Address(Base):
     assert (
         render_span_tree(event)
         == """\
-- op=null: description=null
+- op="test_transaction": description=null
   - op="db": description="SAVEPOINT sa_savepoint_1"
   - op="db": description="SELECT person.id AS person_id, person.name AS person_name \\nFROM person\\n LIMIT ? OFFSET ?"
   - op="db": description="RELEASE SAVEPOINT sa_savepoint_1"
@@ -185,7 +185,7 @@ class Address(Base):
     Session = sessionmaker(bind=engine)  # noqa: N806
     session = Session()
 
-    with start_transaction(name="test_transaction", sampled=True):
+    with sentry_sdk.start_span(name="test_transaction", sampled=True):
         with session.begin_nested():
             session.query(Person).first()
 
@@ -217,7 +217,7 @@ def test_long_sql_query_preserved(sentry_init, capture_events):
     engine = create_engine(
         "sqlite:///:memory:", connect_args={"check_same_thread": False}
     )
-    with start_transaction(name="test"):
+    with sentry_sdk.start_span(name="test"):
         with engine.connect() as con:
             con.execute(text(" UNION ".join("SELECT {}".format(i) for i in range(100))))
 
@@ -246,7 +246,7 @@ def processor(event, hint):
     engine = create_engine(
         "sqlite:///:memory:", connect_args={"check_same_thread": False}
     )
-    with start_transaction(name="test"):
+    with sentry_sdk.start_span(name="test"):
         with engine.connect() as con:
             for _ in range(1500):
                 con.execute(
@@ -306,7 +306,7 @@ def test_query_source_disabled(sentry_init, capture_events):
 
     events = capture_events()
 
-    with start_transaction(name="test_transaction", sampled=True):
+    with sentry_sdk.start_span(name="test_transaction", sampled=True):
         Base = declarative_base()  # noqa: N806
 
         class Person(Base):
@@ -358,7 +358,7 @@ def test_query_source_enabled(sentry_init, capture_events, enable_db_query_sourc
 
     events = capture_events()
 
-    with start_transaction(name="test_transaction", sampled=True):
+    with sentry_sdk.start_span(name="test_transaction", sampled=True):
         Base = declarative_base()  # noqa: N806
 
         class Person(Base):
@@ -405,7 +405,7 @@ def test_query_source(sentry_init, capture_events):
     )
     events = capture_events()
 
-    with start_transaction(name="test_transaction", sampled=True):
+    with sentry_sdk.start_span(name="test_transaction", sampled=True):
         Base = declarative_base()  # noqa: N806
 
         class Person(Base):
@@ -475,7 +475,7 @@ def test_query_source_with_module_in_search_path(sentry_init, capture_events):
         query_first_model_from_session,
     )
 
-    with start_transaction(name="test_transaction", sampled=True):
+    with sentry_sdk.start_span(name="test_transaction", sampled=True):
         Base = declarative_base()  # noqa: N806
 
         class Person(Base):
@@ -533,7 +533,7 @@ def test_no_query_source_if_duration_too_short(sentry_init, capture_events):
     )
     events = capture_events()
 
-    with start_transaction(name="test_transaction", sampled=True):
+    with sentry_sdk.start_span(name="test_transaction", sampled=True):
         Base = declarative_base()  # noqa: N806
 
         class Person(Base):
@@ -601,7 +601,7 @@ def test_query_source_if_duration_over_threshold(sentry_init, capture_events):
     )
     events = capture_events()
 
-    with start_transaction(name="test_transaction", sampled=True):
+    with sentry_sdk.start_span(name="test_transaction", sampled=True):
         Base = declarative_base()  # noqa: N806
 
         class Person(Base):
@@ -620,21 +620,15 @@ class Person(Base):
         bob = Person(name="Bob")
         session.add(bob)
 
-        class fake_record_sql_queries:  # noqa: N801
-            def __init__(self, *args, **kwargs):
-                with freeze_time(datetime(2024, 1, 1, microsecond=0)):
-                    with record_sql_queries(*args, **kwargs) as span:
-                        self.span = span
-                        freezer = freeze_time(datetime(2024, 1, 1, microsecond=99999))
-                        freezer.start()
+        @contextlib.contextmanager
+        def fake_record_sql_queries(*args, **kwargs):  # noqa: N801
+            with freeze_time(datetime(2024, 1, 1, second=0)):
+                with record_sql_queries(*args, **kwargs) as span:
+                    freezer = freeze_time(datetime(2024, 1, 1, second=1))
+                    freezer.start()
+                    yield span
 
-                    freezer.stop()
-
-            def __enter__(self):
-                return self.span
-
-            def __exit__(self, type, value, traceback):
-                pass
+                freezer.stop()
 
         with mock.patch(
             "sentry_sdk.integrations.sqlalchemy.record_sql_queries",
@@ -687,7 +681,7 @@ def test_span_origin(sentry_init, capture_events):
     engine = create_engine(
         "sqlite:///:memory:", connect_args={"check_same_thread": False}
     )
-    with start_transaction(name="foo"):
+    with sentry_sdk.start_span(name="foo"):
         with engine.connect() as con:
             con.execute(text("SELECT 0"))