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

airflow/providers/amazon/aws/transfers/s3_to_redshift.py

@@ -178,7 +178,7 @@ class S3ToRedshiftOperator(BaseOperator):
             where_statement = " AND ".join([f"{self.table}.{k} = {copy_destination}.{k}" for k in keys])
 
             sql = [
-                f"CREATE TABLE {copy_destination} (LIKE {destination});",
+                f"CREATE TABLE {copy_destination} (LIKE {destination} INCLUDING DEFAULTS);",
                 copy_statement,
                 "BEGIN;",
                 f"DELETE FROM {destination} USING {copy_destination} WHERE {where_statement};",

airflow/providers/amazon/aws/triggers/batch.py

@@ -209,7 +209,7 @@ class BatchJobTrigger(AwsBaseWaiterTrigger):
     def __init__(
         self,
         job_id: str | None,
-        region_name: str | None,
+        region_name: str | None = None,
         aws_conn_id: str | None = "aws_default",
         waiter_delay: int = 5,
         waiter_max_attempts: int = 720,

airflow/providers/amazon/aws/triggers/ecs.py

@@ -22,6 +22,7 @@ from typing import Any, AsyncIterator
 
 from botocore.exceptions import ClientError, WaiterError
 
+from airflow import AirflowException
 from airflow.providers.amazon.aws.hooks.base_aws import AwsGenericHook
 from airflow.providers.amazon.aws.hooks.ecs import EcsHook
 from airflow.providers.amazon.aws.hooks.logs import AwsLogsHook
@@ -48,7 +49,7 @@ class ClusterActiveTrigger(AwsBaseWaiterTrigger):
         waiter_delay: int,
         waiter_max_attempts: int,
         aws_conn_id: str | None,
-        region_name: str | None,
+        region_name: str | None = None,
     ):
         super().__init__(
             serialized_fields={"cluster_arn": cluster_arn},
@@ -87,7 +88,7 @@ class ClusterInactiveTrigger(AwsBaseWaiterTrigger):
         waiter_delay: int,
         waiter_max_attempts: int,
         aws_conn_id: str | None,
-        region_name: str | None,
+        region_name: str | None = None,
     ):
         super().__init__(
             serialized_fields={"cluster_arn": cluster_arn},
@@ -170,7 +171,9 @@ class TaskDoneTrigger(BaseTrigger):
                     await waiter.wait(
                         cluster=self.cluster, tasks=[self.task_arn], WaiterConfig={"MaxAttempts": 1}
                     )
-                    break  # we reach this point only if the waiter met a success criteria
+                    # we reach this point only if the waiter met a success criteria
+                    yield TriggerEvent({"status": "success", "task_arn": self.task_arn})
+                    return
                 except WaiterError as error:
                     if "terminal failure" in str(error):
                         raise
@@ -179,8 +182,7 @@ class TaskDoneTrigger(BaseTrigger):
                 finally:
                     if self.log_group and self.log_stream:
                         logs_token = await self._forward_logs(logs_client, logs_token)
-
-        yield TriggerEvent({"status": "success", "task_arn": self.task_arn})
+        raise AirflowException("Waiter error: max attempts reached")
 
     async def _forward_logs(self, logs_client, next_token: str | None = None) -> str | None:
         """

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

@@ -17,11 +17,178 @@
 from __future__ import annotations
 
 import warnings
+from typing import Any
 
 from airflow.exceptions import AirflowProviderDeprecationWarning
 from airflow.providers.amazon.aws.hooks.base_aws import AwsGenericHook
 from airflow.providers.amazon.aws.hooks.eks import EksHook
 from airflow.providers.amazon.aws.triggers.base import AwsBaseWaiterTrigger
+from airflow.providers.amazon.aws.utils.waiter_with_logging import async_wait
+from airflow.triggers.base import TriggerEvent
+
+
+class EksCreateClusterTrigger(AwsBaseWaiterTrigger):
+    """
+    Trigger for EksCreateClusterOperator.
+
+    The trigger will asynchronously wait for the cluster to be created.
+
+    :param cluster_name: The name of the EKS cluster
+    :param waiter_delay: The amount of time in seconds to wait between attempts.
+    :param waiter_max_attempts: The maximum number of attempts to be made.
+    :param aws_conn_id: The Airflow connection used for AWS credentials.
+    :param region_name: Which AWS region the connection should use.
+         If this is None or empty then the default boto3 behaviour is used.
+    """
+
+    def __init__(
+        self,
+        cluster_name: str,
+        waiter_delay: int,
+        waiter_max_attempts: int,
+        aws_conn_id: str,
+        region_name: str | None = None,
+    ):
+        super().__init__(
+            serialized_fields={"cluster_name": cluster_name, "region_name": region_name},
+            waiter_name="cluster_active",
+            waiter_args={"name": cluster_name},
+            failure_message="Error checking Eks cluster",
+            status_message="Eks cluster status is",
+            status_queries=["cluster.status"],
+            return_value=None,
+            waiter_delay=waiter_delay,
+            waiter_max_attempts=waiter_max_attempts,
+            aws_conn_id=aws_conn_id,
+            region_name=region_name,
+        )
+
+    def hook(self) -> AwsGenericHook:
+        return EksHook(aws_conn_id=self.aws_conn_id, region_name=self.region_name)
+
+
+class EksDeleteClusterTrigger(AwsBaseWaiterTrigger):
+    """
+    Trigger for EksDeleteClusterOperator.
+
+    The trigger will asynchronously wait for the cluster to be deleted. If there are
+    any nodegroups or fargate profiles associated with the cluster, they will be deleted
+    before the cluster is deleted.
+
+    :param cluster_name: The name of the EKS cluster
+    :param waiter_delay: The amount of time in seconds to wait between attempts.
+    :param waiter_max_attempts: The maximum number of attempts to be made.
+    :param aws_conn_id: The Airflow connection used for AWS credentials.
+    :param region_name: Which AWS region the connection should use.
+         If this is None or empty then the default boto3 behaviour is used.
+    :param force_delete_compute: If True, any nodegroups or fargate profiles associated
+        with the cluster will be deleted before the cluster is deleted.
+    """
+
+    def __init__(
+        self,
+        cluster_name,
+        waiter_delay: int,
+        waiter_max_attempts: int,
+        aws_conn_id: str,
+        region_name: str | None,
+        force_delete_compute: bool,
+    ):
+        self.cluster_name = cluster_name
+        self.waiter_delay = waiter_delay
+        self.waiter_max_attempts = waiter_max_attempts
+        self.aws_conn_id = aws_conn_id
+        self.region_name = region_name
+        self.force_delete_compute = force_delete_compute
+
+    def serialize(self) -> tuple[str, dict[str, Any]]:
+        return (
+            self.__class__.__module__ + "." + self.__class__.__qualname__,
+            {
+                "cluster_name": self.cluster_name,
+                "waiter_delay": str(self.waiter_delay),
+                "waiter_max_attempts": str(self.waiter_max_attempts),
+                "aws_conn_id": self.aws_conn_id,
+                "region_name": self.region_name,
+                "force_delete_compute": self.force_delete_compute,
+            },
+        )
+
+    def hook(self) -> AwsGenericHook:
+        return EksHook(aws_conn_id=self.aws_conn_id, region_name=self.region_name)
+
+    async def run(self):
+        async with self.hook.async_conn as client:
+            waiter = client.get_waiter("cluster_deleted")
+            if self.force_delete_compute:
+                await self.delete_any_nodegroups(client=client)
+                await self.delete_any_fargate_profiles(client=client)
+                await client.delete_cluster(name=self.cluster_name)
+            await async_wait(
+                waiter=waiter,
+                waiter_delay=int(self.waiter_delay),
+                waiter_max_attempts=int(self.waiter_max_attempts),
+                args={"name": self.cluster_name},
+                failure_message="Error deleting cluster",
+                status_message="Status of cluster is",
+                status_args=["cluster.status"],
+            )
+
+        yield TriggerEvent({"status": "deleted"})
+
+    async def delete_any_nodegroups(self, client) -> None:
+        """
+        Deletes all EKS Nodegroups for a provided Amazon EKS Cluster.
+
+        All the EKS Nodegroups are deleted simultaneously. We wait for
+        all Nodegroups to be deleted before returning.
+        """
+        nodegroups = await client.list_nodegroups(clusterName=self.cluster_name)
+        if nodegroups.get("nodegroups", None):
+            self.log.info("Deleting nodegroups")
+            # ignoring attr-defined here because aws_base hook defines get_waiter for all hooks
+            waiter = self.hook.get_waiter(  # type: ignore[attr-defined]
+                "all_nodegroups_deleted", deferrable=True, client=client
+            )
+            for group in nodegroups["nodegroups"]:
+                await client.delete_nodegroup(clusterName=self.cluster_name, nodegroupName=group)
+            await async_wait(
+                waiter=waiter,
+                waiter_delay=int(self.waiter_delay),
+                waiter_max_attempts=int(self.waiter_max_attempts),
+                args={"clusterName": self.cluster_name},
+                failure_message=f"Error deleting nodegroup for cluster {self.cluster_name}",
+                status_message="Deleting nodegroups associated with the cluster",
+                status_args=["nodegroups"],
+            )
+            self.log.info("All nodegroups deleted")
+        else:
+            self.log.info("No nodegroups associated with cluster %s", self.cluster_name)
+
+    async def delete_any_fargate_profiles(self, client) -> None:
+        """
+        Deletes all EKS Fargate profiles for a provided Amazon EKS Cluster.
+
+        EKS Fargate profiles must be deleted one at a time, so we must wait
+        for one to be deleted before sending the next delete command.
+        """
+        fargate_profiles = await client.list_fargate_profiles(clusterName=self.cluster_name)
+        if fargate_profiles.get("fargateProfileNames"):
+            self.log.info("Waiting for Fargate profiles to delete.  This will take some time.")
+            for profile in fargate_profiles["fargateProfileNames"]:
+                await client.delete_fargate_profile(clusterName=self.cluster_name, fargateProfileName=profile)
+                await async_wait(
+                    waiter=client.get_waiter("fargate_profile_deleted"),
+                    waiter_delay=int(self.waiter_delay),
+                    waiter_max_attempts=int(self.waiter_max_attempts),
+                    args={"clusterName": self.cluster_name, "fargateProfileName": profile},
+                    failure_message=f"Error deleting fargate profile for cluster {self.cluster_name}",
+                    status_message="Status of fargate profile is",
+                    status_args=["fargateProfile.status"],
+                )
+            self.log.info("All Fargate profiles deleted")
+        else:
+            self.log.info(f"No Fargate profiles associated with cluster {self.cluster_name}")
 
 
 class EksCreateFargateProfileTrigger(AwsBaseWaiterTrigger):
@@ -142,10 +309,14 @@ class EksCreateNodegroupTrigger(AwsBaseWaiterTrigger):
         waiter_delay: int,
         waiter_max_attempts: int,
         aws_conn_id: str,
-        region_name: str | None,
+        region_name: str | None = None,
     ):
         super().__init__(
-            serialized_fields={"cluster_name": cluster_name, "nodegroup_name": nodegroup_name},
+            serialized_fields={
+                "cluster_name": cluster_name,
+                "nodegroup_name": nodegroup_name,
+                "region_name": region_name,
+            },
             waiter_name="nodegroup_active",
             waiter_args={"clusterName": cluster_name, "nodegroupName": nodegroup_name},
             failure_message="Error creating nodegroup",
@@ -186,7 +357,7 @@ class EksDeleteNodegroupTrigger(AwsBaseWaiterTrigger):
         waiter_delay: int,
         waiter_max_attempts: int,
         aws_conn_id: str,
-        region_name: str | None,
+        region_name: str | None = None,
     ):
         super().__init__(
             serialized_fields={"cluster_name": cluster_name, "nodegroup_name": nodegroup_name},

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

@@ -24,7 +24,7 @@ from botocore.exceptions import WaiterError
 
 from airflow.exceptions import AirflowProviderDeprecationWarning
 from airflow.providers.amazon.aws.hooks.base_aws import AwsGenericHook
-from airflow.providers.amazon.aws.hooks.emr import EmrContainerHook, EmrHook
+from airflow.providers.amazon.aws.hooks.emr import EmrContainerHook, EmrHook, EmrServerlessHook
 from airflow.providers.amazon.aws.triggers.base import AwsBaseWaiterTrigger
 from airflow.triggers.base import BaseTrigger, TriggerEvent
 
@@ -283,3 +283,217 @@ class EmrStepSensorTrigger(AwsBaseWaiterTrigger):
 
     def hook(self) -> AwsGenericHook:
         return EmrHook(self.aws_conn_id)
+
+
+class EmrServerlessCreateApplicationTrigger(AwsBaseWaiterTrigger):
+    """
+    Poll an Emr Serverless application and wait for it to be created.
+
+    :param application_id: The ID of the application being polled.
+    :waiter_delay: polling period in seconds to check for the status
+    :param waiter_max_attempts: The maximum number of attempts to be made
+    :param aws_conn_id: Reference to AWS connection id
+    """
+
+    def __init__(
+        self,
+        application_id: str,
+        waiter_delay: int = 30,
+        waiter_max_attempts: int = 60,
+        aws_conn_id: str = "aws_default",
+    ) -> None:
+        super().__init__(
+            serialized_fields={"application_id": application_id},
+            waiter_name="serverless_app_created",
+            waiter_args={"applicationId": application_id},
+            failure_message="Application creation failed",
+            status_message="Application status is",
+            status_queries=["application.state", "application.stateDetails"],
+            return_key="application_id",
+            return_value=application_id,
+            waiter_delay=waiter_delay,
+            waiter_max_attempts=waiter_max_attempts,
+            aws_conn_id=aws_conn_id,
+        )
+
+    def hook(self) -> AwsGenericHook:
+        return EmrServerlessHook(self.aws_conn_id)
+
+
+class EmrServerlessStartApplicationTrigger(AwsBaseWaiterTrigger):
+    """
+    Poll an Emr Serverless application and wait for it to be started.
+
+    :param application_id: The ID of the application being polled.
+    :waiter_delay: polling period in seconds to check for the status
+    :param waiter_max_attempts: The maximum number of attempts to be made
+    :param aws_conn_id: Reference to AWS connection id
+    """
+
+    def __init__(
+        self,
+        application_id: str,
+        waiter_delay: int = 30,
+        waiter_max_attempts: int = 60,
+        aws_conn_id: str = "aws_default",
+    ) -> None:
+        super().__init__(
+            serialized_fields={"application_id": application_id},
+            waiter_name="serverless_app_started",
+            waiter_args={"applicationId": application_id},
+            failure_message="Application failed to start",
+            status_message="Application status is",
+            status_queries=["application.state", "application.stateDetails"],
+            return_key="application_id",
+            return_value=application_id,
+            waiter_delay=waiter_delay,
+            waiter_max_attempts=waiter_max_attempts,
+            aws_conn_id=aws_conn_id,
+        )
+
+    def hook(self) -> AwsGenericHook:
+        return EmrServerlessHook(self.aws_conn_id)
+
+
+class EmrServerlessStopApplicationTrigger(AwsBaseWaiterTrigger):
+    """
+    Poll an Emr Serverless application and wait for it to be stopped.
+
+    :param application_id: The ID of the application being polled.
+    :waiter_delay: polling period in seconds to check for the status
+    :param waiter_max_attempts: The maximum number of attempts to be made
+    :param aws_conn_id: Reference to AWS connection id.
+    """
+
+    def __init__(
+        self,
+        application_id: str,
+        waiter_delay: int = 30,
+        waiter_max_attempts: int = 60,
+        aws_conn_id: str = "aws_default",
+    ) -> None:
+        super().__init__(
+            serialized_fields={"application_id": application_id},
+            waiter_name="serverless_app_stopped",
+            waiter_args={"applicationId": application_id},
+            failure_message="Application failed to start",
+            status_message="Application status is",
+            status_queries=["application.state", "application.stateDetails"],
+            return_key="application_id",
+            return_value=application_id,
+            waiter_delay=waiter_delay,
+            waiter_max_attempts=waiter_max_attempts,
+            aws_conn_id=aws_conn_id,
+        )
+
+    def hook(self) -> AwsGenericHook:
+        return EmrServerlessHook(self.aws_conn_id)
+
+
+class EmrServerlessStartJobTrigger(AwsBaseWaiterTrigger):
+    """
+    Poll an Emr Serverless job run and wait for it to be completed.
+
+    :param application_id: The ID of the application the job in being run on.
+    :param job_id: The ID of the job run.
+    :waiter_delay: polling period in seconds to check for the status
+    :param waiter_max_attempts: The maximum number of attempts to be made
+    :param aws_conn_id: Reference to AWS connection id
+    """
+
+    def __init__(
+        self,
+        application_id: str,
+        job_id: str | None,
+        waiter_delay: int = 30,
+        waiter_max_attempts: int = 60,
+        aws_conn_id: str = "aws_default",
+    ) -> None:
+        super().__init__(
+            serialized_fields={"application_id": application_id, "job_id": job_id},
+            waiter_name="serverless_job_completed",
+            waiter_args={"applicationId": application_id, "jobRunId": job_id},
+            failure_message="Serverless Job failed",
+            status_message="Serverless Job status is",
+            status_queries=["jobRun.state", "jobRun.stateDetails"],
+            return_key="job_id",
+            return_value=job_id,
+            waiter_delay=waiter_delay,
+            waiter_max_attempts=waiter_max_attempts,
+            aws_conn_id=aws_conn_id,
+        )
+
+    def hook(self) -> AwsGenericHook:
+        return EmrServerlessHook(self.aws_conn_id)
+
+
+class EmrServerlessDeleteApplicationTrigger(AwsBaseWaiterTrigger):
+    """
+    Poll an Emr Serverless application and wait for it to be deleted.
+
+    :param application_id: The ID of the application being polled.
+    :waiter_delay: polling period in seconds to check for the status
+    :param waiter_max_attempts: The maximum number of attempts to be made
+    :param aws_conn_id: Reference to AWS connection id
+    """
+
+    def __init__(
+        self,
+        application_id: str,
+        waiter_delay: int = 30,
+        waiter_max_attempts: int = 60,
+        aws_conn_id: str = "aws_default",
+    ) -> None:
+        super().__init__(
+            serialized_fields={"application_id": application_id},
+            waiter_name="serverless_app_terminated",
+            waiter_args={"applicationId": application_id},
+            failure_message="Application failed to start",
+            status_message="Application status is",
+            status_queries=["application.state", "application.stateDetails"],
+            return_key="application_id",
+            return_value=application_id,
+            waiter_delay=waiter_delay,
+            waiter_max_attempts=waiter_max_attempts,
+            aws_conn_id=aws_conn_id,
+        )
+
+    def hook(self) -> AwsGenericHook:
+        return EmrServerlessHook(self.aws_conn_id)
+
+
+class EmrServerlessCancelJobsTrigger(AwsBaseWaiterTrigger):
+    """
+    Trigger for canceling a list of jobs in an EMR Serverless application.
+
+    :param application_id: EMR Serverless application ID
+    :param aws_conn_id: Reference to AWS connection id
+    :param waiter_delay: Delay in seconds between each attempt to check the status
+    :param waiter_max_attempts: Maximum number of attempts to check the status
+    """
+
+    def __init__(
+        self,
+        application_id: str,
+        aws_conn_id: str,
+        waiter_delay: int,
+        waiter_max_attempts: int,
+    ) -> None:
+        self.hook_instance = EmrServerlessHook(aws_conn_id)
+        states = list(self.hook_instance.JOB_INTERMEDIATE_STATES.union({"CANCELLING"}))
+        super().__init__(
+            serialized_fields={"application_id": application_id},
+            waiter_name="no_job_running",
+            waiter_args={"applicationId": application_id, "states": states},
+            failure_message="Error while waiting for jobs to cancel",
+            status_message="Currently running jobs",
+            status_queries=["jobRuns[*].applicationId", "jobRuns[*].state"],
+            return_key="application_id",
+            return_value=application_id,
+            waiter_delay=waiter_delay,
+            waiter_max_attempts=waiter_max_attempts,
+            aws_conn_id=aws_conn_id,
+        )
+
+    def hook(self) -> AwsGenericHook:
+        return self.hook_instance