feat(auth): Added support for Custom Cluster Names (CNAME) for Amazon Redshift Serverless

Author: Brooke-white

redshift_connector/iam_helper.py

@@ -223,7 +223,7 @@ def set_iam_credentials(info: RedshiftProperty) -> None:
         if not isinstance(provider, INativePlugin):
             # If the Redshift instance has been identified as using a custom domain name, the hostname must
             # be determined using the redshift client from boto3 API
-            if info.is_cname is True:
+            if info.is_cname is True and not info.is_serverless:
                 IamHelper.set_cluster_identifier(provider, info)
 
             # Redshift database credentials  will be determined using the redshift client from boto3 API
@@ -442,9 +442,18 @@ def set_cluster_credentials(
             if get_creds_api_version == IamHelper.GetClusterCredentialsAPIType.SERVERLESS_V1:
                 # https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/redshift-serverless/client/get_credentials.html#
                 get_cred_args: typing.Dict[str, str] = {"dbName": info.db_name}
+                # if a connection parameter for serverless workgroup is provided it will
+                # be preferred over providing the CustomDomainName. The reason for this
+                # is backwards compatibility with the following cases:
+                # 0/ Serverless with NLB
+                # 1/ Serverless with Custom Domain Name
+                # Providing the CustomDomainName parameter to getCredentials will lead to
+                # failure if the custom domain name is not registered with Redshift. Hence,
+                # the ordering of these conditions is important.
                 if info.serverless_work_group:
                     get_cred_args["workgroupName"] = info.serverless_work_group
-
+                elif info.is_cname:
+                    get_cred_args["customDomainName"] = info.host
                 _logger.debug("Calling get_credentials with parameters %s", get_cred_args)
                 cred = typing.cast(
                     typing.Dict[str, typing.Union[str, datetime.datetime]],

redshift_connector/redshift_property.py

@@ -280,6 +280,10 @@ def set_region_from_endpoint_lookup(self: "RedshiftProperty") -> None:
             ec2_instance_host: str = host_response[0]
             _logger.debug("underlying ec2 instance host %s", ec2_instance_host)
             ec2_region: str = ec2_instance_host.split(".")[1]
+
+            # https://docs.aws.amazon.com/vpc/latest/userguide/vpc-dns.html#vpc-dns-hostnames
+            if ec2_region == "compute-1":
+                ec2_region = "us-east-1"
             self.put(key="region", value=ec2_region)
         except:
             msg: str = "Unable to automatically determine AWS region from host {} port {}. Please check host and port connection parameters are correct.".format(

test/conftest.py

@@ -118,6 +118,19 @@ def provisioned_cname_db_kwargs() -> typing.Dict[str, str]:
     return db_connect
 
 
+@pytest.fixture(scope="class")
+def serverless_cname_db_kwargs() -> typing.Dict[str, typing.Union[str, bool]]:
+    db_connect = {
+        "database": conf.get("redshift-serverless-cname", "database", fallback="mock_database"),
+        "host": conf.get("redshift-serverless-cname", "host", fallback="cname.mytest.com"),
+        "db_user": conf.get("redshift-serverless-cname", "db_user", fallback="mock_user"),
+        "password": conf.get("redshift-serverless-cname", "password", fallback="mock_password"),
+        "is_serverless": conf.getboolean("redshift-serverless-cname", "is_serverless", fallback="mockboolean"),
+    }
+
+    return db_connect
+
+
 @pytest.fixture(scope="class")
 def okta_idp() -> typing.Dict[str, typing.Union[str, bool, int]]:
     db_connect = {

test/manual/test_redshift_custom_domain.py

@@ -76,3 +76,47 @@ def test_nlb_connect() -> None:
     }
     with redshift_connector.connect(**args):  # type: ignore
         pass
+
+
+@pytest.mark.skip(reason="manual")
+@pytest.mark.parametrize("sslmode", (SupportedSSLMode.VERIFY_CA, SupportedSSLMode.VERIFY_FULL))
+def test_serverless_iam_cname_connect(sslmode, serverless_cname_db_kwargs):
+    serverless_cname_db_kwargs["iam"] = True
+    serverless_cname_db_kwargs["profile"] = "default"
+    serverless_cname_db_kwargs["auto_create"] = True
+    serverless_cname_db_kwargs["ssl"] = True
+    serverless_cname_db_kwargs["sslmode"] = sslmode.value
+
+    with redshift_connector.connect(**serverless_cname_db_kwargs) as conn:
+        with conn.cursor() as cursor:
+            cursor.execute("select current_user")
+            print(cursor.fetchone())
+
+
+@pytest.mark.skip(reason="manual")
+@pytest.mark.parametrize("sslmode", (SupportedSSLMode.VERIFY_CA, SupportedSSLMode.VERIFY_FULL))
+def test_serverless_cname_connect(sslmode, serverless_cname_db_kwargs):
+    # this test requires aws default profile contains valid credentials that provide permissions for
+    # redshift-serverless:GetCredentials ( Only called from this test method)
+    import boto3
+
+    profile = "default"
+    client = boto3.client(
+        service_name="redshift-serverless",
+        region_name="us-east-1",
+    )
+    # fetch cluster credentials and pass them as driver connect parameters
+    response = client.get_credentials(
+        customDomainName=serverless_cname_db_kwargs["host"], dbName=serverless_cname_db_kwargs["database"]
+    )
+
+    serverless_cname_db_kwargs["sslmode"] = sslmode.value
+    serverless_cname_db_kwargs["ssl"] = True
+    serverless_cname_db_kwargs["user"] = response["dbUser"]
+    serverless_cname_db_kwargs["password"] = response["dbPassword"]
+    serverless_cname_db_kwargs["profile"] = profile
+
+    with redshift_connector.connect(**serverless_cname_db_kwargs) as conn:
+        with conn.cursor() as cursor:
+            cursor.execute("select current_user")
+            print(cursor.fetchone())

test/unit/test_iam_helper.py

@@ -520,6 +520,68 @@ def test_set_cluster_credentials_uses_custom_domain_name_if_custom_domain_name_c
     )
 
 
+@mock.patch("boto3.client.get_credentials")
+@mock.patch("boto3.client")
+def test_set_cluster_credentials_uses_custom_domain_name_if_custom_domain_name_serverless(
+    mock_boto_client, mock_get_cluster_credentials
+):
+    mock_cred_provider = MagicMock()
+    mock_cred_holder = MagicMock()
+    mock_cred_provider.get_credentials.return_value = mock_cred_holder
+    mock_cred_provider.get_cache_key.return_value = "mocked"
+    mock_cred_holder.has_associated_session = False
+
+    rp: RedshiftProperty = make_redshift_property()
+    rp.put("cluster_identifier", None)
+    rp.put("host", "mycustom.domain.name")
+    rp.put("is_serverless", True)
+    rp.put("is_cname", True)
+
+    IamHelper.credentials_cache.clear()
+
+    IamHelper.set_cluster_credentials(mock_cred_provider, rp)
+
+    assert mock_boto_client.called is True
+    mock_boto_client.assert_has_calls(
+        [
+            call().get_credentials(
+                customDomainName=rp.host,
+                dbName=rp.db_name,
+            )
+        ]
+    )
+
+
+@mock.patch("boto3.client.get_credentials")
+@mock.patch("boto3.client")
+def test_set_cluster_credentials_uses_workgroup_if_nlb_serverless(mock_boto_client, mock_get_cluster_credentials):
+    mock_cred_provider = MagicMock()
+    mock_cred_holder = MagicMock()
+    mock_cred_provider.get_credentials.return_value = mock_cred_holder
+    mock_cred_provider.get_cache_key.return_value = "mocked"
+    mock_cred_holder.has_associated_session = False
+
+    rp: RedshiftProperty = make_redshift_property()
+    rp.put("cluster_identifier", None)
+    rp.put("host", "mycustom.domain.name")
+    rp.put("is_serverless", True)
+    rp.put("serverless_work_group", "xyz")
+    rp.put("serverless_acct_id", "012345678901")
+    IamHelper.credentials_cache.clear()
+
+    IamHelper.set_cluster_credentials(mock_cred_provider, rp)
+
+    assert mock_boto_client.called is True
+    mock_boto_client.assert_has_calls(
+        [
+            call().get_credentials(
+                workgroupName=rp.serverless_work_group,
+                dbName=rp.db_name,
+            )
+        ]
+    )
+
+
 @mock.patch("boto3.client.get_cluster_credentials")
 @mock.patch("boto3.client")
 def test_set_cluster_credentials_caches_credentials(mock_boto_client, mock_get_cluster_credentials) -> None: