Move annotation logic into Query Runner

redash/query_runner/__init__.py

@@ -56,6 +56,7 @@ class NotSupported(Exception):
 
 class BaseQueryRunner(object):
     deprecated = False
+    should_annotate_query = True
     noop_query = None
 
     def __init__(self, configuration):
@@ -74,14 +75,18 @@ def type(cls):
     def enabled(cls):
         return True
 
-    @classmethod
-    def annotate_query(cls):
-        return True
-
     @classmethod
     def configuration_schema(cls):
         return {}
 
+    def annotate_query(self, query, metadata):
+        if not self.should_annotate_query:
+            return query
+
+        annotation = u", ".join([u"{}: {}".format(k, v) for k, v in metadata.iteritems()])
+        annotated_query = u"/* {} */ {}".format(annotation, query)
+        return annotated_query
+
     def test_connection(self):
         if self.noop_query is None:
             raise NotImplementedError()
@@ -150,6 +155,7 @@ def _get_tables_stats(self, tables_dict):
 
 
 class BaseHTTPQueryRunner(BaseQueryRunner):
+    should_annotate_query = False
     response_error = "Endpoint returned unexpected status code"
     requires_authentication = False
     requires_url = True

redash/query_runner/athena.py

@@ -132,9 +132,10 @@ def configuration_schema(cls):
     def enabled(cls):
         return enabled
 
-    @classmethod
-    def annotate_query(cls):
-        return ANNOTATE_QUERY
+    def annotate_query(self, query, metadata):
+        if ANNOTATE_QUERY:
+            return super(Athena, self).annotate_query(query, metadata)
+        return query
 
     @classmethod
     def type(cls):

redash/query_runner/azure_kusto.py

@@ -26,6 +26,7 @@
 
 
 class AzureKusto(BaseQueryRunner):
+    should_annotate_query = False
     noop_query = "let noop = datatable (Noop:string)[1]; noop"
 
     def __init__(self, configuration):
@@ -71,10 +72,6 @@ def configuration_schema(cls):
     def enabled(cls):
         return enabled
 
-    @classmethod
-    def annotate_query(cls):
-        return False
-
     @classmethod
     def type(cls):
         return "azure_kusto"

redash/query_runner/big_query.py

@@ -83,6 +83,7 @@ def _get_query_results(jobs, project_id, location, job_id, start_index):
 
 
 class BigQuery(BaseQueryRunner):
+    should_annotate_query = False
     noop_query = "SELECT 1"
 
     @classmethod
@@ -133,10 +134,6 @@ def configuration_schema(cls):
             'secret': ['jsonKeyFile']
         }
 
-    @classmethod
-    def annotate_query(cls):
-        return False
-
     def _get_bigquery_service(self):
         scope = [
             "https://www.googleapis.com/auth/bigquery",

redash/query_runner/couchbase.py

@@ -70,7 +70,7 @@ def parse_results(results):
 
 
 class Couchbase(BaseQueryRunner):
-
+    should_annotate_query = False
     noop_query = 'Select 1'
 
     @classmethod
@@ -109,10 +109,6 @@ def __init__(self, configuration):
     def enabled(cls):
         return True
 
-    @classmethod
-    def annotate_query(cls):
-        return False
-
     def test_connection(self):
         result = self.call_service(self.noop_query, '')
 

redash/query_runner/dgraph.py

@@ -29,6 +29,7 @@ def reduce_item(reduced_item, key, value):
 
 
 class Dgraph(BaseQueryRunner):
+    should_annotate_query = False
     noop_query = """
     {
       test() {
@@ -64,13 +65,7 @@ def type(cls):
     def enabled(cls):
         return enabled
 
-    @classmethod
-    def annotate_query(cls):
-        """Dgraph uses '#' as a comment delimiter, not '/* */'"""
-        return False
-
     def run_dgraph_query_raw(self, query):
-
         servers = self.configuration.get('servers')
 
         client_stub = pydgraph.DgraphClientStub(servers)

redash/query_runner/dynamodb_sql.py

@@ -33,6 +33,8 @@
 
 
 class DynamoDBSQL(BaseSQLQueryRunner):
+    should_annotate_query = False
+
     @classmethod
     def configuration_schema(cls):
         return {
@@ -57,10 +59,6 @@ def test_connection(self):
         engine = self._connect()
         list(engine.connection.list_tables())
 
-    @classmethod
-    def annotate_query(cls):
-        return False
-
     @classmethod
     def type(cls):
         return "dynamodb_sql"

redash/query_runner/elasticsearch.py

@@ -46,6 +46,7 @@
 
 
 class BaseElasticSearch(BaseQueryRunner):
+    should_annotate_query = False
     DEBUG_ENABLED = False
 
     @classmethod
@@ -288,15 +289,10 @@ def test_connection(self):
 
 
 class Kibana(BaseElasticSearch):
-
     @classmethod
     def enabled(cls):
         return True
 
-    @classmethod
-    def annotate_query(cls):
-        return False
-
     def _execute_simple_query(self, url, auth, _from, mappings, result_fields, result_columns, result_rows):
         url += "&from={0}".format(_from)
         r = requests.get(url, auth=self.auth)
@@ -379,15 +375,10 @@ def run_query(self, query, user):
 
 
 class ElasticSearch(BaseElasticSearch):
-
     @classmethod
     def enabled(cls):
         return True
 
-    @classmethod
-    def annotate_query(cls):
-        return False
-
     @classmethod
     def name(cls):
         return 'Elasticsearch'

redash/query_runner/google_analytics.py

@@ -78,9 +78,7 @@ def parse_ga_response(response):
 
 
 class GoogleAnalytics(BaseSQLQueryRunner):
-    @classmethod
-    def annotate_query(cls):
-        return False
+    should_annotate_query = False
 
     @classmethod
     def type(cls):

redash/query_runner/google_spreadsheets.py

@@ -139,14 +139,12 @@ def request(self, *args, **kwargs):
 
 
 class GoogleSpreadsheet(BaseQueryRunner):
+    should_annotate_query = False
+
     def __init__(self, configuration):
         super(GoogleSpreadsheet, self).__init__(configuration)
         self.syntax = 'custom'
 
-    @classmethod
-    def annotate_query(cls):
-        return False
-
     @classmethod
     def name(cls):
         return "Google Sheets"

redash/query_runner/graphite.py

@@ -26,6 +26,8 @@ def _transform_result(response):
 
 
 class Graphite(BaseQueryRunner):
+    should_annotate_query = False
+
     @classmethod
     def configuration_schema(cls):
         return {
@@ -49,10 +51,6 @@ def configuration_schema(cls):
             'secret': ['password']
         }
 
-    @classmethod
-    def annotate_query(cls):
-        return False
-
     def __init__(self, configuration):
         super(Graphite, self).__init__(configuration)
         self.syntax = 'custom'

redash/query_runner/hive_ds.py

@@ -36,6 +36,7 @@
 
 
 class Hive(BaseSQLQueryRunner):
+    should_annotate_query = False
     noop_query = "SELECT 1"
 
     @classmethod
@@ -60,10 +61,6 @@ def configuration_schema(cls):
             "required": ["host"]
         }
 
-    @classmethod
-    def annotate_query(cls):
-        return False
-
     @classmethod
     def type(cls):
         return "hive"

redash/query_runner/influx_db.py

@@ -48,6 +48,7 @@ def _transform_result(results):
 
 
 class InfluxDB(BaseQueryRunner):
+    should_annotate_query = False
     noop_query = "show measurements limit 1"
 
     @classmethod
@@ -66,10 +67,6 @@ def configuration_schema(cls):
     def enabled(cls):
         return enabled
 
-    @classmethod
-    def annotate_query(cls):
-        return False
-
     @classmethod
     def type(cls):
         return "influxdb"

redash/query_runner/jql.py

@@ -150,10 +150,6 @@ class JiraJQL(BaseHTTPQueryRunner):
     def name(cls):
         return "JIRA (JQL)"
 
-    @classmethod
-    def annotate_query(cls):
-        return False
-
     def __init__(self, configuration):
         super(JiraJQL, self).__init__(configuration)
         self.syntax = 'json'

redash/query_runner/json_ds.py

@@ -159,10 +159,6 @@ def configuration_schema(cls):
             'order': ['username', 'password']
         }
 
-    @classmethod
-    def annotate_query(cls):
-        return False
-
     def __init__(self, configuration):
         super(JSON, self).__init__(configuration)
         self.syntax = 'yaml'

redash/query_runner/memsql_ds.py

@@ -37,6 +37,7 @@
 
 
 class MemSQL(BaseSQLQueryRunner):
+    should_annotate_query = False
     noop_query = 'SELECT 1'
 
     @classmethod
@@ -62,10 +63,6 @@ def configuration_schema(cls):
             "secret": ["password"]
         }
 
-    @classmethod
-    def annotate_query(cls):
-        return False
-
     @classmethod
     def type(cls):
         return "memsql"

redash/query_runner/mongodb.py

@@ -121,6 +121,8 @@ def parse_results(results):
 
 
 class MongoDB(BaseQueryRunner):
+    should_annotate_query = False
+
     @classmethod
     def configuration_schema(cls):
         return {
@@ -146,10 +148,6 @@ def configuration_schema(cls):
     def enabled(cls):
         return enabled
 
-    @classmethod
-    def annotate_query(cls):
-        return False
-
     def __init__(self, configuration):
         super(MongoDB, self).__init__(configuration)
 

redash/query_runner/mssql.py

@@ -26,6 +26,7 @@
 
 
 class SqlServer(BaseSQLQueryRunner):
+    should_annotate_query = False
     noop_query = "SELECT 1"
 
     @classmethod
@@ -78,10 +79,6 @@ def name(cls):
     def type(cls):
         return "mssql"
 
-    @classmethod
-    def annotate_query(cls):
-        return False
-
     def _get_tables(self, schema):
         query = """
         SELECT table_schema, table_name, column_name

redash/query_runner/mssql_odbc.py

@@ -16,6 +16,7 @@
 
 
 class SQLServerODBC(BaseSQLQueryRunner):
+    should_annotate_query = False
     noop_query = "SELECT 1"
 
     @classmethod
@@ -68,10 +69,6 @@ def name(cls):
     def type(cls):
         return "mssql_odbc"
 
-    @classmethod
-    def annotate_query(cls):
-        return False
-
     def _get_tables(self, schema):
         query = """
         SELECT table_schema, table_name, column_name