Feature (Redshift): option to set query group for adhoc/scheduled que…

…ries (#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.
getredash · Sep 2, 2019 · dd36e35 · dd36e35
1 parent f82dedb
commit dd36e35
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 2 deletions.
diff --git 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
@@ -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):