feat(connection): support redshift serverless

Author: Brooke-white

redshift_connector/iam_helper.py

@@ -51,14 +51,14 @@ def set_iam_properties(info: RedshiftProperty) -> RedshiftProperty:
         Helper function to handle IAM connection properties and ensure required parameters are specified.
         Parameters
         """
+        import pkg_resources
+        from packaging.version import Version
+
         if info is None:
             raise InterfaceError("Invalid connection property setting. info must be specified")
 
         # Check for IAM keys and AuthProfile first
         if info.auth_profile is not None:
-            import pkg_resources
-            from packaging.version import Version
-
             if Version(pkg_resources.get_distribution("boto3").version) < Version("1.17.111"):
                 raise pkg_resources.VersionConflict(
                     "boto3 >= 1.17.111 required for authentication via Amazon Redshift authentication profile. "
@@ -83,6 +83,18 @@ def set_iam_properties(info: RedshiftProperty) -> RedshiftProperty:
                 )
                 info.put_all(resp)
 
+        if info.is_serverless_host and info.iam:
+            raise ProgrammingError("This feature is not yet available")
+            # if Version(pkg_resources.get_distribution("boto3").version) <= Version("1.20.22"):
+            #     raise pkg_resources.VersionConflict(
+            #         "boto3 >= XXX required for authentication with Amazon Redshift serverless. "
+            #         "Please upgrade the installed version of boto3 to use this functionality."
+            #     )
+
+        if info.is_serverless_host:
+            info.set_account_id_from_host()
+            info.set_region_from_host()
+
         # IAM requires an SSL connection to work.
         # Make sure that is set to SSL level VERIFY_CA or higher.
         if info.ssl is True:
@@ -109,8 +121,10 @@ def set_iam_properties(info: RedshiftProperty) -> RedshiftProperty:
                 "AWS credentials, Amazon Redshift authentication profile, or AWS profile"
             )
         elif info.iam is True:
+            _logger.debug("boto3 version: {}".format(Version(pkg_resources.get_distribution("boto3").version)))
+            _logger.debug("botocore version: {}".format(Version(pkg_resources.get_distribution("botocore").version)))
 
-            if info.cluster_identifier is None:
+            if info.cluster_identifier is None and not info.is_serverless_host:
                 raise InterfaceError(
                     "Invalid connection property setting. cluster_identifier must be provided when IAM is enabled"
                 )
@@ -174,7 +188,7 @@ def set_iam_properties(info: RedshiftProperty) -> RedshiftProperty:
             info.put("db_groups", [group.lower() for group in info.db_groups])
 
         if info.iam is True:
-            if info.cluster_identifier is None:
+            if info.cluster_identifier is None and not info.is_serverless_host:
                 raise InterfaceError(
                     "Invalid connection property setting. cluster_identifier must be provided when IAM is enabled"
                 )
@@ -305,20 +319,42 @@ def set_iam_credentials(info: RedshiftProperty) -> None:
         IamHelper.set_cluster_credentials(provider, info)
 
     @staticmethod
-    def get_credentials_cache_key(info: RedshiftProperty):
+    def get_credentials_cache_key(
+        info: RedshiftProperty, cred_provider: typing.Union[SamlCredentialsProvider, AWSCredentialsProvider]
+    ):
         db_groups: str = ""
 
         if len(info.db_groups) > 0:
             info.put("db_groups", sorted(info.db_groups))
             db_groups = ",".join(info.db_groups)
 
+        cred_key: str = ""
+
+        if cred_provider:
+            cred_key = str(cred_provider.get_cache_key())
+
         return ";".join(
-            (
-                typing.cast(str, info.db_user),
-                info.db_name,
-                db_groups,
-                typing.cast(str, info.cluster_identifier),
-                str(info.auto_create),
+            filter(
+                None,
+                (
+                    cred_key,
+                    typing.cast(str, info.db_user if info.db_user else info.user_name),
+                    info.db_name,
+                    db_groups,
+                    typing.cast(str, info.account_id if info.is_serverless_host else info.cluster_identifier),
+                    str(info.auto_create),
+                    str(info.duration),
+                    # v2 api parameters
+                    info.preferred_role,
+                    info.web_identity_token,
+                    info.role_arn,
+                    info.role_session_name,
+                    # providers
+                    info.profile,
+                    info.access_key_id,
+                    info.secret_access_key,
+                    info.session_token,
+                ),
             )
         )
 
@@ -339,6 +375,9 @@ def set_cluster_credentials(
             ] = cred_provider.get_credentials()
             session_credentials: typing.Dict[str, str] = credentials_holder.get_session_credentials()
 
+            redshift_client: str = "redshift-serverless" if info.is_serverless_host else "redshift"
+            _logger.debug("boto3.client(service_name={}) being used for IAM auth".format(redshift_client))
+
             for opt_key, opt_val in (("region_name", info.region), ("endpoint_url", info.endpoint_url)):
                 if opt_val is not None:
                     session_credentials[opt_key] = opt_val
@@ -348,21 +387,28 @@ def set_cluster_credentials(
                 cached_session: boto3.Session = typing.cast(
                     ABCAWSCredentialsHolder, credentials_holder
                 ).get_boto_session()
-                client = cached_session.client(service_name="redshift", region_name=info.region)
+                client = cached_session.client(service_name=redshift_client, region_name=info.region)
             else:
-                client = boto3.client(service_name="redshift", **session_credentials)
+                client = boto3.client(service_name=redshift_client, **session_credentials)
 
             if info.host is None or info.host == "" or info.port is None or info.port == "":
-                response = client.describe_clusters(ClusterIdentifier=info.cluster_identifier)
+                response: dict
 
-                info.put("host", response["Clusters"][0]["Endpoint"]["Address"])
-                info.put("port", response["Clusters"][0]["Endpoint"]["Port"])
+                if info.is_serverless_host:
+                    response = client.describe_configuration()
+                    info.put("host", response["endpoint"]["address"])
+                    info.put("port", response["endpoint"]["port"])
+                else:
+                    response = client.describe_clusters(ClusterIdentifier=info.cluster_identifier)
+                    info.put("host", response["Clusters"][0]["Endpoint"]["Address"])
+                    info.put("port", response["Clusters"][0]["Endpoint"]["Port"])
 
             cred: typing.Optional[typing.Dict[str, typing.Union[str, datetime.datetime]]] = None
 
             if info.iam_disable_cache is False:
+                _logger.debug("iam_disable_cache=False")
                 # temporary credentials are cached by redshift_connector and will be used if they have not expired
-                cache_key: str = IamHelper.get_credentials_cache_key(info)
+                cache_key: str = IamHelper.get_credentials_cache_key(info, cred_provider)
                 cred = IamHelper.credentials_cache.get(cache_key, None)
 
                 _logger.debug(
@@ -375,26 +421,42 @@ def set_cluster_credentials(
             if cred is None or typing.cast(datetime.datetime, cred["Expiration"]) < datetime.datetime.now(tz=tzutc()):
                 # retries will occur by default ref:
                 # https://boto3.amazonaws.com/v1/documentation/api/latest/guide/retries.html#legacy-retry-mode
-                cred = typing.cast(
-                    typing.Dict[str, typing.Union[str, datetime.datetime]],
-                    client.get_cluster_credentials(
-                        DbUser=info.db_user,
-                        DbName=info.db_name,
-                        DbGroups=info.db_groups,
-                        ClusterIdentifier=info.cluster_identifier,
-                        AutoCreate=info.auto_create,
-                    ),
-                )
+                _logger.debug("Credentials expired or not found...requesting from boto")
+                if info.is_serverless_host:
+                    cred = typing.cast(
+                        typing.Dict[str, typing.Union[str, datetime.datetime]],
+                        client.get_credentials(
+                            dbName=info.db_name,
+                        ),
+                    )
+                    # re-map expiration for compatibility with redshift credential response
+                    cred["Expiration"] = cred["expiration"]
+                    del cred["expiration"]
+                else:
+                    cred = typing.cast(
+                        typing.Dict[str, typing.Union[str, datetime.datetime]],
+                        client.get_cluster_credentials(
+                            DbUser=info.db_user,
+                            DbName=info.db_name,
+                            DbGroups=info.db_groups,
+                            ClusterIdentifier=info.cluster_identifier,
+                            AutoCreate=info.auto_create,
+                        ),
+                    )
 
                 if info.iam_disable_cache is False:
                     IamHelper.credentials_cache[cache_key] = typing.cast(
                         typing.Dict[str, typing.Union[str, datetime.datetime]], cred
                     )
+            # redshift-serverless api json response payload slightly differs
+            if info.is_serverless_host:
+                info.put("user_name", typing.cast(str, cred["dbUser"]))
+                info.put("password", typing.cast(str, cred["dbPassword"]))
+            else:
+                info.put("user_name", typing.cast(str, cred["DbUser"]))
+                info.put("password", typing.cast(str, cred["DbPassword"]))
 
-            info.put("user_name", typing.cast(str, cred["DbUser"]))
-            info.put("password", typing.cast(str, cred["DbPassword"]))
-
-            _logger.debug("Using temporary aws credentials with expiration: {}".format(cred["Expiration"]))
+            _logger.debug("Using temporary aws credentials with expiration: {}".format(cred.get("Expiration")))
 
         except botocore.exceptions.ClientError as e:
             _logger.error("ClientError: %s", e)

redshift_connector/redshift_property.py

@@ -1,6 +1,9 @@
 import typing
 
 from redshift_connector.config import DEFAULT_PROTOCOL_VERSION
+from redshift_connector.error import ProgrammingError
+
+SERVERLESS_HOST_PATTERN: str = r"(.+)\.(.+).redshift-serverless(-dev)?\.amazonaws\.com(.)*"
 
 
 class RedshiftProperty:
@@ -105,13 +108,17 @@ def __init__(self: "RedshiftProperty", **kwargs):
             # The user name.
             self.user_name: str = ""
             self.web_identity_token: typing.Optional[str] = None
+            # The AWS Account Id
+            self.account_id: typing.Optional[str] = None
 
         else:
             for k, v in kwargs.items():
                 setattr(self, k, v)
 
     def __str__(self: "RedshiftProperty") -> str:
-        return str(self.__dict__)
+        rp = self.__dict__
+        rp["is_serverless_host"] = self.is_serverless_host
+        return str(rp)
 
     def put_all(self, other):
         """
@@ -128,3 +135,37 @@ def put(self: "RedshiftProperty", key: str, value: typing.Any):
         """
         if value is not None:
             setattr(self, key, value)
+
+    @property
+    def is_serverless_host(self: "RedshiftProperty") -> bool:
+        """
+        If the host indicate Redshift serverless will be used for connection.
+        """
+        if not self.host:
+            return False
+
+        import re
+
+        return bool(re.fullmatch(pattern=SERVERLESS_HOST_PATTERN, string=str(self.host)))
+
+    def set_account_id_from_host(self: "RedshiftProperty") -> None:
+        """
+        Returns the AWS account id as parsed from the Redshift serverless endpoint.
+        """
+        import re
+
+        m2 = re.fullmatch(pattern=SERVERLESS_HOST_PATTERN, string=self.host)
+
+        if m2:
+            self.put(key="account_id", value=m2.group(1))
+
+    def set_region_from_host(self: "RedshiftProperty") -> None:
+        """
+        Returns the AWS region as parsed from the Redshift serverless endpoint.
+        """
+        import re
+
+        m2 = re.fullmatch(pattern=SERVERLESS_HOST_PATTERN, string=self.host)
+
+        if m2:
+            self.put(key="region", value=m2.group(2))

test/conftest.py

@@ -76,6 +76,37 @@ def perf_db_kwargs() -> typing.Dict[str, typing.Union[str, bool]]:
     return db_connect
 
 
+@pytest.fixture(scope="class")
+def serverless_native_db_kwargs() -> typing.Dict[str, str]:
+    db_connect = {
+        "database": conf.get("redshift-serverless", "database", fallback="mock_database"),
+        "host": conf.get(
+            "redshift-serverless", "host", fallback="012345678901.us-east-2.redshift-serverless.amazonaws.com"
+        ),
+        "user": conf.get("redshift-serverless", "user", fallback="mock_user"),
+        "password": conf.get("redshift-serverless", "password", fallback="mock_password"),
+    }
+
+    return db_connect
+
+
+@pytest.fixture(scope="class")
+def serverless_iam_db_kwargs() -> typing.Dict[str, typing.Union[str, bool]]:
+    db_connect = {
+        "database": conf.get("redshift-serverless", "database", fallback="mock_database"),
+        "iam": conf.getboolean("redshift-serverless", "iam", fallback=True),
+        "access_key_id": conf.get("redshift-serverless", "access_key_id", fallback="mock_access_key_id"),
+        "secret_access_key": conf.get("redshift-serverless", "secret_access_key", fallback="mock_secret_access_key"),
+        "session_token": conf.get("redshift-serverless", "session_token", fallback="mock_session_token"),
+        "region": conf.get("redshift-serverless", "region", fallback="mock_region"),
+        "host": conf.get(
+            "redshift-serverless", "host", fallback="012345678901.us-east-2.redshift-serverless.amazonaws.com"
+        ),
+    }
+
+    return db_connect  # type: ignore
+
+
 @pytest.fixture(scope="class")
 def okta_idp() -> typing.Dict[str, typing.Union[str, bool, int]]:
     db_connect = {

test/manual/test_redshift_serverless.py

@@ -0,0 +1,39 @@
+import pytest
+
+import redshift_connector
+
+"""
+These functional tests ensure connections to Redshift serverless can be established when
+using various authentication methods.
+
+Please note the pre-requisites were documented while this feature is under public preview,
+and are subject to change.
+
+Pre-requisites:
+1) Redshift serverless configuration
+2) EC2 instance configured for accessing Redshift serverless (i.e. in compatible VPC, subnet)
+3) Perform a sanity check using psql to ensure Redshift serverless connection can be established
+3) EC2 instance has Python installed
+4) Clone redshift_connector on EC2 instance and install
+
+How to use:
+1) Populate config.ini with the Redshift serverless endpoint and user authentication information
+2) Run this file with pytest
+"""
+
+
+def test_native_auth(serverless_native_db_kwargs):
+    with redshift_connector.connect(**serverless_native_db_kwargs):
+        pass
+
+
+def test_iam_auth(serverless_iam_db_kwargs):
+    with redshift_connector.connect(**serverless_iam_db_kwargs):
+        pass
+
+
+def test_idp_auth(okta_idp):
+    okta_idp["host"] = "my_redshift_serverless_endpoint"
+
+    with redshift_connector.connect(**okta_idp):
+        pass