feat(django): Instrument database rollbacks

Author: alexander-alderman-webb

sentry_sdk/consts.py

@@ -116,6 +116,7 @@ class INSTRUMENTER:
 
 class DBOPERATION:
     COMMIT = "COMMIT"
+    ROLLBACK = "ROLLBACK"
 
 
 class SPANDATA:

sentry_sdk/integrations/django/__init__.py

@@ -634,6 +634,7 @@ def install_sql_hook():
         real_executemany = CursorWrapper.executemany
         real_connect = BaseDatabaseWrapper.connect
         real_commit = BaseDatabaseWrapper.commit
+        real_rollback = BaseDatabaseWrapper.rollback
     except AttributeError:
         # This won't work on Django versions < 1.6
         return
@@ -702,11 +703,23 @@ def commit(self):
             _set_db_data(span, self, DBOPERATION.COMMIT)
             return real_commit(self)
 
+    @ensure_integration_enabled(DjangoIntegration, real_rollback)
+    def rollback(self):
+        # type: (BaseDatabaseWrapper) -> None
+        with sentry_sdk.start_span(
+            op=OP.DB,
+            name=DBOPERATION.ROLLBACK,
+            origin=DjangoIntegration.origin_db,
+        ) as span:
+            _set_db_data(span, self, DBOPERATION.ROLLBACK)
+            return real_rollback(self)
+
     CursorWrapper.execute = execute
     CursorWrapper.executemany = executemany
     BaseDatabaseWrapper.connect = connect
     BaseDatabaseWrapper.commit = commit
     ignore_logger("django.db.backends")
+    BaseDatabaseWrapper.rollback = rollback
 
 
 def _set_db_data(span, cursor_or_db, db_operation=None):

tests/integrations/django/myapp/urls.py

@@ -66,11 +66,21 @@ def path(path, *args, **kwargs):
         views.postgres_insert_orm_no_autocommit,
         name="postgres_insert_orm_no_autocommit",
     ),
+    path(
+        "postgres-insert-no-autocommit-rollback",
+        views.postgres_insert_orm_no_autocommit_rollback,
+        name="postgres_insert_orm_no_autocommit_rollback",
+    ),
     path(
         "postgres-insert-atomic",
         views.postgres_insert_orm_atomic,
         name="postgres_insert_orm_atomic",
     ),
+    path(
+        "postgres-insert-atomic-rollback",
+        views.postgres_insert_orm_atomic_rollback,
+        name="postgres_insert_orm_atomic_rollback",
+    ),
     path(
         "postgres-select-slow-from-supplement",
         helper_views.postgres_select_orm,

tests/integrations/django/myapp/views.py

@@ -259,6 +259,18 @@ def postgres_insert_orm_no_autocommit(request, *args, **kwargs):
     return HttpResponse("ok {}".format(user))
 
 
+@csrf_exempt
+def postgres_insert_orm_no_autocommit_rollback(request, *args, **kwargs):
+    transaction.set_autocommit(False, using="postgres")
+    user = User.objects.db_manager("postgres").create_user(
+        username="user1",
+    )
+    transaction.rollback(using="postgres")
+    transaction.set_autocommit(True, using="postgres")
+
+    return HttpResponse("ok {}".format(user))
+
+
 @csrf_exempt
 def postgres_insert_orm_atomic(request, *args, **kwargs):
     with transaction.atomic(using="postgres"):
@@ -268,6 +280,16 @@ def postgres_insert_orm_atomic(request, *args, **kwargs):
     return HttpResponse("ok {}".format(user))
 
 
+@csrf_exempt
+def postgres_insert_orm_atomic_rollback(request, *args, **kwargs):
+    with transaction.atomic(using="postgres"):
+        user = User.objects.db_manager("postgres").create_user(
+            username="user1",
+        )
+        transaction.set_rollback(True, using="postgres")
+    return HttpResponse("ok {}".format(user))
+
+
 @csrf_exempt
 def permission_denied_exc(*args, **kwargs):
     raise PermissionDenied("bye")

tests/integrations/django/test_db_query_data.py

@@ -640,6 +640,118 @@ def test_db_no_autocommit_executemany(sentry_init, client, capture_events):
     assert commit_span["origin"] == "auto.db.django"
 
 
+@pytest.mark.forked
+@pytest_mark_django_db_decorator(transaction=True, databases=["postgres"])
+def test_db_no_autocommit_rollback_execute(sentry_init, client, capture_events):
+    sentry_init(
+        integrations=[DjangoIntegration()],
+        traces_sample_rate=1.0,
+    )
+
+    if "postgres" not in connections:
+        pytest.skip("postgres tests disabled")
+
+    # trigger Django to open a new connection by marking the existing one as None.
+    connections["postgres"].connection = None
+
+    events = capture_events()
+
+    client.get(reverse("postgres_insert_orm_no_autocommit_rollback"))
+
+    (event,) = events
+
+    # Ensure operation is rolled back
+    assert not User.objects.using("postgres").exists()
+
+    assert event["contexts"]["trace"]["origin"] == "auto.http.django"
+
+    for span in event["spans"]:
+        if span["op"] == "db":
+            assert span["origin"] == "auto.db.django"
+        else:
+            assert span["origin"] == "auto.http.django"
+
+    commit_spans = [
+        span
+        for span in event["spans"]
+        if span["data"].get(SPANDATA.DB_OPERATION) == DBOPERATION.ROLLBACK
+    ]
+    assert len(commit_spans) == 1
+    commit_span = commit_spans[0]
+    assert commit_span["origin"] == "auto.db.django"
+
+
+@pytest.mark.forked
+@pytest_mark_django_db_decorator(transaction=True)
+def test_db_no_autocommit_rollback_executemany(sentry_init, client, capture_events):
+    sentry_init(
+        integrations=[DjangoIntegration()],
+        traces_sample_rate=1.0,
+    )
+
+    events = capture_events()
+
+    if "postgres" not in connections:
+        pytest.skip("postgres tests disabled")
+
+    with start_transaction(name="test_transaction"):
+        from django.db import connection, transaction
+
+        cursor = connection.cursor()
+
+        query = """INSERT INTO auth_user (
+    password,
+    is_superuser,
+    username,
+    first_name,
+    last_name,
+    email,
+    is_staff,
+    is_active,
+    date_joined
+)
+VALUES ('password', false, %s, %s, %s, %s, false, true, %s);"""
+
+        query_list = (
+            (
+                "user1",
+                "John",
+                "Doe",
+                "[email protected]",
+                datetime(1970, 1, 1),
+            ),
+            (
+                "user2",
+                "Max",
+                "Mustermann",
+                "[email protected]",
+                datetime(1970, 1, 1),
+            ),
+        )
+
+        transaction.set_autocommit(False)
+        cursor.executemany(query, query_list)
+        transaction.rollback()
+        transaction.set_autocommit(True)
+
+    (event,) = events
+
+    # Ensure operation is rolled back
+    assert not User.objects.exists()
+
+    assert event["contexts"]["trace"]["origin"] == "manual"
+    assert event["spans"][0]["origin"] == "auto.db.django"
+
+    commit_spans = [
+        span
+        for span in event["spans"]
+        if span["data"].get(SPANDATA.DB_OPERATION) == DBOPERATION.ROLLBACK
+    ]
+    assert len(commit_spans) == 1
+    commit_span = commit_spans[0]
+    assert commit_span["origin"] == "auto.db.django"
+
+
 @pytest.mark.forked
 @pytest_mark_django_db_decorator(transaction=True, databases=["postgres"])
 def test_db_atomic_execute(sentry_init, client, capture_events):
@@ -743,3 +855,109 @@ def test_db_atomic_executemany(sentry_init, client, capture_events):
     assert len(commit_spans) == 1
     commit_span = commit_spans[0]
     assert commit_span["origin"] == "auto.db.django"
+
+
+@pytest.mark.forked
+@pytest_mark_django_db_decorator(transaction=True, databases=["postgres"])
+def test_db_atomic_rollback_execute(sentry_init, client, capture_events):
+    sentry_init(
+        integrations=[DjangoIntegration()],
+        send_default_pii=True,
+        traces_sample_rate=1.0,
+    )
+
+    if "postgres" not in connections:
+        pytest.skip("postgres tests disabled")
+
+    # trigger Django to open a new connection by marking the existing one as None.
+    connections["postgres"].connection = None
+
+    events = capture_events()
+
+    client.get(reverse("postgres_insert_orm_atomic_rollback"))
+
+    (event,) = events
+
+    # Ensure operation is rolled back
+    assert not User.objects.using("postgres").exists()
+
+    assert event["contexts"]["trace"]["origin"] == "auto.http.django"
+
+    commit_spans = [
+        span
+        for span in event["spans"]
+        if span["data"].get(SPANDATA.DB_OPERATION) == DBOPERATION.ROLLBACK
+    ]
+    assert len(commit_spans) == 1
+    commit_span = commit_spans[0]
+    assert commit_span["origin"] == "auto.db.django"
+
+
+@pytest.mark.forked
+@pytest_mark_django_db_decorator(transaction=True)
+def test_db_atomic_rollback_executemany(sentry_init, client, capture_events):
+    sentry_init(
+        integrations=[DjangoIntegration()],
+        send_default_pii=True,
+        traces_sample_rate=1.0,
+    )
+
+    if "postgres" not in connections:
+        pytest.skip("postgres tests disabled")
+
+    # trigger Django to open a new connection by marking the existing one as None.
+    connections["postgres"].connection = None
+
+    events = capture_events()
+
+    with start_transaction(name="test_transaction"):
+        with transaction.atomic():
+            cursor = connection.cursor()
+
+            query = """INSERT INTO auth_user (
+    password,
+    is_superuser,
+    username,
+    first_name,
+    last_name,
+    email,
+    is_staff,
+    is_active,
+    date_joined
+)
+VALUES ('password', false, %s, %s, %s, %s, false, true, %s);"""
+
+            query_list = (
+                (
+                    "user1",
+                    "John",
+                    "Doe",
+                    "[email protected]",
+                    datetime(1970, 1, 1),
+                ),
+                (
+                    "user2",
+                    "Max",
+                    "Mustermann",
+                    "[email protected]",
+                    datetime(1970, 1, 1),
+                ),
+            )
+            cursor.executemany(query, query_list)
+            transaction.set_rollback(True)
+
+    (event,) = events
+
+    # Ensure operation is rolled back
+    assert not User.objects.exists()
+
+    assert event["contexts"]["trace"]["origin"] == "manual"
+
+    commit_spans = [
+        span
+        for span in event["spans"]
+        if span["data"].get(SPANDATA.DB_OPERATION) == DBOPERATION.ROLLBACK
+    ]
+    assert len(commit_spans) == 1
+    commit_span = commit_spans[0]
+    assert commit_span["origin"] == "auto.db.django"