Skip to content

Commit db071d1

Browse files
committed
feat(profiling): Add reusable SQL profiling middleware
This commit refactors the SQL query profiling logic out of the AWX controller's API view and into a reusable middleware in the django-ansible-base library. This makes the feature available to all AAP components, such as the gateway and EDA. SQLProfilingMiddleware adds `X-API-Query-Count` and `X-API-Query-Time` headers to API responses. Changes include: - A new `SQLProfilingMiddleware` class in the profiling module. - Configuration is controlled by a new, namespaced setting, `ANSIBLE_BASE_SQL_PROFILING`, with a fallback to the legacy `SQL_DEBUG` setting for backward compatibility. - The middleware logs a warning if it is enabled while Django's `DEBUG` setting is `False`, as query logging is disabled in that state. - Unit tests have been added to ensure correct functionality and prevent regressions. - The profiling documentation has been updated to include usage instructions for the new middleware.
1 parent 1b85c18 commit db071d1

File tree

3 files changed

+125
-3
lines changed

3 files changed

+125
-3
lines changed

ansible_base/lib/middleware/profiling/README.md

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,38 @@ To enable cProfile support, set the following in your Django settings:
3030
ANSIBLE_BASE_CPROFILE_REQUESTS = True
3131
```
3232

33+
## `SQLProfilingMiddleware`
34+
35+
This middleware provides insights into the database queries executed during a request. When enabled, it adds the following headers to the response:
36+
37+
* `X-API-Query-Count`: The total number of database queries executed during the request.
38+
* `X-API-Query-Time`: The total time spent on database queries, in seconds.
39+
40+
To use it, add it to your `MIDDLEWARE` list in your Django settings:
41+
42+
```python
43+
# settings.py
44+
MIDDLEWARE = [
45+
...
46+
'ansible_base.lib.middleware.profiling.profile_request.SQLProfilingMiddleware',
47+
...
48+
]
49+
```
50+
51+
### Enabling SQL Profiling
52+
53+
**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.
54+
55+
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`.
56+
57+
To enable SQL profiling, set the following in your Django settings:
58+
59+
```python
60+
# settings.py
61+
ANSIBLE_BASE_SQL_PROFILING = True
62+
DEBUG = True
63+
```
64+
3365
## `DABProfiler`
3466

3567
The core profiling logic is encapsulated in the `DABProfiler` class. This class can be imported and used directly for profiling non-HTTP contexts, such as background tasks or gRPC services.

ansible_base/lib/middleware/profiling/profile_request.py

Lines changed: 26 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,10 +7,13 @@
77
import uuid
88
from typing import Optional, Union
99

10+
from django.db import connection
11+
from django.conf import settings
1012
from django.utils.translation import gettext_lazy as _
1113

1214
from ansible_base.lib.utils.settings import get_function_from_setting, get_setting
1315

16+
1417
logger = logging.getLogger(__name__)
1518

1619

@@ -79,4 +82,26 @@ def __call__(self, request):
7982
extra=dict(python_objects=dict(request=request, response=response, X_API_CPROFILE_FILE=response["X-API-CProfile-File"])),
8083
)
8184

82-
return response
85+
return response
86+
87+
88+
class SQLProfilingMiddleware:
89+
def __init__(self, get_response):
90+
self.get_response = get_response
91+
92+
def __call__(self, request):
93+
sql_profiling_enabled = get_setting('ANSIBLE_BASE_SQL_PROFILING', get_setting('SQL_DEBUG', False))
94+
if sql_profiling_enabled:
95+
if not settings.DEBUG:
96+
logger.warning("ANSIBLE_BASE_SQL_PROFILING is enabled, but DEBUG is False. No SQL queries will be logged or counted.")
97+
return self.get_response(request)
98+
99+
queries_before = len(connection.queries)
100+
response = self.get_response(request)
101+
q_times = [float(q['time']) for q in connection.queries[queries_before:]]
102+
response['X-API-Query-Count'] = len(q_times)
103+
response['X-API-Query-Time'] = '%0.3fs' % sum(q_times)
104+
else:
105+
response = self.get_response(request)
106+
107+
return response

test_app/tests/lib/middleware/test_profiling_middleware.py

Lines changed: 67 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
1+
import uuid
22
import os
33
import tempfile
44
from unittest.mock import patch
@@ -7,15 +7,24 @@
77
from django.test import TestCase, override_settings
88
from django.urls import path
99

10-
from ansible_base.lib.middleware.profiling.profile_request import ProfileRequestMiddleware
10+
from ansible_base.lib.middleware.profiling.profile_request import ProfileRequestMiddleware, SQLProfilingMiddleware
11+
from ansible_base.lib.utils.settings import get_setting
12+
from test_app.models import User
1113

1214
# A simple view for testing middleware
1315
def simple_view(request):
1416
return HttpResponse("OK")
1517

18+
# A view that performs a database query
19+
def db_view(request):
20+
# Create a user with a unique username to guarantee a query.
21+
User.objects.create(username=f"test-{uuid.uuid4()}")
22+
return HttpResponse("OK")
23+
1624
# Define URL patterns for the test
1725
urlpatterns = [
1826
path('test/', simple_view),
27+
path('test-db/', db_view),
1928
]
2029

2130
@override_settings(ROOT_URLCONF=__name__)
@@ -65,3 +74,59 @@ def test_profile_request_middleware_cprofile_disabled(self):
6574
self.assertNotIn('X-API-CProfile-File', response)
6675

6776

77+
@override_settings(
78+
ROOT_URLCONF=__name__,
79+
MIDDLEWARE=['ansible_base.lib.middleware.profiling.profile_request.SQLProfilingMiddleware']
80+
)
81+
class SQLProfilingMiddlewareTest(TestCase):
82+
def test_sql_profiling_disabled_by_default(self):
83+
"""
84+
Test that the SQLProfilingMiddleware does not add headers when disabled.
85+
"""
86+
response = self.client.get('/test-db/')
87+
self.assertNotIn('X-API-Query-Count', response)
88+
self.assertNotIn('X-API-Query-Time', response)
89+
90+
@override_settings(ANSIBLE_BASE_SQL_PROFILING=True, DEBUG=True)
91+
def test_sql_profiling_enabled_with_new_setting(self):
92+
"""
93+
Test that the SQLProfilingMiddleware adds headers when ANSIBLE_BASE_SQL_PROFILING is True.
94+
"""
95+
response = self.client.get('/test-db/')
96+
self.assertIn('X-API-Query-Count', response)
97+
self.assertGreaterEqual(int(response['X-API-Query-Count']), 1)
98+
self.assertIn('X-API-Query-Time', response)
99+
self.assertTrue(response['X-API-Query-Time'].endswith('s'))
100+
try:
101+
float(response['X-API-Query-Time'][:-1])
102+
except ValueError:
103+
self.fail("X-API-Query-Time value is not a valid float")
104+
105+
@override_settings(SQL_DEBUG=True, DEBUG=True)
106+
def test_sql_profiling_enabled_with_fallback_setting(self):
107+
"""
108+
Test that the SQLProfilingMiddleware adds headers when SQL_DEBUG is True as a fallback.
109+
"""
110+
response = self.client.get('/test-db/')
111+
self.assertIn('X-API-Query-Count', response)
112+
self.assertGreaterEqual(int(response['X-API-Query-Count']), 1)
113+
self.assertIn('X-API-Query-Time', response)
114+
self.assertTrue(response['X-API-Query-Time'].endswith('s'))
115+
try:
116+
float(response['X-API-Query-Time'][:-1])
117+
except ValueError:
118+
self.fail("X-API-Query-Time value is not a valid float")
119+
120+
@override_settings(ANSIBLE_BASE_SQL_PROFILING=True, DEBUG=False)
121+
def test_sql_profiling_logs_warning_if_debug_is_false(self):
122+
"""
123+
Test that the SQLProfilingMiddleware logs a warning and does not add headers
124+
if profiling is enabled but DEBUG is False.
125+
"""
126+
with self.assertLogs('ansible_base.lib.middleware.profiling.profile_request', level='WARNING') as cm:
127+
response = self.client.get('/test-db/')
128+
self.assertIn("ANSIBLE_BASE_SQL_PROFILING is enabled, but DEBUG is False", cm.output[0])
129+
self.assertNotIn('X-API-Query-Count', response)
130+
self.assertNotIn('X-API-Query-Time', response)
131+
132+

0 commit comments

Comments
 (0)