Fix tests and add token authentication method to auth flow

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

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:

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
+        )