apache-airflow-providers-amazon 9.4.0rc1__py3-none-any.whl → 9.5.0__py3-none-any.whl

airflow/providers/amazon/aws/sensors/s3.py

@@ -23,7 +23,6 @@ import os
 import re
 from collections.abc import Sequence
 from datetime import datetime, timedelta
-from functools import cached_property
 from typing import TYPE_CHECKING, Any, Callable, cast
 
 from airflow.configuration import conf
@@ -34,11 +33,13 @@ if TYPE_CHECKING:
 
 from airflow.exceptions import AirflowException
 from airflow.providers.amazon.aws.hooks.s3 import S3Hook
+from airflow.providers.amazon.aws.sensors.base_aws import AwsBaseSensor
 from airflow.providers.amazon.aws.triggers.s3 import S3KeysUnchangedTrigger, S3KeyTrigger
-from airflow.sensors.base import BaseSensorOperator, poke_mode_only
+from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
+from airflow.sensors.base import poke_mode_only
 
 
-class S3KeySensor(BaseSensorOperator):
+class S3KeySensor(AwsBaseSensor[S3Hook]):
     """
     Waits for one or multiple keys (a file-like instance on S3) to be present in a S3 bucket.
 
@@ -65,17 +66,6 @@ class S3KeySensor(BaseSensorOperator):
 
             def check_fn(files: List, **kwargs) -> bool:
                 return any(f.get('Size', 0) > 1048576 for f in files)
-    :param aws_conn_id: a reference to the s3 connection
-    :param verify: Whether to verify SSL certificates for S3 connection.
-        By default, SSL certificates are verified.
-        You can provide the following values:
-
-        - ``False``: do not validate SSL certificates. SSL will still be used
-                 (unless use_ssl is False), but SSL certificates will not be
-                 verified.
-        - ``path/to/cert/bundle.pem``: A filename of the CA cert bundle to uses.
-                 You can specify this argument if you want to use a different
-                 CA cert bundle than the one used by botocore.
     :param deferrable: Run operator in the deferrable mode
     :param use_regex: whether to use regex to check bucket
     :param metadata_keys: List of head_object attributes to gather and send to ``check_fn``.
@@ -83,9 +73,18 @@ class S3KeySensor(BaseSensorOperator):
         all available attributes.
         Default value: "Size".
         If the requested attribute is not found, the key is still included and the value is None.
+    :param aws_conn_id: The Airflow connection used for AWS credentials.
+        If this is ``None`` or empty then the default boto3 behaviour is used. If
+        running Airflow in a distributed manner and aws_conn_id is None or
+        empty, then default boto3 configuration would be used (and must be
+        maintained on each worker node).
+    :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+    :param verify: Whether or not to verify SSL certificates. See:
+        https://boto3.amazonaws.com/v1/documentation/api/latest/reference/core/session.html
     """
 
-    template_fields: Sequence[str] = ("bucket_key", "bucket_name")
+    template_fields: Sequence[str] = aws_template_fields("bucket_key", "bucket_name")
+    aws_hook_class = S3Hook
 
     def __init__(
         self,
@@ -94,7 +93,6 @@ class S3KeySensor(BaseSensorOperator):
         bucket_name: str | None = None,
         wildcard_match: bool = False,
         check_fn: Callable[..., bool] | None = None,
-        aws_conn_id: str | None = "aws_default",
         verify: str | bool | None = None,
         deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
         use_regex: bool = False,
@@ -106,14 +104,13 @@ class S3KeySensor(BaseSensorOperator):
         self.bucket_key = bucket_key
         self.wildcard_match = wildcard_match
         self.check_fn = check_fn
-        self.aws_conn_id = aws_conn_id
         self.verify = verify
         self.deferrable = deferrable
         self.use_regex = use_regex
         self.metadata_keys = metadata_keys if metadata_keys else ["Size"]
 
     def _check_key(self, key, context: Context):
-        bucket_name, key = S3Hook.get_s3_bucket_key(self.bucket_name, key, "bucket_name", "bucket_key")
+        bucket_name, key = self.hook.get_s3_bucket_key(self.bucket_name, key, "bucket_name", "bucket_key")
         self.log.info("Poking for key : s3://%s/%s", bucket_name, key)
 
         """
@@ -195,11 +192,13 @@ class S3KeySensor(BaseSensorOperator):
         self.defer(
             timeout=timedelta(seconds=self.timeout),
             trigger=S3KeyTrigger(
-                bucket_name=cast(str, self.bucket_name),
+                bucket_name=cast("str", self.bucket_name),
                 bucket_key=self.bucket_key,
                 wildcard_match=self.wildcard_match,
                 aws_conn_id=self.aws_conn_id,
+                region_name=self.region_name,
                 verify=self.verify,
+                botocore_config=self.botocore_config,
                 poke_interval=self.poke_interval,
                 should_check_fn=bool(self.check_fn),
                 use_regex=self.use_regex,
@@ -220,13 +219,9 @@ class S3KeySensor(BaseSensorOperator):
         elif event["status"] == "error":
             raise AirflowException(event["message"])
 
-    @cached_property
-    def hook(self) -> S3Hook:
-        return S3Hook(aws_conn_id=self.aws_conn_id, verify=self.verify)
-
 
 @poke_mode_only
-class S3KeysUnchangedSensor(BaseSensorOperator):
+class S3KeysUnchangedSensor(AwsBaseSensor[S3Hook]):
     """
     Return True if inactivity_period has passed with no increase in the number of objects matching prefix.
 
@@ -239,17 +234,7 @@ class S3KeysUnchangedSensor(BaseSensorOperator):
 
     :param bucket_name: Name of the S3 bucket
     :param prefix: The prefix being waited on. Relative path from bucket root level.
-    :param aws_conn_id: a reference to the s3 connection
-    :param verify: Whether or not to verify SSL certificates for S3 connection.
-        By default SSL certificates are verified.
-        You can provide the following values:
-
-        - ``False``: do not validate SSL certificates. SSL will still be used
-                 (unless use_ssl is False), but SSL certificates will not be
-                 verified.
-        - ``path/to/cert/bundle.pem``: A filename of the CA cert bundle to uses.
-                 You can specify this argument if you want to use a different
-                 CA cert bundle than the one used by botocore.
+        https://boto3.amazonaws.com/v1/documentation/api/latest/reference/core/session.html
     :param inactivity_period: The total seconds of inactivity to designate
         keys unchanged. Note, this mechanism is not real time and
         this operator may not return until a poke_interval after this period
@@ -261,16 +246,24 @@ class S3KeysUnchangedSensor(BaseSensorOperator):
         between pokes valid behavior. If true a warning message will be logged
         when this happens. If false an error will be raised.
     :param deferrable: Run sensor in the deferrable mode
+    :param aws_conn_id: The Airflow connection used for AWS credentials.
+        If this is ``None`` or empty then the default boto3 behaviour is used. If
+        running Airflow in a distributed manner and aws_conn_id is None or
+        empty, then default boto3 configuration would be used (and must be
+        maintained on each worker node).
+    :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+    :param verify: Whether or not to verify SSL certificates. See:
+        https://boto3.amazonaws.com/v1/documentation/api/latest/reference/core/session.html
     """
 
-    template_fields: Sequence[str] = ("bucket_name", "prefix")
+    template_fields: Sequence[str] = aws_template_fields("bucket_name", "prefix")
+    aws_hook_class = S3Hook
 
     def __init__(
         self,
         *,
         bucket_name: str,
         prefix: str,
-        aws_conn_id: str | None = "aws_default",
         verify: bool | str | None = None,
         inactivity_period: float = 60 * 60,
         min_objects: int = 1,
@@ -291,15 +284,9 @@ class S3KeysUnchangedSensor(BaseSensorOperator):
         self.inactivity_seconds = 0
         self.allow_delete = allow_delete
         self.deferrable = deferrable
-        self.aws_conn_id = aws_conn_id
         self.verify = verify
         self.last_activity_time: datetime | None = None
 
-    @cached_property
-    def hook(self):
-        """Returns S3Hook."""
-        return S3Hook(aws_conn_id=self.aws_conn_id, verify=self.verify)
-
     def is_keys_unchanged(self, current_objects: set[str]) -> bool:
         """
         Check for new objects after the inactivity_period and update the sensor state accordingly.
@@ -382,7 +369,9 @@ class S3KeysUnchangedSensor(BaseSensorOperator):
                         inactivity_seconds=self.inactivity_seconds,
                         allow_delete=self.allow_delete,
                         aws_conn_id=self.aws_conn_id,
+                        region_name=self.region_name,
                         verify=self.verify,
+                        botocore_config=self.botocore_config,
                         last_activity_time=self.last_activity_time,
                     ),
                     method_name="execute_complete",

airflow/providers/amazon/aws/sensors/sagemaker_unified_studio.py

@@ -0,0 +1,73 @@
+# 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.
+
+"""This module contains the Amazon SageMaker Unified Studio Notebook sensor."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from airflow.exceptions import AirflowException
+from airflow.providers.amazon.aws.hooks.sagemaker_unified_studio import (
+    SageMakerNotebookHook,
+)
+from airflow.sensors.base import BaseSensorOperator
+
+if TYPE_CHECKING:
+    from airflow.utils.context import Context
+
+
+class SageMakerNotebookSensor(BaseSensorOperator):
+    """
+    Waits for a Sagemaker Workflows Notebook execution to reach any of the status below.
+
+    'FAILED', 'STOPPED', 'COMPLETED'
+
+    :param execution_id: The Sagemaker Workflows Notebook running execution identifier
+    :param execution_name: The Sagemaker Workflows Notebook unique execution name
+    """
+
+    def __init__(self, *, execution_id: str, execution_name: str, **kwargs):
+        super().__init__(**kwargs)
+        self.execution_id = execution_id
+        self.execution_name = execution_name
+        self.success_state = ["COMPLETED"]
+        self.in_progress_states = ["PENDING", "RUNNING"]
+
+    def hook(self):
+        return SageMakerNotebookHook(execution_name=self.execution_name)
+
+    # override from base sensor
+    def poke(self, context=None):
+        status = self.hook().get_execution_status(execution_id=self.execution_id)
+
+        if status in self.success_state:
+            log_info_message = f"Exiting Execution {self.execution_id} State: {status}"
+            self.log.info(log_info_message)
+            return True
+        elif status in self.in_progress_states:
+            return False
+        else:
+            error_message = f"Exiting Execution {self.execution_id} State: {status}"
+            self.log.info(error_message)
+            raise AirflowException(error_message)
+
+    def execute(self, context: Context):
+        # This will invoke poke method in the base sensor
+        log_info_message = f"Polling Sagemaker Workflows Artifact execution: {self.execution_name} and execution id: {self.execution_id}"
+        self.log.info(log_info_message)
+        super().execute(context=context)

airflow/providers/amazon/aws/sensors/step_function.py

@@ -81,5 +81,6 @@ class StepFunctionExecutionSensor(AwsBaseSensor[StepFunctionHook]):
             return False
 
         self.log.info("Doing xcom_push of output")
-        self.xcom_push(context, "output", output)
+
+        context["ti"].xcom_push(key="output", value=output)
         return True

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

@@ -103,7 +103,7 @@ class MongoToS3Operator(BaseOperator):
         if self.is_pipeline:
             results: CommandCursor[Any] | Cursor = MongoHook(self.mongo_conn_id).aggregate(
                 mongo_collection=self.mongo_collection,
-                aggregate_query=cast(list, self.mongo_query),
+                aggregate_query=cast("list", self.mongo_query),
                 mongo_db=self.mongo_db,
                 allowDiskUse=self.allow_disk_use,
             )
@@ -111,7 +111,7 @@ class MongoToS3Operator(BaseOperator):
         else:
             results = MongoHook(self.mongo_conn_id).find(
                 mongo_collection=self.mongo_collection,
-                query=cast(dict, self.mongo_query),
+                query=cast("dict", self.mongo_query),
                 projection=self.mongo_projection,
                 mongo_db=self.mongo_db,
                 find_one=False,

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

@@ -29,6 +29,7 @@ from airflow.providers.amazon.aws.hooks.redshift_data import RedshiftDataHook
 from airflow.providers.amazon.aws.hooks.redshift_sql import RedshiftSQLHook
 from airflow.providers.amazon.aws.hooks.s3 import S3Hook
 from airflow.providers.amazon.aws.utils.redshift import build_credentials_block
+from airflow.utils.types import NOTSET, ArgNotSet
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context
@@ -102,7 +103,7 @@ class RedshiftToS3Operator(BaseOperator):
         table: str | None = None,
         select_query: str | None = None,
         redshift_conn_id: str = "redshift_default",
-        aws_conn_id: str | None = "aws_default",
+        aws_conn_id: str | None | ArgNotSet = NOTSET,
         verify: bool | str | None = None,
         unload_options: list | None = None,
         autocommit: bool = False,
@@ -118,7 +119,6 @@ class RedshiftToS3Operator(BaseOperator):
         self.schema = schema
         self.table = table
         self.redshift_conn_id = redshift_conn_id
-        self.aws_conn_id = aws_conn_id
         self.verify = verify
         self.unload_options = unload_options or []
         self.autocommit = autocommit
@@ -127,6 +127,16 @@ class RedshiftToS3Operator(BaseOperator):
         self.table_as_file_name = table_as_file_name
         self.redshift_data_api_kwargs = redshift_data_api_kwargs or {}
         self.select_query = select_query
+        # In execute() we attempt to fetch this aws connection to check for extras. If the user didn't
+        # actually provide a connection note that, because we don't want to let the exception bubble up in
+        # that case (since we're silently injecting a connection on their behalf).
+        self._aws_conn_id: str | None
+        if isinstance(aws_conn_id, ArgNotSet):
+            self.conn_set = False
+            self._aws_conn_id = "aws_default"
+        else:
+            self.conn_set = True
+            self._aws_conn_id = aws_conn_id
 
     def _build_unload_query(
         self, credentials_block: str, select_query: str, s3_key: str, unload_options: str
@@ -176,11 +186,16 @@ class RedshiftToS3Operator(BaseOperator):
                     raise AirflowException(f"Cannot include param '{arg}' in Redshift Data API kwargs")
         else:
             redshift_sql_hook = RedshiftSQLHook(redshift_conn_id=self.redshift_conn_id)
-        conn = S3Hook.get_connection(conn_id=self.aws_conn_id) if self.aws_conn_id else None
+        conn = (
+            S3Hook.get_connection(conn_id=self._aws_conn_id)
+            # Only fetch the connection if it was set by the user and it is not None
+            if self.conn_set and self._aws_conn_id
+            else None
+        )
         if conn and conn.extra_dejson.get("role_arn", False):
             credentials_block = f"aws_iam_role={conn.extra_dejson['role_arn']}"
         else:
-            s3_hook = S3Hook(aws_conn_id=self.aws_conn_id, verify=self.verify)
+            s3_hook = S3Hook(aws_conn_id=self._aws_conn_id, verify=self.verify)
             credentials = s3_hook.get_credentials()
             credentials_block = build_credentials_block(credentials)
 

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

@@ -25,6 +25,7 @@ from airflow.providers.amazon.aws.hooks.redshift_data import RedshiftDataHook
 from airflow.providers.amazon.aws.hooks.redshift_sql import RedshiftSQLHook
 from airflow.providers.amazon.aws.hooks.s3 import S3Hook
 from airflow.providers.amazon.aws.utils.redshift import build_credentials_block
+from airflow.utils.types import NOTSET, ArgNotSet
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context
@@ -93,7 +94,7 @@ class S3ToRedshiftOperator(BaseOperator):
         s3_key: str,
         schema: str | None = None,
         redshift_conn_id: str = "redshift_default",
-        aws_conn_id: str | None = "aws_default",
+        aws_conn_id: str | None | ArgNotSet = NOTSET,
         verify: bool | str | None = None,
         column_list: list[str] | None = None,
         copy_options: list | None = None,
@@ -117,6 +118,16 @@ class S3ToRedshiftOperator(BaseOperator):
         self.method = method
         self.upsert_keys = upsert_keys
         self.redshift_data_api_kwargs = redshift_data_api_kwargs or {}
+        # In execute() we attempt to fetch this aws connection to check for extras. If the user didn't
+        # actually provide a connection note that, because we don't want to let the exception bubble up in
+        # that case (since we're silently injecting a connection on their behalf).
+        self._aws_conn_id: str | None
+        if isinstance(aws_conn_id, ArgNotSet):
+            self.conn_set = False
+            self._aws_conn_id = "aws_default"
+        else:
+            self.conn_set = True
+            self._aws_conn_id = aws_conn_id
 
         if self.redshift_data_api_kwargs:
             for arg in ["sql", "parameters"]:
@@ -149,14 +160,19 @@ class S3ToRedshiftOperator(BaseOperator):
         else:
             redshift_sql_hook = RedshiftSQLHook(redshift_conn_id=self.redshift_conn_id)
 
-        conn = S3Hook.get_connection(conn_id=self.aws_conn_id) if self.aws_conn_id else None
+        conn = (
+            S3Hook.get_connection(conn_id=self._aws_conn_id)
+            # Only fetch the connection if it was set by the user and it is not None
+            if self.conn_set and self._aws_conn_id
+            else None
+        )
         region_info = ""
         if conn and conn.extra_dejson.get("region", False):
             region_info = f"region '{conn.extra_dejson['region']}'"
         if conn and conn.extra_dejson.get("role_arn", False):
             credentials_block = f"aws_iam_role={conn.extra_dejson['role_arn']}"
         else:
-            s3_hook = S3Hook(aws_conn_id=self.aws_conn_id, verify=self.verify)
+            s3_hook = S3Hook(aws_conn_id=self._aws_conn_id, verify=self.verify)
             credentials = s3_hook.get_credentials()
             credentials_block = build_credentials_block(credentials)
 

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

@@ -223,7 +223,7 @@ class SqlToS3Operator(BaseOperator):
             return
         for group_label in (grouped_df := df.groupby(**self.groupby_kwargs)).groups:
             yield (
-                cast(str, group_label),
+                cast("str", group_label),
                 grouped_df.get_group(group_label)
                 .drop(random_column_name, axis=1, errors="ignore")
                 .reset_index(drop=True),

airflow/providers/amazon/aws/triggers/README.md

@@ -65,10 +65,10 @@ To call the asynchronous `wait` function, first create a hook for the particular
 self.redshift_hook = RedshiftHook(aws_conn_id=self.aws_conn_id)
 ```
 
-With this hook, we can use the async_conn property to get access to the aiobotocore client:
+With this hook, we can use the asynchronous get_async_conn method to get access to the aiobotocore client:
 
 ```python
-async with self.redshift_hook.async_conn as client:
+async with await self.redshift_hook.get_async_conn() as client:
     await client.get_waiter("cluster_available").wait(
         ClusterIdentifier=self.cluster_identifier,
         WaiterConfig={
@@ -81,7 +81,7 @@ async with self.redshift_hook.async_conn as client:
 In this case, we are using the built-in cluster_available waiter. If we wanted to use a custom waiter, we would change the code slightly to use the `get_waiter` function from the hook, rather than the aiobotocore client:
 
 ```python
-async with self.redshift_hook.async_conn as client:
+async with await self.redshift_hook.get_async_conn() as client:
     waiter = self.redshift_hook.get_waiter("cluster_paused", deferrable=True, client=client)
     await waiter.wait(
         ClusterIdentifier=self.cluster_identifier,
@@ -131,7 +131,7 @@ For more information about writing custom waiter, see the [README.md](https://gi
 In some cases, a built-in or custom waiter may not be able to solve the problem. In such cases, the asynchronous method used to poll the boto3 API would need to be defined in the hook of the service being used. This method is essentially the same as the synchronous version of the method, except that it will use the aiobotocore client, and will be awaited. For the Redshift example, the async `describe_clusters` method would look as follows:
 
 ```python
-async with self.async_conn as client:
+async with await self.get_async_conn() as client:
     response = client.describe_clusters(ClusterIdentifier=self.cluster_identifier)
 ```
 

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

@@ -55,6 +55,8 @@ class AwsBaseWaiterTrigger(BaseTrigger):
 
     :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 waiter_config_overrides: A dict to update waiter's default configuration. Only specified keys will
+        be updated.
     :param aws_conn_id: The Airflow connection used for AWS credentials. To be used to build the hook.
     :param region_name: The AWS region where the resources to watch are. To be used to build the hook.
     :param verify: Whether or not to verify SSL certificates. To be used to build the hook.
@@ -77,6 +79,7 @@ class AwsBaseWaiterTrigger(BaseTrigger):
         return_value: Any,
         waiter_delay: int,
         waiter_max_attempts: int,
+        waiter_config_overrides: dict[str, Any] | None = None,
         aws_conn_id: str | None,
         region_name: str | None = None,
         verify: bool | str | None = None,
@@ -91,6 +94,7 @@ class AwsBaseWaiterTrigger(BaseTrigger):
         self.failure_message = failure_message
         self.status_message = status_message
         self.status_queries = status_queries
+        self.waiter_config_overrides = waiter_config_overrides
 
         self.return_key = return_key
         self.return_value = return_value
@@ -139,8 +143,13 @@ class AwsBaseWaiterTrigger(BaseTrigger):
 
     async def run(self) -> AsyncIterator[TriggerEvent]:
         hook = self.hook()
-        async with hook.async_conn as client:
-            waiter = hook.get_waiter(self.waiter_name, deferrable=True, client=client)
+        async with await hook.get_async_conn() as client:
+            waiter = hook.get_waiter(
+                self.waiter_name,
+                deferrable=True,
+                client=client,
+                config_overrides=self.waiter_config_overrides,
+            )
             await async_wait(
                 waiter,
                 self.waiter_delay,

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

@@ -167,8 +167,12 @@ class TaskDoneTrigger(BaseTrigger):
 
     async def run(self) -> AsyncIterator[TriggerEvent]:
         async with (
-            EcsHook(aws_conn_id=self.aws_conn_id, region_name=self.region).async_conn as ecs_client,
-            AwsLogsHook(aws_conn_id=self.aws_conn_id, region_name=self.region).async_conn as logs_client,
+            await EcsHook(
+                aws_conn_id=self.aws_conn_id, region_name=self.region
+            ).get_async_conn() as ecs_client,
+            await AwsLogsHook(
+                aws_conn_id=self.aws_conn_id, region_name=self.region
+            ).get_async_conn() as logs_client,
         ):
             waiter = ecs_client.get_waiter("tasks_stopped")
             logs_token = None

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

@@ -70,7 +70,7 @@ class EksCreateClusterTrigger(AwsBaseWaiterTrigger):
         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:
+        async with await self.hook().get_async_conn() as client:
             waiter = client.get_waiter(self.waiter_name)
             try:
                 await async_wait(
@@ -140,7 +140,7 @@ class EksDeleteClusterTrigger(AwsBaseWaiterTrigger):
         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:
+        async with await self.hook().get_async_conn() as client:
             waiter = client.get_waiter("cluster_deleted")
             if self.force_delete_compute:
                 await self.delete_any_nodegroups(client=client)

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

@@ -157,7 +157,7 @@ class GlueCatalogPartitionTrigger(BaseTrigger):
         return bool(partitions)
 
     async def run(self) -> AsyncIterator[TriggerEvent]:
-        async with self.hook.async_conn as client:
+        async with await self.hook.get_async_conn() as client:
             while True:
                 result = await self.poke(client=client)
                 if result:

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

@@ -0,0 +1,128 @@
+# 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 collections.abc import Collection
+from typing import TYPE_CHECKING
+
+from airflow.providers.amazon.aws.hooks.mwaa import MwaaHook
+from airflow.providers.amazon.aws.triggers.base import AwsBaseWaiterTrigger
+from airflow.utils.state import DagRunState
+
+if TYPE_CHECKING:
+    from airflow.providers.amazon.aws.hooks.base_aws import AwsGenericHook
+
+
+class MwaaDagRunCompletedTrigger(AwsBaseWaiterTrigger):
+    """
+    Trigger when an MWAA Dag Run is complete.
+
+    :param external_env_name: The external MWAA environment name that contains the DAG Run you want to wait for
+        (templated)
+    :param external_dag_id: The DAG ID in the external MWAA environment that contains the DAG Run you want to wait for
+        (templated)
+    :param external_dag_run_id: The DAG Run ID in the external MWAA environment that you want to wait for (templated)
+    :param success_states: Collection of DAG Run states that would make this task marked as successful, default is
+        ``{airflow.utils.state.DagRunState.SUCCESS}`` (templated)
+    :param failure_states: Collection of DAG Run states that would make this task marked as failed and raise an
+        AirflowException, default is ``{airflow.utils.state.DagRunState.FAILED}`` (templated)
+    :param waiter_delay: The amount of time in seconds to wait between attempts. (default: 60)
+    :param waiter_max_attempts: The maximum number of attempts to be made. (default: 720)
+    :param aws_conn_id: The Airflow connection used for AWS credentials.
+    """
+
+    def __init__(
+        self,
+        *,
+        external_env_name: str,
+        external_dag_id: str,
+        external_dag_run_id: str,
+        success_states: Collection[str] | None = None,
+        failure_states: Collection[str] | None = None,
+        waiter_delay: int = 60,
+        waiter_max_attempts: int = 720,
+        aws_conn_id: str | None = None,
+    ) -> None:
+        self.success_states = set(success_states) if success_states else {DagRunState.SUCCESS.value}
+        self.failure_states = set(failure_states) if failure_states else {DagRunState.FAILED.value}
+
+        if len(self.success_states & self.failure_states):
+            raise ValueError("success_states and failure_states must not have any values in common")
+
+        in_progress_states = {s.value for s in DagRunState} - self.success_states - self.failure_states
+
+        super().__init__(
+            serialized_fields={
+                "external_env_name": external_env_name,
+                "external_dag_id": external_dag_id,
+                "external_dag_run_id": external_dag_run_id,
+                "success_states": success_states,
+                "failure_states": failure_states,
+            },
+            waiter_name="mwaa_dag_run_complete",
+            waiter_args={
+                "Name": external_env_name,
+                "Path": f"/dags/{external_dag_id}/dagRuns/{external_dag_run_id}",
+                "Method": "GET",
+            },
+            failure_message=f"The DAG run {external_dag_run_id} of DAG {external_dag_id} in MWAA environment {external_env_name} failed with state",
+            status_message="State of DAG run",
+            status_queries=["RestApiResponse.state"],
+            return_key="dag_run_id",
+            return_value=external_dag_run_id,
+            waiter_delay=waiter_delay,
+            waiter_max_attempts=waiter_max_attempts,
+            aws_conn_id=aws_conn_id,
+            waiter_config_overrides={
+                "acceptors": _build_waiter_acceptors(
+                    success_states=self.success_states,
+                    failure_states=self.failure_states,
+                    in_progress_states=in_progress_states,
+                )
+            },
+        )
+
+    def hook(self) -> AwsGenericHook:
+        return MwaaHook(
+            aws_conn_id=self.aws_conn_id,
+            region_name=self.region_name,
+            verify=self.verify,
+            config=self.botocore_config,
+        )
+
+
+def _build_waiter_acceptors(
+    success_states: set[str], failure_states: set[str], in_progress_states: set[str]
+) -> list:
+    acceptors = []
+    for state_set, state_waiter_category in (
+        (success_states, "success"),
+        (failure_states, "failure"),
+        (in_progress_states, "retry"),
+    ):
+        for dag_run_state in state_set:
+            acceptors.append(
+                {
+                    "matcher": "path",
+                    "argument": "RestApiResponse.state",
+                    "expected": dag_run_state,
+                    "state": state_waiter_category,
+                }
+            )
+
+    return acceptors