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 00c53d6.

* 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 <[email protected]>

.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"

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

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)