Make sure to add data before span is closed. some cleanup

Author: antonpirker

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

@@ -148,6 +148,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

@@ -11,7 +11,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 +53,7 @@ class Address(Base):
 
     assert session.query(Person).first() == bob
 
-    capture_message("hi")
+    sentry_sdk.capture_message("hi")
 
     (event,) = events
 
@@ -111,7 +110,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 +134,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 +184,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 +216,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 +245,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 +305,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 +357,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 +404,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 +474,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 +532,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 +600,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):
@@ -687,7 +686,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"))