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}"'