diff --git a/.changes/unreleased/Features-20241217-181340.yaml b/.changes/unreleased/Features-20241217-181340.yaml new file mode 100644 index 000000000..a2c1c523f --- /dev/null +++ b/.changes/unreleased/Features-20241217-181340.yaml @@ -0,0 +1,6 @@ +kind: Features +body: Add IdpTokenAuthPlugin authentication method. +time: 2024-12-17T18:13:40.281494-08:00 +custom: + Author: versusfacit + Issue: "898" diff --git a/dbt/adapters/redshift/connections.py b/dbt/adapters/redshift/connections.py index d93847634..b84c4a109 100644 --- a/dbt/adapters/redshift/connections.py +++ b/dbt/adapters/redshift/connections.py @@ -41,16 +41,26 @@ class IdentityCenterTokenType(StrEnum): ACCESS_TOKEN = "ACCESS_TOKEN" EXT_JWT = "EXT_JWT" + @classmethod + def validate(cls, token_type: str): + try: + cls(token_type) + except ValueError: + raise FailedToConnectError( + f"'token_type' must be set to one of {[token.value for token in iter(cls)]}" + ) + class RedshiftConnectionMethod(StrEnum): DATABASE = "database" IAM = "iam" IAM_ROLE = "iam_role" IAM_IDENTITY_CENTER_BROWSER = "browser_identity_center" + IAM_IDENTITY_CENTER_TOKEN = "iam_idc_token" @classmethod def uses_identity_center(cls, method: str) -> bool: - return method in (cls.IAM_IDENTITY_CENTER_BROWSER,) + return method in (cls.IAM_IDENTITY_CENTER_BROWSER, cls.IAM_IDENTITY_CENTER_TOKEN) @classmethod def is_iam(cls, method: str) -> bool: @@ -153,6 +163,10 @@ class RedshiftCredentials(Credentials): idc_client_display_name: Optional[str] = "Amazon Redshift driver" idp_response_timeout: Optional[int] = None + # token + token: Optional[str] = None + token_type: Optional[str] = None + _ALIASES = {"dbname": "database", "pass": "password"} @property @@ -323,6 +337,18 @@ def __iam_idc_browser_kwargs(credentials) -> Dict[str, Any]: return __iam_kwargs(credentials) | idc_kwargs + def __iam_idc_token_kwargs(credentials) -> Dict[str, Any]: + logger.debug("Connecting to Redshift with '{credentials.method}' credentials method") + + __validate_required_fields("iam_idc_token", ("method", "token", "token_type")) + IdentityCenterTokenType.validate(credentials.token_type) + + return __iam_kwargs(credentials) | { + "credentials_provider": "IdpTokenAuthPlugin", + "token": credentials.token, + "token_type": credentials.token_type, + } + # # Head of function execution # @@ -333,6 +359,7 @@ def __iam_idc_browser_kwargs(credentials) -> Dict[str, Any]: RedshiftConnectionMethod.IAM: __iam_user_kwargs, RedshiftConnectionMethod.IAM_ROLE: __iam_role_kwargs, RedshiftConnectionMethod.IAM_IDENTITY_CENTER_BROWSER: __iam_idc_browser_kwargs, + RedshiftConnectionMethod.IAM_IDENTITY_CENTER_TOKEN: __iam_idc_token_kwargs, } try: diff --git a/tests/unit/test_auth_method.py b/tests/unit/test_auth_method.py index 16d13268f..f81d90ca2 100644 --- a/tests/unit/test_auth_method.py +++ b/tests/unit/test_auth_method.py @@ -673,3 +673,50 @@ def test_invalid_adapter_missing_fields(self): "'idc_region', 'issuer_url' field(s) are required for 'browser_identity_center' credentials method" in context.exception.msg ) + + +class TestIAMIdcToken(AuthMethod): + @mock.patch("redshift_connector.connect", MagicMock()) + def test_profile_idc_token_all_required_fields(self): + """Same as all possible fields""" + self.config.credentials = self.config.credentials.replace( + method="iam_idc_token", + token="token", + token_type="ACCESS_TOKEN", + host="doesnotexist.1235.us-east-2.redshift-serverless.amazonaws.com", + ) + connection = self.adapter.acquire_connection("dummy") + connection.handle + redshift_connector.connect.assert_called_once_with( + iam=False, + host="doesnotexist.1235.us-east-2.redshift-serverless.amazonaws.com", + database="redshift", + cluster_identifier=None, + region=None, + auto_create=False, + db_groups=[], + password="", + user="", + timeout=None, + port=5439, + **DEFAULT_SSL_CONFIG, + credentials_provider="IdpTokenAuthPlugin", + token="token", + token_type="ACCESS_TOKEN", + ) + + @mock.patch("redshift_connector.connect", MagicMock()) + def test_invalid_idc_token_missing_field(self): + # Successful test + self.config.credentials = self.config.credentials.replace( + method="iam_idc_token", + token_type="ACCESS_TOKEN", + host="doesnotexist.1235.us-east-2.redshift-serverless.amazonaws.com", + ) + with self.assertRaises(FailedToConnectError) as context: + connection = self.adapter.acquire_connection("dummy") + connection.handle + assert ( + "'token' field(s) are required for 'iam_idc_token' credentials method" + in context.exception.msg + )