add a single middleware for apps to include that orchestrates things the right way

Author: kdelee

ansible_base/lib/middleware/observability.py

@@ -0,0 +1,27 @@
+"""
+A single middleware to provide a unified observability layer, ensuring that context,
+profiling, and SQL metrics are captured in the correct order.
+"""
+
+from .profiling.profile_request import _ProfileRequestMiddleware, _SQLProfilingMiddleware
+from .request_context import _TraceContextMiddleware
+
+
+class ObservabilityMiddleware:
+    """
+    A single entry point for observability middleware.
+
+    This middleware composes the trace context, request profiling, and SQL
+    profiling middleware in the correct order. Instead of listing all three
+    in your settings, you can now just add this one.
+    """
+
+    def __init__(self, get_response):
+        # Chain the middleware in the desired order. The request will flow
+        # from _TraceContextMiddleware -> _ProfileRequestMiddleware -> _SQLProfilingMiddleware.
+        handler = _SQLProfilingMiddleware(get_response)
+        handler = _ProfileRequestMiddleware(handler)
+        self.handler = _TraceContextMiddleware(handler)
+
+    def __call__(self, request):
+        return self.handler(request)

ansible_base/lib/middleware/profiling/README.md

@@ -1,27 +1,28 @@
-# Request Profiling
+# Request Profiling and Observability
 
-The `ProfileRequestMiddleware` and `DABProfiler` class provide a way to profile requests and other code in your Django application. This functionality is a generalization of the profiling tools found in AWX and can be used by any `django-ansible-base` consumer.
+The `ObservabilityMiddleware` provides a simple way to gain performance and debugging insights into your Django application. It acts as a single entry point for several underlying middleware components, ensuring they are always used in the correct order.
 
-## `ProfileRequestMiddleware`
+## `ObservabilityMiddleware`
 
-This middleware provides performance insights for API requests. To use it, add it to your `MIDDLEWARE` list in your Django settings. For the most accurate and reliable timing, it is recommended to add this middleware to the top of the `MIDDLEWARE` list.
+This single middleware bundles tracing, request profiling, and SQL query analysis. To use it, add it to the top of your `MIDDLEWARE` list in your Django settings.
 
 ```python
 # settings.py
 MIDDLEWARE = [
-    'ansible_base.lib.middleware.profiling.profile_request.ProfileRequestMiddleware',
+    'ansible_base.lib.middleware.observability.ObservabilityMiddleware',
     ...
 ]
 ```
 
 The middleware always adds the following headers to the response:
 
+*   `X-Request-ID`: A unique identifier for the request. If the incoming request includes an `X-Request-ID` header, that value will be used; otherwise, a new UUID will be generated.
 *   `X-API-Time`: The total time taken to process the request, in seconds.
-*   `X-API-Node`: The cluster host ID of the node that served the request. This header is only added if it is not already present in the response.
+*   `X-API-Node`: The cluster host ID of the node that served the request.
 
 ### cProfile Support
 
-When the `ANSIBLE_BASE_CPROFILE_REQUESTS` setting is enabled, the middleware will also perform a cProfile analysis for each request. The resulting `.prof` file is saved to a temporary directory on the node that served the request, and its path is returned in the `X-API-CProfile-File` response header.
+When the `ANSIBLE_BASE_CPROFILE_REQUESTS` setting is enabled, the middleware will also perform a cProfile analysis for each request. The resulting `.prof` file is saved to a temporary directory on the node that served the request, and its path is returned in the `X-API-CProfile-File` response header. The filename will include the request's `X-Request-ID`.
 
 To enable cProfile support, set the following in your Django settings:
 
@@ -30,32 +31,16 @@ To enable cProfile support, set the following in your Django settings:
 ANSIBLE_BASE_CPROFILE_REQUESTS = True
 ```
 
-## `SQLProfilingMiddleware`
+### SQL Profiling Support
 
-This middleware provides insights into the database queries executed during a request. When enabled, it adds the following headers to the response:
+When the `ANSIBLE_BASE_SQL_PROFILING` setting is enabled, the middleware provides insights into the database queries executed during a request. It adds the following headers to the response:
 
 *   `X-API-Query-Count`: The total number of database queries executed during the request.
 *   `X-API-Query-Time`: The total time spent on database queries, in seconds.
 
 It also injects contextual information as a comment into each SQL query, which is invaluable for debugging and tracing. For example:
 `/* trace_id=b71696ed-c483-408d-9740-2e7935b4f2d9, route=api/v2/users/{pk}/, origin=request */ SELECT ...`
 
-To use it, add both the `TraceContextMiddleware` and the `SQLProfilingMiddleware` to your `MIDDLEWARE` list in your Django settings. The `TraceContextMiddleware` should come before the `SQLProfilingMiddleware`.
-
-```python
-# settings.py
-MIDDLEWARE = [
-    ...
-    'ansible_base.lib.middleware.request_context.TraceContextMiddleware',
-    'ansible_base.lib.middleware.profiling.profile_request.SQLProfilingMiddleware',
-    ...
-]
-```
-
-### Enabling SQL Profiling
-
-The middleware is controlled by the `ANSIBLE_BASE_SQL_PROFILING` setting.
-
 To enable SQL profiling, set the following in your Django settings:
 
 ```python
@@ -65,7 +50,7 @@ ANSIBLE_BASE_SQL_PROFILING = True
 
 ## `DABProfiler`
 
-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.
+For profiling non-HTTP contexts, such as background tasks or gRPC services, the `DABProfiler` class can be used directly.
 
 The profiler's cProfile functionality is controlled by the `ANSIBLE_BASE_CPROFILE_REQUESTS` setting.
 
@@ -119,3 +104,4 @@ import pstats
 p = pstats.Stats('/path/to/your/profile.prof')
 p.sort_stats('cumulative').print_stats(10)
 ```
+

ansible_base/lib/middleware/profiling/profile_request.py

@@ -51,15 +51,15 @@ def stop(self, profile_id: Optional[Union[str, uuid.UUID]] = None):
         return elapsed, cprofile_filename
 
 
-class ProfileRequestMiddleware(threading.local):
+class _ProfileRequestMiddleware(threading.local):
     def __init__(self, get_response=None):
         self.get_response = get_response
         self.profiler = DABProfiler()
 
     def __call__(self, request):
         # Logic before the view (formerly process_request)
         self.profiler.start()
-        request_id = request.headers.get('X-Request-ID')
+        request_id = trace_id_var.get()
 
         # Call the next middleware or the view
         response = self.get_response(request)
@@ -111,7 +111,7 @@ def __call__(self, execute, sql, params, many, context):
             self.query_time += time.time() - start_time
 
 
-class SQLProfilingMiddleware:
+class _SQLProfilingMiddleware:
     def __init__(self, get_response):
         self.get_response = get_response
 
@@ -123,7 +123,7 @@ def __call__(self, request):
         if trace_id_var.get() is None:
             logger.warning(
                 "ANSIBLE_BASE_SQL_PROFILING is enabled, but the trace context is not set. "
-                "Please ensure that TraceContextMiddleware is included in your MIDDLEWARE settings before this middleware."
+                "Please use the ObservabilityMiddleware instead of including profiling middleware individually."
             )
 
         metrics = SQLQueryMetrics()

ansible_base/lib/middleware/request_context.py

@@ -3,14 +3,15 @@
 from ansible_base.lib.logging.context import origin_var, route_var, trace_id_var
 
 
-class TraceContextMiddleware:
+class _TraceContextMiddleware:
     def __init__(self, get_response):
         self.get_response = get_response
 
     def __call__(self, request):
         # Set the context for the request and store the tokens
         origin_token = origin_var.set('request')
-        trace_id = request.headers.get('X-Request-ID', str(uuid.uuid4()))
+        # .get is case-insensitive, but we'll use lowercase for consistency
+        trace_id = request.headers.get('x-request-id', str(uuid.uuid4()))
         trace_id_token = trace_id_var.set(trace_id)
 
         route_token = None
@@ -19,6 +20,7 @@ def __call__(self, request):
 
         try:
             response = self.get_response(request)
+            response['X-Request-ID'] = trace_id
         finally:
             # Reset the context variables to their previous state
             origin_var.reset(origin_token)