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

1import json 

2import logging 

3import threading 

4 

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 

10 

11from .core import QueryStats 

12 

13try: 

14 try: 

15 import psycopg 

16 except ImportError: 

17 import psycopg2 as psycopg 

18except ImportError: 

19 psycopg = None 

20 

21logger = logging.getLogger(__name__) 

22_local = threading.local() 

23 

24 

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 

34 

35 

36class QueryStatsMiddleware: 

37 def __init__(self, get_response): 

38 self.get_response = get_response 

39 

40 def __call__(self, request): 

41 if request.GET.get("querystats") == "disable": 

42 return self.get_response(request) 

43 

44 querystats = QueryStats( 

45 # Only want these if we're getting ready to show it 

46 include_tracebacks=request.GET.get("querystats") == "store" 

47 ) 

48 

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) 

53 

54 if settings.DEBUG or is_staff: 

55 # Persist it on the thread 

56 _local.querystats = querystats 

57 

58 with connection.execute_wrapper(_local.querystats): 

59 response = self.get_response(request) 

60 

61 if settings.DEBUG: 

62 # TODO logging settings 

63 logger.debug("Querystats: %s", _local.querystats) 

64 

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() 

68 

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")) 

74 

75 del _local.querystats 

76 

77 return response 

78 

79 else: 

80 return self.get_response(request) 

81 

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 

87 

88 return hasattr(request, "user") and request.user and request.user.is_staff