implement db cursor wrapper so don't need DEBUG=True

Author: kdelee

ansible_base/lib/middleware/profiling/README.md

@@ -50,16 +50,13 @@ MIDDLEWARE = [
 
 ### Enabling SQL Profiling
 
-**Important:** This middleware relies on Django's `connection.queries` list, which is only populated when `settings.DEBUG` is set to `True`. Therefore, you must have `DEBUG = True` in your Django settings for this middleware to have any effect.
-
-The middleware is controlled by the `ANSIBLE_BASE_SQL_PROFILING` setting. For backwards compatibility, it will also be enabled if the standard Django `SQL_DEBUG` setting is `True`.
+The middleware is controlled by the `ANSIBLE_BASE_SQL_PROFILING` setting.
 
 To enable SQL profiling, set the following in your Django settings:
 
 ```python
 # settings.py
 ANSIBLE_BASE_SQL_PROFILING = True
-DEBUG = True
 ```
 
 ## `DABProfiler`

ansible_base/lib/middleware/profiling/profile_request.py

@@ -7,13 +7,12 @@
 import uuid
 from typing import Optional, Union
 
-from django.db import connection
 from django.conf import settings
+from django.db import connection
 from django.utils.translation import gettext_lazy as _
 
 from ansible_base.lib.utils.settings import get_function_from_setting, get_setting
 
-
 logger = logging.getLogger(__name__)
 
 
@@ -85,23 +84,32 @@ def __call__(self, request):
         return response
 
 
+class SQLQueryMetrics:
+    def __init__(self):
+        self.query_count = 0
+        self.query_time = 0.0
+
+    def __call__(self, execute, sql, params, many, context):
+        start_time = time.time()
+        try:
+            return execute(sql, params, many, context)
+        finally:
+            self.query_count += 1
+            self.query_time += time.time() - start_time
+
+
 class SQLProfilingMiddleware:
     def __init__(self, get_response):
         self.get_response = get_response
 
     def __call__(self, request):
-        sql_profiling_enabled = get_setting('ANSIBLE_BASE_SQL_PROFILING', get_setting('SQL_DEBUG', False))
-        if sql_profiling_enabled:
-            if not settings.DEBUG:
-                logger.warning("ANSIBLE_BASE_SQL_PROFILING is enabled, but DEBUG is False. No SQL queries will be logged or counted.")
-                return self.get_response(request)
+        if not get_setting('ANSIBLE_BASE_SQL_PROFILING', False):
+            return self.get_response(request)
 
-            queries_before = len(connection.queries)
-            response = self.get_response(request)
-            q_times = [float(q['time']) for q in connection.queries[queries_before:]]
-            response['X-API-Query-Count'] = len(q_times)
-            response['X-API-Query-Time'] = '%0.3fs' % sum(q_times)
-        else:
+        metrics = SQLQueryMetrics()
+        with connection.execute_wrapper(metrics):
             response = self.get_response(request)
 
+        response['X-API-Query-Count'] = metrics.query_count
+        response['X-API-Query-Time'] = f'{metrics.query_time:.3f}s'
         return response

test_app/tests/lib/middleware/test_profiling_middleware.py

@@ -1,6 +1,6 @@
-import uuid
 import os
 import tempfile
+import uuid
 from unittest.mock import patch
 
 from django.http import HttpResponse
@@ -11,22 +11,26 @@
 from ansible_base.lib.utils.settings import get_setting
 from test_app.models import User
 
+
 # A simple view for testing middleware
 def simple_view(request):
     return HttpResponse("OK")
 
+
 # A view that performs a database query
 def db_view(request):
     # Create a user with a unique username to guarantee a query.
     User.objects.create(username=f"test-{uuid.uuid4()}")
     return HttpResponse("OK")
 
+
 # Define URL patterns for the test
 urlpatterns = [
     path('test/', simple_view),
     path('test-db/', db_view),
 ]
 
+
 @override_settings(ROOT_URLCONF=__name__)
 class ProfileRequestMiddlewareTest(TestCase):
     @override_settings(CLUSTER_HOST_ID='test-node')
@@ -36,7 +40,7 @@ def test_profile_request_middleware_headers(self):
         """
         middleware = ProfileRequestMiddleware(simple_view)
         response = middleware(self.client.get('/test/').wsgi_request)
-        
+
         # Test X-API-Time
         self.assertIn('X-API-Time', response)
         self.assertTrue(response['X-API-Time'].endswith('s'))
@@ -74,10 +78,7 @@ def test_profile_request_middleware_cprofile_disabled(self):
         self.assertNotIn('X-API-CProfile-File', response)
 
 
-@override_settings(
-    ROOT_URLCONF=__name__,
-    MIDDLEWARE=['ansible_base.lib.middleware.profiling.profile_request.SQLProfilingMiddleware']
-)
+@override_settings(ROOT_URLCONF=__name__, MIDDLEWARE=['ansible_base.lib.middleware.profiling.profile_request.SQLProfilingMiddleware'])
 class SQLProfilingMiddlewareTest(TestCase):
     def test_sql_profiling_disabled_by_default(self):
         """
@@ -87,25 +88,10 @@ def test_sql_profiling_disabled_by_default(self):
         self.assertNotIn('X-API-Query-Count', response)
         self.assertNotIn('X-API-Query-Time', response)
 
-    @override_settings(ANSIBLE_BASE_SQL_PROFILING=True, DEBUG=True)
+    @override_settings(ANSIBLE_BASE_SQL_PROFILING=True)
     def test_sql_profiling_enabled_with_new_setting(self):
         """
-        Test that the SQLProfilingMiddleware adds headers when ANSIBLE_BASE_SQL_PROFILING is True.
-        """
-        response = self.client.get('/test-db/')
-        self.assertIn('X-API-Query-Count', response)
-        self.assertGreaterEqual(int(response['X-API-Query-Count']), 1)
-        self.assertIn('X-API-Query-Time', response)
-        self.assertTrue(response['X-API-Query-Time'].endswith('s'))
-        try:
-            float(response['X-API-Query-Time'][:-1])
-        except ValueError:
-            self.fail("X-API-Query-Time value is not a valid float")
-
-    @override_settings(SQL_DEBUG=True, DEBUG=True)
-    def test_sql_profiling_enabled_with_fallback_setting(self):
-        """
-        Test that the SQLProfilingMiddleware adds headers when SQL_DEBUG is True as a fallback.
+        Test that the SQLProfilingMiddleware adds headers when ANSIBLE_BASE_SQL_PROFILING is True
         """
         response = self.client.get('/test-db/')
         self.assertIn('X-API-Query-Count', response)
@@ -116,17 +102,3 @@ def test_sql_profiling_enabled_with_fallback_setting(self):
             float(response['X-API-Query-Time'][:-1])
         except ValueError:
             self.fail("X-API-Query-Time value is not a valid float")
-
-    @override_settings(ANSIBLE_BASE_SQL_PROFILING=True, DEBUG=False)
-    def test_sql_profiling_logs_warning_if_debug_is_false(self):
-        """
-        Test that the SQLProfilingMiddleware logs a warning and does not add headers
-        if profiling is enabled but DEBUG is False.
-        """
-        with self.assertLogs('ansible_base.lib.middleware.profiling.profile_request', level='WARNING') as cm:
-            response = self.client.get('/test-db/')
-            self.assertIn("ANSIBLE_BASE_SQL_PROFILING is enabled, but DEBUG is False", cm.output[0])
-        self.assertNotIn('X-API-Query-Count', response)
-        self.assertNotIn('X-API-Query-Time', response)
-
-