Support query_group for WLM

CHANGES.rst

@@ -1,6 +1,15 @@
 CHANGES
 =======
 
+1.0.1
+-----
+
+New Features:
+
+* Support `query_group` options for `Work Load Management`_
+
+.. _Work Load Management: https://docs.aws.amazon.com/redshift/latest/dg/cm-c-implementing-workload-management.html
+
 1.0.0 (2019/01/29)
 ------------------
 

django_redshift_backend/base.py

@@ -567,6 +567,22 @@ def __init__(self, *args, **kwargs):
         self.introspection = DatabaseIntrospection(self)
         self.validation = BaseDatabaseValidation(self)
 
+    def get_connection_params(self):
+        conn_params = super(DatabaseWrapper, self).get_connection_params()
+        self.query_group = conn_params.pop('query_group', None)
+        return conn_params
+
+    def init_connection_state(self):
+        super(DatabaseWrapper, self).init_connection_state()
+        assignments = []
+        if self.query_group:
+            assignments.append('SET query_group TO \'{query_group}\''.format(
+                query_group=self.query_group))
+
+        if assignments:
+            with self.cursor() as cursor:
+                cursor.execute('; '.join(assignments))
+
     def check_constraints(self, table_names=None):
         """
         No constraints to check in Redshift.

doc/basic.rst

@@ -25,9 +25,21 @@ ENGINE for DATABASES is 'django_redshift_backend'. You can set the name in your
            'PASSWORD': '<your database password>',
            'HOST': '<your database hostname>',
            'PORT': '5439',
+           'OPTIONS': {
+               'query_group': 'webapp',
+           },
        }
    }
 
+
+OPTIONS
+-------
+
+- ``query_group``: Set query_group_ to use `Work Load Management`_
+
+.. _query_group: https://docs.aws.amazon.com/redshift/latest/dg/r_query_group.html
+.. _Work Load Management: https://docs.aws.amazon.com/redshift/latest/dg/cm-c-implementing-workload-management.html
+
 For more information, please refer :doc:`refs`.
 
 

tests/settings.py

@@ -9,6 +9,9 @@
         'PASSWORD': 'password',
         'HOST': 'localhost',
         'PORT': '5439',
+        'OPTIONS': {
+            'query_group': 'webapp',
+        },
     }
 }
 

tests/test_redshift_backend.py

@@ -1,5 +1,6 @@
 # -*- coding: utf-8 -*-
 
+import mock
 import unittest
 
 from django.db import connections
@@ -20,6 +21,20 @@ def test_load_redshift_backend(self):
         db = connections['default']
         self.assertIsNotNone(db)
 
+    def test_query_group(self):
+        db = connections['default']
+        self.assertIsNotNone(db)
+
+        db.get_connection_params()
+        self.assertEqual('webapp', db.query_group)
+
+        db.connection = mock.Mock()
+        db.ensure_timezone = mock.Mock()
+        db.force_debug_cursor = True
+
+        db.init_connection_state()
+        self.assertEqual(1, len(db.queries_log))
+
 
 expected_ddl_normal = norm_sql(
     u'''CREATE TABLE "testapp_testmodel" (