Coverage for /Users/davegaeddert/Developer/dropseed/plain/plain-staff/plain/staff/querystats/middleware.py: 79%
53 statements
« prev ^ index » next coverage.py v7.6.9, created at 2024-12-23 11:16 -0600
« prev ^ index » next coverage.py v7.6.9, created at 2024-12-23 11:16 -0600
1import json
2import logging
3import threading
5from plain.http import ResponseRedirect
6from plain.json import PlainJSONEncoder
7from plain.models import connection
8from plain.runtime import settings
9from plain.urls import reverse
11from .core import QueryStats
13try:
14 try:
15 import psycopg
16 except ImportError:
17 import psycopg2 as psycopg
18except ImportError:
19 psycopg = None
21logger = logging.getLogger(__name__)
22_local = threading.local()
25class QueryStatsJSONEncoder(PlainJSONEncoder):
26 def default(self, obj):
27 try:
28 return super().default(obj)
29 except TypeError:
30 if psycopg and isinstance(obj, psycopg._json.Json):
31 return obj.adapted
32 else:
33 raise
36class QueryStatsMiddleware:
37 def __init__(self, get_response):
38 self.get_response = get_response
40 def __call__(self, request):
41 if request.GET.get("querystats") == "disable":
42 return self.get_response(request)
44 querystats = QueryStats(
45 # Only want these if we're getting ready to show it
46 include_tracebacks=request.GET.get("querystats") == "store"
47 )
49 with connection.execute_wrapper(querystats):
50 # Have to wrap this first call so it is included in the querystats,
51 # but we don't have to wrap everything else unless we are staff or debug
52 is_staff = self.is_staff_request(request)
54 if settings.DEBUG or is_staff:
55 # Persist it on the thread
56 _local.querystats = querystats
58 with connection.execute_wrapper(_local.querystats):
59 response = self.get_response(request)
61 if settings.DEBUG:
62 # TODO logging settings
63 logger.debug("Querystats: %s", _local.querystats)
65 # Make current querystats available on the current page
66 # by using the server timing API which can be parsed client-side
67 response["Server-Timing"] = _local.querystats.as_server_timing()
69 if request.GET.get("querystats") == "store":
70 request.session["querystats"] = json.dumps(
71 _local.querystats.as_context_dict(), cls=QueryStatsJSONEncoder
72 )
73 return ResponseRedirect(reverse("querystats:querystats"))
75 del _local.querystats
77 return response
79 else:
80 return self.get_response(request)
82 @staticmethod
83 def is_staff_request(request):
84 if getattr(request, "impersonator", None):
85 # Support for impersonation (still want the real staff user to see the querystats)
86 return request.impersonator and request.impersonator.is_staff
88 return hasattr(request, "user") and request.user and request.user.is_staff