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

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

@@ -0,0 +1,155 @@
+# 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 operator."""
+
+from __future__ import annotations
+
+from functools import cached_property
+from typing import TYPE_CHECKING
+
+from airflow.configuration import conf
+from airflow.exceptions import AirflowException
+from airflow.models import BaseOperator
+from airflow.providers.amazon.aws.hooks.sagemaker_unified_studio import (
+    SageMakerNotebookHook,
+)
+from airflow.providers.amazon.aws.links.sagemaker_unified_studio import (
+    SageMakerUnifiedStudioLink,
+)
+from airflow.providers.amazon.aws.triggers.sagemaker_unified_studio import (
+    SageMakerNotebookJobTrigger,
+)
+
+if TYPE_CHECKING:
+    from airflow.utils.context import Context
+
+
+class SageMakerNotebookOperator(BaseOperator):
+    """
+    Provides Artifact execution functionality for Sagemaker Unified Studio Workflows.
+
+    Examples:
+     .. code-block:: python
+
+        from airflow.providers.amazon.aws.operators.sagemaker_unified_studio import SageMakerNotebookOperator
+
+        notebook_operator = SageMakerNotebookOperator(
+            task_id="notebook_task",
+            input_config={"input_path": "path/to/notebook.ipynb", "input_params": ""},
+            output_config={"output_format": "ipynb"},
+            wait_for_completion=True,
+            waiter_delay=10,
+            waiter_max_attempts=1440,
+        )
+
+    :param task_id: A unique, meaningful id for the task.
+    :param input_config: Configuration for the input file. Input path should be specified as a relative path.
+        The provided relative path will be automatically resolved to an absolute path within
+        the context of the user's home directory in the IDE. Input params should be a dict.
+        Example: {'input_path': 'folder/input/notebook.ipynb', 'input_params':{'key': 'value'}}
+    :param output_config:  Configuration for the output format. It should include an output_format parameter to control
+        the format of the notebook execution output.
+        Example: {"output_formats": ["NOTEBOOK"]}
+    :param compute: compute configuration to use for the artifact execution. This is a required attribute
+        if the execution is on a remote compute.
+        Example: { "InstanceType": "ml.m5.large", "VolumeSizeInGB": 30, "VolumeKmsKeyId": "", "ImageUri": "string", "ContainerEntrypoint": [ "string" ]}
+    :param termination_condition: conditions to match to terminate the remote execution.
+        Example: { "MaxRuntimeInSeconds": 3600 }
+    :param tags: tags to be associated with the remote execution runs.
+        Example: { "md_analytics": "logs" }
+    :param wait_for_completion: Indicates whether to wait for the notebook execution to complete. If True, wait for completion; if False, don't wait.
+    :param waiter_delay: Interval in seconds to check the notebook execution status.
+    :param waiter_max_attempts: Number of attempts to wait before returning FAILED.
+    :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)
+
+    .. seealso::
+        For more information on how to use this operator, take a look at the guide:
+        :ref:`howto/operator:SageMakerNotebookOperator`
+    """
+
+    operator_extra_links = (SageMakerUnifiedStudioLink(),)
+
+    def __init__(
+        self,
+        task_id: str,
+        input_config: dict,
+        output_config: dict | None = None,
+        compute: dict | None = None,
+        termination_condition: dict | None = None,
+        tags: dict | None = None,
+        wait_for_completion: bool = True,
+        waiter_delay: int = 10,
+        waiter_max_attempts: int = 1440,
+        deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+        **kwargs,
+    ):
+        super().__init__(task_id=task_id, **kwargs)
+        self.execution_name = task_id
+        self.input_config = input_config
+        self.output_config = output_config or {"output_formats": ["NOTEBOOK"]}
+        self.compute = compute or {}
+        self.termination_condition = termination_condition or {}
+        self.tags = tags or {}
+        self.wait_for_completion = wait_for_completion
+        self.waiter_delay = waiter_delay
+        self.waiter_max_attempts = waiter_max_attempts
+        self.deferrable = deferrable
+        self.input_kwargs = kwargs
+
+    @cached_property
+    def notebook_execution_hook(self):
+        if not self.input_config:
+            raise AirflowException("input_config is required")
+
+        if "input_path" not in self.input_config:
+            raise AirflowException("input_path is a required field in the input_config")
+
+        return SageMakerNotebookHook(
+            input_config=self.input_config,
+            output_config=self.output_config,
+            execution_name=self.execution_name,
+            compute=self.compute,
+            termination_condition=self.termination_condition,
+            tags=self.tags,
+            waiter_delay=self.waiter_delay,
+            waiter_max_attempts=self.waiter_max_attempts,
+        )
+
+    def execute(self, context: Context):
+        notebook_execution = self.notebook_execution_hook.start_notebook_execution()
+        execution_id = notebook_execution["execution_id"]
+
+        if self.deferrable:
+            self.defer(
+                trigger=SageMakerNotebookJobTrigger(
+                    execution_id=execution_id,
+                    execution_name=self.execution_name,
+                    waiter_delay=self.waiter_delay,
+                    waiter_max_attempts=self.waiter_max_attempts,
+                ),
+                method_name="execute_complete",
+            )
+        elif self.wait_for_completion:
+            response = self.notebook_execution_hook.wait_for_execution_completion(execution_id, context)
+            status = response["Status"]
+            log_info_message = (
+                f"Notebook Execution: {self.execution_name} Status: {status}. Run Id: {execution_id}"
+            )
+            self.log.info(log_info_message)

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

@@ -311,7 +311,7 @@ class EmrContainerSensor(BaseSensorOperator):
         )
 
         if state in self.FAILURE_STATES:
-            raise AirflowException("EMR Containers sensor failed")
+            raise AirflowException(f"EMR Containers sensor failed due to state: {state}")
 
         if state in self.INTERMEDIATE_STATES:
             return False

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

@@ -0,0 +1,113 @@
+#
+# 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, Sequence
+from typing import TYPE_CHECKING
+
+from airflow.exceptions import AirflowException
+from airflow.providers.amazon.aws.hooks.mwaa import MwaaHook
+from airflow.providers.amazon.aws.sensors.base_aws import AwsBaseSensor
+from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
+from airflow.utils.state import State
+
+if TYPE_CHECKING:
+    from airflow.utils.context import Context
+
+
+class MwaaDagRunSensor(AwsBaseSensor[MwaaHook]):
+    """
+    Waits for a DAG Run in an MWAA Environment to complete.
+
+    If the DAG Run fails, an AirflowException is thrown.
+
+    .. seealso::
+        For more information on how to use this sensor, take a look at the guide:
+        :ref:`howto/sensor:MwaaDagRunSensor`
+
+    :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.State.success_states`` (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.State.failed_states`` (templated)
+    """
+
+    aws_hook_class = MwaaHook
+    template_fields: Sequence[str] = aws_template_fields(
+        "external_env_name",
+        "external_dag_id",
+        "external_dag_run_id",
+        "success_states",
+        "failure_states",
+    )
+
+    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,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        self.success_states = set(success_states if success_states else State.success_states)
+        self.failure_states = set(failure_states if failure_states else State.failed_states)
+
+        if len(self.success_states & self.failure_states):
+            raise AirflowException("allowed_states and failed_states must not have any values in common")
+
+        self.external_env_name = external_env_name
+        self.external_dag_id = external_dag_id
+        self.external_dag_run_id = external_dag_run_id
+
+    def poke(self, context: Context) -> bool:
+        self.log.info(
+            "Poking for DAG run %s of DAG %s in MWAA environment %s",
+            self.external_dag_run_id,
+            self.external_dag_id,
+            self.external_env_name,
+        )
+        response = self.hook.invoke_rest_api(
+            env_name=self.external_env_name,
+            path=f"/dags/{self.external_dag_id}/dagRuns/{self.external_dag_run_id}",
+            method="GET",
+        )
+
+        # If RestApiStatusCode == 200, the RestApiResponse must have the "state" key, otherwise something terrible has
+        # happened in the API and KeyError would be raised
+        # If RestApiStatusCode >= 300, a botocore exception would've already been raised during the
+        # self.hook.invoke_rest_api call
+        # The scope of this sensor is going to only be raising AirflowException due to failure of the DAGRun
+
+        state = response["RestApiResponse"]["state"]
+        if state in self.success_states:
+            return True
+
+        if state in self.failure_states:
+            raise AirflowException(
+                f"The DAG run {self.external_dag_run_id} of DAG {self.external_dag_id} in MWAA environment {self.external_env_name} "
+                f"failed with state {state}."
+            )
+        return False

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

@@ -20,7 +20,7 @@ from collections.abc import Sequence
 from functools import cached_property
 from typing import TYPE_CHECKING
 
-from airflow.exceptions import AirflowNotFoundException
+from airflow.exceptions import AirflowException, AirflowNotFoundException
 from airflow.providers.amazon.aws.hooks.rds import RdsHook
 from airflow.providers.amazon.aws.utils.rds import RdsDbType
 from airflow.sensors.base import BaseSensorOperator
@@ -104,18 +104,17 @@ class RdsExportTaskExistenceSensor(RdsBaseSensor):
 
     :param export_task_identifier: A unique identifier for the snapshot export task.
     :param target_statuses: Target status of export task
+    :param error_statuses: Target error status of export task to fail the sensor
     """
 
-    template_fields: Sequence[str] = (
-        "export_task_identifier",
-        "target_statuses",
-    )
+    template_fields: Sequence[str] = ("export_task_identifier", "target_statuses", "error_statuses")
 
     def __init__(
         self,
         *,
         export_task_identifier: str,
         target_statuses: list[str] | None = None,
+        error_statuses: list[str] | None = None,
         aws_conn_id: str | None = "aws_default",
         **kwargs,
     ):
@@ -129,6 +128,7 @@ class RdsExportTaskExistenceSensor(RdsBaseSensor):
             "canceling",
             "canceled",
         ]
+        self.error_statuses = error_statuses or ["failed"]
 
     def poke(self, context: Context):
         self.log.info(
@@ -136,6 +136,11 @@ class RdsExportTaskExistenceSensor(RdsBaseSensor):
         )
         try:
             state = self.hook.get_export_task_state(self.export_task_identifier)
+            if state in self.error_statuses:
+                raise AirflowException(
+                    f"Export task {self.export_task_identifier} failed with status {state}"
+                )
+
         except AirflowNotFoundException:
             return False
         return state in self.target_statuses

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)
 
         """
@@ -199,7 +196,9 @@ class S3KeySensor(BaseSensorOperator):
                 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/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

@@ -139,7 +139,7 @@ class AwsBaseWaiterTrigger(BaseTrigger):
 
     async def run(self) -> AsyncIterator[TriggerEvent]:
         hook = self.hook()
-        async with hook.async_conn as client:
+        async with await hook.get_async_conn() as client:
             waiter = hook.get_waiter(self.waiter_name, deferrable=True, client=client)
             await async_wait(
                 waiter,

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: