feat(django): Instrument database rollbacks

sentry_sdk/consts.py

@@ -114,6 +114,11 @@ class INSTRUMENTER:
     OTEL = "otel"
 
 
+class DBOPERATION:
+    COMMIT = "COMMIT"
+    ROLLBACK = "ROLLBACK"
+
+
 class SPANDATA:
     """
     Additional information describing the type of the span.

sentry_sdk/integrations/django/__init__.py

@@ -5,7 +5,7 @@
 from importlib import import_module
 
 import sentry_sdk
-from sentry_sdk.consts import OP, SPANDATA
+from sentry_sdk.consts import OP, SPANDATA, DBOPERATION
 from sentry_sdk.scope import add_global_event_processor, should_send_default_pii
 from sentry_sdk.serializer import add_global_repr_processor, add_repr_sequence_type
 from sentry_sdk.tracing import SOURCE_FOR_STYLE, TransactionSource
@@ -132,6 +132,7 @@ def __init__(
         middleware_spans=True,  # type: bool
         signals_spans=True,  # type: bool
         cache_spans=False,  # type: bool
+        database_transaction_spans=False,  # type: bool
         signals_denylist=None,  # type: Optional[list[signals.Signal]]
         http_methods_to_capture=DEFAULT_HTTP_METHODS_TO_CAPTURE,  # type: tuple[str, ...]
     ):
@@ -148,6 +149,7 @@ def __init__(
         self.signals_denylist = signals_denylist or []
 
         self.cache_spans = cache_spans
+        self.database_transaction_spans = database_transaction_spans
 
         self.http_methods_to_capture = tuple(map(str.upper, http_methods_to_capture))
 
@@ -633,6 +635,8 @@ def install_sql_hook():
         real_execute = CursorWrapper.execute
         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
@@ -690,18 +694,55 @@ def connect(self):
             _set_db_data(span, self)
             return real_connect(self)
 
+    @ensure_integration_enabled(DjangoIntegration, real_commit)
+    def _commit(self):
+        # type: (BaseDatabaseWrapper) -> None
+        integration = sentry_sdk.get_client().get_integration(DjangoIntegration)
+
+        if integration is None or not integration.database_transaction_spans:
+            return real_commit(self)
+
+        with sentry_sdk.start_span(
+            op=OP.DB,
+            name=DBOPERATION.COMMIT,
+            origin=DjangoIntegration.origin_db,
+        ) as span:
+            _set_db_data(span, self, DBOPERATION.COMMIT)
+            return real_commit(self)
+
+    @ensure_integration_enabled(DjangoIntegration, real_rollback)
+    def _rollback(self):
+        # type: (BaseDatabaseWrapper) -> None
+        integration = sentry_sdk.get_client().get_integration(DjangoIntegration)
+
+        if integration is None or not integration.database_transaction_spans:
+            return real_rollback(self)
+
+        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
+    BaseDatabaseWrapper._rollback = _rollback
     ignore_logger("django.db.backends")
 
 
-def _set_db_data(span, cursor_or_db):
-    # type: (Span, Any) -> None
+def _set_db_data(span, cursor_or_db, db_operation=None):
+    # type: (Span, Any, Optional[str]) -> None
     db = cursor_or_db.db if hasattr(cursor_or_db, "db") else cursor_or_db
     vendor = db.vendor
     span.set_data(SPANDATA.DB_SYSTEM, vendor)
 
+    if db_operation is not None:
+        span.set_data(SPANDATA.DB_OPERATION, db_operation)
+
     # Some custom backends override `__getattr__`, making it look like `cursor_or_db`
     # actually has a `connection` and the `connection` has a `get_dsn_parameters`
     # attribute, only to throw an error once you actually want to call it.

tests/integrations/django/myapp/urls.py

@@ -61,6 +61,31 @@ def path(path, *args, **kwargs):
     path("template-test4", views.template_test4, name="template_test4"),
     path("postgres-select", views.postgres_select, name="postgres_select"),
     path("postgres-select-slow", views.postgres_select_orm, name="postgres_select_orm"),
+    path(
+        "postgres-insert-no-autocommit",
+        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

@@ -2,6 +2,7 @@
 import json
 import threading
 
+from django.db import transaction
 from django.contrib.auth import login
 from django.contrib.auth.models import User
 from django.core.exceptions import PermissionDenied
@@ -246,6 +247,73 @@ def postgres_select_orm(request, *args, **kwargs):
     return HttpResponse("ok {}".format(user))
 
 
+@csrf_exempt
+def postgres_insert_orm_no_autocommit(request, *args, **kwargs):
+    transaction.set_autocommit(False, using="postgres")
+    try:
+        user = User.objects.db_manager("postgres").create_user(
+            username="user1",
+        )
+        transaction.commit(using="postgres")
+    except Exception:
+        transaction.set_autocommit(True, using="postgres")
+        transaction.rollback()
+        raise
+
+    transaction.set_autocommit(True, using="postgres")
+    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.set_autocommit(True, using="postgres")
+        transaction.rollback()
+        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"):
+        user = User.objects.db_manager("postgres").create_user(
+            username="user1",
+        )
+    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")