added and implemented query restrictions

redash/query_runner/snowflake.py

@@ -15,6 +15,7 @@
     TYPE_BOOLEAN,
 )
 from redash.utils import json_dumps, json_loads
+import re
 
 TYPES_MAP = {
     0: TYPE_INTEGER,
@@ -30,6 +31,31 @@
 }
 
 
+def _query_restrictions(query):
+    query_without_comments = ''
+    for line in query.split('\n'):
+        line = line.strip()
+        if line.find('--') != -1:
+            line = line[:line.find('--')]
+        query_without_comments += ' ' + line  # creates one line query
+    query = query_without_comments
+    query = query.lower()
+    # replace multiple spaces with one space
+    query = re.sub(' +', ' ', query)
+    # get rid of prefix like bigbrain. or final.
+    query = re.sub('bigbrain.', '', re.sub('final.', '', re.sub('raw.', '', query)))
+    occurrences = re.findall(" from events ", query) + re.findall(" join events ", query)
+    # print("num of occurrences : ", len(occurrences))
+    if len(occurrences) > 1:
+        return False, f'Querying events table multiple times is forbidden.The query contains {len(occurrences)} occurrences of the table events.'
+
+    if occurrences:
+        if query.find("created_at") + query.find("ingestion_time") == -2:
+            return False, 'Querying events table should always be with time constraint (by created_at for ' \
+                          'FINAL.events & ingestion_time for RAW.events) '
+    return True, None
+
+
 class Snowflake(BaseQueryRunner):
     noop_query = "SELECT 1"
 
@@ -128,13 +154,18 @@ def run_query(self, query, user, query_id=None):
         connection = self._get_connection()
         cursor = connection.cursor()
 
+        passed, message = _query_restrictions(query)
+
+        if not passed:
+            return {}, message
+
         try:
             cursor.execute("USE WAREHOUSE {}".format(self.configuration["warehouse"]))
             cursor.execute("USE {}".format(self.configuration["database"]))
 
             user_id = "redash" if user is None else user.email
             query_id = str(query_id) if query_id else ''
-            query += "-- REDASH USER: " + user_id + " QUERY ID: " + query_id
+            query += '-- {"REDASH USER": "' + user_id + '"  , "QUERY ID": "' + query_id + '"}'
 
             cursor.execute(query)