Do query time tracking at the psycopg2 level instead of the Django level

This allows us to track the query time of SQLAlchemy and raw queries.

(imported from commit 818a4ee41786ffc57b80d7ed1cfba075f29b6ee5)

zerver/lib/db.py

@@ -0,0 +1,44 @@
+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 = []

zerver/middleware.py

@@ -149,15 +149,16 @@ def write_log_line(log_data, path, method, remote_ip, email, client_name,
 
     # Get the amount of time spent doing database queries
     db_time_output = ""
-    if len(connection.queries) > 0:
-        query_time = sum(float(query.get('time', 0)) for query in connection.queries)
+    queries = connection.connection.queries if connection.connection is not None else []
+    if len(queries) > 0:
+        query_time = sum(float(query.get('time', 0)) for query in queries)
         db_time_output = " (db: %s/%sq)" % (format_timedelta(query_time),
-                                            len(connection.queries))
+                                            len(queries))
 
         if not suppress_statsd:
             # Log ms, db ms, and num queries to statsd
             statsd.timing("%s.dbtime" % (statsd_path,), timedelta_ms(query_time))
-            statsd.incr("%s.dbq" % (statsd_path, ), len(connection.queries))
+            statsd.incr("%s.dbq" % (statsd_path, ), len(queries))
             statsd.timing("%s.total" % (statsd_path,), timedelta_ms(time_delta))
 
     if 'extra' in log_data:
@@ -195,7 +196,8 @@ class LogRequests(object):
     def process_request(self, request):
         request._log_data = dict()
         record_request_start_data(request._log_data)
-        connection.queries = []
+        if connection.connection is not None:
+            connection.connection.queries = []
 
     def process_view(self, request, view_func, args, kwargs):
         # process_request was already run; we save the initialization
@@ -206,7 +208,8 @@ class LogRequests(object):
         # And then completely reset our tracking to only cover work
         # done as part of this request
         record_request_start_data(request._log_data)
-        connection.queries = []
+        if connection.connection is not None:
+            connection.connection.queries = []
 
     def process_response(self, request, response):
         # The reverse proxy might have sent us the real external IP
@@ -253,30 +256,6 @@ def csrf_failure(request, reason=""):
     else:
         return html_csrf_failure(request, reason)
 
-# Monkeypatch in time tracking to the Django non-debug cursor
-# Code comes from CursorDebugWrapper
-def wrapper_execute(self, action, sql, params=()):
-    start = time.time()
-    try:
-        return action(sql, params)
-    finally:
-        stop = time.time()
-        duration = stop - start
-        self.db.queries.append({
-                'time': "%.3f" % duration,
-                })
-
-from django.db.backends.util import CursorWrapper
-def cursor_execute(self, sql, params=()):
-    return wrapper_execute(self, self.cursor.execute, sql, params)
-def cursor_executemany(self, sql, params=()):
-    return wrapper_execute(self, self.cursor.executemany, sql, params)
-
-if not settings.DEBUG:
-    # If settings.DEBUG, the default cursor will do the appropriate logging already
-    CursorWrapper.execute = cursor_execute
-    CursorWrapper.executemany = cursor_executemany
-
 class RateLimitMiddleware(object):
     def process_response(self, request, response):
         if not settings.RATE_LIMITING:

zerver/tests.py

@@ -6,7 +6,6 @@ from django.test.simple import DjangoTestSuiteRunner
 from django.utils.timezone import now
 from django.core.exceptions import ValidationError
 from django.db.models import Q
-from django.db.backends.util import CursorDebugWrapper
 from guardian.shortcuts import assign_perm, remove_perm
 
 from zilencer.models import Deployment
@@ -33,6 +32,7 @@ from zerver.lib.rate_limiter import clear_user_history
 from zerver.lib.alert_words import alert_words_in_realm, user_alert_words, \
     add_user_alert_words, remove_user_alert_words
 from zerver.lib.digest import send_digest_email
+from zerver.lib.db import TimeTrackingCursor
 from zerver.forms import not_mit_mailing_list
 from zerver.lib.validator import check_string, check_list, check_dict, \
     check_bool, check_int
@@ -114,25 +114,21 @@ def queries_captured():
                     'time': "%.3f" % duration,
                     })
 
-    old_settings = settings.DEBUG
-    settings.DEBUG = True
-
-    old_execute = CursorDebugWrapper.execute
-    old_executemany = CursorDebugWrapper.executemany
+    old_execute = TimeTrackingCursor.execute
+    old_executemany = TimeTrackingCursor.executemany
 
     def cursor_execute(self, sql, params=()):
-        return wrapper_execute(self, self.cursor.execute, sql, params)
-    CursorDebugWrapper.execute = cursor_execute
+        return wrapper_execute(self, super(TimeTrackingCursor, self).execute, sql, params)
+    TimeTrackingCursor.execute = cursor_execute
 
     def cursor_executemany(self, sql, params=()):
-        return wrapper_execute(self, self.cursor.executemany, sql, params)
-    CursorDebugWrapper.executemany = cursor_executemany
+        return wrapper_execute(self, super(TimeTrackingCursor, self).executemany, sql, params)
+    TimeTrackingCursor.executemany = cursor_executemany
 
     yield queries
 
-    settings.DEBUG = old_settings
-    CursorDebugWrapper.execute = old_execute
-    CursorDebugWrapper.executemany = old_executemany
+    TimeTrackingCursor.execute = old_execute
+    TimeTrackingCursor.executemany = old_executemany
 
 
 def bail(msg):

zerver/worker/queue_processors.py

@@ -18,7 +18,7 @@ from zerver.lib.digest import handle_digest_email
 from zerver.decorator import JsonableError
 from zerver.lib.socket import req_redis_key
 from confirmation.models import Confirmation
-from django.db import reset_queries
+from zerver.lib.db import reset_queries
 from django.core.mail import EmailMessage
 
 import os

zproject/settings.py

@@ -9,6 +9,7 @@ import sys
 import ConfigParser
 
 from zerver.openid import openid_failure_handler
+from zerver.lib.db import TimeTrackingConnection
 
 config_file = ConfigParser.RawConfigParser()
 config_file.read("/etc/zulip/zulip.conf")
@@ -62,6 +63,7 @@ DATABASES = {"default": {
     'OPTIONS': {
         'sslmode': 'verify-full',
         'autocommit': True,
+        'connection_factory': TimeTrackingConnection
         },
     },
 }
@@ -72,6 +74,7 @@ if ENTERPRISE:
             'HOST': '',
             'OPTIONS': {
                 'autocommit': True,
+                'connection_factory': TimeTrackingConnection
             }
             })
 elif not DEPLOYED:
@@ -80,6 +83,7 @@ elif not DEPLOYED:
             'HOST': 'localhost',
             'OPTIONS': {
                 'autocommit': True,
+                'connection_factory': TimeTrackingConnection
             }
             })
     INTERNAL_ZULIP_USERS = []

zproject/test_settings.py

@@ -8,7 +8,7 @@ DATABASES["default"] = {"NAME": "zulip_test",
                         "SCHEMA": "zulip",
                         "ENGINE": "django.db.backends.postgresql_psycopg2",
                         "TEST_NAME": "django_zulip_tests",
-                        "OPTIONS": { },}
+                        "OPTIONS": {"connection_factory": TimeTrackingConnection },}
 
 
 if "TORNADO_SERVER" in os.environ: