debugsqlshell.py 1.0 KB

123456789101112131415161718192021222324252627282930313233
  1. from time import perf_counter
  2. import sqlparse
  3. from django.core.management.commands.shell import Command
  4. from django.db import connection
  5. if connection.vendor == "postgresql":
  6. from django.db.backends.postgresql import base as base_module
  7. else:
  8. from django.db.backends import utils as base_module
  9. # 'debugsqlshell' is the same as the 'shell'.
  10. # Command is required to exist to be loaded via
  11. # django.core.managementload_command_class
  12. __all__ = ["Command", "PrintQueryWrapper"]
  13. class PrintQueryWrapper(base_module.CursorDebugWrapper):
  14. def execute(self, sql, params=()):
  15. start_time = perf_counter()
  16. try:
  17. return self.cursor.execute(sql, params)
  18. finally:
  19. raw_sql = self.db.ops.last_executed_query(self.cursor, sql, params)
  20. end_time = perf_counter()
  21. duration = (end_time - start_time) * 1000
  22. formatted_sql = sqlparse.format(raw_sql, reindent=True)
  23. print(f"{formatted_sql} [{duration:.2f}ms]")
  24. base_module.CursorDebugWrapper = PrintQueryWrapper