apache-airflow-providers-amazon 8.12.0__tar.gz → 8.13.0rc1__tar.gz

{apache_airflow_providers_amazon-8.12.0 → apache_airflow_providers_amazon-8.13.0rc1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: apache-airflow-providers-amazon
-Version: 8.12.0
+Version: 8.13.0rc1
 Summary: Provider package apache-airflow-providers-amazon for Apache Airflow
 Keywords: airflow-provider,amazon,airflow,integration
 Author-email: Apache Software Foundation <dev@airflow.apache.org>
@@ -20,9 +20,9 @@ Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Topic :: System :: Monitoring
-Requires-Dist: apache-airflow-providers-common-sql>=1.3.1
+Requires-Dist: apache-airflow-providers-common-sql>=1.3.1.dev0
 Requires-Dist: apache-airflow-providers-http
-Requires-Dist: apache-airflow>=2.5.0
+Requires-Dist: apache-airflow>=2.6.0.dev0
 Requires-Dist: asgiref
 Requires-Dist: boto3>=1.28.0
 Requires-Dist: botocore>=1.31.0
@@ -45,12 +45,12 @@ Requires-Dist: apache-airflow-providers-mongo ; extra == "mongo"
 Requires-Dist: apache-airflow-providers-openlineage ; extra == "openlineage"
 Requires-Dist: pandas>=0.17.1 ; extra == "pandas"
 Requires-Dist: python3-saml>=1.16.0 ; extra == "python3-saml"
-Requires-Dist: s3fs>=2023.9.2 ; extra == "s3fs"
+Requires-Dist: s3fs>=2023.10.0 ; extra == "s3fs"
 Requires-Dist: apache-airflow-providers-salesforce ; extra == "salesforce"
 Requires-Dist: apache-airflow-providers-ssh ; extra == "ssh"
 Project-URL: Bug Tracker, https://github.com/apache/airflow/issues
-Project-URL: Changelog, https://airflow.apache.org/docs/apache-airflow-providers-amazon/8.12.0/changelog.html
-Project-URL: Documentation, https://airflow.apache.org/docs/apache-airflow-providers-amazon/8.12.0
+Project-URL: Changelog, https://airflow.apache.org/docs/apache-airflow-providers-amazon/8.13.0/changelog.html
+Project-URL: Documentation, https://airflow.apache.org/docs/apache-airflow-providers-amazon/8.13.0
 Project-URL: Slack Chat, https://s.apache.org/airflow-slack
 Project-URL: Source Code, https://github.com/apache/airflow
 Project-URL: Twitter, https://twitter.com/ApacheAirflow
@@ -117,7 +117,7 @@ Provides-Extra: ssh
 
 Package ``apache-airflow-providers-amazon``
 
-Release: ``8.12.0``
+Release: ``8.13.0.rc1``
 
 
 Amazon integration (including `Amazon Web Services (AWS) <https://aws.amazon.com/>`__).
@@ -130,7 +130,7 @@ This is a provider package for ``amazon`` provider. All classes for this provide
 are in ``airflow.providers.amazon`` python package.
 
 You can find package information and changelog for the provider
-in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-amazon/8.12.0/>`_.
+in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-amazon/8.13.0/>`_.
 
 Installation
 ------------
@@ -147,7 +147,7 @@ Requirements
 =======================================  ==================
 PIP package                              Version required
 =======================================  ==================
-``apache-airflow``                       ``>=2.5.0``
+``apache-airflow``                       ``>=2.6.0``
 ``apache-airflow-providers-common-sql``  ``>=1.3.1``
 ``apache-airflow-providers-http``
 ``boto3``                                ``>=1.28.0``
@@ -192,4 +192,4 @@ Dependent package
 ======================================================================================================================  ===================
 
 The changelog for the provider package can be found in the
-`changelog <https://airflow.apache.org/docs/apache-airflow-providers-amazon/8.12.0/changelog.html>`_.
+`changelog <https://airflow.apache.org/docs/apache-airflow-providers-amazon/8.13.0/changelog.html>`_.

{apache_airflow_providers_amazon-8.12.0 → apache_airflow_providers_amazon-8.13.0rc1}/README.rst

@@ -42,7 +42,7 @@
 
 Package ``apache-airflow-providers-amazon``
 
-Release: ``8.12.0``
+Release: ``8.13.0.rc1``
 
 
 Amazon integration (including `Amazon Web Services (AWS) <https://aws.amazon.com/>`__).
@@ -55,7 +55,7 @@ This is a provider package for ``amazon`` provider. All classes for this provide
 are in ``airflow.providers.amazon`` python package.
 
 You can find package information and changelog for the provider
-in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-amazon/8.12.0/>`_.
+in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-amazon/8.13.0/>`_.
 
 Installation
 ------------
@@ -72,7 +72,7 @@ Requirements
 =======================================  ==================
 PIP package                              Version required
 =======================================  ==================
-``apache-airflow``                       ``>=2.5.0``
+``apache-airflow``                       ``>=2.6.0``
 ``apache-airflow-providers-common-sql``  ``>=1.3.1``
 ``apache-airflow-providers-http``
 ``boto3``                                ``>=1.28.0``
@@ -117,4 +117,4 @@ Dependent package
 ======================================================================================================================  ===================
 
 The changelog for the provider package can be found in the
-`changelog <https://airflow.apache.org/docs/apache-airflow-providers-amazon/8.12.0/changelog.html>`_.
+`changelog <https://airflow.apache.org/docs/apache-airflow-providers-amazon/8.13.0/changelog.html>`_.

{apache_airflow_providers_amazon-8.12.0 → apache_airflow_providers_amazon-8.13.0rc1}/airflow/providers/amazon/__init__.py

@@ -27,7 +27,7 @@ import packaging.version
 
 __all__ = ["__version__"]
 
-__version__ = "8.12.0"
+__version__ = "8.13.0"
 
 try:
     from airflow import __version__ as airflow_version
@@ -35,8 +35,8 @@ except ImportError:
     from airflow.version import version as airflow_version
 
 if packaging.version.parse(packaging.version.parse(airflow_version).base_version) < packaging.version.parse(
-    "2.5.0"
+    "2.6.0"
 ):
     raise RuntimeError(
-        f"The package `apache-airflow-providers-amazon:{__version__}` needs Apache Airflow 2.5.0+"
+        f"The package `apache-airflow-providers-amazon:{__version__}` needs Apache Airflow 2.6.0+"
     )

apache_airflow_providers_amazon-8.13.0rc1/airflow/providers/amazon/aws/auth_manager/avp/entities.py

@@ -0,0 +1,64 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from enum import Enum
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from airflow.auth.managers.base_auth_manager import ResourceMethod
+
+AVP_PREFIX_ENTITIES = "Airflow::"
+
+
+class AvpEntities(Enum):
+    """Enum of Amazon Verified Permissions entities."""
+
+    ACTION = "Action"
+    ROLE = "Role"
+    USER = "User"
+
+    # Resource types
+    CONFIGURATION = "Configuration"
+    CONNECTION = "Connection"
+    DATASET = "Dataset"
+    POOL = "Pool"
+    VARIABLE = "Variable"
+    VIEW = "View"
+
+
+def get_entity_type(resource_type: AvpEntities) -> str:
+    """
+    Return entity type.
+
+    :param resource_type: Resource type.
+
+    Example: Airflow::Action, Airflow::Role, Airflow::Variable, Airflow::User.
+    """
+    return AVP_PREFIX_ENTITIES + resource_type.value
+
+
+def get_action_id(resource_type: AvpEntities, method: ResourceMethod):
+    """
+    Return action id.
+
+    Convention for action ID is <resource_type>::<method>. Example: Variable::GET.
+
+    :param resource_type: Resource type.
+    :param method: Resource method.
+    """
+    return f"{resource_type.value}::{method}"

apache_airflow_providers_amazon-8.13.0rc1/airflow/providers/amazon/aws/auth_manager/avp/facade.py

@@ -0,0 +1,126 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from functools import cached_property
+from typing import TYPE_CHECKING, Callable
+
+from airflow.configuration import conf
+from airflow.exceptions import AirflowException
+from airflow.providers.amazon.aws.auth_manager.avp.entities import AvpEntities, get_action_id, get_entity_type
+from airflow.providers.amazon.aws.auth_manager.constants import (
+    CONF_AVP_POLICY_STORE_ID_KEY,
+    CONF_CONN_ID_KEY,
+    CONF_SECTION_NAME,
+)
+from airflow.providers.amazon.aws.hooks.verified_permissions import VerifiedPermissionsHook
+from airflow.utils.log.logging_mixin import LoggingMixin
+
+if TYPE_CHECKING:
+    from airflow.auth.managers.base_auth_manager import ResourceMethod
+    from airflow.providers.amazon.aws.auth_manager.user import AwsAuthManagerUser
+
+
+class AwsAuthManagerAmazonVerifiedPermissionsFacade(LoggingMixin):
+    """
+    Facade for Amazon Verified Permissions.
+
+    Used as an intermediate layer between AWS auth manager and Amazon Verified Permissions.
+    """
+
+    @cached_property
+    def avp_client(self):
+        """Build Amazon Verified Permissions client."""
+        aws_conn_id = conf.get(CONF_SECTION_NAME, CONF_CONN_ID_KEY)
+        return VerifiedPermissionsHook(aws_conn_id=aws_conn_id).conn
+
+    @cached_property
+    def avp_policy_store_id(self):
+        """Get the Amazon Verified Permission policy store ID from config."""
+        return conf.get_mandatory_value(CONF_SECTION_NAME, CONF_AVP_POLICY_STORE_ID_KEY)
+
+    def is_authorized(
+        self,
+        *,
+        method: ResourceMethod,
+        entity_type: AvpEntities,
+        user: AwsAuthManagerUser,
+        entity_id: str | None = None,
+        entity_fetcher: Callable | None = None,
+    ) -> bool:
+        """
+        Make an authorization decision against Amazon Verified Permissions.
+
+        Check whether the user has permissions to access given resource.
+
+        :param method: the method to perform
+        :param entity_type: the entity type the user accesses
+        :param user: the user
+        :param entity_id: the entity ID the user accesses. If not provided, all entities of the type will be
+            considered.
+        :param entity_fetcher: function that returns list of entities to be passed to Amazon Verified
+            Permissions. Only needed if some resource properties are used in the policies (e.g. DAG folder).
+        """
+        entity_list = self._get_user_role_entities(user)
+        if entity_fetcher and entity_id:
+            # If no entity ID is provided, there is no need to fetch entities.
+            # We just need to know whether the user has permissions to access all resources from this type
+            entity_list += entity_fetcher()
+
+        self.log.debug(
+            "Making authorization request for user=%s, method=%s, entity_type=%s, entity_id=%s",
+            user.get_id(),
+            method,
+            entity_type,
+            entity_id,
+        )
+
+        resp = self.avp_client.is_authorized(
+            policyStoreId=self.avp_policy_store_id,
+            principal={"entityType": get_entity_type(AvpEntities.USER), "entityId": user.get_id()},
+            action={
+                "actionType": get_entity_type(AvpEntities.ACTION),
+                "actionId": get_action_id(entity_type, method),
+            },
+            resource={"entityType": get_entity_type(entity_type), "entityId": entity_id or "*"},
+            entities={"entityList": entity_list},
+        )
+
+        self.log.debug("Authorization response: %s", resp)
+
+        if len(resp.get("errors", [])) > 0:
+            self.log.error(
+                "Error occurred while making an authorization decision. Errors: %s", resp["errors"]
+            )
+            raise AirflowException("Error occurred while making an authorization decision.")
+
+        return resp["decision"] == "ALLOW"
+
+    @staticmethod
+    def _get_user_role_entities(user: AwsAuthManagerUser) -> list[dict]:
+        user_entity = {
+            "identifier": {"entityType": get_entity_type(AvpEntities.USER), "entityId": user.get_id()},
+            "parents": [
+                {"entityType": get_entity_type(AvpEntities.ROLE), "entityId": group}
+                for group in user.get_groups()
+            ],
+        }
+        role_entities = [
+            {"identifier": {"entityType": get_entity_type(AvpEntities.ROLE), "entityId": group}}
+            for group in user.get_groups()
+        ]
+        return [user_entity, *role_entities]

{apache_airflow_providers_amazon-8.12.0 → apache_airflow_providers_amazon-8.13.0rc1}/airflow/providers/amazon/aws/auth_manager/aws_auth_manager.py

@@ -23,6 +23,8 @@ from flask import session, url_for
 
 from airflow.configuration import conf
 from airflow.exceptions import AirflowOptionalProviderFeatureException
+from airflow.providers.amazon.aws.auth_manager.avp.entities import AvpEntities
+from airflow.providers.amazon.aws.auth_manager.avp.facade import AwsAuthManagerAmazonVerifiedPermissionsFacade
 from airflow.providers.amazon.aws.auth_manager.constants import (
     CONF_ENABLE_KEY,
     CONF_SECTION_NAME,
@@ -72,6 +74,10 @@ class AwsAuthManager(BaseAuthManager):
                 "The AWS auth manager is currently being built. It is not finalized. It is not intended to be used yet."
             )
 
+    @cached_property
+    def avp_facade(self):
+        return AwsAuthManagerAmazonVerifiedPermissionsFacade()
+
     def get_user(self) -> AwsAuthManagerUser | None:
         return session["aws_user"] if self.is_logged_in() else None
 
@@ -85,7 +91,13 @@ class AwsAuthManager(BaseAuthManager):
         details: ConfigurationDetails | None = None,
         user: BaseUser | None = None,
     ) -> bool:
-        return self.is_logged_in()
+        config_section = details.section if details else None
+        return self.avp_facade.is_authorized(
+            method=method,
+            entity_type=AvpEntities.CONFIGURATION,
+            user=user or self.get_user(),
+            entity_id=config_section,
+        )
 
     def is_authorized_cluster_activity(self, *, method: ResourceMethod, user: BaseUser | None = None) -> bool:
         return self.is_logged_in()
@@ -97,7 +109,13 @@ class AwsAuthManager(BaseAuthManager):
         details: ConnectionDetails | None = None,
         user: BaseUser | None = None,
     ) -> bool:
-        return self.is_logged_in()
+        connection_id = details.conn_id if details else None
+        return self.avp_facade.is_authorized(
+            method=method,
+            entity_type=AvpEntities.CONNECTION,
+            user=user or self.get_user(),
+            entity_id=connection_id,
+        )
 
     def is_authorized_dag(
         self,
@@ -112,17 +130,35 @@ class AwsAuthManager(BaseAuthManager):
     def is_authorized_dataset(
         self, *, method: ResourceMethod, details: DatasetDetails | None = None, user: BaseUser | None = None
     ) -> bool:
-        return self.is_logged_in()
+        dataset_uri = details.uri if details else None
+        return self.avp_facade.is_authorized(
+            method=method,
+            entity_type=AvpEntities.DATASET,
+            user=user or self.get_user(),
+            entity_id=dataset_uri,
+        )
 
     def is_authorized_pool(
         self, *, method: ResourceMethod, details: PoolDetails | None = None, user: BaseUser | None = None
     ) -> bool:
-        return self.is_logged_in()
+        pool_name = details.name if details else None
+        return self.avp_facade.is_authorized(
+            method=method,
+            entity_type=AvpEntities.POOL,
+            user=user or self.get_user(),
+            entity_id=pool_name,
+        )
 
     def is_authorized_variable(
         self, *, method: ResourceMethod, details: VariableDetails | None = None, user: BaseUser | None = None
     ) -> bool:
-        return self.is_logged_in()
+        variable_key = details.key if details else None
+        return self.avp_facade.is_authorized(
+            method=method,
+            entity_type=AvpEntities.VARIABLE,
+            user=user or self.get_user(),
+            entity_id=variable_key,
+        )
 
     def is_authorized_view(
         self,
@@ -130,7 +166,12 @@ class AwsAuthManager(BaseAuthManager):
         access_view: AccessView,
         user: BaseUser | None = None,
     ) -> bool:
-        return self.is_logged_in()
+        return self.avp_facade.is_authorized(
+            method="GET",
+            entity_type=AvpEntities.VIEW,
+            user=user or self.get_user(),
+            entity_id=access_view.value,
+        )
 
     def get_url_login(self, **kwargs) -> str:
         return url_for("AwsAuthManagerAuthenticationViews.login")

{apache_airflow_providers_amazon-8.12.0 → apache_airflow_providers_amazon-8.13.0rc1}/airflow/providers/amazon/aws/auth_manager/constants.py

@@ -18,6 +18,8 @@
 # Configuration keys
 from __future__ import annotations
 
+CONF_ENABLE_KEY = "enable"
 CONF_SECTION_NAME = "aws_auth_manager"
+CONF_CONN_ID_KEY = "conn_id"
 CONF_SAML_METADATA_URL_KEY = "saml_metadata_url"
-CONF_ENABLE_KEY = "enable"
+CONF_AVP_POLICY_STORE_ID_KEY = "avp_policy_store_id"

{apache_airflow_providers_amazon-8.12.0 → apache_airflow_providers_amazon-8.13.0rc1}/airflow/providers/amazon/aws/auth_manager/user.py

@@ -49,3 +49,6 @@ class AwsAuthManagerUser(BaseUser):
 
     def get_name(self) -> str:
         return self.username or self.email or self.user_id
+
+    def get_groups(self):
+        return self.groups

{apache_airflow_providers_amazon-8.12.0 → apache_airflow_providers_amazon-8.13.0rc1}/airflow/providers/amazon/aws/fs/s3.py

@@ -25,7 +25,7 @@ import requests
 from botocore import UNSIGNED
 from requests import HTTPError
 
-from airflow.providers.amazon.aws.hooks.base_aws import AwsGenericHook
+from airflow.providers.amazon.aws.hooks.s3 import S3Hook
 
 if TYPE_CHECKING:
     from botocore.awsrequest import AWSRequest
@@ -55,14 +55,14 @@ def get_fs(conn_id: str | None) -> AbstractFileSystem:
             "pip install apache-airflow-providers-amazon[s3fs]"
         )
 
-    aws: AwsGenericHook = AwsGenericHook(aws_conn_id=conn_id, client_type="s3")
-    session = aws.get_session(deferrable=True)
-    endpoint_url = aws.conn_config.get_service_endpoint_url(service_name="s3")
+    s3_hook = S3Hook(aws_conn_id=conn_id)
+    session = s3_hook.get_session(deferrable=True)
+    endpoint_url = s3_hook.conn_config.get_service_endpoint_url(service_name="s3")
 
-    config_kwargs: dict[str, Any] = aws.conn_config.extra_config.get("config_kwargs", {})
+    config_kwargs: dict[str, Any] = s3_hook.conn_config.extra_config.get("config_kwargs", {})
     register_events: dict[str, Callable[[Properties], None]] = {}
 
-    s3_service_config = aws.service_config
+    s3_service_config = s3_hook.service_config
     if signer := s3_service_config.get("signer", None):
         log.info("Loading signer %s", signer)
         if singer_func := SIGNERS.get(signer):

{apache_airflow_providers_amazon-8.12.0 → apache_airflow_providers_amazon-8.13.0rc1}/airflow/providers/amazon/aws/hooks/athena.py

@@ -292,24 +292,17 @@ class AthenaHook(AwsBaseHook):
 
         :param query_execution_id: Id of submitted athena query
         """
-        output_location = None
-        if query_execution_id:
-            response = self.get_query_info(query_execution_id=query_execution_id, use_cache=True)
-
-            if response:
-                try:
-                    output_location = response["QueryExecution"]["ResultConfiguration"]["OutputLocation"]
-                except KeyError:
-                    self.log.error(
-                        "Error retrieving OutputLocation. Query execution id: %s", query_execution_id
-                    )
-                    raise
-            else:
-                raise
-        else:
-            raise ValueError("Invalid Query execution id. Query execution id: %s", query_execution_id)
+        if not query_execution_id:
+            raise ValueError(f"Invalid Query execution id. Query execution id: {query_execution_id}")
+
+        if not (response := self.get_query_info(query_execution_id=query_execution_id, use_cache=True)):
+            raise ValueError(f"Unable to get query information for execution id: {query_execution_id}")
 
-        return output_location
+        try:
+            return response["QueryExecution"]["ResultConfiguration"]["OutputLocation"]
+        except KeyError:
+            self.log.error("Error retrieving OutputLocation. Query execution id: %s", query_execution_id)
+            raise
 
     def stop_query(self, query_execution_id: str) -> dict:
         """Cancel the submitted query.

{apache_airflow_providers_amazon-8.12.0 → apache_airflow_providers_amazon-8.13.0rc1}/airflow/providers/amazon/aws/hooks/ec2.py

@@ -19,16 +19,21 @@ from __future__ import annotations
 
 import functools
 import time
+from typing import Callable, TypeVar
 
 from airflow.exceptions import AirflowException
 from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
+from airflow.typing_compat import ParamSpec
 
+PS = ParamSpec("PS")
+RT = TypeVar("RT")
 
-def only_client_type(func):
+
+def only_client_type(func: Callable[PS, RT]) -> Callable[PS, RT]:
     @functools.wraps(func)
-    def checker(self, *args, **kwargs):
-        if self._api_type == "client_type":
-            return func(self, *args, **kwargs)
+    def checker(*args, **kwargs) -> RT:
+        if args[0]._api_type == "client_type":
+            return func(*args, **kwargs)
 
         ec2_doc_link = "https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/ec2.html"
         raise AirflowException(
@@ -85,7 +90,7 @@ class EC2Hook(AwsBaseHook):
         :return: Instance object
         """
         if self._api_type == "client_type":
-            return self.get_instances(filters=filters, instance_ids=[instance_id])
+            return self.get_instances(filters=filters, instance_ids=[instance_id])[0]
 
         return self.conn.Instance(id=instance_id)
 

{apache_airflow_providers_amazon-8.12.0 → apache_airflow_providers_amazon-8.13.0rc1}/airflow/providers/amazon/aws/hooks/emr.py

@@ -437,8 +437,6 @@ class EmrContainerHook(AwsBaseHook):
 
         :param job_id: The ID of the job run request.
         """
-        reason = None  # We absorb any errors if we can't retrieve the job status
-
         try:
             response = self.conn.describe_job_run(
                 virtualClusterId=self.virtual_cluster_id,
@@ -446,13 +444,13 @@ class EmrContainerHook(AwsBaseHook):
             )
             failure_reason = response["jobRun"]["failureReason"]
             state_details = response["jobRun"]["stateDetails"]
-            reason = f"{failure_reason} - {state_details}"
+            return f"{failure_reason} - {state_details}"
         except KeyError:
             self.log.error("Could not get status of the EMR on EKS job")
         except ClientError as ex:
             self.log.error("AWS request failed, check logs for more info: %s", ex)
 
-        return reason
+        return None
 
     def check_query_status(self, job_id: str) -> str | None:
         """
@@ -491,26 +489,21 @@ class EmrContainerHook(AwsBaseHook):
         :param max_polling_attempts: Number of times to poll for query state before function exits
         """
         try_number = 1
-        final_query_state = None  # Query state when query reaches final state or max_polling_attempts reached
-
         while True:
             query_state = self.check_query_status(job_id)
+            if query_state in self.TERMINAL_STATES:
+                self.log.info("Try %s: Query execution completed. Final state is %s", try_number, query_state)
+                return query_state
             if query_state is None:
                 self.log.info("Try %s: Invalid query state. Retrying again", try_number)
-            elif query_state in self.TERMINAL_STATES:
-                self.log.info("Try %s: Query execution completed. Final state is %s", try_number, query_state)
-                final_query_state = query_state
-                break
             else:
                 self.log.info("Try %s: Query is still in non-terminal state - %s", try_number, query_state)
             if (
                 max_polling_attempts and try_number >= max_polling_attempts
             ):  # Break loop if max_polling_attempts reached
-                final_query_state = query_state
-                break
+                return query_state
             try_number += 1
             time.sleep(poll_interval)
-        return final_query_state
 
     def stop_query(self, job_id: str) -> dict:
         """

{apache_airflow_providers_amazon-8.12.0 → apache_airflow_providers_amazon-8.13.0rc1}/airflow/providers/amazon/aws/hooks/redshift_sql.py

@@ -102,24 +102,47 @@ class RedshiftSQLHook(DbApiHook):
         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,
-        )
-        token = cluster_creds["DbPassword"]
-        login = cluster_creds["DbUser"]
+        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_serverless_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_serverless_client.get_credentials(
+                dbName=conn.schema,
+                workgroupName=serverless_work_group,
+                durationSeconds=serverless_token_duration_seconds,
+            )
+            token = cluster_creds["dbPassword"]
+            login = cluster_creds["dbUser"]
+        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
 
     def get_uri(self) -> str: