Skip to content

Commit

Permalink
Remove group_federation for serverless (#836)
Browse files Browse the repository at this point in the history
* 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]>
  • Loading branch information
Fleid and colin-rogers-dbt authored Jun 13, 2024
1 parent a83553b commit 885db2d
Show file tree
Hide file tree
Showing 3 changed files with 135 additions and 7 deletions.
6 changes: 6 additions & 0 deletions .changes/unreleased/Fixes-20240531-113620.yaml
Original file line number Diff line number Diff line change
@@ -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"
19 changes: 12 additions & 7 deletions dbt/adapters/redshift/connections.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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:"
Expand Down
117 changes: 117 additions & 0 deletions tests/unit/test_auth_method.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)

0 comments on commit 885db2d

Please sign in to comment.