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):