mirror of https://github.com/zulip/zulip.git
45 lines
1.5 KiB
Python
45 lines
1.5 KiB
Python
|
from __future__ import absolute_import
|
||
|
|
|
||
|
|
import time
|
||
|
|
from psycopg2.extensions import cursor, connection
|
||
|
|
|
||
|
|
# Similar to the tracking done in Django's CursorDebugWrapper, but done at the
|
||
|
|
# psycopg2 cursor level so it works with SQLAlchemy.
|
||
|
|
def wrapper_execute(self, action, sql, params=()):
|
||
|
|
start = time.time()
|
||
|
|
try:
|
||
|
|
return action(sql, params)
|
||
|
|
finally:
|
||
|
|
stop = time.time()
|
||
|
|
duration = stop - start
|
||
|
|
self.connection.queries.append({
|
||
|
|
'time': "%.3f" % duration,
|
||
|
|
})
|
||
|
|
|
||
|
|
class TimeTrackingCursor(cursor):
|
||
|
|
"""A psycopg2 cursor class that tracks the time spent executing queries."""
|
||
|
|
|
||
|
|
def execute(self, query, vars=None):
|
||
|
|
return wrapper_execute(self, super(TimeTrackingCursor, self).execute, query, vars)
|
||
|
|
|
||
|
|
def executemany(self, query, vars):
|
||
|
|
return wrapper_execute(self, super(TimeTrackingCursor, self).executemany, query, vars)
|
||
|
|
|
||
|
|
class TimeTrackingConnection(connection):
|
||
|
|
"""A psycopg2 connection class that uses TimeTrackingCursors."""
|
||
|
|
|
||
|
|
def __init__(self, *args, **kwargs):
|
||
|
|
self.queries = []
|
||
|
|
super(TimeTrackingConnection, self).__init__(*args, **kwargs)
|
||
|
|
|
||
|
|
def cursor(self, name=None):
|
||
|
|
if name is None:
|
||
|
|
return super(TimeTrackingConnection, self).cursor(cursor_factory=TimeTrackingCursor)
|
||
|
|
else:
|
||
|
|
return super(TimeTrackingConnection, self).cursor(name, cursor_factory=TimeTrackingCursor)
|
||
|
|
|
||
|
|
def reset_queries():
|
||
|
|
from django.db import connections
|
||
|
|
for conn in connections.all():
|
||
|
|
conn.connection.queries = []
|