From 885db2db7c049308340038efdaa0e83214219ded Mon Sep 17 00:00:00 2001 From: Florian Eiden Date: Wed, 12 Jun 2024 21:30:33 -0700 Subject: [PATCH] Remove group_federation for serverless (#836) * Remove grou_federation for serverless * Add changie log * Fix test on cluster_id * Add explicit role_arn support * Revert "Add explicit role_arn support" This reverts commit 00c53d6e5d1d4bb5433ea0fae4cebcb1fd7319a6. * Add tests * Linting * Linting * Tests should ignore group_federation, and not expect false * Pick the correct method for the test * db_user should line up with IAMRole on Provisionned * Align tests to spec * Explicitely set group federation for serverless * Solved the cluster id issue - should be null on serverless * Update changie desc --------- Co-authored-by: Colin Rogers <111200756+colin-rogers-dbt@users.noreply.github.com> --- .../unreleased/Fixes-20240531-113620.yaml | 6 + dbt/adapters/redshift/connections.py | 19 +-- tests/unit/test_auth_method.py | 117 ++++++++++++++++++ 3 files changed, 135 insertions(+), 7 deletions(-) create mode 100644 .changes/unreleased/Fixes-20240531-113620.yaml diff --git a/.changes/unreleased/Fixes-20240531-113620.yaml b/.changes/unreleased/Fixes-20240531-113620.yaml new file mode 100644 index 000000000..a020ca045 --- /dev/null +++ b/.changes/unreleased/Fixes-20240531-113620.yaml @@ -0,0 +1,6 @@ +kind: Fixes +body: Support IAM Role authentication for Redshift Serverless +time: 2024-05-31T11:36:20.397521-07:00 +custom: + Author: fleid + Issue: "835" diff --git a/dbt/adapters/redshift/connections.py b/dbt/adapters/redshift/connections.py index dfd7ab30b..d3fbcafea 100644 --- a/dbt/adapters/redshift/connections.py +++ b/dbt/adapters/redshift/connections.py @@ -237,10 +237,15 @@ def _iam_user_kwargs(self) -> Dict[str, Any]: def _iam_role_kwargs(self) -> Dict[str, Optional[Any]]: logger.debug("Connecting to redshift with 'iam_role' credentials method") kwargs = self._iam_kwargs - kwargs.update( - group_federation=True, - db_user=None, - ) + + # It's a role, we're ignoring the user + kwargs.update(db_user=None) + + # Serverless shouldn't get group_federation, Provisoned clusters should + if "serverless" in self.credentials.host: + kwargs.update(group_federation=False) + else: + kwargs.update(group_federation=True) if iam_profile := self.credentials.iam_profile: kwargs.update(profile=iam_profile) @@ -256,10 +261,10 @@ def _iam_kwargs(self) -> Dict[str, Any]: password="", ) - if cluster_id := self.credentials.cluster_id: - kwargs.update(cluster_identifier=cluster_id) - elif "serverless" in self.credentials.host: + if "serverless" in self.credentials.host: kwargs.update(cluster_identifier=None) + elif cluster_id := self.credentials.cluster_id: + kwargs.update(cluster_identifier=cluster_id) else: raise FailedToConnectError( "Failed to use IAM method:" diff --git a/tests/unit/test_auth_method.py b/tests/unit/test_auth_method.py index bd9912d0c..55b1aad74 100644 --- a/tests/unit/test_auth_method.py +++ b/tests/unit/test_auth_method.py @@ -456,3 +456,120 @@ def test_profile(self): group_federation=True, **DEFAULT_SSL_CONFIG, ) + + +class TestIAMRoleMethodServerless(AuthMethod): + # Should behave like IAM Role provisioned, with the exception of not having group_federation set + + @mock.patch("redshift_connector.connect", MagicMock()) + def test_profile_default_region(self): + self.config.credentials = self.config.credentials.replace( + method="iam_role", + iam_profile="iam_profile_test", + host="doesnotexist.1233.us-east-2.redshift-serverless.amazonaws.com", + ) + connection = self.adapter.acquire_connection("dummy") + connection.handle + redshift_connector.connect.assert_called_once_with( + iam=True, + host="doesnotexist.1233.us-east-2.redshift-serverless.amazonaws.com", + database="redshift", + cluster_identifier=None, + region=None, + auto_create=False, + db_groups=[], + db_user=None, + password="", + user="", + profile="iam_profile_test", + timeout=None, + port=5439, + group_federation=False, + **DEFAULT_SSL_CONFIG, + ) + + @mock.patch("redshift_connector.connect", MagicMock()) + def test_profile_ignore_cluster(self): + self.config.credentials = self.config.credentials.replace( + method="iam_role", + iam_profile="iam_profile_test", + host="doesnotexist.1233.us-east-2.redshift-serverless.amazonaws.com", + cluster_id="my_redshift", + ) + connection = self.adapter.acquire_connection("dummy") + connection.handle + redshift_connector.connect.assert_called_once_with( + iam=True, + host="doesnotexist.1233.us-east-2.redshift-serverless.amazonaws.com", + database="redshift", + cluster_identifier=None, + region=None, + auto_create=False, + db_groups=[], + db_user=None, + password="", + user="", + profile="iam_profile_test", + timeout=None, + port=5439, + group_federation=False, + **DEFAULT_SSL_CONFIG, + ) + + @mock.patch("redshift_connector.connect", MagicMock()) + def test_profile_explicit_region(self): + # Successful test + self.config.credentials = self.config.credentials.replace( + method="iam_role", + iam_profile="iam_profile_test", + host="doesnotexist.1233.redshift-serverless.amazonaws.com", + region="us-east-2", + ) + connection = self.adapter.acquire_connection("dummy") + connection.handle + redshift_connector.connect.assert_called_once_with( + iam=True, + host="doesnotexist.1233.redshift-serverless.amazonaws.com", + database="redshift", + cluster_identifier=None, + region="us-east-2", + auto_create=False, + db_groups=[], + db_user=None, + password="", + user="", + profile="iam_profile_test", + timeout=None, + port=5439, + group_federation=False, + **DEFAULT_SSL_CONFIG, + ) + + @mock.patch("redshift_connector.connect", MagicMock()) + def test_profile_invalid_serverless(self): + self.config.credentials = self.config.credentials.replace( + method="iam_role", + iam_profile="iam_profile_test", + host="doesnotexist.1233.us-east-2.redshift-srvrlss.amazonaws.com", + ) + with self.assertRaises(FailedToConnectError) as context: + connection = self.adapter.acquire_connection("dummy") + connection.handle + redshift_connector.connect.assert_called_once_with( + iam=True, + host="doesnotexist.1233.us-east-2.redshift-srvrlss.amazonaws.com", + database="redshift", + cluster_identifier=None, + region=None, + auto_create=False, + db_groups=[], + db_user=None, + password="", + user="", + profile="iam_profile_test", + port=5439, + timeout=None, + group_federation=False, + **DEFAULT_SSL_CONFIG, + ) + self.assertTrue("'host' must be provided" in context.exception.msg)