record span and breadcrumb when Django opens db connection (#1250)

Author: chdsbd

sentry_sdk/integrations/django/__init__.py

@@ -481,9 +481,17 @@ def install_sql_hook():
     except ImportError:
         from django.db.backends.util import CursorWrapper
 
+    try:
+        # django 1.6 and 1.7 compatability
+        from django.db.backends import BaseDatabaseWrapper
+    except ImportError:
+        # django 1.8 or later
+        from django.db.backends.base.base import BaseDatabaseWrapper
+
     try:
         real_execute = CursorWrapper.execute
         real_executemany = CursorWrapper.executemany
+        real_connect = BaseDatabaseWrapper.connect
     except AttributeError:
         # This won't work on Django versions < 1.6
         return
@@ -510,6 +518,19 @@ def executemany(self, sql, param_list):
         ):
             return real_executemany(self, sql, param_list)
 
+    def connect(self):
+        # type: (BaseDatabaseWrapper) -> None
+        hub = Hub.current
+        if hub.get_integration(DjangoIntegration) is None:
+            return real_connect(self)
+
+        with capture_internal_exceptions():
+            hub.add_breadcrumb(message="connect", category="query")
+
+        with hub.start_span(op="db", description="connect"):
+            return real_connect(self)
+
     CursorWrapper.execute = execute
     CursorWrapper.executemany = executemany
+    BaseDatabaseWrapper.connect = connect
     ignore_logger("django.db.backends")

tests/integrations/django/myapp/urls.py

@@ -47,6 +47,7 @@ def path(path, *args, **kwargs):
     path("template-exc", views.template_exc, name="template_exc"),
     path("template-test", views.template_test, name="template_test"),
     path("template-test2", views.template_test2, name="template_test2"),
+    path("postgres-select", views.postgres_select, name="postgres_select"),
     path(
         "permission-denied-exc",
         views.permission_denied_exc,

tests/integrations/django/myapp/views.py

@@ -127,6 +127,15 @@ def template_test2(request, *args, **kwargs):
     )
 
 
+@csrf_exempt
+def postgres_select(request, *args, **kwargs):
+    from django.db import connections
+
+    cursor = connections["postgres"].cursor()
+    cursor.execute("SELECT 1;")
+    return HttpResponse("ok")
+
+
 @csrf_exempt
 def permission_denied_exc(*args, **kwargs):
     raise PermissionDenied("bye")

tests/integrations/django/test_basic.py

@@ -19,19 +19,24 @@
 
 from sentry_sdk import capture_message, capture_exception, configure_scope
 from sentry_sdk.integrations.django import DjangoIntegration
+from functools import partial
 
 from tests.integrations.django.myapp.wsgi import application
 
 # Hack to prevent from experimental feature introduced in version `4.3.0` in `pytest-django` that
 # requires explicit database allow from failing the test
-pytest_mark_django_db_decorator = pytest.mark.django_db
+pytest_mark_django_db_decorator = partial(pytest.mark.django_db)
 try:
     pytest_version = tuple(map(int, pytest_django.__version__.split(".")))
     if pytest_version > (4, 2, 0):
-        pytest_mark_django_db_decorator = pytest.mark.django_db(databases="__all__")
+        pytest_mark_django_db_decorator = partial(
+            pytest.mark.django_db, databases="__all__"
+        )
 except ValueError:
     if "dev" in pytest_django.__version__:
-        pytest_mark_django_db_decorator = pytest.mark.django_db(databases="__all__")
+        pytest_mark_django_db_decorator = partial(
+            pytest.mark.django_db, databases="__all__"
+        )
 except AttributeError:
     pass
 
@@ -259,7 +264,7 @@ def test_sql_queries(sentry_init, capture_events, with_integration):
 
 
 @pytest.mark.forked
-@pytest_mark_django_db_decorator
+@pytest_mark_django_db_decorator()
 def test_sql_dict_query_params(sentry_init, capture_events):
     sentry_init(
         integrations=[DjangoIntegration()],
@@ -304,7 +309,7 @@ def test_sql_dict_query_params(sentry_init, capture_events):
     ],
 )
 @pytest.mark.forked
-@pytest_mark_django_db_decorator
+@pytest_mark_django_db_decorator()
 def test_sql_psycopg2_string_composition(sentry_init, capture_events, query):
     sentry_init(
         integrations=[DjangoIntegration()],
@@ -337,7 +342,7 @@ def test_sql_psycopg2_string_composition(sentry_init, capture_events, query):
 
 
 @pytest.mark.forked
-@pytest_mark_django_db_decorator
+@pytest_mark_django_db_decorator()
 def test_sql_psycopg2_placeholders(sentry_init, capture_events):
     sentry_init(
         integrations=[DjangoIntegration()],
@@ -397,6 +402,72 @@ def test_sql_psycopg2_placeholders(sentry_init, capture_events):
     ]
 
 
+@pytest.mark.forked
+@pytest_mark_django_db_decorator(transaction=True)
+def test_django_connect_trace(sentry_init, client, capture_events, render_span_tree):
+    """
+    Verify we record a span when opening a new database.
+    """
+    sentry_init(
+        integrations=[DjangoIntegration()],
+        send_default_pii=True,
+        traces_sample_rate=1.0,
+    )
+
+    from django.db import connections
+
+    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()
+
+    content, status, headers = client.get(reverse("postgres_select"))
+    assert status == "200 OK"
+
+    assert '- op="db": description="connect"' in render_span_tree(events[0])
+
+
+@pytest.mark.forked
+@pytest_mark_django_db_decorator(transaction=True)
+def test_django_connect_breadcrumbs(
+    sentry_init, client, capture_events, render_span_tree
+):
+    """
+    Verify we record a breadcrumb when opening a new database.
+    """
+    sentry_init(
+        integrations=[DjangoIntegration()],
+        send_default_pii=True,
+    )
+
+    from django.db import connections
+
+    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()
+
+    cursor = connections["postgres"].cursor()
+    cursor.execute("select 1")
+
+    # trigger recording of event.
+    capture_message("HI")
+    (event,) = events
+    for crumb in event["breadcrumbs"]["values"]:
+        del crumb["timestamp"]
+
+    assert event["breadcrumbs"]["values"][-2:] == [
+        {"message": "connect", "category": "query", "type": "default"},
+        {"message": "select 1", "category": "query", "data": {}, "type": "default"},
+    ]
+
+
 @pytest.mark.parametrize(
     "transaction_style,expected_transaction",
     [