fix(sqlalchemy): Fix broken nesting under begin_nested after rollback (#697)

Author: untitaker

sentry_sdk/integrations/sqlalchemy.py

@@ -84,3 +84,14 @@ def _handle_error(context, *args):
 
     if span is not None:
         span.set_status("internal_error")
+
+    # _after_cursor_execute does not get called for crashing SQL stmts. Judging
+    # from SQLAlchemy codebase it does seem like any error coming into this
+    # handler is going to be fatal.
+    ctx_mgr = getattr(
+        conn, "_sentry_sql_span_manager", None
+    )  # type: ContextManager[Any]
+
+    if ctx_mgr is not None:
+        conn._sentry_sql_span_manager = None
+        ctx_mgr.__exit__(None, None, None)

tests/conftest.py

@@ -313,3 +313,28 @@ def maybe_monkeypatched_threading(request):
         assert request.param is None
 
     return request.param
+
+
+@pytest.fixture
+def render_span_tree():
+    def inner(event):
+        assert event["type"] == "transaction"
+
+        by_parent = {}
+        for span in event["spans"]:
+            by_parent.setdefault(span["parent_span_id"], []).append(span)
+
+        def render_span(span):
+            yield "- op={!r}: description={!r}".format(
+                span.get("op"), span.get("description")
+            )
+            for subspan in by_parent.get(span["span_id"]) or ():
+                for line in render_span(subspan):
+                    yield "  {}".format(line)
+
+        root_span = event["contexts"]["trace"]
+
+        # Return a list instead of a multiline string because black will know better how to format that
+        return "\n".join(render_span(root_span))
+
+    return inner

tests/integrations/sqlalchemy/test_sqlalchemy.py

@@ -1,8 +1,12 @@
-from sqlalchemy import Column, ForeignKey, Integer, String
+import sys
+import pytest
+
+from sqlalchemy import Column, ForeignKey, Integer, String, create_engine
+from sqlalchemy.exc import IntegrityError
 from sqlalchemy.ext.declarative import declarative_base
 from sqlalchemy.orm import relationship, sessionmaker
-from sqlalchemy import create_engine
 
+import sentry_sdk
 from sentry_sdk import capture_message
 from sentry_sdk.integrations.sqlalchemy import SqlalchemyIntegration
 
@@ -63,3 +67,70 @@ class Address(Base):
             "type": "default",
         },
     ]
+
+
+@pytest.mark.skipif(
+    sys.version_info < (3,), reason="This sqla usage seems to be broken on Py2"
+)
+def test_transactions(sentry_init, capture_events, render_span_tree):
+
+    sentry_init(
+        integrations=[SqlalchemyIntegration()], _experiments={"record_sql_params": True}
+    )
+    events = capture_events()
+
+    Base = declarative_base()  # noqa: N806
+
+    class Person(Base):
+        __tablename__ = "person"
+        id = Column(Integer, primary_key=True)
+        name = Column(String(250), nullable=False)
+
+    class Address(Base):
+        __tablename__ = "address"
+        id = Column(Integer, primary_key=True)
+        street_name = Column(String(250))
+        street_number = Column(String(250))
+        post_code = Column(String(250), nullable=False)
+        person_id = Column(Integer, ForeignKey("person.id"))
+        person = relationship(Person)
+
+    engine = create_engine("sqlite:///:memory:")
+    Base.metadata.create_all(engine)
+
+    Session = sessionmaker(bind=engine)  # noqa: N806
+    session = Session()
+
+    with sentry_sdk.start_span(transaction="test_transaction", sampled=True):
+        with session.begin_nested():
+            session.query(Person).first()
+
+        for _ in range(2):
+            with pytest.raises(IntegrityError):
+                with session.begin_nested():
+                    session.add(Person(id=1, name="bob"))
+                    session.add(Person(id=1, name="bob"))
+
+        with session.begin_nested():
+            session.query(Person).first()
+
+    (event,) = events
+
+    assert (
+        render_span_tree(event)
+        == """\
+- op=None: description=None
+  - 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'
+  - op='db': description='SAVEPOINT sa_savepoint_2'
+  - op='db': description='INSERT INTO person (id, name) VALUES (?, ?)'
+  - op='db': description='ROLLBACK TO SAVEPOINT sa_savepoint_2'
+  - op='db': description='SAVEPOINT sa_savepoint_3'
+  - op='db': description='INSERT INTO person (id, name) VALUES (?, ?)'
+  - op='db': description='ROLLBACK TO SAVEPOINT sa_savepoint_3'
+  - op='db': description='SAVEPOINT sa_savepoint_4'
+  - 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_4'\
+"""
+    )