feat(serverless): support nlb connection

Author: Brooke-white

redshift_connector/__init__.py

@@ -180,6 +180,9 @@ def connect(
     provider_name: typing.Optional[str] = None,
     scope: typing.Optional[str] = None,
     numeric_to_float: typing.Optional[bool] = False,
+    is_serverless: typing.Optional[bool] = False,
+    serverless_acct_id: typing.Optional[str] = None,
+    serverless_work_group: typing.Optional[str] = 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.
@@ -275,6 +278,12 @@ def connect(
         Scope for BrowserAzureOauth2CredentialsProvider authentication.
     numeric_to_float: Optional[str]
         Specifies if NUMERIC datatype values will be converted from ``decimal.Decimal`` to ``float``. By default NUMERIC values are received as ``decimal.Decimal``.
+    is_serverless: Optional[bool]
+        Redshift end-point is serverless or provisional. Default value false.
+    serverless_acct_id: Optional[str]
+        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.
     Returns
     -------
     A Connection object associated with the specified Amazon Redshift cluster: :class:`Connection`
@@ -304,6 +313,7 @@ def connect(
     info.put("idp_host", idp_host)
     info.put("idp_response_timeout", idp_response_timeout)
     info.put("idp_tenant", idp_tenant)
+    info.put("is_serverless", is_serverless)
     info.put("listen_port", listen_port)
     info.put("login_url", login_url)
     info.put("max_prepared_statements", max_prepared_statements)
@@ -321,6 +331,8 @@ def connect(
     info.put("role_session_name", role_session_name)
     info.put("scope", scope)
     info.put("secret_access_key", secret_access_key)
+    info.put("serverless_acct_id", serverless_acct_id)
+    info.put("serverless_work_group", serverless_work_group)
     info.put("session_token", session_token)
     info.put("source_address", source_address)
     info.put("ssl", ssl)

redshift_connector/iam_helper.py

@@ -2,7 +2,9 @@
 import logging
 import typing
 
+import pkg_resources
 from dateutil.tz import tzutc
+from packaging.version import Version
 
 from redshift_connector.auth.aws_credentials_provider import AWSCredentialsProvider
 from redshift_connector.credentials_holder import (
@@ -35,21 +37,24 @@ def set_iam_properties(info: RedshiftProperty) -> RedshiftProperty:
         # set properties present for both IAM, Native authentication
         IamHelper.set_auth_properties(info)
 
-        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 and info.iam:
+            if Version(pkg_resources.get_distribution("boto3").version) < Version("1.24.5"):
+                raise pkg_resources.VersionConflict(
+                    "boto3 >= 1.24.5 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()
-            info.set_work_group_from_host()
+            # consider overridden connection parameters
+            if not info.region:
+                info.set_region_from_host()
+            if not info.serverless_acct_id:
+                info.set_serverless_acct_id()
+            if not info.serverless_work_group:
+                info.set_serverless_work_group_from_host()
 
         if info.iam is True:
-            if info.cluster_identifier is None and not info.is_serverless_host:
+            if info.cluster_identifier is None and not info._is_serverless:
                 raise InterfaceError(
                     "Invalid connection property setting. cluster_identifier must be provided when IAM is enabled"
                 )
@@ -136,8 +141,10 @@ def get_credentials_cache_key(info: RedshiftProperty, cred_provider: typing.Unio
                     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),
-                    typing.cast(str, info.work_group if info.is_serverless_host and info.work_group else ""),
+                    typing.cast(str, info.serverless_acct_id if info._is_serverless else info.cluster_identifier),
+                    typing.cast(
+                        str, info.serverless_work_group if info._is_serverless and info.serverless_work_group else ""
+                    ),
                     str(info.auto_create),
                     str(info.duration),
                     # v2 api parameters
@@ -171,7 +178,7 @@ def set_cluster_credentials(
             ] = cred_provider.get_credentials()  # type: ignore
             session_credentials: typing.Dict[str, str] = credentials_holder.get_session_credentials()
 
-            redshift_client: str = "redshift-serverless" if info.is_serverless_host else "redshift"
+            redshift_client: str = "redshift-serverless" if info._is_serverless 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)):
@@ -190,7 +197,7 @@ def set_cluster_credentials(
             if info.host is None or info.host == "" or info.port is None or info.port == "":
                 response: dict
 
-                if info.is_serverless_host:
+                if info._is_serverless:
                     response = client.describe_configuration()
                     info.put("host", response["endpoint"]["address"])
                     info.put("port", response["endpoint"]["port"])
@@ -218,10 +225,10 @@ 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_host:
+                if info._is_serverless:
                     get_cred_args: typing.Dict[str, str] = {"dbName": info.db_name}
-                    if info.work_group:
-                        get_cred_args["workgroupName"] = info.work_group
+                    if info.serverless_work_group:
+                        get_cred_args["workgroupName"] = info.serverless_work_group
 
                     cred = typing.cast(
                         typing.Dict[str, typing.Union[str, datetime.datetime]],
@@ -247,7 +254,7 @@ def set_cluster_credentials(
                         typing.Dict[str, typing.Union[str, datetime.datetime]], cred
                     )
             # redshift-serverless api json response payload slightly differs
-            if info.is_serverless_host:
+            if info._is_serverless:
                 info.put("user_name", typing.cast(str, cred["dbUser"]))
                 info.put("password", typing.cast(str, cred["dbPassword"]))
             else:

redshift_connector/idp_auth_helper.py

@@ -68,7 +68,7 @@ def set_auth_properties(info: RedshiftProperty):
             _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 and not info.is_serverless_host:
+            if info.cluster_identifier is None and not info._is_serverless:
                 raise InterfaceError(
                     "Invalid connection property setting. cluster_identifier must be provided when IAM is enabled"
                 )

redshift_connector/redshift_property.py

@@ -110,14 +110,13 @@ 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
             # The name of the Redshift Native Auth Provider
             self.provider_name: typing.Optional[str] = None
             self.scope: str = ""
             self.numeric_to_float: bool = False
-            # The work group used for Amazon serverless
-            self.work_group: typing.Optional[str] = None
+            self.is_serverless: bool = False
+            self.serverless_acct_id: typing.Optional[str] = None
+            self.serverless_work_group: typing.Optional[str] = None
 
         else:
             for k, v in kwargs.items():
@@ -126,6 +125,7 @@ def __init__(self: "RedshiftProperty", **kwargs):
     def __str__(self: "RedshiftProperty") -> str:
         rp = self.__dict__
         rp["is_serverless_host"] = self.is_serverless_host
+        rp["_is_serverless"] = self._is_serverless
         return str(rp)
 
     def put_all(self, other):
@@ -158,7 +158,14 @@ def is_serverless_host(self: "RedshiftProperty") -> bool:
             re.fullmatch(pattern=SERVERLESS_WITH_WORKGROUP_HOST_PATTERN, string=str(self.host))
         )
 
-    def set_account_id_from_host(self: "RedshiftProperty") -> None:
+    @property
+    def _is_serverless(self):
+        """
+        Returns True if host patches serverless pattern or if is_serverless flag set by user
+        """
+        return self.is_serverless_host or self.is_serverless
+
+    def set_serverless_acct_id(self: "RedshiftProperty") -> None:
         """
         Sets the AWS account id as parsed from the Redshift serverless endpoint.
         """
@@ -168,7 +175,7 @@ def set_account_id_from_host(self: "RedshiftProperty") -> None:
             m2 = re.fullmatch(pattern=serverless_pattern, string=self.host)
 
             if m2:
-                self.put(key="account_id", value=m2.group(typing.cast(int, m2.lastindex) - 1))
+                self.put(key="serverless_acct_id", value=m2.group(typing.cast(int, m2.lastindex) - 1))
                 break
 
     def set_region_from_host(self: "RedshiftProperty") -> None:
@@ -184,7 +191,7 @@ def set_region_from_host(self: "RedshiftProperty") -> None:
                 self.put(key="region", value=m2.group(typing.cast(int, m2.lastindex)))
                 break
 
-    def set_work_group_from_host(self: "RedshiftProperty") -> None:
+    def set_serverless_work_group_from_host(self: "RedshiftProperty") -> None:
         """
         Sets the work_group as parsed from the Redshift serverless endpoint.
         """
@@ -193,4 +200,4 @@ def set_work_group_from_host(self: "RedshiftProperty") -> None:
         m2 = re.fullmatch(pattern=SERVERLESS_WITH_WORKGROUP_HOST_PATTERN, string=self.host)
 
         if m2:
-            self.put(key="work_group", value=m2.group(1))
+            self.put(key="serverless_work_group", value=m2.group(1))

test/unit/test_iam_helper.py

@@ -364,9 +364,9 @@ def test_set_iam_credentials_for_serverless_calls_get_credentials(
     for k, v in serverless_iam_db_kwargs.items():
         rp.put(k, v)
 
-    rp.set_account_id_from_host()
+    rp.set_serverless_acct_id()
     rp.set_region_from_host()
-    rp.set_work_group_from_host()
+    rp.set_serverless_work_group_from_host()
 
     mock_cred_provider = MagicMock()
     mock_cred_holder = MagicMock()
@@ -378,8 +378,10 @@ def test_set_iam_credentials_for_serverless_calls_get_credentials(
     IamHelper.set_cluster_credentials(mock_cred_provider, rp)
 
     # ensure describe_configuration is called
-    if rp.work_group:
-        mock_boto_client.assert_has_calls([call().get_credentials(dbName=rp.db_name, workgroupName=rp.work_group)])
+    if rp.serverless_work_group:
+        mock_boto_client.assert_has_calls(
+            [call().get_credentials(dbName=rp.db_name, workgroupName=rp.serverless_work_group)]
+        )
     else:
         mock_boto_client.assert_has_calls([call().get_credentials(dbName=rp.db_name)])
 
@@ -391,6 +393,48 @@ def test_set_iam_credentials_for_serverless_calls_get_credentials(
     assert "password" in [c[0][0] for c in spy.call_args_list]
 
 
+def test_serverless_properties_used_when_is_serverless_true(mocker):
+    rp: RedshiftProperty = RedshiftProperty()
+    rp.host = "test-endpoint-xxxx.123456789123.us-east-2.redshift-serverless.amazonaws.com"
+    rp.is_serverless = True
+    rp.serverless_work_group = "something"
+    rp.serverless_acct_id = "111111111111"
+
+    mocker.patch(
+        "redshift_connector.native_plugin_helper.NativeAuthPluginHelper.get_native_auth_plugin_credentials",
+        return_value=None,
+    )
+
+    result = IamHelper.set_iam_properties(rp)
+
+    assert result.is_serverless == True
+    assert result.serverless_work_group == rp.serverless_work_group
+    assert result.serverless_acct_id == rp.serverless_acct_id
+
+
+def test_internal_is_serverless_prop_true_for_nlb_host():
+    rp: RedshiftProperty = RedshiftProperty()
+    rp.is_serverless = True
+
+    assert rp._is_serverless is True
+
+
+@pytest.mark.parametrize(
+    "serverless_host",
+    (
+        "testwg1.012345678901.us-east-2.redshift-serverless.amazonaws.com",
+        "012345678901.us-east-2.redshift-serverless.amazonaws.com",
+    ),
+)
+def test_internal_is_serverless_prop_true_for_serverless_host(serverless_host):
+    rp: RedshiftProperty = RedshiftProperty()
+    rp.host = serverless_host
+    rp.serverless_work_group = "something"
+    rp.serverless_acct_id = "111111111111"
+
+    assert rp._is_serverless is True
+
+
 def test_dynamically_loading_credential_holder(mocker):
     external_class_name: str = "test.unit.MockCredentialsProvider"
     mocker.patch("{}.get_credentials".format(external_class_name))
@@ -812,7 +856,7 @@ def test_set_iam_properties_calls_set_auth_props(mocker):
     spy = mocker.spy(IdpAuthHelper, "set_auth_properties")
     mock_rp: MagicMock = MagicMock()
     mock_rp.credentials_provider = None
-    mock_rp.is_serverless_host = False
+    mock_rp._is_serverless = False
     IamHelper.set_iam_properties(mock_rp)
 
     assert spy.called is True

test/unit/test_redshift_property.py

@@ -25,11 +25,11 @@ def test_is_serverless_host(host, exp_is_serverless):
         ("012345678901.ap-northeast-3.redshift-serverless.amazonaws.com", "012345678901"),
     ],
 )
-def test_set_account_id_from_host(host, exp_account_id):
+def test_set_serverless_acct_id_from_host(host, exp_account_id):
     info: RedshiftProperty = RedshiftProperty()
     info.host = host
-    info.set_account_id_from_host()
-    assert info.account_id == exp_account_id
+    info.set_serverless_acct_id()
+    assert info.serverless_acct_id == exp_account_id
 
 
 @pytest.mark.parametrize(
@@ -55,8 +55,8 @@ def test_set_region_from_host(host, exp_region):
         ("testwg2.012345678901.ap-northeast-3.redshift-serverless.amazonaws.com", "testwg2"),
     ],
 )
-def test_set_work_group_from_host(host, exp_work_group):
+def test_set_serverless_work_group_from_host(host, exp_work_group):
     info: RedshiftProperty = RedshiftProperty()
     info.host = host
-    info.set_work_group_from_host()
-    assert info.work_group == exp_work_group
+    info.set_serverless_work_group_from_host()
+    assert info.serverless_work_group == exp_work_group

tutorials/001 - Connecting to Amazon Redshift.ipynb

@@ -495,7 +495,7 @@
    "source": [
     "# Connecting to a Redshift Serverless Endpoint\n",
     "\n",
-    "Authentication methods discussed below are supported for Redshift serverless endpoints.\n",
+    "Authentication methods discussed below are supported for Redshift serverless endpoints. If connecting using a network load balancer (NLB) or Redshift-managed VPC endpoint please set ``is_serverless=True`` and specify workgroup and serverless account id information using ``serverless_acct_id`` and ``serverless_work_group``, respectively.\n",
     "\n",
     "### Using Database credentials (Native authentication)\n",
     "\n",
@@ -584,4 +584,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 1
-}
+}