Repository: scoutapp/scout_apm_python

Issue:
Use Django 2.0+ DB instrumentation hook
Rather than monkey-patching, use the built-in instrumentation on Django 2.0+:

https://docs.djangoproject.com/en/2.0/topics/db/instrumentation/


Discussion:
(suggested by Aymeric)

Resolving patch:
diff --git a/src/scout_apm/django/instruments/sql.py b/src/scout_apm/django/instruments/sql.py
--- a/src/scout_apm/django/instruments/sql.py
+++ b/src/scout_apm/django/instruments/sql.py
@@ -3,6 +3,10 @@
 
 import logging
 
+import django
+from django.db import connections
+from django.db.backends.signals import connection_created
+
 from scout_apm.core.monkey import monkeypatch_method
 from scout_apm.core.tracked_request import TrackedRequest
 
@@ -21,31 +25,74 @@
 logger = logging.getLogger(__name__)
 
 
+def db_execute_hook(execute, sql, params, many, context):
+    """
+    Database instrumentation hook for Django 2.0+
+    https://docs.djangoproject.com/en/2.0/topics/db/instrumentation/
+    """
+    if many:
+        operation = "SQL/Many"
+    else:
+        operation = "SQL/Query"
+
+    tracked_request = TrackedRequest.instance()
+    span = tracked_request.start_span(operation=operation)
+    span.tag("db.statement", sql)
+
+    try:
+        return execute(sql, params, many, context)
+    finally:
+        tracked_request.stop_span()
+        tracked_request.callset.update(sql, 1, span.duration())
+        if tracked_request.callset.should_capture_backtrace(sql):
+            span.capture_backtrace()
+
+
+def install_db_execute_hook(connection, **kwargs):
+    """
+    Install db_execute_hook on the given database connection.
+
+    Rather than use the documented API of the `execute_wrapper()` context
+    manager, directly insert the hook. This is done because:
+    1. Deleting the context manager closes it, so it's not possible to enter
+       it here and not exit it, unless we store it forever in some variable
+    2. We want to be the first hook, so we can capture every query (although
+       potentially later hooks will change the SQL)
+    3. We want to be idempotent and only install the hook once
+    """
+    if db_execute_hook not in connection.execute_wrappers:
+        connection.execute_wrappers.insert(0, db_execute_hook)
+
+
 class _DetailedTracingCursorWrapper(CursorWrapper):
+    """
+    Monkey-patched-in cursor wrapper for Django < 2.0
+    """
+
     def execute(self, sql, params=None):
-        tr = TrackedRequest.instance()
-        span = tr.start_span(operation="SQL/Query")
+        tracked_request = TrackedRequest.instance()
+        span = tracked_request.start_span(operation="SQL/Query")
         span.tag("db.statement", sql)
 
         try:
             return self.cursor.execute(sql, params)
         finally:
-            tr.stop_span()
-            tr.callset.update(sql, 1, span.duration())
-            if tr.callset.should_capture_backtrace(sql):
+            tracked_request.stop_span()
+            tracked_request.callset.update(sql, 1, span.duration())
+            if tracked_request.callset.should_capture_backtrace(sql):
                 span.capture_backtrace()
 
     def executemany(self, sql, param_list):
-        tr = TrackedRequest.instance()
-        span = tr.start_span(operation="SQL/Many")
+        tracked_request = TrackedRequest.instance()
+        span = tracked_request.start_span(operation="SQL/Many")
         span.tag("db.statement", sql)
 
         try:
             return self.cursor.executemany(sql, param_list)
         finally:
-            tr.stop_span()
-            tr.callset.update(sql, 1, span.duration())
-            if tr.callset.should_capture_backtrace(sql):
+            tracked_request.stop_span()
+            tracked_request.callset.update(sql, 1, span.duration())
+            if tracked_request.callset.should_capture_backtrace(sql):
                 span.capture_backtrace()
 
 
@@ -54,9 +101,16 @@ def install_sql_instrumentation():
         return
     install_sql_instrumentation.installed = True
 
-    @monkeypatch_method(BaseDatabaseWrapper)
-    def cursor(original, self, *args, **kwargs):
-        result = original(*args, **kwargs)
-        return _DetailedTracingCursorWrapper(result, self)
+    if django.VERSION >= (2, 0):
+        for connection in connections.all():
+            install_db_execute_hook(connection=connection)
+        connection_created.connect(install_db_execute_hook)
+        logger.debug("Installed DB connection created signal handler")
+    else:
+
+        @monkeypatch_method(BaseDatabaseWrapper)
+        def cursor(original, self, *args, **kwargs):
+            result = original(*args, **kwargs)
+            return _DetailedTracingCursorWrapper(result, self)
 
-    logger.debug("Monkey patched SQL")
+        logger.debug("Monkey patched SQL")
diff --git a/src/scout_apm/django/middleware.py b/src/scout_apm/django/middleware.py
--- a/src/scout_apm/django/middleware.py
+++ b/src/scout_apm/django/middleware.py
@@ -85,16 +85,15 @@ def __call__(self, request):
         be recorded.  This can happen if a middleware further along the stack
         doesn't call onward, and instead returns a response directly.
         """
-
-        tr = TrackedRequest.instance()
+        tracked_request = TrackedRequest.instance()
 
         # This operation name won't be recorded unless changed later in
         # process_view
-        tr.start_span(operation="Unknown")
+        tracked_request.start_span(operation="Unknown")
         try:
             response = self.get_response(request)
         finally:
-            tr.stop_span()
+            tracked_request.stop_span()
         return response
 
     def process_view(self, request, view_func, view_args, view_kwargs):


Test patch: