From 503bb5cdb6539b9dbde5589fbf348d30c8650def Mon Sep 17 00:00:00 2001 From: Arik Fraimovich Date: Sun, 1 Sep 2019 13:44:05 +0300 Subject: [PATCH 1/5] Code formatting --- redash/query_runner/pg.py | 72 +++++++++++++++++++++------------------ 1 file changed, 39 insertions(+), 33 deletions(-) diff --git a/redash/query_runner/pg.py b/redash/query_runner/pg.py index df5dacfba1..9042138e2a 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,9 +263,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'], @@ -300,7 +307,6 @@ def _get_tables(self, schema): class CockroachDB(PostgreSQL): - @classmethod def type(cls): return "cockroach" From 8548b8f61d00c2d71c19aaad08f8770c49dd2315 Mon Sep 17 00:00:00 2001 From: Arik Fraimovich Date: Sun, 1 Sep 2019 16:18:23 +0300 Subject: [PATCH 2/5] Move annotation logic into query runner, so it can be overriden in the query runner. --- redash/query_runner/__init__.py | 15 +++++++++++---- redash/query_runner/athena.py | 7 ++++--- redash/query_runner/azure_kusto.py | 8 ++------ redash/query_runner/big_query.py | 6 +----- redash/query_runner/couchbase.py | 6 +----- redash/query_runner/dgraph.py | 10 ++-------- redash/query_runner/dynamodb_sql.py | 6 +----- redash/query_runner/elasticsearch.py | 12 +----------- redash/query_runner/google_analytics.py | 6 +----- redash/query_runner/google_spreadsheets.py | 6 +----- redash/query_runner/graphite.py | 6 +----- redash/query_runner/hive_ds.py | 6 +----- redash/query_runner/influx_db.py | 6 +----- redash/query_runner/jql.py | 6 +----- redash/query_runner/json_ds.py | 8 ++------ redash/query_runner/memsql_ds.py | 6 +----- redash/query_runner/mongodb.py | 6 +----- redash/query_runner/mssql.py | 6 +----- redash/query_runner/mssql_odbc.py | 6 +----- redash/query_runner/prometheus.py | 9 ++------- redash/query_runner/python.py | 6 +----- redash/query_runner/qubole.py | 9 ++------- redash/query_runner/query_results.py | 8 ++------ redash/query_runner/salesforce.py | 9 ++------- redash/query_runner/script.py | 6 +----- redash/query_runner/treasuredata.py | 6 +----- redash/query_runner/uptycs.py | 6 +----- redash/query_runner/url.py | 8 ++------ redash/query_runner/yandex_metrica.py | 6 +----- redash/tasks/queries.py | 14 ++++---------- 30 files changed, 54 insertions(+), 171 deletions(-) diff --git a/redash/query_runner/__init__.py b/redash/query_runner/__init__.py index dcd4ccdf80..fd9b062e2f 100644 --- a/redash/query_runner/__init__.py +++ b/redash/query_runner/__init__.py @@ -54,6 +54,12 @@ class NotSupported(Exception): pass +class NoAnnotationMixin(object): + """Simple Mixin to disable query annotation.""" + def annotate_query(self, query, metadata): + return query + + class BaseQueryRunner(object): deprecated = False noop_query = None @@ -74,14 +80,15 @@ 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): + 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() 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..56a014b162 100644 --- a/redash/query_runner/azure_kusto.py +++ b/redash/query_runner/azure_kusto.py @@ -1,5 +1,5 @@ from redash.query_runner import BaseQueryRunner, register -from redash.query_runner import TYPE_STRING, TYPE_DATE, TYPE_DATETIME, TYPE_INTEGER, TYPE_FLOAT, TYPE_BOOLEAN +from redash.query_runner import TYPE_STRING, TYPE_DATE, TYPE_DATETIME, TYPE_INTEGER, TYPE_FLOAT, TYPE_BOOLEAN, NoAnnotationMixin from redash.utils import json_dumps, json_loads @@ -25,7 +25,7 @@ } -class AzureKusto(BaseQueryRunner): +class AzureKusto(BaseQueryRunner, NoAnnotationMixin): noop_query = "let noop = datatable (Noop:string)[1]; noop" def __init__(self, configuration): @@ -71,10 +71,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..ae06c45889 100644 --- a/redash/query_runner/big_query.py +++ b/redash/query_runner/big_query.py @@ -82,7 +82,7 @@ def _get_query_results(jobs, project_id, location, job_id, start_index): return query_reply -class BigQuery(BaseQueryRunner): +class BigQuery(BaseQueryRunner, NoAnnotationMixin): noop_query = "SELECT 1" @classmethod @@ -133,10 +133,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..bd912637b4 100644 --- a/redash/query_runner/couchbase.py +++ b/redash/query_runner/couchbase.py @@ -69,7 +69,7 @@ def parse_results(results): return rows, columns -class Couchbase(BaseQueryRunner): +class Couchbase(BaseQueryRunner, NoAnnotationMixin): noop_query = 'Select 1' @@ -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..8355dedd6a 100644 --- a/redash/query_runner/dgraph.py +++ b/redash/query_runner/dgraph.py @@ -6,7 +6,7 @@ except ImportError: enabled = False -from redash.query_runner import BaseQueryRunner, register +from redash.query_runner import BaseQueryRunner, NoAnnotationMixin, register from redash.utils import json_dumps @@ -28,7 +28,7 @@ def reduce_item(reduced_item, key, value): reduced_item[key] = value -class Dgraph(BaseQueryRunner): +class Dgraph(BaseQueryRunner, NoAnnotationMixin): noop_query = """ { test() { @@ -64,13 +64,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..977750e428 100644 --- a/redash/query_runner/dynamodb_sql.py +++ b/redash/query_runner/dynamodb_sql.py @@ -32,7 +32,7 @@ } -class DynamoDBSQL(BaseSQLQueryRunner): +class DynamoDBSQL(BaseSQLQueryRunner, NoAnnotationMixin): @classmethod def configuration_schema(cls): return { @@ -57,10 +57,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..fd7f97ef05 100644 --- a/redash/query_runner/elasticsearch.py +++ b/redash/query_runner/elasticsearch.py @@ -45,7 +45,7 @@ } -class BaseElasticSearch(BaseQueryRunner): +class BaseElasticSearch(BaseQueryRunner, NoAnnotationMixin): DEBUG_ENABLED = False @classmethod @@ -288,15 +288,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 +374,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..06c20ab82b 100644 --- a/redash/query_runner/google_analytics.py +++ b/redash/query_runner/google_analytics.py @@ -77,11 +77,7 @@ def parse_ga_response(response): return {'columns': columns, 'rows': rows} -class GoogleAnalytics(BaseSQLQueryRunner): - @classmethod - def annotate_query(cls): - return False - +class GoogleAnalytics(BaseSQLQueryRunner, NoAnnotationMixin): @classmethod def type(cls): return "google_analytics" diff --git a/redash/query_runner/google_spreadsheets.py b/redash/query_runner/google_spreadsheets.py index 5b144f4459..623c5a00b8 100644 --- a/redash/query_runner/google_spreadsheets.py +++ b/redash/query_runner/google_spreadsheets.py @@ -138,15 +138,11 @@ def request(self, *args, **kwargs): return super(TimeoutSession, self).request(*args, **kwargs) -class GoogleSpreadsheet(BaseQueryRunner): +class GoogleSpreadsheet(BaseQueryRunner, NoAnnotationMixin): 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..6d5e65fcc3 100644 --- a/redash/query_runner/graphite.py +++ b/redash/query_runner/graphite.py @@ -25,7 +25,7 @@ def _transform_result(response): return json_dumps(data) -class Graphite(BaseQueryRunner): +class Graphite(BaseQueryRunner, NoAnnotationMixin): @classmethod def configuration_schema(cls): return { @@ -49,10 +49,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..aba7546498 100644 --- a/redash/query_runner/hive_ds.py +++ b/redash/query_runner/hive_ds.py @@ -35,7 +35,7 @@ } -class Hive(BaseSQLQueryRunner): +class Hive(BaseSQLQueryRunner, NoAnnotationMixin): noop_query = "SELECT 1" @classmethod @@ -60,10 +60,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..1f990ec84f 100644 --- a/redash/query_runner/influx_db.py +++ b/redash/query_runner/influx_db.py @@ -47,7 +47,7 @@ def _transform_result(results): }) -class InfluxDB(BaseQueryRunner): +class InfluxDB(BaseQueryRunner, NoAnnotationMixin): noop_query = "show measurements limit 1" @classmethod @@ -66,10 +66,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..026e9cb520 100644 --- a/redash/query_runner/jql.py +++ b/redash/query_runner/jql.py @@ -138,7 +138,7 @@ def get_dict_output_field_name(cls, field_name, member_name): return None -class JiraJQL(BaseHTTPQueryRunner): +class JiraJQL(BaseHTTPQueryRunner, NoAnnotationMixin): noop_query = '{"queryType": "count"}' response_error = "JIRA returned unexpected status code" requires_authentication = True @@ -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..137d9252b7 100644 --- a/redash/query_runner/json_ds.py +++ b/redash/query_runner/json_ds.py @@ -10,7 +10,7 @@ from redash.utils.compat import long from redash.query_runner import (BaseHTTPQueryRunner, register, TYPE_BOOLEAN, TYPE_DATETIME, TYPE_FLOAT, - TYPE_INTEGER, TYPE_STRING) + TYPE_INTEGER, TYPE_STRING, NoAnnotationMixin) class QueryParseError(Exception): @@ -138,7 +138,7 @@ def parse_json(data, path, fields): return {'rows': rows, 'columns': columns} -class JSON(BaseHTTPQueryRunner): +class JSON(BaseHTTPQueryRunner, NoAnnotationMixin): requires_url = False @classmethod @@ -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..ca6d875098 100644 --- a/redash/query_runner/memsql_ds.py +++ b/redash/query_runner/memsql_ds.py @@ -36,7 +36,7 @@ } -class MemSQL(BaseSQLQueryRunner): +class MemSQL(BaseSQLQueryRunner, NoAnnotationMixin): noop_query = 'SELECT 1' @classmethod @@ -62,10 +62,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..ea06c6a6e3 100644 --- a/redash/query_runner/mongodb.py +++ b/redash/query_runner/mongodb.py @@ -120,7 +120,7 @@ def parse_results(results): return rows, columns -class MongoDB(BaseQueryRunner): +class MongoDB(BaseQueryRunner, NoAnnotationMixin): @classmethod def configuration_schema(cls): return { @@ -146,10 +146,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..f7d03500b4 100644 --- a/redash/query_runner/mssql.py +++ b/redash/query_runner/mssql.py @@ -25,7 +25,7 @@ } -class SqlServer(BaseSQLQueryRunner): +class SqlServer(BaseSQLQueryRunner, NoAnnotationMixin): noop_query = "SELECT 1" @classmethod @@ -78,10 +78,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..4c9a5e5290 100644 --- a/redash/query_runner/mssql_odbc.py +++ b/redash/query_runner/mssql_odbc.py @@ -15,7 +15,7 @@ enabled = False -class SQLServerODBC(BaseSQLQueryRunner): +class SQLServerODBC(BaseSQLQueryRunner, NoAnnotationMixin): noop_query = "SELECT 1" @classmethod @@ -68,10 +68,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/prometheus.py b/redash/query_runner/prometheus.py index 088291df86..66aef49a5c 100644 --- a/redash/query_runner/prometheus.py +++ b/redash/query_runner/prometheus.py @@ -3,7 +3,7 @@ from datetime import datetime from dateutil import parser from urlparse import parse_qs -from redash.query_runner import BaseQueryRunner, register, TYPE_DATETIME, TYPE_STRING +from redash.query_runner import BaseQueryRunner, register, TYPE_DATETIME, TYPE_STRING, NoAnnotationMixin from redash.utils import json_dumps @@ -63,8 +63,7 @@ def convert_query_range(payload): payload.update(query_range) -class Prometheus(BaseQueryRunner): - +class Prometheus(BaseQueryRunner, NoAnnotationMixin): @classmethod def configuration_schema(cls): return { @@ -78,10 +77,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..ef4ed6bafd 100644 --- a/redash/query_runner/python.py +++ b/redash/query_runner/python.py @@ -35,7 +35,7 @@ def __call__(self): return self -class Python(BaseQueryRunner): +class Python(BaseQueryRunner, NoAnnotationMixin): safe_builtins = ( 'sorted', 'reversed', 'map', 'reduce', 'any', 'all', 'slice', 'filter', 'len', 'next', 'enumerate', @@ -63,10 +63,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..e143daf45a 100644 --- a/redash/query_runner/qubole.py +++ b/redash/query_runner/qubole.py @@ -4,7 +4,7 @@ import logging from cStringIO import StringIO -from redash.query_runner import BaseQueryRunner, register +from redash.query_runner import BaseQueryRunner, NoAnnotationMixin, register from redash.query_runner import TYPE_STRING from redash.utils import json_dumps @@ -18,8 +18,7 @@ enabled = False -class Qubole(BaseQueryRunner): - +class Qubole(BaseQueryRunner, NoAnnotationMixin): @classmethod def configuration_schema(cls): return { @@ -62,10 +61,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 910df7c9c1..bb950188f9 100644 --- a/redash/query_runner/query_results.py +++ b/redash/query_runner/query_results.py @@ -4,7 +4,7 @@ from redash import models from redash.permissions import has_access, not_view_only -from redash.query_runner import BaseQueryRunner, TYPE_STRING, guess_type, register +from redash.query_runner import BaseQueryRunner, NoAnnotationMixin, TYPE_STRING, guess_type, register from redash.utils import json_dumps, json_loads logger = logging.getLogger(__name__) @@ -103,7 +103,7 @@ def create_table(connection, table_name, query_results): connection.execute(insert_template, values) -class Results(BaseQueryRunner): +class Results(BaseQueryRunner, NoAnnotationMixin): noop_query = 'SELECT 1' @classmethod @@ -114,10 +114,6 @@ def configuration_schema(cls): } } - @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..f3120af737 100644 --- a/redash/query_runner/salesforce.py +++ b/redash/query_runner/salesforce.py @@ -3,7 +3,7 @@ import re import logging from collections import OrderedDict -from redash.query_runner import BaseQueryRunner, register +from redash.query_runner import BaseQueryRunner, NoAnnotationMixin, register from redash.query_runner import TYPE_STRING, TYPE_DATE, TYPE_DATETIME, TYPE_INTEGER, TYPE_FLOAT, TYPE_BOOLEAN from redash.utils import json_dumps logger = logging.getLogger(__name__) @@ -49,16 +49,11 @@ # https://developer.salesforce.com/docs/atlas.en-us.soql_sosl.meta/soql_sosl/sforce_api_calls_soql_select_examples.htm -class Salesforce(BaseQueryRunner): - +class Salesforce(BaseQueryRunner, NoAnnotationMixin): @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..72126d0cd8 100644 --- a/redash/query_runner/script.py +++ b/redash/query_runner/script.py @@ -28,11 +28,7 @@ def run_script(script, shell): return output, None -class Script(BaseQueryRunner): - @classmethod - def annotate_query(cls): - return False - +class Script(BaseQueryRunner, NoAnnotationMixin): @classmethod def enabled(cls): return "check_output" in subprocess.__dict__ diff --git a/redash/query_runner/treasuredata.py b/redash/query_runner/treasuredata.py index 320f4e3457..7036eb4f40 100644 --- a/redash/query_runner/treasuredata.py +++ b/redash/query_runner/treasuredata.py @@ -33,7 +33,7 @@ } -class TreasureData(BaseQueryRunner): +class TreasureData(BaseQueryRunner, NoAnnotationMixin): noop_query = "SELECT 1" @classmethod @@ -67,10 +67,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..03dcc75853 100644 --- a/redash/query_runner/uptycs.py +++ b/redash/query_runner/uptycs.py @@ -9,7 +9,7 @@ logger = logging.getLogger(__name__) -class Uptycs(BaseSQLQueryRunner): +class Uptycs(BaseSQLQueryRunner, NoAnnotationMixin): noop_query = "SELECT 1" @classmethod @@ -40,10 +40,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..58abeed17b 100644 --- a/redash/query_runner/url.py +++ b/redash/query_runner/url.py @@ -1,15 +1,11 @@ -from redash.query_runner import BaseHTTPQueryRunner, register +from redash.query_runner import BaseHTTPQueryRunner, NoAnnotationMixin, register from redash.utils import deprecated @deprecated() -class Url(BaseHTTPQueryRunner): +class Url(BaseHTTPQueryRunner, NoAnnotationMixin): 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..6816ff909b 100644 --- a/redash/query_runner/yandex_metrica.py +++ b/redash/query_runner/yandex_metrica.py @@ -61,11 +61,7 @@ def parse_ym_response(response): return {'columns': columns, 'rows': rows} -class YandexMetrica(BaseSQLQueryRunner): - @classmethod - def annotate_query(cls): - return False - +class YandexMetrica(BaseSQLQueryRunner, NoAnnotationMixin): @classmethod def type(cls): # This is written with a "k" for backward-compatibility. See #2874. diff --git a/redash/tasks/queries.py b/redash/tasks/queries.py index 1ba0c8fc79..b76286c849 100644 --- a/redash/tasks/queries.py +++ b/redash/tasks/queries.py @@ -400,16 +400,10 @@ 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'] + return query_runner.annotate_query(self.query, self.metadata) def _log_progress(self, state): logger.info( From fd09c41812532ddde4d5f800e4b383d251c2dde2 Mon Sep 17 00:00:00 2001 From: Arik Fraimovich Date: Sun, 1 Sep 2019 17:04:26 +0300 Subject: [PATCH 3/5] Add mixin to __all__ --- redash/query_runner/__init__.py | 1 + 1 file changed, 1 insertion(+) diff --git a/redash/query_runner/__init__.py b/redash/query_runner/__init__.py index fd9b062e2f..2074cf4be0 100644 --- a/redash/query_runner/__init__.py +++ b/redash/query_runner/__init__.py @@ -12,6 +12,7 @@ __all__ = [ 'BaseQueryRunner', + 'NoAnnotationMixin', 'BaseHTTPQueryRunner', 'InterruptException', 'BaseSQLQueryRunner', From bec6132102cc2ad04250af0d56be11118ec5cb1e Mon Sep 17 00:00:00 2001 From: Arik Fraimovich Date: Mon, 2 Sep 2019 15:19:04 +0300 Subject: [PATCH 4/5] Switch to flag instead of mixin --- redash/query_runner/__init__.py | 12 +++++------- redash/query_runner/azure_kusto.py | 5 +++-- redash/query_runner/big_query.py | 3 ++- redash/query_runner/couchbase.py | 4 ++-- redash/query_runner/dgraph.py | 5 +++-- redash/query_runner/dynamodb_sql.py | 4 +++- redash/query_runner/elasticsearch.py | 3 ++- redash/query_runner/google_analytics.py | 4 +++- redash/query_runner/google_spreadsheets.py | 4 +++- redash/query_runner/graphite.py | 4 +++- redash/query_runner/hive_ds.py | 3 ++- redash/query_runner/influx_db.py | 3 ++- redash/query_runner/jql.py | 2 +- redash/query_runner/json_ds.py | 4 ++-- redash/query_runner/memsql_ds.py | 3 ++- redash/query_runner/mongodb.py | 4 +++- redash/query_runner/mssql.py | 3 ++- redash/query_runner/mssql_odbc.py | 3 ++- redash/query_runner/prometheus.py | 6 ++++-- redash/query_runner/python.py | 4 +++- redash/query_runner/qubole.py | 6 ++++-- redash/query_runner/query_results.py | 5 +++-- redash/query_runner/salesforce.py | 6 ++++-- redash/query_runner/script.py | 4 +++- redash/query_runner/treasuredata.py | 3 ++- redash/query_runner/uptycs.py | 3 ++- redash/query_runner/url.py | 4 ++-- redash/query_runner/yandex_metrica.py | 4 +++- 28 files changed, 75 insertions(+), 43 deletions(-) diff --git a/redash/query_runner/__init__.py b/redash/query_runner/__init__.py index 2074cf4be0..52175a93b7 100644 --- a/redash/query_runner/__init__.py +++ b/redash/query_runner/__init__.py @@ -12,7 +12,6 @@ __all__ = [ 'BaseQueryRunner', - 'NoAnnotationMixin', 'BaseHTTPQueryRunner', 'InterruptException', 'BaseSQLQueryRunner', @@ -55,14 +54,9 @@ class NotSupported(Exception): pass -class NoAnnotationMixin(object): - """Simple Mixin to disable query annotation.""" - def annotate_query(self, query, metadata): - return query - - class BaseQueryRunner(object): deprecated = False + should_annotate_query = True noop_query = None def __init__(self, configuration): @@ -86,6 +80,9 @@ 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 @@ -158,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/azure_kusto.py b/redash/query_runner/azure_kusto.py index 56a014b162..d045f54d71 100644 --- a/redash/query_runner/azure_kusto.py +++ b/redash/query_runner/azure_kusto.py @@ -1,5 +1,5 @@ from redash.query_runner import BaseQueryRunner, register -from redash.query_runner import TYPE_STRING, TYPE_DATE, TYPE_DATETIME, TYPE_INTEGER, TYPE_FLOAT, TYPE_BOOLEAN, NoAnnotationMixin +from redash.query_runner import TYPE_STRING, TYPE_DATE, TYPE_DATETIME, TYPE_INTEGER, TYPE_FLOAT, TYPE_BOOLEAN from redash.utils import json_dumps, json_loads @@ -25,7 +25,8 @@ } -class AzureKusto(BaseQueryRunner, NoAnnotationMixin): +class AzureKusto(BaseQueryRunner): + should_annotate_query = False noop_query = "let noop = datatable (Noop:string)[1]; noop" def __init__(self, configuration): diff --git a/redash/query_runner/big_query.py b/redash/query_runner/big_query.py index ae06c45889..782a600409 100644 --- a/redash/query_runner/big_query.py +++ b/redash/query_runner/big_query.py @@ -82,7 +82,8 @@ def _get_query_results(jobs, project_id, location, job_id, start_index): return query_reply -class BigQuery(BaseQueryRunner, NoAnnotationMixin): +class BigQuery(BaseQueryRunner): + should_annotate_query = False noop_query = "SELECT 1" @classmethod diff --git a/redash/query_runner/couchbase.py b/redash/query_runner/couchbase.py index bd912637b4..093cd387a6 100644 --- a/redash/query_runner/couchbase.py +++ b/redash/query_runner/couchbase.py @@ -69,8 +69,8 @@ def parse_results(results): return rows, columns -class Couchbase(BaseQueryRunner, NoAnnotationMixin): - +class Couchbase(BaseQueryRunner): + should_annotate_query = False noop_query = 'Select 1' @classmethod diff --git a/redash/query_runner/dgraph.py b/redash/query_runner/dgraph.py index 8355dedd6a..d5342e163e 100644 --- a/redash/query_runner/dgraph.py +++ b/redash/query_runner/dgraph.py @@ -6,7 +6,7 @@ except ImportError: enabled = False -from redash.query_runner import BaseQueryRunner, NoAnnotationMixin, register +from redash.query_runner import BaseQueryRunner, register from redash.utils import json_dumps @@ -28,7 +28,8 @@ def reduce_item(reduced_item, key, value): reduced_item[key] = value -class Dgraph(BaseQueryRunner, NoAnnotationMixin): +class Dgraph(BaseQueryRunner): + should_annotate_query = False noop_query = """ { test() { diff --git a/redash/query_runner/dynamodb_sql.py b/redash/query_runner/dynamodb_sql.py index 977750e428..12be5a7f65 100644 --- a/redash/query_runner/dynamodb_sql.py +++ b/redash/query_runner/dynamodb_sql.py @@ -32,7 +32,9 @@ } -class DynamoDBSQL(BaseSQLQueryRunner, NoAnnotationMixin): +class DynamoDBSQL(BaseSQLQueryRunner): + should_annotate_query = False + @classmethod def configuration_schema(cls): return { diff --git a/redash/query_runner/elasticsearch.py b/redash/query_runner/elasticsearch.py index fd7f97ef05..2bb338a767 100644 --- a/redash/query_runner/elasticsearch.py +++ b/redash/query_runner/elasticsearch.py @@ -45,7 +45,8 @@ } -class BaseElasticSearch(BaseQueryRunner, NoAnnotationMixin): +class BaseElasticSearch(BaseQueryRunner): + should_annotate_query = False DEBUG_ENABLED = False @classmethod diff --git a/redash/query_runner/google_analytics.py b/redash/query_runner/google_analytics.py index 06c20ab82b..479403d6de 100644 --- a/redash/query_runner/google_analytics.py +++ b/redash/query_runner/google_analytics.py @@ -77,7 +77,9 @@ def parse_ga_response(response): return {'columns': columns, 'rows': rows} -class GoogleAnalytics(BaseSQLQueryRunner, NoAnnotationMixin): +class GoogleAnalytics(BaseSQLQueryRunner): + should_annotate_query = False + @classmethod def type(cls): return "google_analytics" diff --git a/redash/query_runner/google_spreadsheets.py b/redash/query_runner/google_spreadsheets.py index 623c5a00b8..5c369d3534 100644 --- a/redash/query_runner/google_spreadsheets.py +++ b/redash/query_runner/google_spreadsheets.py @@ -138,7 +138,9 @@ def request(self, *args, **kwargs): return super(TimeoutSession, self).request(*args, **kwargs) -class GoogleSpreadsheet(BaseQueryRunner, NoAnnotationMixin): +class GoogleSpreadsheet(BaseQueryRunner): + should_annotate_query = False + def __init__(self, configuration): super(GoogleSpreadsheet, self).__init__(configuration) self.syntax = 'custom' diff --git a/redash/query_runner/graphite.py b/redash/query_runner/graphite.py index 6d5e65fcc3..711584c70d 100644 --- a/redash/query_runner/graphite.py +++ b/redash/query_runner/graphite.py @@ -25,7 +25,9 @@ def _transform_result(response): return json_dumps(data) -class Graphite(BaseQueryRunner, NoAnnotationMixin): +class Graphite(BaseQueryRunner): + should_annotate_query = False + @classmethod def configuration_schema(cls): return { diff --git a/redash/query_runner/hive_ds.py b/redash/query_runner/hive_ds.py index aba7546498..361bdf2105 100644 --- a/redash/query_runner/hive_ds.py +++ b/redash/query_runner/hive_ds.py @@ -35,7 +35,8 @@ } -class Hive(BaseSQLQueryRunner, NoAnnotationMixin): +class Hive(BaseSQLQueryRunner): + should_annotate_query = False noop_query = "SELECT 1" @classmethod diff --git a/redash/query_runner/influx_db.py b/redash/query_runner/influx_db.py index 1f990ec84f..bec53c8d27 100644 --- a/redash/query_runner/influx_db.py +++ b/redash/query_runner/influx_db.py @@ -47,7 +47,8 @@ def _transform_result(results): }) -class InfluxDB(BaseQueryRunner, NoAnnotationMixin): +class InfluxDB(BaseQueryRunner): + should_annotate_query = False noop_query = "show measurements limit 1" @classmethod diff --git a/redash/query_runner/jql.py b/redash/query_runner/jql.py index 026e9cb520..76e707e3a3 100644 --- a/redash/query_runner/jql.py +++ b/redash/query_runner/jql.py @@ -138,7 +138,7 @@ def get_dict_output_field_name(cls, field_name, member_name): return None -class JiraJQL(BaseHTTPQueryRunner, NoAnnotationMixin): +class JiraJQL(BaseHTTPQueryRunner): noop_query = '{"queryType": "count"}' response_error = "JIRA returned unexpected status code" requires_authentication = True diff --git a/redash/query_runner/json_ds.py b/redash/query_runner/json_ds.py index 137d9252b7..9cf3226ed8 100644 --- a/redash/query_runner/json_ds.py +++ b/redash/query_runner/json_ds.py @@ -10,7 +10,7 @@ from redash.utils.compat import long from redash.query_runner import (BaseHTTPQueryRunner, register, TYPE_BOOLEAN, TYPE_DATETIME, TYPE_FLOAT, - TYPE_INTEGER, TYPE_STRING, NoAnnotationMixin) + TYPE_INTEGER, TYPE_STRING) class QueryParseError(Exception): @@ -138,7 +138,7 @@ def parse_json(data, path, fields): return {'rows': rows, 'columns': columns} -class JSON(BaseHTTPQueryRunner, NoAnnotationMixin): +class JSON(BaseHTTPQueryRunner): requires_url = False @classmethod diff --git a/redash/query_runner/memsql_ds.py b/redash/query_runner/memsql_ds.py index ca6d875098..917e4962cb 100644 --- a/redash/query_runner/memsql_ds.py +++ b/redash/query_runner/memsql_ds.py @@ -36,7 +36,8 @@ } -class MemSQL(BaseSQLQueryRunner, NoAnnotationMixin): +class MemSQL(BaseSQLQueryRunner): + should_annotate_query = False noop_query = 'SELECT 1' @classmethod diff --git a/redash/query_runner/mongodb.py b/redash/query_runner/mongodb.py index ea06c6a6e3..b6dad02747 100644 --- a/redash/query_runner/mongodb.py +++ b/redash/query_runner/mongodb.py @@ -120,7 +120,9 @@ def parse_results(results): return rows, columns -class MongoDB(BaseQueryRunner, NoAnnotationMixin): +class MongoDB(BaseQueryRunner): + should_annotate_query = False + @classmethod def configuration_schema(cls): return { diff --git a/redash/query_runner/mssql.py b/redash/query_runner/mssql.py index f7d03500b4..4349acebf3 100644 --- a/redash/query_runner/mssql.py +++ b/redash/query_runner/mssql.py @@ -25,7 +25,8 @@ } -class SqlServer(BaseSQLQueryRunner, NoAnnotationMixin): +class SqlServer(BaseSQLQueryRunner): + should_annotate_query = False noop_query = "SELECT 1" @classmethod diff --git a/redash/query_runner/mssql_odbc.py b/redash/query_runner/mssql_odbc.py index 4c9a5e5290..7736c56fba 100644 --- a/redash/query_runner/mssql_odbc.py +++ b/redash/query_runner/mssql_odbc.py @@ -15,7 +15,8 @@ enabled = False -class SQLServerODBC(BaseSQLQueryRunner, NoAnnotationMixin): +class SQLServerODBC(BaseSQLQueryRunner): + should_annotate_query = False noop_query = "SELECT 1" @classmethod diff --git a/redash/query_runner/prometheus.py b/redash/query_runner/prometheus.py index 66aef49a5c..6279d90a69 100644 --- a/redash/query_runner/prometheus.py +++ b/redash/query_runner/prometheus.py @@ -3,7 +3,7 @@ from datetime import datetime from dateutil import parser from urlparse import parse_qs -from redash.query_runner import BaseQueryRunner, register, TYPE_DATETIME, TYPE_STRING, NoAnnotationMixin +from redash.query_runner import BaseQueryRunner, register, TYPE_DATETIME, TYPE_STRING from redash.utils import json_dumps @@ -63,7 +63,9 @@ def convert_query_range(payload): payload.update(query_range) -class Prometheus(BaseQueryRunner, NoAnnotationMixin): +class Prometheus(BaseQueryRunner): + should_annotate_query = False + @classmethod def configuration_schema(cls): return { diff --git a/redash/query_runner/python.py b/redash/query_runner/python.py index ef4ed6bafd..8a516965d3 100644 --- a/redash/query_runner/python.py +++ b/redash/query_runner/python.py @@ -35,7 +35,9 @@ def __call__(self): return self -class Python(BaseQueryRunner, NoAnnotationMixin): +class Python(BaseQueryRunner): + should_annotate_query = False + safe_builtins = ( 'sorted', 'reversed', 'map', 'reduce', 'any', 'all', 'slice', 'filter', 'len', 'next', 'enumerate', diff --git a/redash/query_runner/qubole.py b/redash/query_runner/qubole.py index e143daf45a..82276ca139 100644 --- a/redash/query_runner/qubole.py +++ b/redash/query_runner/qubole.py @@ -4,7 +4,7 @@ import logging from cStringIO import StringIO -from redash.query_runner import BaseQueryRunner, NoAnnotationMixin, register +from redash.query_runner import BaseQueryRunner, register from redash.query_runner import TYPE_STRING from redash.utils import json_dumps @@ -18,7 +18,9 @@ enabled = False -class Qubole(BaseQueryRunner, NoAnnotationMixin): +class Qubole(BaseQueryRunner): + should_annotate_query = False + @classmethod def configuration_schema(cls): return { diff --git a/redash/query_runner/query_results.py b/redash/query_runner/query_results.py index bb950188f9..35faee8201 100644 --- a/redash/query_runner/query_results.py +++ b/redash/query_runner/query_results.py @@ -4,7 +4,7 @@ from redash import models from redash.permissions import has_access, not_view_only -from redash.query_runner import BaseQueryRunner, NoAnnotationMixin, TYPE_STRING, guess_type, register +from redash.query_runner import BaseQueryRunner, TYPE_STRING, guess_type, register from redash.utils import json_dumps, json_loads logger = logging.getLogger(__name__) @@ -103,7 +103,8 @@ def create_table(connection, table_name, query_results): connection.execute(insert_template, values) -class Results(BaseQueryRunner, NoAnnotationMixin): +class Results(BaseQueryRunner): + should_annotate_query = False noop_query = 'SELECT 1' @classmethod diff --git a/redash/query_runner/salesforce.py b/redash/query_runner/salesforce.py index f3120af737..8cc72910ff 100644 --- a/redash/query_runner/salesforce.py +++ b/redash/query_runner/salesforce.py @@ -3,7 +3,7 @@ import re import logging from collections import OrderedDict -from redash.query_runner import BaseQueryRunner, NoAnnotationMixin, register +from redash.query_runner import BaseQueryRunner, register from redash.query_runner import TYPE_STRING, TYPE_DATE, TYPE_DATETIME, TYPE_INTEGER, TYPE_FLOAT, TYPE_BOOLEAN from redash.utils import json_dumps logger = logging.getLogger(__name__) @@ -49,7 +49,9 @@ # https://developer.salesforce.com/docs/atlas.en-us.soql_sosl.meta/soql_sosl/sforce_api_calls_soql_select_examples.htm -class Salesforce(BaseQueryRunner, NoAnnotationMixin): +class Salesforce(BaseQueryRunner): + should_annotate_query = False + @classmethod def enabled(cls): return enabled diff --git a/redash/query_runner/script.py b/redash/query_runner/script.py index 72126d0cd8..6c529e9e39 100644 --- a/redash/query_runner/script.py +++ b/redash/query_runner/script.py @@ -28,7 +28,9 @@ def run_script(script, shell): return output, None -class Script(BaseQueryRunner, NoAnnotationMixin): +class Script(BaseQueryRunner): + should_annotate_query = False + @classmethod def enabled(cls): return "check_output" in subprocess.__dict__ diff --git a/redash/query_runner/treasuredata.py b/redash/query_runner/treasuredata.py index 7036eb4f40..895becaac3 100644 --- a/redash/query_runner/treasuredata.py +++ b/redash/query_runner/treasuredata.py @@ -33,7 +33,8 @@ } -class TreasureData(BaseQueryRunner, NoAnnotationMixin): +class TreasureData(BaseQueryRunner): + should_annotate_query = False noop_query = "SELECT 1" @classmethod diff --git a/redash/query_runner/uptycs.py b/redash/query_runner/uptycs.py index 03dcc75853..c2573a26d0 100644 --- a/redash/query_runner/uptycs.py +++ b/redash/query_runner/uptycs.py @@ -9,7 +9,8 @@ logger = logging.getLogger(__name__) -class Uptycs(BaseSQLQueryRunner, NoAnnotationMixin): +class Uptycs(BaseSQLQueryRunner): + should_annotate_query = False noop_query = "SELECT 1" @classmethod diff --git a/redash/query_runner/url.py b/redash/query_runner/url.py index 58abeed17b..d53cf1a9f0 100644 --- a/redash/query_runner/url.py +++ b/redash/query_runner/url.py @@ -1,9 +1,9 @@ -from redash.query_runner import BaseHTTPQueryRunner, NoAnnotationMixin, register +from redash.query_runner import BaseHTTPQueryRunner, register from redash.utils import deprecated @deprecated() -class Url(BaseHTTPQueryRunner, NoAnnotationMixin): +class Url(BaseHTTPQueryRunner): requires_url = False def test_connection(self): diff --git a/redash/query_runner/yandex_metrica.py b/redash/query_runner/yandex_metrica.py index 6816ff909b..d008b8a505 100644 --- a/redash/query_runner/yandex_metrica.py +++ b/redash/query_runner/yandex_metrica.py @@ -61,7 +61,9 @@ def parse_ym_response(response): return {'columns': columns, 'rows': rows} -class YandexMetrica(BaseSQLQueryRunner, NoAnnotationMixin): +class YandexMetrica(BaseSQLQueryRunner): + should_annotate_query = False + @classmethod def type(cls): # This is written with a "k" for backward-compatibility. See #2874. From dd36e35cb8bdcd8f8c2fd7a00e11532c46998ddd Mon Sep 17 00:00:00 2001 From: Arik Fraimovich Date: Mon, 2 Sep 2019 15:32:27 +0300 Subject: [PATCH 5/5] Feature (Redshift): option to set query group for adhoc/scheduled queries (#4114) * Add scheduled status to query job metadata. * Add: option to set query group for adhoc/scheduled Redshift queries * Scheduled might not be set for already enqueued queries. --- redash/query_runner/pg.py | 28 ++++++++++++++++++++++++++-- redash/tasks/queries.py | 2 ++ 2 files changed, 28 insertions(+), 2 deletions(-) diff --git a/redash/query_runner/pg.py b/redash/query_runner/pg.py index 9042138e2a..e9e4cc5431 100644 --- a/redash/query_runner/pg.py +++ b/redash/query_runner/pg.py @@ -266,12 +266,36 @@ def configuration_schema(cls): "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 diff --git a/redash/tasks/queries.py b/redash/tasks/queries.py index b76286c849..0ee9cd0ab0 100644 --- a/redash/tasks/queries.py +++ b/redash/tasks/queries.py @@ -403,6 +403,8 @@ def _annotate_query(self, query_runner): 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):