Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Remove group_federation for serverless #836

Merged
merged 16 commits into from
Jun 13, 2024
Merged
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: Remove group_federation when using serverless for iam authentication
Fleid marked this conversation as resolved.
Show resolved Hide resolved
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 @@ -234,10 +234,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:
Fleid marked this conversation as resolved.
Show resolved Hide resolved
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 @@ -253,10 +258,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)
Loading