diff --git a/CHANGES.rst b/CHANGES.rst index 337194b..6efe9dc 100644 --- a/CHANGES.rst +++ b/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) ------------------ diff --git a/django_redshift_backend/base.py b/django_redshift_backend/base.py index 4a72fa5..6c3eac3 100644 --- a/django_redshift_backend/base.py +++ b/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. diff --git a/doc/basic.rst b/doc/basic.rst index 4224272..273b5a8 100644 --- a/doc/basic.rst +++ b/doc/basic.rst @@ -25,9 +25,21 @@ ENGINE for DATABASES is 'django_redshift_backend'. You can set the name in your 'PASSWORD': '', 'HOST': '', '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`. diff --git a/tests/settings.py b/tests/settings.py index 73816ee..c46594c 100644 --- a/tests/settings.py +++ b/tests/settings.py @@ -9,6 +9,9 @@ 'PASSWORD': 'password', 'HOST': 'localhost', 'PORT': '5439', + 'OPTIONS': { + 'query_group': 'webapp', + }, } } diff --git a/tests/test_redshift_backend.py b/tests/test_redshift_backend.py index 0b10915..36e9eb5 100644 --- a/tests/test_redshift_backend.py +++ b/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" (