Support IAM authentication for Redshift serverless

airflow/providers/amazon/aws/hooks/redshift_sql.py

@@ -102,22 +102,43 @@ def get_iam_token(self, conn: Connection) -> tuple[str, str, int]:
         Port is required. If none is provided, default is used for each service.
         """
         port = conn.port or 5439
-        # Pull the custer-identifier from the beginning of the Redshift URL
-        # ex. my-cluster.ccdre4hpd39h.us-east-1.redshift.amazonaws.com returns my-cluster
-        cluster_identifier = conn.extra_dejson.get("cluster_identifier")
-        if not cluster_identifier:
-            if conn.host:
-                cluster_identifier = conn.host.split(".", 1)[0]
-            else:
-                raise AirflowException("Please set cluster_identifier or host in redshift connection.")
-        redshift_client = AwsBaseHook(aws_conn_id=self.aws_conn_id, client_type="redshift").conn
-        # https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/redshift.html#Redshift.Client.get_cluster_credentials
-        cluster_creds = redshift_client.get_cluster_credentials(
-            DbUser=conn.login,
-            DbName=conn.schema,
-            ClusterIdentifier=cluster_identifier,
-            AutoCreate=False,
-        )
+        is_serverless = conn.extra_dejson.get("is_serverless", False)
+        if is_serverless:
+            serverless_work_group = conn.extra_dejson.get("serverless_work_group")
+            if not serverless_work_group:
+                raise AirflowException(
+                    "Please set serverless_work_group in redshift connection to use IAM with"
+                    " Redshift Serverless."
+                )
+            serverless_token_duration_seconds = conn.extra_dejson.get(
+                "serverless_token_duration_seconds", 3600
+            )
+            redshift_client = AwsBaseHook(
+                aws_conn_id=self.aws_conn_id, client_type="redshift-serverless"
+            ).conn
+            # https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/redshift-serverless/client/get_credentials.html#get-credentials
+            cluster_creds = redshift_client.get_cluster_credentials(
+                DbName=conn.schema,
+                workgroupName=serverless_work_group,
+                durationSeconds=serverless_token_duration_seconds,
+            )
+        else:
+            # Pull the custer-identifier from the beginning of the Redshift URL
+            # ex. my-cluster.ccdre4hpd39h.us-east-1.redshift.amazonaws.com returns my-cluster
+            cluster_identifier = conn.extra_dejson.get("cluster_identifier")
+            if not cluster_identifier:
+                if conn.host:
+                    cluster_identifier = conn.host.split(".", 1)[0]
+                else:
+                    raise AirflowException("Please set cluster_identifier or host in redshift connection.")
+            redshift_client = AwsBaseHook(aws_conn_id=self.aws_conn_id, client_type="redshift").conn
+            # https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/redshift.html#Redshift.Client.get_cluster_credentials
+            cluster_creds = redshift_client.get_cluster_credentials(
+                DbUser=conn.login,
+                DbName=conn.schema,
+                ClusterIdentifier=cluster_identifier,
+                AutoCreate=False,
+            )
         token = cluster_creds["DbPassword"]
         login = cluster_creds["DbUser"]
         return login, token, port

docs/apache-airflow-providers-amazon/connections/redshift.rst

@@ -97,3 +97,23 @@ inferred by the **Host** field in Connection.
       "database": "dev",
       "profile": "default"
     }
+
+If you want to use IAM with Amazon Redshift Serverless, you need to set **is_serverless** to true and provide
+**serverless_work_group**. You can also set **serverless_token_duration_seconds** to specify the number of seconds
+until the returned temporary password expires; the minimum is 900 seconds, the maximum is 3600 seconds and by default
+it's 3600 seconds.
+
+* **Extra**:
+
+.. code-block:: json
+
+    {
+      "iam": true,
+      "is_serverless": true,
+      "serverless_work_group": "default",
+      "serverless_token_duration_seconds": 3600,
+      "port": 5439,
+      "region": "us-east-1",
+      "database": "dev",
+      "profile": "default"
+    }

tests/providers/amazon/aws/hooks/test_redshift_sql.py

@@ -120,6 +120,51 @@ def test_get_conn_iam(self, mock_connect, mock_aws_hook_conn, aws_conn_id):
             iam=True,
         )
 
+    @mock.patch("airflow.providers.amazon.aws.hooks.base_aws.AwsBaseHook.conn")
+    @mock.patch("airflow.providers.amazon.aws.hooks.redshift_sql.redshift_connector.connect")
+    @pytest.mark.parametrize("aws_conn_id", [NOTSET, None, "mock_aws_conn"])
+    def test_get_conn_iam_serverless_redshift(self, mock_connect, mock_aws_hook_conn, aws_conn_id):
+        mock_work_group = "my-test-workgroup"
+        mock_conn_extra = {
+            "iam": True,
+            "is_serverless": True,
+            "profile": "default",
+            "serverless_work_group": mock_work_group,
+        }
+        if aws_conn_id is not NOTSET:
+            self.db_hook.aws_conn_id = aws_conn_id
+        self.connection.extra = json.dumps(mock_conn_extra)
+
+        mock_db_user = f"IAM:{self.connection.login}"
+        mock_db_pass = "aws_token"
+
+        # Mock AWS Connection
+        mock_aws_hook_conn.get_cluster_credentials.return_value = {
+            "DbPassword": mock_db_pass,
+            "DbUser": mock_db_user,
+        }
+
+        self.db_hook.get_conn()
+
+        # Check boto3 'redshift' client method `get_cluster_credentials` call args
+        mock_aws_hook_conn.get_cluster_credentials.assert_called_once_with(
+            DbName=LOGIN_SCHEMA,
+            workgroupName=mock_work_group,
+            durationSeconds=3600,
+        )
+
+        mock_connect.assert_called_once_with(
+            user=mock_db_user,
+            password=mock_db_pass,
+            host=LOGIN_HOST,
+            port=LOGIN_PORT,
+            serverless_work_group=mock_work_group,
+            profile="default",
+            database=LOGIN_SCHEMA,
+            iam=True,
+            is_serverless=True,
+        )
+
     @pytest.mark.parametrize(
         "conn_params, conn_extra, expected_call_args",
         [