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/db",
 20    "plain/utils/functional",
 21    "plain/core/servers",
 22    "plain/core/handlers",
 23]
 24
 25
 26def pretty_print_sql(sql):
 27    return sqlparse.format(sql, reindent=True, keyword_case="upper")
 28
 29
 30def get_stack():
 31    return "".join(tidy_stack(traceback.format_stack()))
 32
 33
 34def tidy_stack(stack):
 35    lines = []
 36
 37    skip_next = False
 38
 39    for line in stack:
 40        if skip_next:
 41            skip_next = False
 42            continue
 43
 44        if line.startswith('  File "') and any(
 45            ignore in line for ignore in IGNORE_STACK_FILES
 46        ):
 47            skip_next = True
 48            continue
 49
 50        lines.append(line)
 51
 52    return lines
 53
 54
 55class QueryStats:
 56    def __init__(self, include_tracebacks):
 57        self.queries = []
 58        self.include_tracebacks = include_tracebacks
 59
 60    def __str__(self):
 61        s = f"{self.num_queries} queries in {self.total_time_display}"
 62        if self.duplicate_queries:
 63            s += f" ({self.num_duplicate_queries} duplicates)"
 64        return s
 65
 66    def __call__(self, execute, sql, params, many, context):
 67        current_query = {"sql": sql, "params": params, "many": many}
 68        start = time.monotonic()
 69
 70        result = execute(sql, params, many, context)
 71
 72        if self.include_tracebacks:
 73            current_query["tb"] = get_stack()
 74
 75        # if many, then X times is len(params)
 76
 77        current_query["result"] = result
 78
 79        current_query["duration"] = time.monotonic() - start
 80
 81        self.queries.append(current_query)
 82        return result
 83
 84    @cached_property
 85    def total_time(self):
 86        return sum(q["duration"] for q in self.queries)
 87
 88    @staticmethod
 89    def get_time_display(seconds):
 90        if seconds < 0.01:
 91            return f"{seconds * 1000:.0f} ms"
 92        return f"{seconds:.2f} seconds"
 93
 94    @cached_property
 95    def total_time_display(self):
 96        return self.get_time_display(self.total_time)
 97
 98    @cached_property
 99    def num_queries(self):
100        return len(self.queries)
101
102    # @cached_property
103    # def models(self):
104    #     # parse table names from self.queries sql
105    #     table_names = [x for x in [q['sql'].split(' ')[2] for q in self.queries] if x]
106    #     models = connection.introspection.installed_models(table_names)
107    #     return models
108
109    @cached_property
110    def duplicate_queries(self):
111        sqls = [q["sql"] for q in self.queries]
112        duplicates = {k: v for k, v in Counter(sqls).items() if v > 1}
113        return duplicates
114
115    @cached_property
116    def num_duplicate_queries(self):
117        # Count the number of "excess" queries by getting how many there
118        # are minus the initial one (and potentially only one required)
119        return sum(self.duplicate_queries.values()) - len(self.duplicate_queries)
120
121    def as_summary_dict(self):
122        return {
123            "summary": str(self),
124            "total_time": self.total_time,
125            "num_queries": self.num_queries,
126            "num_duplicate_queries": self.num_duplicate_queries,
127        }
128
129    def as_context_dict(self):
130        # If we don't create a dict, the instance of this class
131        # is lost before we can use it in the template
132        for query in self.queries:
133            # Add some useful display info
134            query["duration_display"] = self.get_time_display(query["duration"])
135            query["sql_display"] = pretty_print_sql(query["sql"])
136            duplicates = self.duplicate_queries.get(query["sql"], 0)
137            if duplicates:
138                query["duplicate_count"] = duplicates
139
140        summary = self.as_summary_dict()
141
142        return {
143            **summary,
144            "total_time_display": self.total_time_display,
145            "queries": self.queries,
146        }
147
148    def as_server_timing(self):
149        duration = self.total_time * 1000  # put in ms
150        duration = round(duration, 2)
151        description = str(self)
152        return f'querystats;dur={duration};desc="{description}"'