From e524db0215f70d4dab0a1f998fa9852ce217d690 Mon Sep 17 00:00:00 2001
From: Allen Short <ashort@mozilla.com>
Date: Fri, 9 Dec 2016 14:40:00 -0600
Subject: [PATCH] Measure query time with statsd

---
 redash/metrics/database.py | 124 ++++++++++---------------------------
 redash/models.py           |   1 +
 2 files changed, 32 insertions(+), 93 deletions(-)

diff --git a/redash/metrics/database.py b/redash/metrics/database.py
index a4ce6e312f..4cfb1ad83e 100644
--- a/redash/metrics/database.py
+++ b/redash/metrics/database.py
@@ -1,99 +1,37 @@
-from functools import wraps
 import time
 import logging
-from playhouse.gfk import Model
-from playhouse.postgres_ext import PostgresqlExtDatabase
-from redash import statsd_client
-
-metrics_logger = logging.getLogger("metrics")
-
-
-class MeteredPostgresqlExtDatabase(PostgresqlExtDatabase):
-    def __init__(self, *args, **kwargs):
-        self.query_count = 0
-        self.query_duration = 0
-        return super(MeteredPostgresqlExtDatabase, self).__init__(*args, **kwargs)
-
-    def execute_sql(self, *args, **kwargs):
-        start_time = time.time()
-
-        try:
-            result = super(MeteredPostgresqlExtDatabase, self).execute_sql(*args, **kwargs)
-            return result
-        finally:
-            self.query_count += 1
-            # TODO: there is a noticeable few ms discrepancy between the duration here and the one calculated in
-            # metered_execute. Need to think what to do about it.
-            duration = (time.time() - start_time) * 1000
-            self.query_duration += duration
-
-    def reset_metrics(self):
-        # TODO: instead of manually managing reset of metrics, we should store them in a LocalProxy based object, that
-        # is guaranteed to be "replaced" when the current request is done.
-        self.query_count = 0
-        self.query_duration = 0
-
-
-def patch_query_execute():
-    real_execute = peewee.Query._execute
-    real_clone = peewee.Query.clone
-
-    @wraps(real_execute)
-    def metered_execute(self, *args, **kwargs):
-        name = self.model_class.__name__
 
-        action = getattr(self, 'model_action', 'unknown')
+from sqlalchemy.engine import Engine
 
-        start_time = time.time()
-        try:
-            result = real_execute(self, *args, **kwargs)
-            return result
-        finally:
-            duration = (time.time() - start_time) * 1000
-            statsd_client.timing('db.{}.{}'.format(name, action), duration)
-            metrics_logger.debug("model=%s query=%s duration=%.2f", name, action, duration)
-
-    @wraps(real_clone)
-    def extended_clone(self):
-        cloned = real_clone(self)
-        setattr(cloned, 'model_action', getattr(self, 'model_action', 'unknown'))
-        return cloned
-
-    peewee.Query._execute = metered_execute
-    peewee.Query.clone = extended_clone
-
-
-class MeteredModel(Model):
-    @classmethod
-    def select(cls, *args, **kwargs):
-        return cls._execute_and_measure('select', args, kwargs)
-
-    @classmethod
-    def update(cls, *args, **kwargs):
-        return cls._execute_and_measure('update', args, kwargs)
-
-    @classmethod
-    def insert(cls, *args, **kwargs):
-        return cls._execute_and_measure('insert', args, kwargs)
-
-    @classmethod
-    def insert_many(cls, *args, **kwargs):
-        return cls._execute_and_measure('insert_many', args, kwargs)
-
-    @classmethod
-    def insert_from(cls, *args, **kwargs):
-        return cls._execute_and_measure('insert_from', args, kwargs)
-
-    @classmethod
-    def delete(cls, *args, **kwargs):
-        return cls._execute_and_measure('delete', args, kwargs)
+from redash import statsd_client
 
-    @classmethod
-    def raw(cls, *args, **kwargs):
-        return cls._execute_and_measure('raw', args, kwargs)
+metrics_logger = logging.getLogger("metrics")
 
-    @classmethod
-    def _execute_and_measure(cls, action, args, kwargs):
-        result = getattr(super(MeteredModel, cls), action)(*args, **kwargs)
-        setattr(result, 'model_action', action)
-        return result
+from sqlalchemy.orm.util import _ORMJoin
+from sqlalchemy.event import listens_for
+
+
+@listens_for(Engine, "before_execute")
+def before_execute(conn, elt, multiparams, params):
+    conn.info.setdefault('query_start_time', []).append(time.time())
+
+
+@listens_for(Engine, "after_execute")
+def after_execute(conn, elt, multiparams, params, result):
+    duration = time.time() - conn.info['query_start_time'].pop(-1)
+    action = elt.__class__.__name__
+
+    if action == 'Select':
+        t = elt.froms[0]
+        while isinstance(t, _ORMJoin):
+            t = t.left
+        name = t.name
+    elif action in ['Update', 'Insert', 'Delete']:
+        name = elt.table.name
+    else:
+        # create/drop tables, sqlalchemy internal schema queries, etc
+        return
+    statsd_client.timing('db.{}.{}'.format(name, action), duration)
+    metrics_logger.debug("model=%s query=%s duration=%.2f", name, action,
+                         duration)
+    return result
diff --git a/redash/models.py b/redash/models.py
index 542252aa7e..84f8a3e773 100644
--- a/redash/models.py
+++ b/redash/models.py
@@ -25,6 +25,7 @@
 from redash.query_runner import get_query_runner, get_configuration_schema_for_query_runner_type
 from redash.utils import generate_token, json_dumps
 from redash.utils.configuration import ConfigurationContainer
+from redash.metrics import database
 
 db = SQLAlchemy()
 Column = functools.partial(db.Column, nullable=False)