diff --git a/redash/query_runner/__init__.py b/redash/query_runner/__init__.py index dcd4ccdf80..52175a93b7 100644 --- a/redash/query_runner/__init__.py +++ b/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 diff --git a/redash/query_runner/athena.py b/redash/query_runner/athena.py index db13297caa..7735c4182d 100644 --- a/redash/query_runner/athena.py +++ b/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): diff --git a/redash/query_runner/azure_kusto.py b/redash/query_runner/azure_kusto.py index a2eb32954d..d045f54d71 100644 --- a/redash/query_runner/azure_kusto.py +++ b/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" diff --git a/redash/query_runner/big_query.py b/redash/query_runner/big_query.py index d85574372a..782a600409 100644 --- a/redash/query_runner/big_query.py +++ b/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", diff --git a/redash/query_runner/couchbase.py b/redash/query_runner/couchbase.py index ba83fa0915..093cd387a6 100644 --- a/redash/query_runner/couchbase.py +++ b/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, '') diff --git a/redash/query_runner/dgraph.py b/redash/query_runner/dgraph.py index 60a12edbd1..d5342e163e 100644 --- a/redash/query_runner/dgraph.py +++ b/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) diff --git a/redash/query_runner/dynamodb_sql.py b/redash/query_runner/dynamodb_sql.py index 014d1bd5f6..12be5a7f65 100644 --- a/redash/query_runner/dynamodb_sql.py +++ b/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" diff --git a/redash/query_runner/elasticsearch.py b/redash/query_runner/elasticsearch.py index 57c0edd500..2bb338a767 100644 --- a/redash/query_runner/elasticsearch.py +++ b/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' diff --git a/redash/query_runner/google_analytics.py b/redash/query_runner/google_analytics.py index 71be522015..479403d6de 100644 --- a/redash/query_runner/google_analytics.py +++ b/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): diff --git a/redash/query_runner/google_spreadsheets.py b/redash/query_runner/google_spreadsheets.py index 5b144f4459..5c369d3534 100644 --- a/redash/query_runner/google_spreadsheets.py +++ b/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" diff --git a/redash/query_runner/graphite.py b/redash/query_runner/graphite.py index 1fb5ec1503..711584c70d 100644 --- a/redash/query_runner/graphite.py +++ b/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' diff --git a/redash/query_runner/hive_ds.py b/redash/query_runner/hive_ds.py index 2107d0d0b9..361bdf2105 100644 --- a/redash/query_runner/hive_ds.py +++ b/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" diff --git a/redash/query_runner/influx_db.py b/redash/query_runner/influx_db.py index 47f3a4201f..bec53c8d27 100644 --- a/redash/query_runner/influx_db.py +++ b/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" diff --git a/redash/query_runner/jql.py b/redash/query_runner/jql.py index d24ee0b9f8..76e707e3a3 100644 --- a/redash/query_runner/jql.py +++ b/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' diff --git a/redash/query_runner/json_ds.py b/redash/query_runner/json_ds.py index a278b8b683..9cf3226ed8 100644 --- a/redash/query_runner/json_ds.py +++ b/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' diff --git a/redash/query_runner/memsql_ds.py b/redash/query_runner/memsql_ds.py index bbec2836d4..917e4962cb 100644 --- a/redash/query_runner/memsql_ds.py +++ b/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" diff --git a/redash/query_runner/mongodb.py b/redash/query_runner/mongodb.py index 9a204c69fb..b6dad02747 100644 --- a/redash/query_runner/mongodb.py +++ b/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) diff --git a/redash/query_runner/mssql.py b/redash/query_runner/mssql.py index c4b4fea1e0..4349acebf3 100644 --- a/redash/query_runner/mssql.py +++ b/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 diff --git a/redash/query_runner/mssql_odbc.py b/redash/query_runner/mssql_odbc.py index a729e037c7..7736c56fba 100644 --- a/redash/query_runner/mssql_odbc.py +++ b/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 diff --git a/redash/query_runner/pg.py b/redash/query_runner/pg.py index df5dacfba1..e9e4cc5431 100644 --- a/redash/query_runner/pg.py +++ b/redash/query_runner/pg.py @@ -38,10 +38,8 @@ def default(self, o): items = [ o._bounds[0], - str(o._lower), - ', ', - str(o._upper), - o._bounds[1] + str(o._lower), ', ', + str(o._upper), o._bounds[1] ] return ''.join(items) @@ -92,9 +90,9 @@ def configuration_schema(cls): "title": "Database Name" }, "sslmode": { - "type": "string", - "title": "SSL Mode", - "default": "prefer" + "type": "string", + "title": "SSL Mode", + "default": "prefer" } }, "order": ['host', 'port', 'user', 'password'], @@ -116,7 +114,8 @@ def _get_definitions(self, schema, query): for row in results['rows']: if row['table_schema'] != 'public': - table_name = u'{}.{}'.format(row['table_schema'], row['table_name']) + table_name = u'{}.{}'.format(row['table_schema'], + row['table_name']) else: table_name = row['table_name'] @@ -168,13 +167,14 @@ def _get_tables(self, schema): return schema.values() def _get_connection(self): - connection = psycopg2.connect(user=self.configuration.get('user'), - password=self.configuration.get('password'), - host=self.configuration.get('host'), - port=self.configuration.get('port'), - dbname=self.configuration.get('dbname'), - sslmode=self.configuration.get('sslmode'), - async_=True) + connection = psycopg2.connect( + user=self.configuration.get('user'), + password=self.configuration.get('password'), + host=self.configuration.get('host'), + port=self.configuration.get('port'), + dbname=self.configuration.get('dbname'), + sslmode=self.configuration.get('sslmode'), + async_=True) return connection @@ -189,12 +189,18 @@ def run_query(self, query, user): _wait(connection) if cursor.description is not None: - columns = self.fetch_columns([(i[0], types_map.get(i[1], None)) for i in cursor.description]) - rows = [dict(zip((c['name'] for c in columns), row)) for row in cursor] + columns = self.fetch_columns([(i[0], types_map.get(i[1], None)) + for i in cursor.description]) + rows = [ + dict(zip((c['name'] for c in columns), row)) + for row in cursor + ] data = {'columns': columns, 'rows': rows} error = None - json_data = json_dumps(data, ignore_nan=True, cls=PostgreSQLJSONEncoder) + json_data = json_dumps(data, + ignore_nan=True, + cls=PostgreSQLJSONEncoder) else: error = 'Query completed but it returned no data.' json_data = None @@ -220,22 +226,23 @@ def type(cls): return "redshift" def _get_connection(self): - sslrootcert_path = os.path.join(os.path.dirname(__file__), './files/redshift-ca-bundle.crt') - - connection = psycopg2.connect(user=self.configuration.get('user'), - password=self.configuration.get('password'), - host=self.configuration.get('host'), - port=self.configuration.get('port'), - dbname=self.configuration.get('dbname'), - sslmode=self.configuration.get('sslmode', 'prefer'), - sslrootcert=sslrootcert_path, - async_=True) + sslrootcert_path = os.path.join(os.path.dirname(__file__), + './files/redshift-ca-bundle.crt') + + connection = psycopg2.connect( + user=self.configuration.get('user'), + password=self.configuration.get('password'), + host=self.configuration.get('host'), + port=self.configuration.get('port'), + dbname=self.configuration.get('dbname'), + sslmode=self.configuration.get('sslmode', 'prefer'), + sslrootcert=sslrootcert_path, + async_=True) return connection @classmethod def configuration_schema(cls): - return { "type": "object", "properties": { @@ -256,15 +263,39 @@ def configuration_schema(cls): "title": "Database Name" }, "sslmode": { - "type": "string", - "title": "SSL Mode", - "default": "prefer" - } + "type": "string", + "title": "SSL Mode", + "default": "prefer" + }, + "adhoc_query_group": { + "type": "string", + "title": "Query Group for Adhoc Queries", + "default": "default" + }, + "scheduled_query_group": { + "type": "string", + "title": "Query Group for Scheduled Queries", + "default": "default" + }, }, - "order": ['host', 'port', 'user', 'password'], + "order": ['host', 'port', 'user', 'password', 'dbname', 'sslmode', 'adhoc_query_group', 'scheduled_query_group'], "required": ["dbname", "user", "password", "host", "port"], "secret": ["password"] } + + def annotate_query(self, query, metadata): + annotated = super(Redshift, self).annotate_query(query, metadata) + + if metadata.get('Scheduled', False): + query_group = self.configuration.get('scheduled_query_group') + else: + query_group = self.configuration.get('adhoc_query_group') + + if query_group: + set_query_group = 'set query_group to {};'.format(query_group) + annotated = '{}\n{}'.format(set_query_group, annotated) + + return annotated def _get_tables(self, schema): # Use svv_columns to include internal & external (Spectrum) tables and views data for Redshift @@ -300,7 +331,6 @@ def _get_tables(self, schema): class CockroachDB(PostgreSQL): - @classmethod def type(cls): return "cockroach" diff --git a/redash/query_runner/prometheus.py b/redash/query_runner/prometheus.py index 088291df86..6279d90a69 100644 --- a/redash/query_runner/prometheus.py +++ b/redash/query_runner/prometheus.py @@ -64,6 +64,7 @@ def convert_query_range(payload): class Prometheus(BaseQueryRunner): + should_annotate_query = False @classmethod def configuration_schema(cls): @@ -78,10 +79,6 @@ def configuration_schema(cls): "required": ["url"] } - @classmethod - def annotate_query(cls): - return False - def test_connection(self): resp = requests.get(self.configuration.get("url", None)) return resp.ok diff --git a/redash/query_runner/python.py b/redash/query_runner/python.py index 36209cd0ea..8a516965d3 100644 --- a/redash/query_runner/python.py +++ b/redash/query_runner/python.py @@ -36,6 +36,8 @@ def __call__(self): class Python(BaseQueryRunner): + should_annotate_query = False + safe_builtins = ( 'sorted', 'reversed', 'map', 'reduce', 'any', 'all', 'slice', 'filter', 'len', 'next', 'enumerate', @@ -63,10 +65,6 @@ def configuration_schema(cls): def enabled(cls): return True - @classmethod - def annotate_query(cls): - return False - def __init__(self, configuration): super(Python, self).__init__(configuration) diff --git a/redash/query_runner/qubole.py b/redash/query_runner/qubole.py index d261f4fcff..82276ca139 100644 --- a/redash/query_runner/qubole.py +++ b/redash/query_runner/qubole.py @@ -19,6 +19,7 @@ class Qubole(BaseQueryRunner): + should_annotate_query = False @classmethod def configuration_schema(cls): @@ -62,10 +63,6 @@ def name(cls): def enabled(cls): return enabled - @classmethod - def annotate_query(cls): - return False - def test_connection(self): headers = self._get_header() r = requests.head("%s/api/latest/users" % self.configuration.get('endpoint'), headers=headers) diff --git a/redash/query_runner/query_results.py b/redash/query_runner/query_results.py index 3183fc8ca8..97e174e398 100644 --- a/redash/query_runner/query_results.py +++ b/redash/query_runner/query_results.py @@ -113,16 +113,13 @@ def create_table(connection, table_name, query_results): class Results(BaseQueryRunner): + should_annotate_query = False noop_query = 'SELECT 1' @classmethod def configuration_schema(cls): return {"type": "object", "properties": {}} - @classmethod - def annotate_query(cls): - return False - @classmethod def name(cls): return "Query Results" diff --git a/redash/query_runner/salesforce.py b/redash/query_runner/salesforce.py index b1187bef58..8cc72910ff 100644 --- a/redash/query_runner/salesforce.py +++ b/redash/query_runner/salesforce.py @@ -50,15 +50,12 @@ class Salesforce(BaseQueryRunner): - + should_annotate_query = False + @classmethod def enabled(cls): return enabled - @classmethod - def annotate_query(cls): - return False - @classmethod def configuration_schema(cls): return { diff --git a/redash/query_runner/script.py b/redash/query_runner/script.py index 38e3ae62c5..6c529e9e39 100644 --- a/redash/query_runner/script.py +++ b/redash/query_runner/script.py @@ -29,9 +29,7 @@ def run_script(script, shell): class Script(BaseQueryRunner): - @classmethod - def annotate_query(cls): - return False + should_annotate_query = False @classmethod def enabled(cls): diff --git a/redash/query_runner/treasuredata.py b/redash/query_runner/treasuredata.py index 320f4e3457..895becaac3 100644 --- a/redash/query_runner/treasuredata.py +++ b/redash/query_runner/treasuredata.py @@ -34,6 +34,7 @@ class TreasureData(BaseQueryRunner): + should_annotate_query = False noop_query = "SELECT 1" @classmethod @@ -67,10 +68,6 @@ def configuration_schema(cls): def enabled(cls): return enabled - @classmethod - def annotate_query(cls): - return False - @classmethod def type(cls): return "treasuredata" diff --git a/redash/query_runner/uptycs.py b/redash/query_runner/uptycs.py index 9e6e7ff989..c2573a26d0 100644 --- a/redash/query_runner/uptycs.py +++ b/redash/query_runner/uptycs.py @@ -10,6 +10,7 @@ class Uptycs(BaseSQLQueryRunner): + should_annotate_query = False noop_query = "SELECT 1" @classmethod @@ -40,10 +41,6 @@ def configuration_schema(cls): "secret": ["secret", "key"] } - @classmethod - def annotate_query(cls): - return False - def generate_header(self, key, secret): header = {} utcnow = datetime.datetime.utcnow() diff --git a/redash/query_runner/url.py b/redash/query_runner/url.py index d32a20ee01..d53cf1a9f0 100644 --- a/redash/query_runner/url.py +++ b/redash/query_runner/url.py @@ -6,10 +6,6 @@ class Url(BaseHTTPQueryRunner): requires_url = False - @classmethod - def annotate_query(cls): - return False - def test_connection(self): pass diff --git a/redash/query_runner/yandex_metrica.py b/redash/query_runner/yandex_metrica.py index 82f47d8565..d008b8a505 100644 --- a/redash/query_runner/yandex_metrica.py +++ b/redash/query_runner/yandex_metrica.py @@ -62,9 +62,7 @@ def parse_ym_response(response): class YandexMetrica(BaseSQLQueryRunner): - @classmethod - def annotate_query(cls): - return False + should_annotate_query = False @classmethod def type(cls): diff --git a/redash/tasks/queries.py b/redash/tasks/queries.py index 1ba0c8fc79..0ee9cd0ab0 100644 --- a/redash/tasks/queries.py +++ b/redash/tasks/queries.py @@ -400,16 +400,12 @@ def run(self): return result def _annotate_query(self, query_runner): - if query_runner.annotate_query(): - self.metadata['Task ID'] = self.task.request.id - self.metadata['Query Hash'] = self.query_hash - self.metadata['Queue'] = self.task.request.delivery_info['routing_key'] - - annotation = u", ".join([u"{}: {}".format(k, v) for k, v in self.metadata.iteritems()]) - annotated_query = u"/* {} */ {}".format(annotation, self.query) - else: - annotated_query = self.query - return annotated_query + self.metadata['Task ID'] = self.task.request.id + self.metadata['Query Hash'] = self.query_hash + self.metadata['Queue'] = self.task.request.delivery_info['routing_key'] + self.metadata['Scheduled'] = self.scheduled_query is not None + + return query_runner.annotate_query(self.query, self.metadata) def _log_progress(self, state): logger.info(