Coverage for /Users/davegaeddert/Developer/dropseed/plain/plain-staff/plain/staff/querystats/core.py: 67%

79 statements  

« prev     ^ index     » next       coverage.py v7.6.9, created at 2024-12-23 11:16 -0600

1import time 

2import traceback 

3from collections import Counter 

4 

5import sqlparse 

6 

7from plain.utils.functional import cached_property 

8 

9IGNORE_STACK_FILES = [ 

10 "threading", 

11 "socketserver", 

12 "wsgiref", 

13 "gunicorn", 

14 "whitenoise", 

15 "sentry_sdk", 

16 "querystats/core", 

17 "plain/template/base", 

18 "plain/utils/decorators", 

19 "plain/utils/deprecation", 

20 "plain/db", 

21 "plain/utils/functional", 

22 "plain/core/servers", 

23 "plain/core/handlers", 

24] 

25 

26 

27def pretty_print_sql(sql): 

28 return sqlparse.format(sql, reindent=True, keyword_case="upper") 

29 

30 

31def get_stack(): 

32 return "".join(tidy_stack(traceback.format_stack())) 

33 

34 

35def tidy_stack(stack): 

36 lines = [] 

37 

38 skip_next = False 

39 

40 for line in stack: 

41 if skip_next: 

42 skip_next = False 

43 continue 

44 

45 if line.startswith(' File "') and any( 

46 ignore in line for ignore in IGNORE_STACK_FILES 

47 ): 

48 skip_next = True 

49 continue 

50 

51 lines.append(line) 

52 

53 return lines 

54 

55 

56class QueryStats: 

57 def __init__(self, include_tracebacks): 

58 self.queries = [] 

59 self.include_tracebacks = include_tracebacks 

60 

61 def __str__(self): 

62 s = f"{self.num_queries} queries in {self.total_time_display}" 

63 if self.duplicate_queries: 

64 s += f" ({self.num_duplicate_queries} duplicates)" 

65 return s 

66 

67 def __call__(self, execute, sql, params, many, context): 

68 current_query = {"sql": sql, "params": params, "many": many} 

69 start = time.monotonic() 

70 

71 result = execute(sql, params, many, context) 

72 

73 if self.include_tracebacks: 

74 current_query["tb"] = get_stack() 

75 

76 # if many, then X times is len(params) 

77 

78 current_query["result"] = result 

79 

80 current_query["duration"] = time.monotonic() - start 

81 

82 self.queries.append(current_query) 

83 return result 

84 

85 @cached_property 

86 def total_time(self): 

87 return sum(q["duration"] for q in self.queries) 

88 

89 @staticmethod 

90 def get_time_display(seconds): 

91 if seconds < 0.01: 

92 return f"{seconds * 1000:.0f} ms" 

93 return f"{seconds:.2f} seconds" 

94 

95 @cached_property 

96 def total_time_display(self): 

97 return self.get_time_display(self.total_time) 

98 

99 @cached_property 

100 def num_queries(self): 

101 return len(self.queries) 

102 

103 # @cached_property 

104 # def models(self): 

105 # # parse table names from self.queries sql 

106 # table_names = [x for x in [q['sql'].split(' ')[2] for q in self.queries] if x] 

107 # models = connection.introspection.installed_models(table_names) 

108 # return models 

109 

110 @cached_property 

111 def duplicate_queries(self): 

112 sqls = [q["sql"] for q in self.queries] 

113 duplicates = {k: v for k, v in Counter(sqls).items() if v > 1} 

114 return duplicates 

115 

116 @cached_property 

117 def num_duplicate_queries(self): 

118 # Count the number of "excess" queries by getting how many there 

119 # are minus the initial one (and potentially only one required) 

120 return sum(self.duplicate_queries.values()) - len(self.duplicate_queries) 

121 

122 def as_summary_dict(self): 

123 return { 

124 "summary": str(self), 

125 "total_time": self.total_time, 

126 "num_queries": self.num_queries, 

127 "num_duplicate_queries": self.num_duplicate_queries, 

128 } 

129 

130 def as_context_dict(self): 

131 # If we don't create a dict, the instance of this class 

132 # is lost before we can use it in the template 

133 for query in self.queries: 

134 # Add some useful display info 

135 query["duration_display"] = self.get_time_display(query["duration"]) 

136 query["sql_display"] = pretty_print_sql(query["sql"]) 

137 duplicates = self.duplicate_queries.get(query["sql"], 0) 

138 if duplicates: 

139 query["duplicate_count"] = duplicates 

140 

141 summary = self.as_summary_dict() 

142 

143 return { 

144 **summary, 

145 "total_time_display": self.total_time_display, 

146 "queries": self.queries, 

147 } 

148 

149 def as_server_timing(self): 

150 duration = self.total_time * 1000 # put in ms 

151 duration = round(duration, 2) 

152 description = str(self) 

153 return f'querystats;dur={duration};desc="{description}"'