Merge pull request #1091 from whummer/feature/cache_embeds

Add caching for queries used in embeds

redash/handlers/embed.py

@@ -1,14 +1,16 @@
 import json
 import pystache
+import time
+import logging
 
 from funcy import project
 from flask import render_template, request
 from flask_login import login_required, current_user
 from flask_restful import abort
 
-from redash import models, settings
+from redash import models, settings, utils
 from redash import serializers
-from redash.utils import json_dumps, collect_parameters_from_request
+from redash.utils import json_dumps, collect_parameters_from_request, gen_query_hash
 from redash.handlers import routes
 from redash.handlers.base import org_scoped_rule, record_event
 from redash.handlers.query_results import collect_query_parameters
@@ -21,7 +23,7 @@ from authentication import current_org
 #             removed once we refactor the query results API endpoints and handling
 #             on the client side. Please don't reuse in other API handlers.
 #
-def run_query_sync(data_source, parameter_values, query_text):
+def run_query_sync(data_source, parameter_values, query_text, max_age=0):
     query_parameters = set(collect_query_parameters(query_text))
     missing_params = set(query_parameters) - set(parameter_values.keys())
     if missing_params:
@@ -30,13 +32,36 @@ def run_query_sync(data_source, parameter_values, query_text):
     if query_parameters:
         query_text = pystache.render(query_text, parameter_values)
 
+    if max_age <= 0:
+        query_result = None
+    else:
+        query_result = models.QueryResult.get_latest(data_source, query_text, max_age)
+
+    query_hash = gen_query_hash(query_text)
+
+    if query_result:
+        logging.info("Returning cached result for query %s" % query_hash)
+        return query_result.data
+
     try:
+        started_at = time.time()
         data, error = data_source.query_runner.run_query(query_text)
+
         if error:
             return None
+        # update cache
+        if max_age > 0:
+            run_time = time.time() - started_at
+            query_result, updated_query_ids = models.QueryResult.store_result(data_source.org_id, data_source.id,
+                                                                                  query_hash, query_text, data,
+                                                                                  run_time, utils.utcnow())
+
         return data
     except Exception, e:
-        abort(503, message="Unable to get result from the database.")
+        if max_age > 0:
+            abort(404, message="Unable to get result from the database, and no cached query result found.")
+        else:
+            abort(503, message="Unable to get result from the database.")
         return None
 
 
@@ -59,7 +84,8 @@ def embed(query_id, visualization_id, org_slug=None):
             # WARNING: Note that the external query parameters
             #          are a potential risk of SQL injections.
             #
-            results = run_query_sync(query.data_source, parameter_values, query.query)
+            max_age = int(request.args.get('maxAge', 0))
+            results = run_query_sync(query.data_source, parameter_values, query.query, max_age=max_age)
             if results is None:
                 abort(400, message="Unable to get results for this query")
             else:

tests/handlers/test_embed.py

@@ -16,20 +16,28 @@ class TestEmbedVisualization(BaseTestCase):
         # set configuration
         settings.ALLOW_PARAMETERS_IN_EMBEDS = True
 
-        vis = self.factory.create_visualization_with_params()
-        param1_name = "param1"
-        param1_value = "12345"
+        try:
+            vis = self.factory.create_visualization_with_params()
+            param1_name = "param1"
+            param1_value = "12345"
 
-        res = self.make_request("get", "/embed/query/{}/visualization/{}?p_{}={}".format(vis.query.id, vis.id, param1_name, param1_value), is_json=False)
+            res = self.make_request("get", "/embed/query/{}/visualization/{}?p_{}={}".format(vis.query.id, vis.id, param1_name, param1_value), is_json=False)
 
-        # reset configuration
-        settings.ALLOW_PARAMETERS_IN_EMBEDS = previous
+            # Currently we are expecting a 503 error which indicates that
+            # the database is unavailable. This ensures that the code in embed.py
+            # reaches the point where a DB query is made, where we then fail
+            # intentionally (because DB connection is not available in the tests).
+            self.assertEqual(res.status_code, 503)
 
-        # Currently we are expecting a 503 error which indicates that
-        # the database is unavailable. This ensures that the code in embed.py
-        # reaches the point where a DB query is made, where we then fail
-        # intentionally (because DB connection is not available in the tests).
-        self.assertEqual(res.status_code, 503)
+            # run embed query with maxAge to test caching
+            res = self.make_request("get", "/embed/query/{}/visualization/{}?p_{}={}&maxAge=60".format(vis.query.id, vis.id, param1_name, param1_value), is_json=False)
+            # If the 'maxAge' parameter is set and the query fails (because DB connection
+            # is not available in the tests), we're expecting a 404 error here.
+            self.assertEqual(res.status_code, 404)
+
+        finally:
+            # reset configuration
+            settings.ALLOW_PARAMETERS_IN_EMBEDS = previous
 
 
 class TestPublicDashboard(BaseTestCase):