Support additional connection methods

.changes/unreleased/Features-20230206-112433.yaml

@@ -0,0 +1,5 @@
+kind: Features
+time: 2023-02-06T11:24:33.926088-08:00
+custom:
+  Author: jiezhen-chen
+  Issue: "6232"

dbt/adapters/redshift/connections.py

@@ -1,3 +1,4 @@
+import os
 import re
 from multiprocessing import Lock
 from contextlib import contextmanager
@@ -39,13 +40,16 @@ def json_schema(self):
 class RedshiftConnectionMethod(StrEnum):
     DATABASE = "database"
     IAM = "iam"
+    AUTH_PROFILE = "auth_profile"
+    IDP = "IdP"
 
 
 @dataclass
 class RedshiftCredentials(Credentials):
-    host: str
-    user: str
     port: Port
+    region: Optional[str] = None
+    host: Optional[str] = None
+    user: Optional[str] = None
     method: str = RedshiftConnectionMethod.DATABASE  # type: ignore
     password: Optional[str] = None  # type: ignore
     cluster_id: Optional[str] = field(
@@ -59,7 +63,19 @@ class RedshiftCredentials(Credentials):
     connect_timeout: int = 30
     role: Optional[str] = None
     sslmode: Optional[str] = None
+    application_name: Optional[str] = "dbt"
     retries: int = 1
+    auth_profile: Optional[str] = None
+    # Azure identity provider plugin
+    credentials_provider: Optional[str] = None
+    idp_tenant: Optional[str] = None
+    client_id: Optional[str] = None
+    client_secret: Optional[str] = None
+    preferred_role: Optional[str] = None
+    # Okta identity provider plugin
+    idp_host: Optional[str] = None
+    app_id: Optional[str] = None
+    app_name: Optional[str] = None
 
     _ALIASES = {"dbname": "database", "pass": "password"}
 
@@ -68,11 +84,20 @@ def type(self):
         return "redshift"
 
     def _connection_keys(self):
-        return "host", "port", "user", "database", "schema", "method", "cluster_id", "iam_profile"
+        return (
+            "host",
+            "port",
+            "user",
+            "database",
+            "schema",
+            "method",
+            "cluster_id",
+            "iam_profile",
+        )
 
     @property
     def unique_field(self) -> str:
-        return self.host
+        return self.host if self.host else self.database
 
 
 class RedshiftConnectMethodFactory:
@@ -84,14 +109,18 @@ def __init__(self, credentials):
     def get_connect_method(self):
         method = self.credentials.method
         kwargs = {
-            "host": self.credentials.host,
+            "host": "",
+            "region": self.credentials.region,
             "database": self.credentials.database,
             "port": self.credentials.port if self.credentials.port else 5439,
             "auto_create": self.credentials.autocreate,
             "db_groups": self.credentials.db_groups,
-            "region": self.credentials.host.split(".")[2],
             "timeout": self.credentials.connect_timeout,
+            "application_name": self.credentials.application_name,
         }
+        if method == RedshiftConnectionMethod.IAM or method == RedshiftConnectionMethod.DATABASE:
+            kwargs["host"] = self.credentials.host
+            kwargs["region"] = self.credentials.host.split(".")[2]
         if self.credentials.sslmode:
             kwargs["sslmode"] = self.credentials.sslmode
 
@@ -107,14 +136,115 @@ def get_connect_method(self):
             def connect():
                 logger.debug("Connecting to redshift with username/password based auth...")
                 c = redshift_connector.connect(
-                    user=self.credentials.user, password=self.credentials.password, **kwargs
+                    user=self.credentials.user,
+                    password=self.credentials.password,
+                    **kwargs,
+                )
+                if self.credentials.role:
+                    c.cursor().execute("set role {}".format(self.credentials.role))
+                return c
+
+            return connect
+
+        elif method == RedshiftConnectionMethod.AUTH_PROFILE:
+            if not self.credentials.auth_profile:
+                raise dbt.exceptions.FailedToConnectError(
+                    "Failed to use auth profile method. 'auth_profile' must be provided."
+                )
+            if not self.credentials.region:
+                raise dbt.exceptions.FailedToConnectError(
+                    "Failed to use auth profile method. 'region' must be provided."
+                )
+
+            def connect():
+                logger.debug("Connecting to redshift with authentication profile...")
+                c = redshift_connector.connect(
+                    iam=True,
+                    access_key_id=os.environ["AWS_ACCESS_KEY_ID"],
+                    secret_access_key=os.environ["AWS_SECRET_ACCESS_KEY"],
+                    session_token=os.environ["AWS_SESSION_TOKEN"],
+                    db_user=self.credentials.user,
+                    auth_profile=self.credentials.auth_profile,
+                    **kwargs,
                 )
                 if self.credentials.role:
                     c.cursor().execute("set role {}".format(self.credentials.role))
                 return c
 
             return connect
+        elif method == RedshiftConnectionMethod.IDP:
+            if not self.credentials.credentials_provider:
+                raise dbt.exceptions.FailedToConnectError(
+                    "Failed to use IdP credentials. 'credentials_provider' must be provided."
+                )
 
+            if not self.credentials.region:
+                raise dbt.exceptions.FailedToConnectError(
+                    "Failed to use IdP credentials. 'region' must be provided."
+                )
+
+            if not self.credentials.password or not self.credentials.user:
+                raise dbt.exceptions.FailedToConnectError(
+                    "Failed to use IdP credentials. 'password' and 'user' must be provided."
+                )
+
+            if self.credentials.credentials_provider == "AzureCredentialsProvider":
+                if (
+                    not self.credentials.idp_tenant
+                    or not self.credentials.client_id
+                    or not self.credentials.client_secret
+                    or not self.credentials.preferred_role
+                ):
+                    raise dbt.exceptions.FailedToConnectError(
+                        "Failed to use Azure credential. 'idp_tenant', 'client_id', 'client_secret', "
+                        "and 'preferred_role' must be provided"
+                    )
+
+                def connect():
+                    logger.debug("Connecting to redshift with Azure Credentials Provider...")
+                    c = redshift_connector.connect(
+                        iam=True,
+                        region=self.credentials.region,
+                        database=self.credentials.database,
+                        cluster_identifier=self.credentials.cluster_id,
+                        credentials_provider="AzureCredentialsProvider",
+                        user=self.credentials.user,
+                        password=self.credentials.password,
+                        idp_tenant=self.credentials.idp_tenant,
+                        client_id=self.credentials.client_id,
+                        client_secret=self.credentials.client_secret,
+                        preferred_role=self.credentials.preferred_role,
+                    )
+                    return c
+
+                return connect
+            elif self.credentials.credentials_provider == "OktaCredentialsProvider":
+                if (
+                    not self.credentials.idp_host
+                    or not self.credentials.app_id
+                    or not self.credentials.app_name
+                ):
+                    raise dbt.exceptions.FailedToConnectError(
+                        "Failed to use Okta credential. 'idp_host', 'app_id', 'app_name' must be provided."
+                    )
+
+                def connect():
+                    logger.debug("Connecting to redshift with Okta Credentials Provider...")
+                    c = redshift_connector.connect(
+                        iam=True,
+                        region=self.credentials.region,
+                        database=self.credentials.database,
+                        cluster_identifier=self.credentials.cluster_id,
+                        credentials_provider="OktaCredentialsProvider",
+                        user=self.credentials.user,
+                        password=self.credentials.password,
+                        idp_host=self.credentials.idp_host,
+                        app_id=self.credentials.app_id,
+                        app_name=self.credentials.app_name,
+                    )
+                    return c
+
+                return connect
         elif method == RedshiftConnectionMethod.IAM:
             if not self.credentials.cluster_id and "serverless" not in self.credentials.host:
                 raise dbt.exceptions.FailedToConnectError(
@@ -138,6 +268,7 @@ def connect():
                 return c
 
             return connect
+
         else:
             raise dbt.exceptions.FailedToConnectError(
                 "Invalid 'method' in profile: '{}'".format(method)