Remove tree validations and introduce ParameterizedQuery

redash/handlers/embed.py

@@ -1,66 +1,14 @@
 from __future__ import absolute_import
-import logging
-import time
 
 from flask import request
 
 from .authentication import current_org
 from flask_login import current_user, login_required
-from flask_restful import abort
-from redash import models, utils
+from redash import models
 from redash.handlers import routes
 from redash.handlers.base import (get_object_or_404, org_scoped_rule,
                                   record_event)
-from redash.utils import find_missing_params
 from redash.handlers.static import render_index
-from redash.utils import gen_query_hash, mustache_render
-
-
-#
-# Run a parameterized query synchronously and return the result
-# DISCLAIMER: Temporary solution to support parameters in queries. Should be
-#             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, max_age=0):
-    missing_params = find_missing_params(query_text, parameter_values)
-    if missing_params:
-        raise Exception('Missing parameter value for: {}'.format(", ".join(missing_params)))
-
-    query_text = mustache_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, current_user)
-
-        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())
-
-            models.db.session.commit()
-        return data
-    except Exception:
-        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
 
 
 @routes.route(org_scoped_rule('/embed/query/<query_id>/visualization/<visualization_id>'), methods=['GET'])

redash/handlers/query_results.py

@@ -10,90 +10,59 @@
                                 require_permission, view_only)
 from redash.tasks import QueryTask, record_event
 from redash.tasks.queries import enqueue_query
-from redash.utils import (collect_parameters_from_request, find_missing_params, gen_query_hash, json_dumps, utcnow)
-from redash.utils.sql_query import SQLInjectionError, SQLQuery
+from redash.utils import (collect_parameters_from_request, gen_query_hash, json_dumps, utcnow)
+from redash.utils.parameterized_query import SQLInjectionError
 
 
 def error_response(message):
     return {'job': {'status': 4, 'error': message}}, 400
 
 
-def apply_parameters(template, parameters, data_source):
-    query = SQLQuery(template).apply(parameters)
-
-    # for now we only log `SQLInjectionError` to detect false positives
-    try:
-        text = query.text
-    except SQLInjectionError:
-        record_event({
-            'action': 'sql_injection',
-            'object_type': 'query',
-            'query': template,
-            'parameters': parameters,
-            'timestamp': time.time(),
-            'org_id': data_source.org_id
-        })
-    except Exception as e:
-        logging.info(u"Failed applying parameters for query %s: %s", gen_query_hash(query.query), e.message)
-    finally:
-        text = query.query
-
-    return text
-
-
 #
 # Run a parameterized query synchronously and return the result
 # DISCLAIMER: Temporary solution to support parameters in queries. Should be
 #             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, max_age=0):
-    missing_params = find_missing_params(query_text, parameter_values)
-    if missing_params:
-        raise Exception('Missing parameter value for: {}'.format(", ".join(missing_params)))
-
-    query_text = apply_parameters(query_text, parameter_values, data_source)
+    parameterized_query_class = data_source.query_runner.parameterized_query_class
+    parameterized_query = parameterized_query_class(query_text).apply(parameter_values)
 
-    if max_age <= 0:
-        query_result = None
-    else:
-        query_result = models.QueryResult.get_latest(data_source, query_text, max_age)
+    if parameterized_query.missing_params:
+        raise Exception('Missing parameter value for: {}'.format(", ".join(parameterized_query.missing_params)))
 
+    query_text = parameterized_query.query
     query_hash = gen_query_hash(query_text)
 
-    if query_result:
+    if max_age:
+        query_result = models.QueryResult.get_latest(data_source, query_text, max_age)
         logging.info("Returning cached result for query %s" % query_hash)
         return query_result
-
-    try:
-        started_at = time.time()
-        data, error = data_source.query_runner.run_query(query_text, current_user)
-
-        if error:
-            logging.info('got bak error')
-            logging.info(error)
+    else:
+        try:
+            started_at = time.time()
+            data, error = data_source.query_runner.run_query(query_text, current_user)
+
+            if error:
+                logging.info('got bak error')
+                logging.info(error)
+                return None
+
+            run_time = time.time() - started_at
+            query_result, _ = models.QueryResult.store_result(data_source.org_id, data_source, query_hash,
+                                                              query_text, data, run_time, utcnow())
+
+            models.db.session.commit()
+            return query_result
+        except Exception:
+            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
 
-        run_time = time.time() - started_at
-        query_result, updated_query_ids = models.QueryResult.store_result(data_source.org_id, data_source,
-                                                                              query_hash, query_text, data,
-                                                                              run_time, utcnow())
-
-        models.db.session.commit()
-        return query_result
-    except Exception as e:
-        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
-
 
 def run_query(data_source, parameter_values, query_text, query_id, max_age=0):
-    missing_params = find_missing_params(query_text, parameter_values)
-    if missing_params:
-        return error_response(u'Missing parameter value for: {}'.format(u", ".join(missing_params)))
-
     if data_source.paused:
         if data_source.pause_reason:
             message = '{} is paused ({}). Please try later.'.format(data_source.name, data_source.pause_reason)
@@ -102,17 +71,37 @@ def run_query(data_source, parameter_values, query_text, query_id, max_age=0):
 
         return error_response(message)
 
-    query_text = apply_parameters(query_text, parameter_values, data_source)
+    parameterized_query_class = data_source.query_runner.parameterized_query_class
+    parameterized_query = parameterized_query_class(query_text).apply(parameter_values)
 
-    if max_age == 0:
-        query_result = None
-    else:
-        query_result = models.QueryResult.get_latest(data_source, query_text, max_age)
+    if parameterized_query.missing_params:
+        return error_response(u'Missing parameter value for: {}'.format(u", ".join(parameterized_query.missing_params)))
 
-    if query_result:
+    # for now we only log `SQLInjectionError` to detect false positives
+    try:
+        query_text = parameterized_query.text
+    except SQLInjectionError:
+        record_event({
+            'action': 'sql_injection',
+            'object_type': 'query',
+            'query': query_text,
+            'parameters': parameter_values,
+            'timestamp': time.time(),
+            'org_id': data_source.org_id
+        })
+        query_text = parameterized_query.query
+    except Exception as e:
+        logging.info(u"Failed applying parameters for query %s: %s", gen_query_hash(parameterized_query.query), e.message)
+        query_text = parameterized_query.query
+
+    if max_age:
+        query_result = models.QueryResult.get_latest(data_source, query_text, max_age)
         return {'query_result': query_result.to_dict()}
     else:
-        job = enqueue_query(query_text, data_source, current_user.id, metadata={"Username": current_user.email, "Query ID": query_id})
+        job = enqueue_query(query_text, data_source, current_user.id, metadata={
+            "Username": current_user.email,
+            "Query ID": query_id
+        })
         return {'job': job.to_dict()}
 
 

redash/query_runner/__init__.py

@@ -3,6 +3,7 @@
 
 from redash import settings
 from redash.utils import json_loads
+from redash.utils.parameterized_query import ParameterizedSqlQuery, ParameterizedQuery
 
 logger = logging.getLogger(__name__)
 
@@ -115,6 +116,13 @@ def _run_query_internal(self, query):
             raise Exception("Failed running query [%s]." % query)
         return json_loads(results)['rows']
 
+    @property
+    def parameterized_query_class(self):
+        if self.syntax == 'sql':
+            return ParameterizedSqlQuery
+        else:
+            return ParameterizedQuery
+
     @classmethod
     def to_dict(cls):
         return {

redash/utils/__init__.py

@@ -15,7 +15,7 @@
 import pystache
 import pytz
 import simplejson
-from funcy import distinct, select_values
+from funcy import select_values
 from redash import settings
 from sqlalchemy.orm.query import Query
 
@@ -153,41 +153,6 @@ def writerows(self, rows):
             self.writerow(row)
 
 
-def _collect_key_names(nodes):
-    keys = []
-    for node in nodes._parse_tree:
-        if isinstance(node, pystache.parser._EscapeNode):
-            keys.append(node.key)
-        elif isinstance(node, pystache.parser._SectionNode):
-            keys.append(node.key)
-            keys.extend(_collect_key_names(node.parsed))
-
-    return distinct(keys)
-
-
-def collect_query_parameters(query):
-    nodes = pystache.parse(query)
-    keys = _collect_key_names(nodes)
-    return keys
-
-
-def parameter_names(parameter_values):
-    names = []
-    for key, value in parameter_values.iteritems():
-        if isinstance(value, dict):
-            for inner_key in value.keys():
-                names.append(u'{}.{}'.format(key, inner_key))
-        else:
-            names.append(key)
-
-    return names
-
-
-def find_missing_params(query_text, parameter_values):
-    query_parameters = set(collect_query_parameters(query_text))
-    return set(query_parameters) - set(parameter_names(parameter_values))
-
-
 def collect_parameters_from_request(args):
     parameters = {}
 

redash/utils/sql_query.py → redash/utils/parameterized_query.py

@@ -1,7 +1,9 @@
 import re
 
 import sqlparse
+import pystache
 from redash.utils import mustache_render
+from funcy import distinct
 
 
 def _replace_params(template):
@@ -44,15 +46,61 @@ def _same_type(a, b):
         return _equivalent_leaves(a, b)
 
 
-class SQLQuery(object):
+def _collect_key_names(nodes):
+    keys = []
+    for node in nodes._parse_tree:
+        if isinstance(node, pystache.parser._EscapeNode):
+            keys.append(node.key)
+        elif isinstance(node, pystache.parser._SectionNode):
+            keys.append(node.key)
+            keys.extend(_collect_key_names(node.parsed))
+
+    return distinct(keys)
+
+
+def _collect_query_parameters(query):
+    nodes = pystache.parse(query)
+    keys = _collect_key_names(nodes)
+    return keys
+
+
+def _parameter_names(parameter_values):
+    names = []
+    for key, value in parameter_values.iteritems():
+        if isinstance(value, dict):
+            for inner_key in value.keys():
+                names.append(u'{}.{}'.format(key, inner_key))
+        else:
+            names.append(key)
+
+    return names
+
+
+class ParameterizedQuery(object):
     def __init__(self, template):
         self.template = template
         self.query = template
+        self.parameters = {}
 
     def apply(self, parameters):
-        self.query = mustache_render(self.template, parameters)
+        self.parameters.update(parameters)
+        self.query = mustache_render(self.template, self.parameters)
         return self
 
+    @property
+    def missing_params(self):
+        query_parameters = set(_collect_query_parameters(self.template))
+        return set(query_parameters) - set(_parameter_names(self.parameters))
+
+    @property
+    def text(self):
+        return self.query
+
+
+class ParameterizedSqlQuery(ParameterizedQuery):
+    def __init__(self, template):
+        ParameterizedQuery.__init__(self, template)
+
     def is_safe(self):
         template_tree = sqlparse.parse(_replace_params(self.template))
         query_tree = sqlparse.parse(self.query)
@@ -64,7 +112,7 @@ def text(self):
         if not self.is_safe():
             raise SQLInjectionError()
         else:
-            return self.query
+            return super(ParameterizedSqlQuery, self).text
 
 
 class SQLInjectionError(Exception):