feat(django): Instrument database rollbacks (#5115)

Add spans for SQL rollbacks issued when Django calls `rollback()` on a
PEP-249 database connection.

Rollback spans are generated for `transaction.atomic` blocks and for
manual `transaction.rollback()` calls when auto-commit is disabled. Tests
cover both cases, for SQLite and PostgreSQL, respectively.

Author: alexander-alderman-webb

sentry_sdk/consts.py

@@ -116,6 +116,7 @@ class INSTRUMENTER:
 
 class SPANNAME:
     DB_COMMIT = "COMMIT"
+    DB_ROLLBACK = "ROLLBACK"
 
 
 class SPANDATA:

sentry_sdk/integrations/django/__init__.py

@@ -636,6 +636,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
@@ -708,10 +709,26 @@ def _commit(self):
             _set_db_data(span, self, SPANNAME.DB_COMMIT)
             return real_commit(self)
 
+    def _rollback(self):
+        # type: (BaseDatabaseWrapper) -> None
+        integration = sentry_sdk.get_client().get_integration(DjangoIntegration)
+
+        if integration is None or not integration.db_transaction_spans:
+            return real_rollback(self)
+
+        with sentry_sdk.start_span(
+            op=OP.DB,
+            name=SPANNAME.DB_ROLLBACK,
+            origin=DjangoIntegration.origin_db,
+        ) as span:
+            _set_db_data(span, self, SPANNAME.DB_ROLLBACK)
+            return real_rollback(self)
+
     CursorWrapper.execute = execute
     CursorWrapper.executemany = executemany
     BaseDatabaseWrapper.connect = connect
     BaseDatabaseWrapper._commit = _commit
+    BaseDatabaseWrapper._rollback = _rollback
     ignore_logger("django.db.backends")
 
 

tests/integrations/django/myapp/urls.py

@@ -66,11 +66,26 @@ 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-insert-atomic-exception",
+        views.postgres_insert_orm_atomic_exception,
+        name="postgres_insert_orm_atomic_exception",
+    ),
     path(
         "postgres-select-slow-from-supplement",
         helper_views.postgres_select_orm,

tests/integrations/django/myapp/views.py

@@ -264,6 +264,23 @@ 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")
+    try:
+        user = User.objects.db_manager("postgres").create_user(
+            username="user1",
+        )
+        transaction.rollback(using="postgres")
+    except Exception:
+        transaction.rollback(using="postgres")
+        transaction.set_autocommit(True, using="postgres")
+        raise
+
+    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"):
@@ -273,6 +290,30 @@ 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 postgres_insert_orm_atomic_exception(request, *args, **kwargs):
+    try:
+        with transaction.atomic(using="postgres"):
+            user = User.objects.db_manager("postgres").create_user(
+                username="user1",
+            )
+            transaction.set_rollback(True, using="postgres")
+            1 / 0
+    except ZeroDivisionError:
+        pass
+    return HttpResponse("ok {}".format(user))
+
+
 @csrf_exempt
 def permission_denied_exc(*args, **kwargs):
     raise PermissionDenied("bye")