From 796012bcfca730a6d6974cff1f12a4c6dc16edf9 Mon Sep 17 00:00:00 2001 From: Aman Pandey Date: Mon, 24 Jun 2024 19:24:26 +0530 Subject: [PATCH] asyn middleware --- debug_toolbar/middleware.py | 79 +++++++++++++++++++++++++++++++++++++ 1 file changed, 79 insertions(+) diff --git a/debug_toolbar/middleware.py b/debug_toolbar/middleware.py index 65b5282c5..05c8465f4 100644 --- a/debug_toolbar/middleware.py +++ b/debug_toolbar/middleware.py @@ -6,6 +6,7 @@ import socket from functools import lru_cache +from asgiref.sync import iscoroutinefunction, markcoroutinefunction from django.conf import settings from django.utils.module_loading import import_string @@ -64,10 +65,30 @@ class DebugToolbarMiddleware: on outgoing response. """ + sync_capable = True + async_capable = True + def __init__(self, get_response): + if get_response is None: + raise ValueError("get_response must be provided.") self.get_response = get_response + # If get_response is a coroutine function, turns us into async mode so + # a thread is not consumed during a whole request. + self.async_mode = iscoroutinefunction(self.get_response) + + if self.async_mode: + # Mark the class as async-capable, but do the actual switch inside + # __call__ to avoid swapping out dunder methods. + markcoroutinefunction(self) + + def _enable_panel_instrumentation(self, panels): + for panel in panels: + panel.enable_instrumentation() def __call__(self, request): + if self.async_mode: + return self.__acall__(request) + # Decide whether the toolbar is active for this request. show_toolbar = get_show_toolbar() if not show_toolbar(request) or DebugToolbar.is_toolbar_request(request): @@ -123,6 +144,64 @@ def __call__(self, request): response["Content-Length"] = len(response.content) return response + async def __acall__(self, request): + # Decide whether the toolbar is active for this request. + show_toolbar = get_show_toolbar() + if not show_toolbar(request) or DebugToolbar.is_toolbar_request(request): + response = await self.get_response(request) + return response + + toolbar = DebugToolbar(request, self.get_response) + + # Activate instrumentation ie. monkey-patch. + for panel in toolbar.enabled_panels: + panel.enable_instrumentation() + try: + # Run panels like Django middleware. + print(request) + response = await toolbar.process_request(request) + finally: + clear_stack_trace_caches() + # Deactivate instrumentation ie. monkey-unpatch. This must run + # regardless of the response. Keep 'return' clauses below. + for panel in reversed(toolbar.enabled_panels): + panel.disable_instrumentation() + + # Generate the stats for all requests when the toolbar is being shown, + # but not necessarily inserted. + for panel in reversed(toolbar.enabled_panels): + panel.generate_stats(request, response) + panel.generate_server_timing(request, response) + + # Always render the toolbar for the history panel, even if it is not + # included in the response. + rendered = toolbar.render_toolbar() + + for header, value in self.get_headers(request, toolbar.enabled_panels).items(): + response.headers[header] = value + + # Check for responses where the toolbar can't be inserted. + content_encoding = response.get("Content-Encoding", "") + content_type = response.get("Content-Type", "").split(";")[0] + if ( + getattr(response, "streaming", False) + or content_encoding != "" + or content_type not in _HTML_TYPES + ): + return response + + # Insert the toolbar in the response. + content = response.content.decode(response.charset) + insert_before = dt_settings.get_config()["INSERT_BEFORE"] + pattern = re.escape(insert_before) + bits = re.split(pattern, content, flags=re.IGNORECASE) + if len(bits) > 1: + bits[-2] += rendered + response.content = insert_before.join(bits) + if "Content-Length" in response: + response["Content-Length"] = len(response.content) + return response + @staticmethod def get_headers(request, panels): headers = {}