feat(auth, iam): support group federation

Author: Brooke-white

redshift_connector/__init__.py

@@ -183,6 +183,7 @@ def connect(
     is_serverless: typing.Optional[bool] = False,
     serverless_acct_id: typing.Optional[str] = None,
     serverless_work_group: typing.Optional[str] = None,
+    group_federation: typing.Optional[bool] = None,
 ) -> Connection:
     """
     Establishes a :class:`Connection` to an Amazon Redshift cluster. This function validates user input, optionally authenticates using an identity provider plugin, then constructs a :class:`Connection` object.
@@ -284,6 +285,8 @@ def connect(
         The account ID of the serverless. Default value None
     serverless_work_group: Optional[str]
         The name of work group for serverless end point. Default value None.
+    group_federation: Optional[bool]
+        Use the IDP Groups in the Redshift. Default value False.
     Returns
     -------
     A Connection object associated with the specified Amazon Redshift cluster: :class:`Connection`
@@ -307,6 +310,7 @@ def connect(
     info.put("db_user", db_user)
     info.put("endpoint_url", endpoint_url)
     info.put("force_lowercase", force_lowercase)
+    info.put("group_federation", group_federation)
     info.put("host", host)
     info.put("iam", iam)
     info.put("iam_disable_cache", iam_disable_cache)

redshift_connector/iam_helper.py

@@ -1,4 +1,5 @@
 import datetime
+import enum
 import logging
 import typing
 
@@ -25,9 +26,33 @@
 
 
 class IamHelper(IdpAuthHelper):
+    class GetClusterCredentialsAPIType(enum.Enum):
+        SERVERLESS_V1 = "get_credentials()"
+        IAM_V1 = "get_cluster_credentials()"
+        IAM_V2 = "get_cluster_credentials_with_iam()"
+
+        @staticmethod
+        def can_support_v2(info: RedshiftProperty):
+            return Version(pkg_resources.get_distribution("boto3").version) >= Version("1.24.5")
 
     credentials_cache: typing.Dict[str, dict] = {}
 
+    @staticmethod
+    def get_cluster_credentials_api_type(info: RedshiftProperty):
+        if not info._is_serverless:
+            if not info.group_federation:
+                return IamHelper.GetClusterCredentialsAPIType.IAM_V1
+            elif (not info.credentials_provider) and IamHelper.GetClusterCredentialsAPIType.can_support_v2(info):
+                return IamHelper.GetClusterCredentialsAPIType.IAM_V2
+            else:
+                raise InterfaceError("Authentication with plugin is not supported for group federation")
+        elif not info.group_federation:
+            return IamHelper.GetClusterCredentialsAPIType.SERVERLESS_V1
+        elif (not info.credentials_provider) and IamHelper.GetClusterCredentialsAPIType.can_support_v2(info):
+            return IamHelper.GetClusterCredentialsAPIType.IAM_V2
+        else:
+            raise InterfaceError("Authentication with plugin is not supported for group federation")
+
     @staticmethod
     def set_iam_properties(info: RedshiftProperty) -> RedshiftProperty:
         """
@@ -225,7 +250,12 @@ def set_cluster_credentials(
                 # retries will occur by default ref:
                 # https://boto3.amazonaws.com/v1/documentation/api/latest/guide/retries.html#legacy-retry-mode
                 _logger.debug("Credentials expired or not found...requesting from boto")
-                if info._is_serverless:
+                get_creds_api_version: IamHelper.GetClusterCredentialsAPIType = (
+                    IamHelper.get_cluster_credentials_api_type(info)
+                )
+                _logger.debug("boto3 get_credentials api version: {} will be used".format(get_creds_api_version.value))
+
+                if get_creds_api_version == IamHelper.GetClusterCredentialsAPIType.SERVERLESS_V1:
                     get_cred_args: typing.Dict[str, str] = {"dbName": info.db_name}
                     if info.serverless_work_group:
                         get_cred_args["workgroupName"] = info.serverless_work_group
@@ -237,6 +267,15 @@ def set_cluster_credentials(
                     # re-map expiration for compatibility with redshift credential response
                     cred["Expiration"] = cred["expiration"]
                     del cred["expiration"]
+                elif get_creds_api_version == IamHelper.GetClusterCredentialsAPIType.IAM_V2:
+                    cred = typing.cast(
+                        typing.Dict[str, typing.Union[str, datetime.datetime]],
+                        client.get_cluster_credentials_with_iam(
+                            DbName=info.db_name,
+                            ClusterIdentifier=info.cluster_identifier,
+                            DurationSeconds=info.duration,
+                        ),
+                    )
                 else:
                     cred = typing.cast(
                         typing.Dict[str, typing.Union[str, datetime.datetime]],

redshift_connector/plugin/i_plugin.py

@@ -42,3 +42,7 @@ def refresh(self: "IPlugin") -> None:
         Refreshes the credentials, stored in :class:NativeTokenHolder, for the current plugin.
         """
         pass  # pragma: no cover
+
+    @abstractmethod
+    def set_group_federation(self: "IPlugin", group_federation: bool):
+        pass

redshift_connector/plugin/jwt_credentials_provider.py

@@ -41,10 +41,14 @@ def add_parameter(
         self.provider_name = info.provider_name
         self.ssl_insecure = info.ssl_insecure
         self.disable_cache = info.iam_disable_cache
+        self.group_federation = False
 
         if info.role_session_name is not None:
             self.role_session_name = info.role_session_name
 
+    def set_group_federation(self: "JwtCredentialsProvider", group_federation: bool):
+        self.group_federation = group_federation
+
     def get_credentials(self: "JwtCredentialsProvider") -> NativeTokenHolder:
         credentials: typing.Optional[NativeTokenHolder] = None
 

redshift_connector/plugin/saml_credentials_provider.py

@@ -35,6 +35,7 @@ def __init__(self: "SamlCredentialsProvider") -> None:
         self.auto_create: typing.Optional[bool] = None
         self.region: typing.Optional[str] = None
         self.principal: typing.Optional[str] = None
+        self.group_federation: bool = False
 
         self.cache: dict = {}
 
@@ -53,6 +54,9 @@ def add_parameter(self: "SamlCredentialsProvider", info: RedshiftProperty) -> No
         self.region = info.region
         self.principal = info.principal
 
+    def set_group_federation(self: "SamlCredentialsProvider", group_federation: bool):
+        self.group_federation = group_federation
+
     def get_sub_type(self) -> int:
         return IdpAuthHelper.SAML_PLUGIN
 

redshift_connector/redshift_property.py

@@ -117,6 +117,7 @@ def __init__(self: "RedshiftProperty", **kwargs):
             self.is_serverless: bool = False
             self.serverless_acct_id: typing.Optional[str] = None
             self.serverless_work_group: typing.Optional[str] = None
+            self.group_federation: bool = False
 
         else:
             for k, v in kwargs.items():

test/manual/auth/test_aws_credentials.py

@@ -31,3 +31,34 @@ def test_use_aws_credentials_default_profile():
     ) as con:
         with con.cursor() as cursor:
             cursor.execute("select 1")
+
+
+"""
+How to use:
+0) Generate credentials using instructions:  https://docs.aws.amazon.com/sdk-for-javascript/v2/developer-guide/getting-your-credentials.html
+1) In the connect method below,  specify the connection parameters
+3) Specify the AWS IAM credentials in the variables above
+4) Update iam_helper.py to include correct min version. line `Version(pkg_resources.get_distribution("boto3").version) > Version("9.99.9999"):`
+5) Manually execute this test
+"""
+
+
+@pytest.mark.skip(reason="manual")
+def test_use_get_cluster_credentials_with_iam(db_kwargs):
+    role_name = "groupFederationTest"
+    with redshift_connector.connect(**db_kwargs) as conn:
+        with conn.cursor() as cursor:
+            # https://docs.aws.amazon.com/redshift/latest/dg/r_CREATE_USER.html
+            cursor.execute('create user "IAMR:{}" with password disable;'.format(role_name))
+    with redshift_connector.connect(
+        iam=True,
+        database="replace_me",
+        cluster_identifier="replace_me",
+        region="replace_me",
+        profile="replace_me",  # contains credentials for AssumeRole groupFederationTest
+        group_federation=True,
+    ) as con:
+        with con.cursor() as cursor:
+            cursor.execute("select 1")
+            cursor.execute("select current_user")
+            assert cursor.fetchone()[0] == role_name

test/unit/test_iam_helper.py

@@ -670,6 +670,35 @@ def test_set_cluster_credentials_refreshes_stale_credentials(
     )
 
 
+@pytest.mark.parametrize(
+    "conn_params, exp_result",
+    (
+        ({"credentials_provider": "BrowserSamlCredentialsProvider"}, IamHelper.GetClusterCredentialsAPIType.IAM_V1),
+        ({"group_federation": True}, IamHelper.GetClusterCredentialsAPIType.IAM_V2),
+        ({"is_serverless": True}, IamHelper.GetClusterCredentialsAPIType.SERVERLESS_V1),
+        ({"is_serverless": True, "group_federation": True}, IamHelper.GetClusterCredentialsAPIType.IAM_V2),
+        (
+            {"group_federation": True, "credentials_provider": "BrowserSamlCredentialsProvider"},
+            "Authentication with plugin is not supported for group federation",
+        ),
+        (
+            {"is_serverless": True, "group_federation": True, "credentials_provider": "BrowserSamlCredentialsProvider"},
+            "Authentication with plugin is not supported for group federation",
+        ),
+    ),
+)
+def test_get_cluster_credentials_api_type_will_use_correct_api(conn_params, exp_result):
+    info = RedshiftProperty()
+    for param in conn_params.items():
+        info.put(param[0], param[1])
+
+    if isinstance(exp_result, IamHelper.GetClusterCredentialsAPIType):
+        assert IamHelper.get_cluster_credentials_api_type(info) == exp_result
+    else:
+        with pytest.raises(InterfaceError, match=exp_result):
+            IamHelper.get_cluster_credentials_api_type(info)
+
+
 @pytest.mark.parametrize(
     "boto3_version",
     (