feat(django): Instrument database commits

Author: alexander-alderman-webb

sentry_sdk/consts.py

@@ -114,6 +114,10 @@ class INSTRUMENTER:
     OTEL = "otel"
 
 
+class DBOPERATION:
+    COMMIT = "COMMIT"
+
+
 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
@@ -633,6 +633,7 @@ def install_sql_hook():
         real_execute = CursorWrapper.execute
         real_executemany = CursorWrapper.executemany
         real_connect = BaseDatabaseWrapper.connect
+        real_commit = BaseDatabaseWrapper.commit
     except AttributeError:
         # This won't work on Django versions < 1.6
         return
@@ -690,14 +691,26 @@ def connect(self):
             _set_db_data(span, self)
             return real_connect(self)
 
+    @ensure_integration_enabled(DjangoIntegration, real_commit)
+    def commit(self):
+        # type: (BaseDatabaseWrapper) -> None
+        print("commiting")
+        with sentry_sdk.start_span(
+            op=OP.DB,
+            name="commit",  # DBOPERATION.COMMIT,
+            origin=DjangoIntegration.origin_db,
+        ) as span:
+            _set_db_data(span, self, DBOPERATION.COMMIT)
+            return real_commit(self)
+
     CursorWrapper.execute = execute
     CursorWrapper.executemany = executemany
     BaseDatabaseWrapper.connect = connect
-    ignore_logger("django.db.backends")
+    BaseDatabaseWrapper.commit = commit
 
 
-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)
@@ -735,6 +748,9 @@ def _set_db_data(span, cursor_or_db):
     if db_name is not None:
         span.set_data(SPANDATA.DB_NAME, db_name)
 
+    if db_operation is not None:
+        span.set_data(SPANDATA.DB_OPERATION, db_operation)
+
     server_address = connection_params.get("host")
     if server_address is not None:
         span.set_data(SPANDATA.SERVER_ADDRESS, server_address)

tests/integrations/django/test_db_query_data.py

@@ -5,7 +5,7 @@
 from unittest import mock
 
 from django import VERSION as DJANGO_VERSION
-from django.db import connections
+from django.db import connection, connections, transaction
 
 try:
     from django.urls import reverse
@@ -15,7 +15,7 @@
 from werkzeug.test import Client
 
 from sentry_sdk import start_transaction
-from sentry_sdk.consts import SPANDATA
+from sentry_sdk.consts import SPANDATA, DBOPERATION
 from sentry_sdk.integrations.django import DjangoIntegration
 from sentry_sdk.tracing_utils import record_sql_queries
 
@@ -481,6 +481,7 @@ def test_db_span_origin_execute(sentry_init, client, capture_events):
     assert event["contexts"]["trace"]["origin"] == "auto.http.django"
 
     for span in event["spans"]:
+        print("span is", span["op"], span["description"])
         if span["op"] == "db":
             assert span["origin"] == "auto.db.django"
         else:
@@ -524,3 +525,190 @@ def test_db_span_origin_executemany(sentry_init, client, capture_events):
 
     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.COMMIT
+    ]
+    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_execute(sentry_init, client, capture_events):
+    """
+    Verify we record a breadcrumb when opening a new database.
+    """
+    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_select_orm_no_autocommit"))
+
+    (event,) = events
+
+    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.COMMIT
+    ]
+    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_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 = """UPDATE auth_user SET username = %s where id = %s;"""
+        query_list = (
+            (
+                "test1",
+                1,
+            ),
+            (
+                "test2",
+                2,
+            ),
+        )
+        cursor.executemany(query, query_list)
+
+        transaction.commit()
+
+    (event,) = events
+
+    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.COMMIT
+    ]
+    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_execute(sentry_init, client, capture_events):
+    """
+    Verify we record a breadcrumb when opening a new database.
+    """
+    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 transaction.atomic():
+        client.get(reverse("postgres_select_orm_atomic"))
+        connections["postgres"].commit()
+
+    (event,) = events
+
+    assert event["contexts"]["trace"]["origin"] == "auto.http.django"
+
+    commit_spans = [
+        span
+        for span in event["spans"]
+        if span["data"].get(SPANDATA.DB_OPERATION) == DBOPERATION.COMMIT
+    ]
+    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_executemany(sentry_init, client, capture_events):
+    """
+    Verify we record a breadcrumb when opening a new database.
+    """
+    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 = """UPDATE auth_user SET username = %s where id = %s;"""
+            query_list = (
+                (
+                    "test1",
+                    1,
+                ),
+                (
+                    "test2",
+                    2,
+                ),
+            )
+            cursor.executemany(query, query_list)
+
+    (event,) = events
+
+    assert event["contexts"]["trace"]["origin"] == "manual"
+
+    commit_spans = [
+        span
+        for span in event["spans"]
+        if span["data"].get(SPANDATA.DB_OPERATION) == DBOPERATION.COMMIT
+    ]
+    assert len(commit_spans) == 1
+    commit_span = commit_spans[0]
+    assert commit_span["origin"] == "auto.db.django"