feat: Added support for Identity Enhanced Credentials authentication

Author: Sungum Sogani

redshift_connector/__init__.py

@@ -303,6 +303,12 @@ def connect(
         The access token to be used with IdC basic credentials provider plugin. Default value is None.
     token_type: Optional[str]
         The token type to be used for authentication using IdP token auth plugin. Default value is None.
+    access_key_id : Optional[str]
+        The access key for the user configured for IAM database authentication. Can also be used with IdpTokenAuthPlugin for identity-enhanced credentials flow.
+    secret_access_key : Optional[str]
+        The secret access key for the IAM role or IAM user configured for IAM database authentication. Can also be used with IdpTokenAuthPlugin for identity-enhanced credentials flow.
+    session_token : Optional[str]
+        The session token for temporary AWS credentials. Required when using temporary credentials with IAM authentication or IdpTokenAuthPlugin identity-enhanced credentials flow.
     Returns
     -------
     A Connection object associated with the specified Amazon Redshift cluster: :class:`Connection`
@@ -458,6 +464,9 @@ def connect(
         identity_namespace=info.identity_namespace,
         token_type=info.token_type,
         idc_client_display_name=info.idc_client_display_name,
+        access_key_id=info.access_key_id,
+        secret_access_key=info.secret_access_key,
+        session_token=info.session_token,
     )
 
 

redshift_connector/core.py

@@ -432,6 +432,9 @@ def __init__(
         identity_namespace: typing.Optional[str] = None,
         token_type: typing.Optional[str] = None,
         idc_client_display_name: typing.Optional[str] = None,
+        access_key_id: typing.Optional[str] = None,
+        secret_access_key: typing.Optional[str] = None,
+        session_token: typing.Optional[str] = None,
     ):
         """
         Creates a :class:`Connection` to an Amazon Redshift cluster. For more information on establishing a connection to an Amazon Redshift cluster using `federated API access <https://aws.amazon.com/blogs/big-data/federated-api-access-to-amazon-redshift-using-an-amazon-redshift-connector-for-python/>`_ see our examples page.
@@ -486,6 +489,12 @@ def __init__(
             The token type to be used for authentication using IdP Token auth plugin
         idc_client_display_name: Optional[str]
             The client display name to be used for user consent in IdC browser auth plugin.
+        access_key_id: Optional[str]
+            The AWS access key ID for identity-enhanced credentials flow with IdpTokenAuthPlugin.
+        secret_access_key: Optional[str]
+            The AWS secret access key for identity-enhanced credentials flow with IdpTokenAuthPlugin.
+        session_token: Optional[str]
+            The AWS session token for identity-enhanced credentials flow with IdpTokenAuthPlugin.
         """
         self.merge_socket_read = True
 
@@ -570,7 +579,10 @@ def get_calling_module() -> str:
                 "BrowserIdcAuthPlugin",
             ):
                 redshift_native_auth = True
-                self.set_idc_plugins_params(init_params, credentials_provider, identity_namespace, token_type)
+                self.set_idc_plugins_params(init_params, credentials_provider,
+                    identity_namespace, token_type, idc_client_display_name,
+                    access_key_id, secret_access_key, session_token
+                )
 
             if redshift_native_auth and provider_name:
                 init_params["provider_name"] = provider_name
@@ -2520,6 +2532,9 @@ def set_idc_plugins_params(
         identity_namespace: typing.Optional[str] = None,
         token_type: typing.Optional[str] = None,
         idc_client_display_name: typing.Optional[str] = None,
+        access_key_id: typing.Optional[str] = None,
+        secret_access_key: typing.Optional[str] = None,
+        session_token: typing.Optional[str] = None,
     ) -> None:
         plugin_name = typing.cast(str, credentials_provider).split(".")[-1]
         init_params["idp_type"] = "AwsIdc"
@@ -2529,6 +2544,9 @@ def set_idc_plugins_params(
 
         if plugin_name == "BrowserIdcAuthPlugin":
             init_params["token_type"] = "ACCESS_TOKEN"
+        elif plugin_name == "IdpTokenAuthPlugin" and access_key_id and secret_access_key and session_token:
+            # Check if using identity-enhanced credentials flow with IdPTokenAuthPlugin
+            init_params["token_type"] = "SUBJECT_TOKEN"
         elif token_type:
             init_params["token_type"] = token_type
 

redshift_connector/iam_helper.py

@@ -137,6 +137,11 @@ def set_iam_properties(info: RedshiftProperty) -> RedshiftProperty:
         if not info.region:
             info.set_region_from_host()
 
+        # Extract cluster_identifier from host for provisioned clusters if not already set
+        # This is needed for IDC plugins (like IdpTokenAuthPlugin) which don't use iam=True
+        if info.iam is not True and info.cluster_identifier is None and info.is_provisioned_host:
+            info.set_cluster_identifier_from_host()
+
         if info.iam is True:
             if info.region is None:
                 _logger.debug("Setting region via DNS lookup as region was not provided in connection parameters")