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

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/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/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

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

@@ -53,6 +53,9 @@ class S3KeyTrigger(BaseTrigger):
         poke_interval: float = 5.0,
         should_check_fn: bool = False,
         use_regex: bool = False,
+        region_name: str | None = None,
+        verify: bool | str | None = None,
+        botocore_config: dict | None = None,
         **hook_params: Any,
     ):
         super().__init__()
@@ -64,6 +67,9 @@ class S3KeyTrigger(BaseTrigger):
         self.poke_interval = poke_interval
         self.should_check_fn = should_check_fn
         self.use_regex = use_regex
+        self.region_name = region_name
+        self.verify = verify
+        self.botocore_config = botocore_config
 
     def serialize(self) -> tuple[str, dict[str, Any]]:
         """Serialize S3KeyTrigger arguments and classpath."""
@@ -78,17 +84,25 @@ class S3KeyTrigger(BaseTrigger):
                 "poke_interval": self.poke_interval,
                 "should_check_fn": self.should_check_fn,
                 "use_regex": self.use_regex,
+                "region_name": self.region_name,
+                "verify": self.verify,
+                "botocore_config": self.botocore_config,
             },
         )
 
     @cached_property
     def hook(self) -> S3Hook:
-        return S3Hook(aws_conn_id=self.aws_conn_id, verify=self.hook_params.get("verify"))
+        return S3Hook(
+            aws_conn_id=self.aws_conn_id,
+            region_name=self.region_name,
+            verify=self.verify,
+            config=self.botocore_config,
+        )
 
     async def run(self) -> AsyncIterator[TriggerEvent]:
         """Make an asynchronous connection using S3HookAsync."""
         try:
-            async with self.hook.async_conn as client:
+            async with await self.hook.get_async_conn() as client:
                 while True:
                     if await self.hook.check_key_async(
                         client, self.bucket_name, self.bucket_key, self.wildcard_match, self.use_regex
@@ -143,7 +157,9 @@ class S3KeysUnchangedTrigger(BaseTrigger):
         allow_delete: bool = True,
         aws_conn_id: str | None = "aws_default",
         last_activity_time: datetime | None = None,
+        region_name: str | None = None,
         verify: bool | str | None = None,
+        botocore_config: dict | None = None,
         **hook_params: Any,
     ):
         super().__init__()
@@ -160,8 +176,10 @@ class S3KeysUnchangedTrigger(BaseTrigger):
         self.allow_delete = allow_delete
         self.aws_conn_id = aws_conn_id
         self.last_activity_time = last_activity_time
-        self.verify = verify
         self.polling_period_seconds = 0
+        self.region_name = region_name
+        self.verify = verify
+        self.botocore_config = botocore_config
         self.hook_params = hook_params
 
     def serialize(self) -> tuple[str, dict[str, Any]]:
@@ -179,19 +197,26 @@ class S3KeysUnchangedTrigger(BaseTrigger):
                 "aws_conn_id": self.aws_conn_id,
                 "last_activity_time": self.last_activity_time,
                 "hook_params": self.hook_params,
-                "verify": self.verify,
                 "polling_period_seconds": self.polling_period_seconds,
+                "region_name": self.region_name,
+                "verify": self.verify,
+                "botocore_config": self.botocore_config,
             },
         )
 
     @cached_property
     def hook(self) -> S3Hook:
-        return S3Hook(aws_conn_id=self.aws_conn_id, verify=self.hook_params.get("verify"))
+        return S3Hook(
+            aws_conn_id=self.aws_conn_id,
+            region_name=self.region_name,
+            verify=self.verify,
+            config=self.botocore_config,
+        )
 
     async def run(self) -> AsyncIterator[TriggerEvent]:
         """Make an asynchronous connection using S3Hook."""
         try:
-            async with self.hook.async_conn as client:
+            async with await self.hook.get_async_conn() as client:
                 while True:
                     result = await self.hook.is_keys_unchanged_async(
                         client=client,

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

@@ -108,7 +108,7 @@ class SageMakerTrigger(BaseTrigger):
 
     async def run(self):
         self.log.info("job name is %s and job type is %s", self.job_name, self.job_type)
-        async with self.hook.async_conn as client:
+        async with await self.hook.get_async_conn() as client:
             waiter = self.hook.get_waiter(
                 self._get_job_type_waiter(self.job_type), deferrable=True, client=client
             )
@@ -166,7 +166,7 @@ class SageMakerPipelineTrigger(BaseTrigger):
 
     async def run(self) -> AsyncIterator[TriggerEvent]:
         hook = SageMakerHook(aws_conn_id=self.aws_conn_id)
-        async with hook.async_conn as conn:
+        async with await hook.get_async_conn() as conn:
             waiter = hook.get_waiter(self._waiter_name[self.waiter_type], deferrable=True, client=conn)
             for _ in range(self.waiter_max_attempts):
                 try:

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

@@ -0,0 +1,66 @@
+# 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 job trigger."""
+
+from __future__ import annotations
+
+from airflow.triggers.base import BaseTrigger
+
+
+class SageMakerNotebookJobTrigger(BaseTrigger):
+    """
+    Watches for a notebook job, triggers when it finishes.
+
+    Examples:
+     .. code-block:: python
+
+        from airflow.providers.amazon.aws.triggers.sagemaker_unified_studio import SageMakerNotebookJobTrigger
+
+        notebook_trigger = SageMakerNotebookJobTrigger(
+            execution_id="notebook_job_1234",
+            execution_name="notebook_task",
+            waiter_delay=10,
+            waiter_max_attempts=1440,
+        )
+
+    :param execution_id: A unique, meaningful id for the task.
+    :param execution_name: A unique, meaningful name for the task.
+    :param waiter_delay: Interval in seconds to check the notebook execution status.
+    :param waiter_max_attempts: Number of attempts to wait before returning FAILED.
+    """
+
+    def __init__(self, execution_id, execution_name, waiter_delay, waiter_max_attempts, **kwargs):
+        super().__init__(**kwargs)
+        self.execution_id = execution_id
+        self.execution_name = execution_name
+        self.waiter_delay = waiter_delay
+        self.waiter_max_attempts = waiter_max_attempts
+
+    def serialize(self):
+        return (
+            # dynamically generate the fully qualified name of the class
+            self.__class__.__module__ + "." + self.__class__.__qualname__,
+            {
+                "execution_id": self.execution_id,
+                "execution_name": self.execution_name,
+                "poll_interval": self.poll_interval,
+            },
+        )
+
+    async def run(self):
+        pass

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

@@ -23,14 +23,22 @@ from typing import TYPE_CHECKING, Any
 from airflow.exceptions import AirflowException
 from airflow.providers.amazon.aws.hooks.sqs import SqsHook
 from airflow.providers.amazon.aws.utils.sqs import process_response
-from airflow.triggers.base import BaseTrigger, TriggerEvent
+from airflow.providers.amazon.version_compat import AIRFLOW_V_3_0_PLUS
+
+if AIRFLOW_V_3_0_PLUS:
+    from airflow.triggers.base import BaseEventTrigger, TriggerEvent
+else:
+    from airflow.triggers.base import (  # type: ignore
+        BaseTrigger as BaseEventTrigger,
+        TriggerEvent,
+    )
 
 if TYPE_CHECKING:
     from airflow.providers.amazon.aws.hooks.base_aws import BaseAwsConnection
     from airflow.providers.amazon.aws.utils.sqs import MessageFilteringType
 
 
-class SqsSensorTrigger(BaseTrigger):
+class SqsSensorTrigger(BaseEventTrigger):
     """
     Asynchronously get messages from an Amazon SQS queue and then delete the messages from the queue.
 
@@ -176,7 +184,7 @@ class SqsSensorTrigger(BaseTrigger):
         while True:
             # This loop will run indefinitely until the timeout, which is set in the self.defer
             # method, is reached.
-            async with self.hook.async_conn as client:
+            async with await self.hook.get_async_conn() as client:
                 result = await self.poke(client=client)
                 if result:
                     yield TriggerEvent({"status": "success", "message_batch": result})

airflow/providers/amazon/aws/{auth_manager/security_manager/__init__.py → utils/sagemaker_unified_studio.py}

@@ -14,3 +14,15 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
+
+"""This module contains utils for the Amazon SageMaker Unified Studio Notebook plugin."""
+
+from __future__ import annotations
+
+import os
+
+workflows_env_key = "WORKFLOWS_ENV"
+
+
+def is_local_runner():
+    return os.getenv(workflows_env_key, "") == "Local"