feat(auth): Support Redshift custom domain name

Author: Brooke-white

.pre-commit-config.yaml

@@ -11,7 +11,7 @@ repos:
       - id: isort
         args: ["--profile", "black", "."]
   - repo: https://github.com/ambv/black
-    rev: 20.8b1
+    rev: 22.3.0
     hooks:
       - id: black
   - repo: https://github.com/pre-commit/mirrors-mypy

redshift_connector/core.py

@@ -394,6 +394,20 @@ def client_os_version(self: "Connection") -> str:
             os_version = "unknown"
         return os_version
 
+    @staticmethod
+    def __get_host_address_info(host: str, port: int):
+        """
+        Returns IPv4 address and port given a host name and port
+        """
+        # https://docs.python.org/3/library/socket.html#socket.getaddrinfo
+        response = socket.getaddrinfo(host=host, port=port, family=socket.AF_INET)
+        _logger.debug("getaddrinfo response {}".format(response))
+
+        if not response:
+            raise InterfaceError("Unable to determine ip for host {} port {}".format(host, port))
+
+        return response[0][4]
+
     def __init__(
         self: "Connection",
         user: str,
@@ -593,7 +607,11 @@ def get_calling_module() -> str:
                 self._usock.settimeout(timeout)
 
             if unix_sock is None and host is not None:
-                self._usock.connect((host, port))
+                hostport: typing.Tuple[str, int] = Connection.__get_host_address_info(host, port)
+                _logger.debug(
+                    "Attempting to create connection socket with address {} {}".format(hostport[0], str(hostport[1]))
+                )
+                self._usock.connect(hostport)
             elif unix_sock is not None:
                 self._usock.connect(unix_sock)
 

redshift_connector/iam_helper.py

@@ -108,7 +108,7 @@ def set_auth_properties(info: RedshiftProperty):
                     )
                     info.put_all(resp)
 
-            if info.cluster_identifier is None and not info._is_serverless:
+            if info.cluster_identifier is None and not info._is_serverless and not info.is_cname:
                 raise InterfaceError(
                     "Invalid connection property setting. cluster_identifier must be provided when IAM is enabled"
                 )

redshift_connector/idp_auth_helper.py

@@ -1,11 +1,14 @@
+import logging
 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(.)*"
 SERVERLESS_WITH_WORKGROUP_HOST_PATTERN: str = r"(.+)\.(.+)\.(.+).redshift-serverless(-dev)?\.amazonaws\.com(.)*"
 IAM_URL_PATTERN: str = r"^(https)://[-a-zA-Z0-9+&@#/%?=~_!:,.']*[-a-zA-Z0-9+&@#/%=~_']"
+PROVISIONED_HOST_PATTERN: str = r"(.+)\.(.+)\.(.+).redshift(-dev)?\.amazonaws\.com(.)*"
+
+_logger: logging.Logger = logging.getLogger(__name__)
 
 
 class RedshiftProperty:
@@ -118,6 +121,8 @@ def __init__(self: "RedshiftProperty", **kwargs):
             self.serverless_acct_id: typing.Optional[str] = None
             self.serverless_work_group: typing.Optional[str] = None
             self.group_federation: bool = False
+            # flag indicating if host name and RedshiftProperty indicate Redshift with custom domain name is used
+            self.is_cname: bool = False
 
         else:
             for k, v in kwargs.items():
@@ -162,9 +167,44 @@ def is_serverless_host(self: "RedshiftProperty") -> bool:
         )
 
     @property
-    def _is_serverless(self):
+    def is_provisioned_host(self: "RedshiftProperty") -> bool:
+        """
+        Returns True if host matches Regex for Redshift provisioned. Otherwise returns False.
+        """
+        if not self.host:
+            return False
+
+        import re
+
+        return bool(re.fullmatch(pattern=PROVISIONED_HOST_PATTERN, string=str(self.host)))
+
+    def set_is_cname(self: "RedshiftProperty") -> None:
+        """
+        Sets RedshiftProperty is_cname attribute based on RedshiftProperty attribute values and host name Regex matching.
+        """
+        is_cname: bool = False
+        _logger.debug("determining if host indicates Redshift instance with custom name")
+
+        if self.is_provisioned_host:
+            _logger.debug("cluster identified as Redshift provisioned")
+        elif self.is_serverless_host:
+            _logger.debug("cluster identified as Redshift serverless")
+        elif self.is_serverless:
+            if self.serverless_work_group is not None:
+                _logger.debug("cluster identified as Redshift serverless with NLB")
+            else:
+                _logger.debug("cluster identified as Redshift serverless with with custom name")
+                is_cname = True
+        else:
+            _logger.debug("cluster identified as Redshift provisioned with with custom name/NLB")
+            is_cname = True
+
+        self.put(key="is_cname", value=is_cname)
+
+    @property
+    def _is_serverless(self: "RedshiftProperty"):
         """
-        Returns True if host patches serverless pattern or if is_serverless flag set by user
+        Returns True if host matches serverless pattern or if is_serverless flag set by user. Otherwise returns False.
         """
         return self.is_serverless_host or self.is_serverless
 

redshift_connector/redshift_property.py

@@ -107,6 +107,18 @@ def serverless_iam_db_kwargs() -> typing.Dict[str, typing.Union[str, bool]]:
     return db_connect  # type: ignore
 
 
+@pytest.fixture(scope="class")
+def provisioned_cname_db_kwargs() -> typing.Dict[str, str]:
+    db_connect = {
+        "database": conf.get("redshift-provisioned-cname", "database", fallback="mock_database"),
+        "host": conf.get("redshift-provisioned-cname", "host", fallback="cname.mytest.com"),
+        "db_user": conf.get("redshift-provisioned-cname", "db_user", fallback="mock_user"),
+        "password": conf.get("redshift-provisioned-cname", "password", fallback="mock_password"),
+    }
+
+    return db_connect
+
+
 @pytest.fixture(scope="class")
 def okta_idp() -> typing.Dict[str, typing.Union[str, bool, int]]:
     db_connect = {

test/conftest.py

@@ -66,16 +66,11 @@ def fin():
     request.addfinalizer(fin)
 
 
-def test_socket_missing():
-    conn_params = {
-        "unix_sock": "/file-does-not-exist",
-        "user": "doesn't-matter",
-        "password": "hunter2",
-        "database": "myDb",
-    }
+def test_socket_missing(db_kwargs):
+    db_kwargs["unix_sock"] = "/file-does-not-exist"
 
     with pytest.raises(redshift_connector.InterfaceError):
-        redshift_connector.connect(**conn_params)
+        redshift_connector.connect(**db_kwargs)
 
 
 def test_database_missing(db_kwargs):
@@ -150,7 +145,7 @@ def test_unicode_database_name(db_kwargs):
 
 
 def test_bytes_database_name(db_kwargs):
-    """ Should only raise an exception saying db doesn't exist """
+    """Should only raise an exception saying db doesn't exist"""
 
     db_kwargs["database"] = bytes("redshift_connector_sn\uFF6Fw", "utf8")
     with pytest.raises(redshift_connector.ProgrammingError, match="3D000"):

test/integration/test_connection.py

@@ -0,0 +1,79 @@
+import pytest
+
+import redshift_connector
+from redshift_connector.idp_auth_helper import SupportedSSLMode
+
+"""
+These functional tests ensure connections to Redshift provisioned customer with custom domain name can be established
+when using various authentication methods.
+
+Pre-requisites:
+1) Redshift provisioned configuration
+2) Existing custom domain association with instance created in step 1
+"""
+
+
+# @pytest.mark.skip(reason="manual")
+@pytest.mark.parametrize("sslmode", (SupportedSSLMode.VERIFY_CA, SupportedSSLMode.VERIFY_FULL))
+def test_native_connect(provisioned_cname_db_kwargs, sslmode):
+    # this test requires aws default profile contains valid credentials that provide permissions for
+    # redshift:GetClusterCredentials ( Only called from this test method)
+    import boto3
+
+    profile = "default"
+    client = boto3.client(
+        service_name="redshift",
+        region_name="eu-north-1",
+    )
+    # fetch cluster credentials and pass them as driver connect parameters
+    response = client.get_cluster_credentials(
+        CustomDomainName=provisioned_cname_db_kwargs["host"], DbUser=provisioned_cname_db_kwargs["db_user"]
+    )
+
+    provisioned_cname_db_kwargs["password"] = response["DbPassword"]
+    provisioned_cname_db_kwargs["user"] = response["DbUser"]
+    provisioned_cname_db_kwargs["profile"] = profile
+    provisioned_cname_db_kwargs["ssl"] = True
+    provisioned_cname_db_kwargs["sslmode"] = sslmode.value
+
+    with redshift_connector.connect(**provisioned_cname_db_kwargs):
+        pass
+
+
+# @pytest.mark.skip(reason="manual")
+@pytest.mark.parametrize("sslmode", (SupportedSSLMode.VERIFY_CA, SupportedSSLMode.VERIFY_FULL))
+def test_iam_connect(provisioned_cname_db_kwargs, sslmode):
+    # this test requires aws default profile contains valid credentials that provide permissions for
+    # redshift:GetClusterCredentials (called from driver)
+    # redshift:DescribeClusters (called from driver)
+    # redshift:DescribeCustomDomainAssociations (called from driver)
+    provisioned_cname_db_kwargs["iam"] = True
+    provisioned_cname_db_kwargs["profile"] = "default"
+    provisioned_cname_db_kwargs["auto_create"] = True
+    provisioned_cname_db_kwargs["region"] = "eu-north-1"
+    provisioned_cname_db_kwargs["ssl"] = True
+    provisioned_cname_db_kwargs["sslmode"] = sslmode.value
+    with redshift_connector.connect(**provisioned_cname_db_kwargs):
+        pass
+
+
+def test_idp_connect(okta_idp, provisioned_cname_db_kwargs):
+    # todo
+    pass
+
+
+# @pytest.mark.skip(reason="manual")
+def test_nlb_connect():
+    args = {
+        "iam": True,
+        # "access_key_id": "xxx",
+        # "secret_access_key": "xxx",
+        "cluster_identifier": "replace-me",
+        "region": "us-east-1",
+        "host": "replace-me",
+        "port": 5439,
+        "database": "dev",
+        "db_user": "replace-me",
+    }
+    with redshift_connector.connect(**args):
+        pass