apache-airflow-providers-amazon 8.3.1__py3-none-any.whl → 8.4.0__py3-none-any.whl

airflow/providers/amazon/__init__.py

@@ -28,14 +28,16 @@ import packaging.version
 
 __all__ = ["__version__"]
 
-__version__ = "8.3.1"
+__version__ = "8.4.0"
 
 try:
     from airflow import __version__ as airflow_version
 except ImportError:
     from airflow.version import version as airflow_version
 
-if packaging.version.parse(airflow_version) < packaging.version.parse("2.4.0"):
+if packaging.version.parse(packaging.version.parse(airflow_version).base_version) < packaging.version.parse(
+    "2.4.0"
+):
     raise RuntimeError(
         f"The package `apache-airflow-providers-amazon:{__version__}` requires Apache Airflow 2.4.0+"  # NOQA: E501
     )

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

@@ -197,22 +197,35 @@ class BaseSessionFactory(LoggingMixin):
     def _create_session_with_assume_role(
         self, session_kwargs: dict[str, Any], deferrable: bool = False
     ) -> boto3.session.Session:
-
         if self.conn.assume_role_method == "assume_role_with_web_identity":
             # Deferred credentials have no initial credentials
             credential_fetcher = self._get_web_identity_credential_fetcher()
-            credentials = botocore.credentials.DeferredRefreshableCredentials(
-                method="assume-role-with-web-identity",
-                refresh_using=credential_fetcher.fetch_credentials,
-                time_fetcher=lambda: datetime.datetime.now(tz=tzlocal()),
-            )
+
+            params = {
+                "method": "assume-role-with-web-identity",
+                "refresh_using": credential_fetcher.fetch_credentials,
+                "time_fetcher": lambda: datetime.datetime.now(tz=tzlocal()),
+            }
+
+            if deferrable:
+                from aiobotocore.credentials import AioDeferredRefreshableCredentials
+
+                credentials = AioDeferredRefreshableCredentials(**params)
+            else:
+                credentials = botocore.credentials.DeferredRefreshableCredentials(**params)
         else:
             # Refreshable credentials do have initial credentials
-            credentials = botocore.credentials.RefreshableCredentials.create_from_metadata(
-                metadata=self._refresh_credentials(),
-                refresh_using=self._refresh_credentials,
-                method="sts-assume-role",
-            )
+            params = {
+                "metadata": self._refresh_credentials(),
+                "refresh_using": self._refresh_credentials,
+                "method": "sts-assume-role",
+            }
+            if deferrable:
+                from aiobotocore.credentials import AioRefreshableCredentials
+
+                credentials = AioRefreshableCredentials.create_from_metadata(**params)
+            else:
+                credentials = botocore.credentials.RefreshableCredentials.create_from_metadata(**params)
 
         if deferrable:
             from aiobotocore.session import get_session as async_get_session
@@ -796,7 +809,11 @@ class AwsGenericHook(BaseHook, Generic[BaseAwsConnection]):
         """
         try:
             session = self.get_session()
-            conn_info = session.client("sts").get_caller_identity()
+            test_endpoint_url = self.conn_config.extra_config.get("test_endpoint_url")
+            conn_info = session.client(
+                "sts",
+                endpoint_url=test_endpoint_url,
+            ).get_caller_identity()
             metadata = conn_info.pop("ResponseMetadata", {})
             if metadata.get("HTTPStatusCode") != 200:
                 try:

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

@@ -256,9 +256,14 @@ class EmrServerlessHook(AwsBaseHook):
         kwargs["client_type"] = "emr-serverless"
         super().__init__(*args, **kwargs)
 
-    def cancel_running_jobs(self, application_id: str, waiter_config: dict = {}):
+    def cancel_running_jobs(
+        self, application_id: str, waiter_config: dict | None = None, wait_for_completion: bool = True
+    ) -> int:
         """
-        List all jobs in an intermediate state, cancel them, then wait for those jobs to reach terminal state.
+        Cancel jobs in an intermediate state, and return the number of cancelled jobs.
+
+        If wait_for_completion is True, then the method will wait until all jobs are
+        cancelled before returning.
 
         Note: if new jobs are triggered while this operation is ongoing,
         it's going to time out and return an error.
@@ -284,13 +289,16 @@ class EmrServerlessHook(AwsBaseHook):
                 )
                 for job_id in job_ids:
                     self.conn.cancel_job_run(applicationId=application_id, jobRunId=job_id)
-        if count > 0:
-            self.log.info("now waiting for the %s cancelled job(s) to terminate", count)
-            self.get_waiter("no_job_running").wait(
-                applicationId=application_id,
-                states=list(self.JOB_INTERMEDIATE_STATES.union({"CANCELLING"})),
-                WaiterConfig=waiter_config,
-            )
+        if wait_for_completion:
+            if count > 0:
+                self.log.info("now waiting for the %s cancelled job(s) to terminate", count)
+                self.get_waiter("no_job_running").wait(
+                    applicationId=application_id,
+                    states=list(self.JOB_INTERMEDIATE_STATES.union({"CANCELLING"})),
+                    WaiterConfig=waiter_config or {},
+                )
+
+        return count
 
 
 class EmrContainerHook(AwsBaseHook):

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

@@ -0,0 +1,27 @@
+# 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 airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
+
+
+class EventBridgeHook(AwsBaseHook):
+    """Amazon EventBridge Hook."""
+
+    def __init__(self, *args, **kwargs):
+        """Creating object."""
+        super().__init__(client_type="events", *args, **kwargs)

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

@@ -60,6 +60,7 @@ class RedshiftDataHook(AwsGenericHook["RedshiftDataAPIServiceClient"]):
         with_event: bool = False,
         wait_for_completion: bool = True,
         poll_interval: int = 10,
+        workgroup_name: str | None = None,
     ) -> str:
         """
         Execute a statement against Amazon Redshift.
@@ -74,6 +75,9 @@ class RedshiftDataHook(AwsGenericHook["RedshiftDataAPIServiceClient"]):
         :param with_event: indicates whether to send an event to EventBridge
         :param wait_for_completion: indicates whether to wait for a result, if True wait, if False don't wait
         :param poll_interval: how often in seconds to check the query status
+        :param workgroup_name: name of the Redshift Serverless workgroup. Mutually exclusive with
+            `cluster_identifier`. Specify this parameter to query Redshift Serverless. More info
+            https://docs.aws.amazon.com/redshift/latest/mgmt/working-with-serverless.html
 
         :returns statement_id: str, the UUID of the statement
         """
@@ -85,6 +89,7 @@ class RedshiftDataHook(AwsGenericHook["RedshiftDataAPIServiceClient"]):
             "WithEvent": with_event,
             "SecretArn": secret_arn,
             "StatementName": statement_name,
+            "WorkgroupName": workgroup_name,
         }
         if isinstance(sql, list):
             kwargs["Sqls"] = sql
@@ -95,6 +100,9 @@ class RedshiftDataHook(AwsGenericHook["RedshiftDataAPIServiceClient"]):
 
         statement_id = resp["Id"]
 
+        if bool(cluster_identifier) is bool(workgroup_name):
+            raise ValueError("Either 'cluster_identifier' or 'workgroup_name' must be specified.")
+
         if wait_for_completion:
             self.wait_for_results(statement_id, poll_interval=poll_interval)
 
@@ -127,6 +135,7 @@ class RedshiftDataHook(AwsGenericHook["RedshiftDataAPIServiceClient"]):
         database: str,
         schema: str | None = "public",
         cluster_identifier: str | None = None,
+        workgroup_name: str | None = None,
         db_user: str | None = None,
         secret_arn: str | None = None,
         statement_name: str | None = None,
@@ -168,6 +177,7 @@ class RedshiftDataHook(AwsGenericHook["RedshiftDataAPIServiceClient"]):
             sql=sql,
             database=database,
             cluster_identifier=cluster_identifier,
+            workgroup_name=workgroup_name,
             db_user=db_user,
             secret_arn=secret_arn,
             statement_name=statement_name,

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

@@ -23,6 +23,7 @@ import re
 import tarfile
 import tempfile
 import time
+import warnings
 from collections import Counter
 from datetime import datetime
 from functools import partial
@@ -30,7 +31,7 @@ from typing import Any, Callable, Generator, cast
 
 from botocore.exceptions import ClientError
 
-from airflow.exceptions import AirflowException
+from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
 from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
 from airflow.providers.amazon.aws.hooks.logs import AwsLogsHook
 from airflow.providers.amazon.aws.hooks.s3 import S3Hook
@@ -1061,7 +1062,7 @@ class SageMakerHook(AwsBaseHook):
         display_name: str = "airflow-triggered-execution",
         pipeline_params: dict | None = None,
         wait_for_completion: bool = False,
-        check_interval: int = 30,
+        check_interval: int | None = None,
         verbose: bool = True,
     ) -> str:
         """Start a new execution for a SageMaker pipeline.
@@ -1073,14 +1074,19 @@ class SageMakerHook(AwsBaseHook):
         :param display_name: The name this pipeline execution will have in the UI. Doesn't need to be unique.
         :param pipeline_params: Optional parameters for the pipeline.
             All parameters supplied need to already be present in the pipeline definition.
-        :param wait_for_completion: Will only return once the pipeline is complete if true.
-        :param check_interval: How long to wait between checks for pipeline status when waiting for
-            completion.
-        :param verbose: Whether to print steps details when waiting for completion.
-            Defaults to true, consider turning off for pipelines that have thousands of steps.
 
         :return: the ARN of the pipeline execution launched.
         """
+        if wait_for_completion or check_interval is not None:
+            warnings.warn(
+                "parameter `wait_for_completion` and `check_interval` are deprecated, "
+                "remove them and call check_status yourself if you want to wait for completion",
+                AirflowProviderDeprecationWarning,
+                stacklevel=2,
+            )
+        if check_interval is None:
+            check_interval = 30
+
         formatted_params = format_tags(pipeline_params, key_label="Name")
 
         try:
@@ -1108,7 +1114,7 @@ class SageMakerHook(AwsBaseHook):
         self,
         pipeline_exec_arn: str,
         wait_for_completion: bool = False,
-        check_interval: int = 10,
+        check_interval: int | None = None,
         verbose: bool = True,
         fail_if_not_running: bool = False,
     ) -> str:
@@ -1119,12 +1125,6 @@ class SageMakerHook(AwsBaseHook):
 
         :param pipeline_exec_arn: Amazon Resource Name (ARN) of the pipeline execution.
             It's the ARN of the pipeline itself followed by "/execution/" and an id.
-        :param wait_for_completion: Whether to wait for the pipeline to reach a final state.
-            (i.e. either 'Stopped' or 'Failed')
-        :param check_interval: How long to wait between checks for pipeline status when waiting for
-            completion.
-        :param verbose: Whether to print steps details when waiting for completion.
-            Defaults to true, consider turning off for pipelines that have thousands of steps.
         :param fail_if_not_running: This method will raise an exception if the pipeline we're trying to stop
             is not in an "Executing" state when the call is sent (which would mean that the pipeline is
             already either stopping or stopped).
@@ -1133,6 +1133,16 @@ class SageMakerHook(AwsBaseHook):
         :return: Status of the pipeline execution after the operation.
             One of 'Executing'|'Stopping'|'Stopped'|'Failed'|'Succeeded'.
         """
+        if wait_for_completion or check_interval is not None:
+            warnings.warn(
+                "parameter `wait_for_completion` and `check_interval` are deprecated, "
+                "remove them and call check_status yourself if you want to wait for completion",
+                AirflowProviderDeprecationWarning,
+                stacklevel=2,
+            )
+        if check_interval is None:
+            check_interval = 10
+
         retries = 2  # i.e. 3 calls max, 1 initial + 2 retries
         while True:
             try:

airflow/providers/amazon/aws/notifications/chime.py

@@ -53,7 +53,7 @@ class ChimeNotifier(BaseNotifier):
         """To reduce overhead cache the hook for the notifier."""
         return ChimeWebhookHook(chime_conn_id=self.chime_conn_id)
 
-    def notify(self, context: Context) -> None:
+    def notify(self, context: Context) -> None:  # type: ignore[override]
         """Send a message to a Chime Chat Room."""
         self.hook.send_message(message=self.message)
 

airflow/providers/amazon/aws/operators/eks.py

@@ -21,7 +21,7 @@ import logging
 import warnings
 from ast import literal_eval
 from datetime import timedelta
-from typing import TYPE_CHECKING, List, Sequence, cast
+from typing import TYPE_CHECKING, Any, List, Sequence, cast
 
 from botocore.exceptions import ClientError, WaiterError
 
@@ -30,8 +30,10 @@ from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarni
 from airflow.models import BaseOperator
 from airflow.providers.amazon.aws.hooks.eks import EksHook
 from airflow.providers.amazon.aws.triggers.eks import (
+    EksCreateClusterTrigger,
     EksCreateFargateProfileTrigger,
     EksCreateNodegroupTrigger,
+    EksDeleteClusterTrigger,
     EksDeleteFargateProfileTrigger,
     EksDeleteNodegroupTrigger,
 )
@@ -187,6 +189,9 @@ class EksCreateClusterOperator(BaseOperator):
          (templated)
     :param waiter_delay: Time (in seconds) to wait between two consecutive calls to check cluster state
     :param waiter_max_attempts: The maximum number of attempts to check cluster state
+    :param deferrable: If True, the operator will wait asynchronously for the job to complete.
+        This implies waiting for completion. This mode requires aiobotocore module to be installed.
+        (default: False)
 
     """
 
@@ -225,6 +230,7 @@ class EksCreateClusterOperator(BaseOperator):
         wait_for_completion: bool = False,
         aws_conn_id: str = DEFAULT_CONN_ID,
         region: str | None = None,
+        deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
         waiter_delay: int = 30,
         waiter_max_attempts: int = 40,
         **kwargs,
@@ -237,7 +243,7 @@ class EksCreateClusterOperator(BaseOperator):
         self.nodegroup_role_arn = nodegroup_role_arn
         self.fargate_pod_execution_role_arn = fargate_pod_execution_role_arn
         self.create_fargate_profile_kwargs = create_fargate_profile_kwargs or {}
-        self.wait_for_completion = wait_for_completion
+        self.wait_for_completion = False if deferrable else wait_for_completion
         self.waiter_delay = waiter_delay
         self.waiter_max_attempts = waiter_max_attempts
         self.aws_conn_id = aws_conn_id
@@ -246,6 +252,7 @@ class EksCreateClusterOperator(BaseOperator):
         self.create_nodegroup_kwargs = create_nodegroup_kwargs or {}
         self.fargate_selectors = fargate_selectors or [{"namespace": DEFAULT_NAMESPACE_NAME}]
         self.fargate_profile_name = fargate_profile_name
+        self.deferrable = deferrable
         super().__init__(
             **kwargs,
         )
@@ -274,12 +281,25 @@ class EksCreateClusterOperator(BaseOperator):
 
         # Short circuit early if we don't need to wait to attach compute
         # and the caller hasn't requested to wait for the cluster either.
-        if not self.compute and not self.wait_for_completion:
+        if not any([self.compute, self.wait_for_completion, self.deferrable]):
             return None
 
-        self.log.info("Waiting for EKS Cluster to provision.  This will take some time.")
+        self.log.info("Waiting for EKS Cluster to provision. This will take some time.")
         client = self.eks_hook.conn
 
+        if self.deferrable:
+            self.defer(
+                trigger=EksCreateClusterTrigger(
+                    cluster_name=self.cluster_name,
+                    aws_conn_id=self.aws_conn_id,
+                    region_name=self.region,
+                    waiter_delay=self.waiter_delay,
+                    waiter_max_attempts=self.waiter_max_attempts,
+                ),
+                method_name="deferrable_create_cluster_next",
+                timeout=timedelta(seconds=self.waiter_max_attempts * self.waiter_delay),
+            )
+
         try:
             client.get_waiter("cluster_active").wait(
                 name=self.cluster_name,
@@ -311,6 +331,89 @@ class EksCreateClusterOperator(BaseOperator):
             subnets=cast(List[str], self.resources_vpc_config.get("subnetIds")),
         )
 
+    def deferrable_create_cluster_next(self, context: Context, event: dict[str, Any] | None = None) -> None:
+        if event is None:
+            self.log.error("Trigger error: event is None")
+            raise AirflowException("Trigger error: event is None")
+        elif event["status"] == "failed":
+            self.log.error("Cluster failed to start and will be torn down.")
+            self.eks_hook.delete_cluster(name=self.cluster_name)
+            self.defer(
+                trigger=EksDeleteClusterTrigger(
+                    cluster_name=self.cluster_name,
+                    waiter_delay=self.waiter_delay,
+                    waiter_max_attempts=self.waiter_max_attempts,
+                    aws_conn_id=self.aws_conn_id,
+                    region_name=self.region,
+                    force_delete_compute=False,
+                ),
+                method_name="execute_failed",
+                timeout=timedelta(seconds=self.waiter_max_attempts * self.waiter_delay),
+            )
+        elif event["status"] == "success":
+            self.log.info("Cluster is ready to provision compute.")
+            _create_compute(
+                compute=self.compute,
+                cluster_name=self.cluster_name,
+                aws_conn_id=self.aws_conn_id,
+                region=self.region,
+                wait_for_completion=self.wait_for_completion,
+                waiter_delay=self.waiter_delay,
+                waiter_max_attempts=self.waiter_max_attempts,
+                nodegroup_name=self.nodegroup_name,
+                nodegroup_role_arn=self.nodegroup_role_arn,
+                create_nodegroup_kwargs=self.create_nodegroup_kwargs,
+                fargate_profile_name=self.fargate_profile_name,
+                fargate_pod_execution_role_arn=self.fargate_pod_execution_role_arn,
+                fargate_selectors=self.fargate_selectors,
+                create_fargate_profile_kwargs=self.create_fargate_profile_kwargs,
+                subnets=cast(List[str], self.resources_vpc_config.get("subnetIds")),
+            )
+            if self.compute == "fargate":
+                self.defer(
+                    trigger=EksCreateFargateProfileTrigger(
+                        cluster_name=self.cluster_name,
+                        fargate_profile_name=self.fargate_profile_name,
+                        waiter_delay=self.waiter_delay,
+                        waiter_max_attempts=self.waiter_max_attempts,
+                        aws_conn_id=self.aws_conn_id,
+                        region=self.region,
+                    ),
+                    method_name="execute_complete",
+                    timeout=timedelta(seconds=self.waiter_max_attempts * self.waiter_delay),
+                )
+            else:
+                self.defer(
+                    trigger=EksCreateNodegroupTrigger(
+                        nodegroup_name=self.nodegroup_name,
+                        cluster_name=self.cluster_name,
+                        aws_conn_id=self.aws_conn_id,
+                        region_name=self.region,
+                        waiter_delay=self.waiter_delay,
+                        waiter_max_attempts=self.waiter_max_attempts,
+                    ),
+                    method_name="execute_complete",
+                    timeout=timedelta(seconds=self.waiter_max_attempts * self.waiter_delay),
+                )
+
+    def execute_failed(self, context: Context, event: dict[str, Any] | None = None) -> None:
+        if event is None:
+            self.log.info("Trigger error: event is None")
+            raise AirflowException("Trigger error: event is None")
+        elif event["status"] == "delteted":
+            self.log.info("Cluster deleted")
+            raise event["exception"]
+
+    def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> None:
+        resource = "fargate profile" if self.compute == "fargate" else self.compute
+        if event is None:
+            self.log.info("Trigger error: event is None")
+            raise AirflowException("Trigger error: event is None")
+        elif event["status"] != "success":
+            raise AirflowException(f"Error creating {resource}: {event}")
+
+        self.log.info("%s created successfully", resource)
+
 
 class EksCreateNodegroupOperator(BaseOperator):
     """
@@ -564,6 +667,11 @@ class EksDeleteClusterOperator(BaseOperator):
          maintained on each worker node).
     :param region: Which AWS region the connection should use. (templated)
         If this is None or empty then the default boto3 behaviour is used.
+    :param waiter_delay: Time (in seconds) to wait between two consecutive calls to check cluster state
+    :param waiter_max_attempts: The maximum number of attempts to check cluster state
+    :param deferrable: If True, the operator will wait asynchronously for the cluster to be deleted.
+        This implies waiting for completion. This mode requires aiobotocore module to be installed.
+        (default: False)
 
     """
 
@@ -582,13 +690,19 @@ class EksDeleteClusterOperator(BaseOperator):
         wait_for_completion: bool = False,
         aws_conn_id: str = DEFAULT_CONN_ID,
         region: str | None = None,
+        deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+        waiter_delay: int = 30,
+        waiter_max_attempts: int = 40,
         **kwargs,
     ) -> None:
         self.cluster_name = cluster_name
         self.force_delete_compute = force_delete_compute
-        self.wait_for_completion = wait_for_completion
+        self.wait_for_completion = False if deferrable else wait_for_completion
         self.aws_conn_id = aws_conn_id
         self.region = region
+        self.deferrable = deferrable
+        self.waiter_delay = waiter_delay
+        self.waiter_max_attempts = waiter_max_attempts
         super().__init__(**kwargs)
 
     def execute(self, context: Context):
@@ -596,8 +710,20 @@ class EksDeleteClusterOperator(BaseOperator):
             aws_conn_id=self.aws_conn_id,
             region_name=self.region,
         )
-
-        if self.force_delete_compute:
+        if self.deferrable:
+            self.defer(
+                trigger=EksDeleteClusterTrigger(
+                    cluster_name=self.cluster_name,
+                    waiter_delay=self.waiter_delay,
+                    waiter_max_attempts=self.waiter_max_attempts,
+                    aws_conn_id=self.aws_conn_id,
+                    region_name=self.region,
+                    force_delete_compute=self.force_delete_compute,
+                ),
+                method_name="execute_complete",
+                timeout=timedelta(seconds=self.waiter_delay * self.waiter_max_attempts),
+            )
+        elif self.force_delete_compute:
             self.delete_any_nodegroups(eks_hook)
             self.delete_any_fargate_profiles(eks_hook)
 
@@ -645,6 +771,13 @@ class EksDeleteClusterOperator(BaseOperator):
                 )
         self.log.info(SUCCESS_MSG.format(compute=FARGATE_FULL_NAME))
 
+    def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> None:
+        if event is None:
+            self.log.error("Trigger error. Event is None")
+            raise AirflowException("Trigger error. Event is None")
+        elif event["status"] == "success":
+            self.log.info("Cluster deleted successfully.")
+
 
 class EksDeleteNodegroupOperator(BaseOperator):
     """