Merge pull request getredash#1846 from getredash/athena-updates

Athena updates: bring back ability to disable annotations and disable formatter

redash/query_runner/athena.py

@@ -1,12 +1,13 @@
-import logging
 import json
+import logging
+import os
 
-from redash.utils import JSONEncoder
 from redash.query_runner import *
-
+from redash.settings import parse_boolean
+from redash.utils import JSONEncoder
 
 logger = logging.getLogger(__name__)
-
+ANNOTATE_QUERY = parse_boolean(os.environ.get('ATHENA_ANNOTATE_QUERY', 'true'))
 
 try:
     import pyathena
@@ -33,6 +34,11 @@
 }
 
 
+class SimpleFormatter(object):
+    def format(self, operation, parameters=None):
+        return operation
+
+
 class Athena(BaseQueryRunner):
     noop_query = 'SELECT 1'
 
@@ -83,6 +89,10 @@ def configuration_schema(cls):
     def enabled(cls):
         return enabled
 
+    @classmethod
+    def annotate_query(cls):
+        return ANNOTATE_QUERY
+
     @classmethod
     def type(cls):
         return "athena"
@@ -119,7 +129,8 @@ def run_query(self, query, user):
             aws_secret_access_key=self.configuration.get('aws_secret_key', None),
             schema_name=self.configuration.get('schema', 'default'),
             encryption_option=self.configuration.get('encryption_option', None),
-            kms_key=self.configuration.get('kms_key', None)).cursor()
+            kms_key=self.configuration.get('kms_key', None),
+            formatter=SimpleFormatter()).cursor()
 
         try:
             cursor.execute(query)