feat: support Redshift native authentication, Add Azure Oauth2 IdP

Author: Brooke-white

redshift_connector/__init__.py

@@ -2,7 +2,7 @@
 import typing
 
 from redshift_connector import plugin
-from redshift_connector.config import DEFAULT_PROTOCOL_VERSION
+from redshift_connector.config import DEFAULT_PROTOCOL_VERSION, ClientProtocolVersion
 from redshift_connector.core import BINARY, Connection, Cursor
 from redshift_connector.error import (
     ArrayContentNotHomogenousError,
@@ -177,6 +177,8 @@ def connect(
     iam_disable_cache: typing.Optional[bool] = None,
     auth_profile: typing.Optional[str] = None,
     endpoint_url: typing.Optional[str] = None,
+    provider_name: typing.Optional[str] = None,
+    scope: 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.
@@ -266,6 +268,10 @@ def connect(
         The name of an Amazon Redshift Authentication profile having connection properties as JSON. See :class:RedshiftProperty to learn how connection properties should be named.
     endpoint_url: Optional[str]
         The Amazon Redshift endpoint url. This option is only used by AWS internal teams.
+    provider_name: Optional[str]
+        The name of the Redshift Native Auth Provider.
+    scope: Optional[str]
+        Scope for BrowserAzureOauth2CredentialsProvider authentication.
     Returns
     -------
     A Connection object associated with the specified Amazon Redshift cluster: :class:`Connection`
@@ -304,10 +310,12 @@ def connect(
     info.put("preferred_role", preferred_role)
     info.put("principal", principal_arn)
     info.put("profile", profile)
+    info.put("provider_name", provider_name)
     info.put("region", region)
     info.put("replication", replication)
     info.put("role_arn", role_arn)
     info.put("role_session_name", role_session_name)
+    info.put("scope", scope)
     info.put("secret_access_key", secret_access_key)
     info.put("session_token", session_token)
     info.put("source_address", source_address)
@@ -326,7 +334,27 @@ def connect(
     _logger.debug(mask_secure_info_in_props(info).__str__())
     _logger.debug(make_divider_block())
 
-    IamHelper.set_iam_properties(info)
+    if (info.ssl is False) and (info.iam is True):
+        raise InterfaceError("Invalid connection property setting. SSL must be enabled when using IAM")
+
+    if (info.iam is False) and (info.ssl_insecure is False):
+        raise InterfaceError("Invalid connection property setting. IAM must be enabled when using ssl_insecure")
+
+    if info.client_protocol_version not in ClientProtocolVersion.list():
+        raise InterfaceError(
+            "Invalid connection property setting. client_protocol_version must be in: {}".format(
+                ClientProtocolVersion.list()
+            )
+        )
+
+    redshift_native_auth: bool = False
+    if info.iam:
+        if info.credentials_provider == "BasicJwtCredentialsProvider":
+            redshift_native_auth = True
+            _logger.debug("redshift_native_auth enabled")
+
+    if not redshift_native_auth:
+        IamHelper.set_iam_properties(info)
 
     _logger.debug(make_divider_block())
     _logger.debug("Connection arguments following validation and IAM auth (if applicable)")
@@ -352,6 +380,8 @@ def connect(
         client_protocol_version=info.client_protocol_version,
         database_metadata_current_db_only=info.database_metadata_current_db_only,
         credentials_provider=info.credentials_provider,
+        provider_name=info.provider_name,
+        web_identity_token=info.web_identity_token,
     )
 
 

redshift_connector/core.py

@@ -417,6 +417,8 @@ def __init__(
         client_protocol_version: int = DEFAULT_PROTOCOL_VERSION,
         database_metadata_current_db_only: bool = True,
         credentials_provider: typing.Optional[str] = None,
+        provider_name: typing.Optional[str] = None,
+        web_identity_token: typing.Optional[str] = None,
     ):
         """
         Creates a :class:`Connection` to an Amazon Redshift cluster. For more information on establishing a connection to an Amazon Redshift cluster using `federated API access <https://aws.amazon.com/blogs/big-data/federated-api-access-to-amazon-redshift-using-an-amazon-redshift-connector-for-python/>`_ see our examples page.
@@ -455,6 +457,10 @@ def __init__(
             Is `datashare <https://docs.aws.amazon.com/redshift/latest/dg/datashare-overview.html>`_ disabled. Default value is True, implying datasharing will not be used.
         credentials_provider : Optional[str]
             The class-path of the IdP plugin used for authentication with Amazon Redshift.
+        provider_name : Optional[str]
+            The name of the Redshift Native Auth Provider.
+        web_identity_token: Optional[str]
+            A web identity token used for authentication via Redshift Native IDP Integration
         """
         self.merge_socket_read = True
 
@@ -483,11 +489,15 @@ def __init__(
         # for receiving some datatypes
         self._enable_protocol_based_conversion_funcs()
 
+        self.web_identity_token = web_identity_token
+
         if user is None:
             raise InterfaceError("The 'user' connection parameter cannot be None")
 
+        redshift_native_auth: bool = False
+
         init_params: typing.Dict[str, typing.Optional[typing.Union[str, bytes]]] = {
-            "user": user,
+            "user": "",
             "database": database,
             "application_name": application_name,
             "replication": replication,
@@ -499,6 +509,19 @@ def __init__(
         if credentials_provider:
             init_params["plugin_name"] = credentials_provider
 
+            if credentials_provider.split(".")[-1] in (
+                "BasicJwtCredentialsProvider",
+                "BrowserAzureOAuth2CredentialsProvider",
+            ):
+                redshift_native_auth = True
+                init_params["idp_type"] = "AzureAD"
+
+                if provider_name:
+                    init_params["provider_name"] = provider_name
+
+        if not redshift_native_auth or user:
+            init_params["user"] = user
+
         _logger.debug(make_divider_block())
         _logger.debug("Establishing a connection")
         _logger.debug(init_params)
@@ -512,7 +535,10 @@ def __init__(
             elif not isinstance(v, (bytes, bytearray)):
                 raise InterfaceError("The parameter " + k + " can't be of type " + str(type(v)) + ".")
 
-        self.user: bytes = typing.cast(bytes, init_params["user"])
+        if "user" in init_params:
+            self.user: bytes = typing.cast(bytes, init_params["user"])
+        else:
+            self.user = b""
 
         if isinstance(password, str):
             self.password: bytes = password.encode("utf8")
@@ -1217,6 +1243,7 @@ def handle_AUTHENTICATION_REQUEST(self: "Connection", data: bytes, cursor: Curso
                   7 = GSSAPI (not supported)
                   8 = GSSAPI data (not supported)
                   9 = SSPI (not supported)
+                  14 = Redshift Native IDP Integration
 
         Please note that some authentication messages have additional data following the authentication code.
         That data is documented in the appropriate conditional branch below.
@@ -1284,6 +1311,22 @@ def handle_AUTHENTICATION_REQUEST(self: "Connection", data: bytes, cursor: Curso
         elif auth_code == 12:
             # AuthenticationSASLFinal
             self.auth.set_server_final(data[4:].decode("utf8"))
+        elif auth_code == 14:
+            # Redshift Native IDP Integration
+            aad_token: str = typing.cast(str, self.web_identity_token)
+            _logger.debug("<=BE Authentication request IDP")
+
+            if not aad_token:
+                raise ConnectionAbortedError(
+                    "The server requested AAD token-based authentication, but no token was provided."
+                )
+
+            _logger.debug("FE=> IDP(AAD Token)")
+
+            token: bytes = aad_token.encode(encoding="utf-8")
+            self._write(create_message(b"i", token))
+            # self._write(NULL_BYTE)
+            self._flush()
 
         elif auth_code in (2, 4, 6, 7, 8, 9):
             raise InterfaceError("Authentication method " + str(auth_code) + " not supported by redshift_connector.")