apache-airflow-providers-amazon 9.11.0rc1__tar.gz → 9.12.0rc1__tar.gz

{apache_airflow_providers_amazon-9.11.0rc1 → apache_airflow_providers_amazon-9.12.0rc1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: apache-airflow-providers-amazon
-Version: 9.11.0rc1
+Version: 9.12.0rc1
 Summary: Provider package apache-airflow-providers-amazon for Apache Airflow
 Keywords: airflow-provider,amazon,airflow,integration
 Author-email: Apache Software Foundation <dev@airflow.apache.org>
@@ -38,7 +38,7 @@ Requires-Dist: marshmallow>=3
 Requires-Dist: aiobotocore[boto3]>=2.21.1 ; extra == "aiobotocore"
 Requires-Dist: apache-airflow-providers-apache-hive ; extra == "apache-hive"
 Requires-Dist: apache-airflow-providers-cncf-kubernetes>=7.2.0rc1 ; extra == "cncf-kubernetes"
-Requires-Dist: apache-airflow-providers-common-messaging>=1.0.1rc1 ; extra == "common-messaging"
+Requires-Dist: apache-airflow-providers-common-messaging>=1.0.3rc1 ; extra == "common-messaging"
 Requires-Dist: apache-airflow-providers-exasol ; extra == "exasol"
 Requires-Dist: apache-airflow-providers-fab>=2.2.0rc1 ; extra == "fab" and ( python_version < '3.13')
 Requires-Dist: apache-airflow-providers-ftp ; extra == "ftp"
@@ -55,8 +55,8 @@ Requires-Dist: apache-airflow-providers-salesforce ; extra == "salesforce"
 Requires-Dist: apache-airflow-providers-ssh ; extra == "ssh"
 Requires-Dist: apache-airflow-providers-standard ; extra == "standard"
 Project-URL: Bug Tracker, https://github.com/apache/airflow/issues
-Project-URL: Changelog, https://airflow.staged.apache.org/docs/apache-airflow-providers-amazon/9.11.0/changelog.html
-Project-URL: Documentation, https://airflow.staged.apache.org/docs/apache-airflow-providers-amazon/9.11.0
+Project-URL: Changelog, https://airflow.staged.apache.org/docs/apache-airflow-providers-amazon/9.12.0/changelog.html
+Project-URL: Documentation, https://airflow.staged.apache.org/docs/apache-airflow-providers-amazon/9.12.0
 Project-URL: Mastodon, https://fosstodon.org/@airflow
 Project-URL: Slack Chat, https://s.apache.org/airflow-slack
 Project-URL: Source Code, https://github.com/apache/airflow
@@ -104,7 +104,7 @@ Provides-Extra: standard
 
 Package ``apache-airflow-providers-amazon``
 
-Release: ``9.11.0``
+Release: ``9.12.0``
 
 Release Date: ``|PypiReleaseDate|``
 
@@ -118,7 +118,7 @@ This is a provider package for ``amazon`` provider. All classes for this provide
 are in ``airflow.providers.amazon`` python package.
 
 You can find package information and changelog for the provider
-in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-amazon/9.11.0/>`_.
+in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-amazon/9.12.0/>`_.
 
 Installation
 ------------
@@ -186,5 +186,5 @@ Dependent package
 ========================================================================================================================  ====================
 
 The changelog for the provider package can be found in the
-`changelog <https://airflow.apache.org/docs/apache-airflow-providers-amazon/9.11.0/changelog.html>`_.
+`changelog <https://airflow.apache.org/docs/apache-airflow-providers-amazon/9.12.0/changelog.html>`_.
 

{apache_airflow_providers_amazon-9.11.0rc1 → apache_airflow_providers_amazon-9.12.0rc1}/README.rst

@@ -23,7 +23,7 @@
 
 Package ``apache-airflow-providers-amazon``
 
-Release: ``9.11.0``
+Release: ``9.12.0``
 
 Release Date: ``|PypiReleaseDate|``
 
@@ -37,7 +37,7 @@ This is a provider package for ``amazon`` provider. All classes for this provide
 are in ``airflow.providers.amazon`` python package.
 
 You can find package information and changelog for the provider
-in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-amazon/9.11.0/>`_.
+in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-amazon/9.12.0/>`_.
 
 Installation
 ------------
@@ -105,4 +105,4 @@ Dependent package
 ========================================================================================================================  ====================
 
 The changelog for the provider package can be found in the
-`changelog <https://airflow.apache.org/docs/apache-airflow-providers-amazon/9.11.0/changelog.html>`_.
+`changelog <https://airflow.apache.org/docs/apache-airflow-providers-amazon/9.12.0/changelog.html>`_.

{apache_airflow_providers_amazon-9.11.0rc1 → apache_airflow_providers_amazon-9.12.0rc1}/docs/changelog.rst

@@ -26,6 +26,41 @@
 Changelog
 ---------
 
+9.12.0
+......
+
+Features
+~~~~~~~~
+
+* ``Add MwaaTaskSensor to Amazon Provider Package (#51719)``
+* ``Support executor_config on Lambda Executor (#53994)``
+
+Bug Fixes
+~~~~~~~~~
+
+* ``Fix AWS Lambda executor error handling for DLQ vs task failures (#53990)``
+
+Misc
+~~~~
+
+* ``Set minimum version for common.messaging to 1.0.3 (#54160)``
+* ``Use timezone from new TaskSDK public API where possible (#53949)``
+* ``Use timezone from new TaskSDK public API where possible (second part) (#53986)``
+* ``Refactor bundle view_url to not instantiate bundle on server components (#52876)``
+
+Doc-only
+~~~~~~~~
+
+* ``Add missing PR number in amazon changelog (#53880)``
+
+.. Below changes are excluded from the changelog. Move them to
+   appropriate section above if needed. Do not delete the lines(!):
+   * ``SageMaker Unified Studios System Tests Update (#54038)``
+   * ``Do not run export operations in 'example_dynamodb_to_s3' (#54158)``
+   * ``Add system test for Lambda executor Dead Letter Queue (DLQ) processing (#54042)``
+   * ``Fix 'importskip' statements in tests (#54135)``
+   * ``Increase timeout to delete tables in 'example_s3_to_dynamodb' (#54096)``
+
 9.11.0
 ......
 
@@ -68,7 +103,7 @@ Misc
    * ``Handle ruff PT028 changes (#53235)``
    * ``Create connection with API instead of directly through Session (#53161)``
    * ``Make dag_version_id in TI non-nullable (#50825)``
-   * ``Changing import path to 'airflow.sdk' in Amazon provider package``
+   * ``Changing import path to 'airflow.sdk' in Amazon provider package (#50659)``
 
 9.10.0
 ......

{apache_airflow_providers_amazon-9.11.0rc1 → apache_airflow_providers_amazon-9.12.0rc1}/docs/executors/batch-executor.rst

@@ -69,7 +69,7 @@ Options <https://airflow.apache.org/docs/apache-airflow/stable/howto/set-config.
    :end-before: .. END CONFIG_OPTIONS_PRECEDENCE
 
 .. note::
-   ``exec_config`` is an optional parameter that can be provided to operators. It is a dictionary type and in the context of the Batch Executor it represents a ``submit_job_kwargs`` configuration which is then updated over-top of the ``submit_job_kwargs`` specified in Airflow config above (if present). It is a recursive update which essentially applies Python update to each nested dictionary in the configuration. Loosely approximated as: ``submit_job_kwargs.update(exec_config)``
+   ``executor_config`` is an optional parameter that can be provided to operators. It is a dictionary type and in the context of the Batch Executor it represents a ``submit_job_kwargs`` configuration which is then updated over-top of the ``submit_job_kwargs`` specified in Airflow config above (if present). It is a recursive update which essentially applies Python update to each nested dictionary in the configuration. Loosely approximated as: ``submit_job_kwargs.update(executor_config)``
 
 Required config options:
 ~~~~~~~~~~~~~~~~~~~~~~~~
@@ -98,7 +98,7 @@ hints and examples, see the ``config_templates`` folder in the Amazon
 provider package.
 
 .. note::
-   ``exec_config`` is an optional parameter that can be provided to operators. It is a dictionary type and in the context of the Batch Executor it represents a ``submit_job_kwargs`` configuration which is then updated over-top of the ``submit_job_kwargs`` specified in Airflow config above (if present). It is a recursive update which essentially applies Python update to each nested dictionary in the configuration. Loosely approximated as: ``submit_job_kwargs.update(exec_config)``
+   ``executor_config`` is an optional parameter that can be provided to operators. It is a dictionary type and in the context of the Batch Executor it represents a ``submit_job_kwargs`` configuration which is then updated over-top of the ``submit_job_kwargs`` specified in Airflow config above (if present). It is a recursive update which essentially applies Python update to each nested dictionary in the configuration. Loosely approximated as: ``submit_job_kwargs.update(executor_config)``
 
 .. _dockerfile_for_batch_executor:
 

{apache_airflow_providers_amazon-9.11.0rc1 → apache_airflow_providers_amazon-9.12.0rc1}/docs/executors/ecs-executor.rst

@@ -74,7 +74,7 @@ In the case of conflicts, the order of precedence from lowest to highest is:
    provided.
 
 .. note::
-   ``exec_config`` is an optional parameter that can be provided to operators. It is a dictionary type and in the context of the ECS Executor it represents a ``run_task_kwargs`` configuration which is then updated over-top of the ``run_task_kwargs`` specified in Airflow config above (if present). It is a recursive update which essentially applies Python update to each nested dictionary in the configuration. Loosely approximated as: ``run_task_kwargs.update(exec_config)``
+   ``executor_config`` is an optional parameter that can be provided to operators. It is a dictionary type and in the context of the ECS Executor it represents a ``run_task_kwargs`` configuration which is then updated over-top of the ``run_task_kwargs`` specified in Airflow config above (if present). It is a recursive update which essentially applies Python update to each nested dictionary in the configuration. Loosely approximated as: ``run_task_kwargs.update(executor_config)``
 
 Required config options:
 ~~~~~~~~~~~~~~~~~~~~~~~~

{apache_airflow_providers_amazon-9.11.0rc1 → apache_airflow_providers_amazon-9.12.0rc1}/docs/executors/lambda-executor.rst

@@ -88,6 +88,9 @@ In the case of conflicts, the order of precedence from lowest to highest is:
    environment variables. These are checked with Airflow's config
    precedence.
 
+.. note::
+   ``executor_config`` is an optional parameter that can be provided to operators. It is a dictionary type and in the context of the Lambda Executor it is passed to each Lambda invocation as part of the payload. This allows you to pass additional context to the Lambda function for any particular task execution. The Lambda function can then access this configuration via the ``executor_config`` key in the payload within the Lambda handler code.
+
 Required config options:
 ~~~~~~~~~~~~~~~~~~~~~~~~
 

{apache_airflow_providers_amazon-9.11.0rc1 → apache_airflow_providers_amazon-9.12.0rc1}/docs/index.rst

@@ -87,7 +87,7 @@ apache-airflow-providers-amazon package
 Amazon integration (including `Amazon Web Services (AWS) <https://aws.amazon.com/>`__).
 
 
-Release: 9.11.0
+Release: 9.12.0
 
 Release Date: ``|PypiReleaseDate|``
 
@@ -168,5 +168,5 @@ Downloading official packages
 You can download officially released packages and verify their checksums and signatures from the
 `Official Apache Download site <https://downloads.apache.org/airflow/providers/>`_
 
-* `The apache-airflow-providers-amazon 9.11.0 sdist package <https://downloads.apache.org/airflow/providers/apache_airflow_providers_amazon-9.11.0.tar.gz>`_ (`asc <https://downloads.apache.org/airflow/providers/apache_airflow_providers_amazon-9.11.0.tar.gz.asc>`__, `sha512 <https://downloads.apache.org/airflow/providers/apache_airflow_providers_amazon-9.11.0.tar.gz.sha512>`__)
-* `The apache-airflow-providers-amazon 9.11.0 wheel package <https://downloads.apache.org/airflow/providers/apache_airflow_providers_amazon-9.11.0-py3-none-any.whl>`_ (`asc <https://downloads.apache.org/airflow/providers/apache_airflow_providers_amazon-9.11.0-py3-none-any.whl.asc>`__, `sha512 <https://downloads.apache.org/airflow/providers/apache_airflow_providers_amazon-9.11.0-py3-none-any.whl.sha512>`__)
+* `The apache-airflow-providers-amazon 9.12.0 sdist package <https://downloads.apache.org/airflow/providers/apache_airflow_providers_amazon-9.12.0.tar.gz>`_ (`asc <https://downloads.apache.org/airflow/providers/apache_airflow_providers_amazon-9.12.0.tar.gz.asc>`__, `sha512 <https://downloads.apache.org/airflow/providers/apache_airflow_providers_amazon-9.12.0.tar.gz.sha512>`__)
+* `The apache-airflow-providers-amazon 9.12.0 wheel package <https://downloads.apache.org/airflow/providers/apache_airflow_providers_amazon-9.12.0-py3-none-any.whl>`_ (`asc <https://downloads.apache.org/airflow/providers/apache_airflow_providers_amazon-9.12.0-py3-none-any.whl.asc>`__, `sha512 <https://downloads.apache.org/airflow/providers/apache_airflow_providers_amazon-9.12.0-py3-none-any.whl.sha512>`__)

{apache_airflow_providers_amazon-9.11.0rc1 → apache_airflow_providers_amazon-9.12.0rc1}/docs/operators/mwaa.rst

@@ -76,6 +76,22 @@ In the following example, the task ``wait_for_dag_run`` waits for the DAG run cr
     :start-after: [START howto_sensor_mwaa_dag_run]
     :end-before: [END howto_sensor_mwaa_dag_run]
 
+.. _howto/sensor:MwaaTaskSensor:
+
+Wait on the state of an AWS MWAA Task
+========================================
+
+To wait for a DAG task instance across MWAA environments until it reaches one of the given states, you can use the
+:class:`~airflow.providers.amazon.aws.sensors.mwaa.MwaaTaskSensor`
+
+In the following example, the task ``wait_for_task`` waits for the DAG run created in the above task to complete.
+
+.. exampleinclude:: /../../amazon/tests/system/amazon/aws/example_mwaa.py
+    :language: python
+    :dedent: 4
+    :start-after: [START howto_sensor_mwaa_task]
+    :end-before: [END howto_sensor_mwaa_task]
+
 References
 ----------
 

{apache_airflow_providers_amazon-9.11.0rc1 → apache_airflow_providers_amazon-9.12.0rc1}/provider.yaml

@@ -22,12 +22,13 @@ description: |
   Amazon integration (including `Amazon Web Services (AWS) <https://aws.amazon.com/>`__).
 
 state: ready
-source-date-epoch: 1753688567
+source-date-epoch: 1754502801
 # Note that those versions are maintained by release manager - do not update them manually
 # with the exception of case where other provider in sources has >= new provider version.
 # In such case adding >= NEW_VERSION and bumping to NEW_VERSION in a provider have
 # to be done in the same PR
 versions:
+  - 9.12.0
   - 9.11.0
   - 9.10.0
   - 9.9.0

{apache_airflow_providers_amazon-9.11.0rc1 → apache_airflow_providers_amazon-9.12.0rc1}/pyproject.toml

@@ -25,7 +25,7 @@ build-backend = "flit_core.buildapi"
 
 [project]
 name = "apache-airflow-providers-amazon"
-version = "9.11.0rc1"
+version = "9.12.0rc1"
 description = "Provider package apache-airflow-providers-amazon for Apache Airflow"
 readme = "README.rst"
 authors = [
@@ -140,7 +140,7 @@ dependencies = [
     "apache-airflow-providers-standard"
 ]
 "common.messaging" = [
-    "apache-airflow-providers-common-messaging>=1.0.1rc1"
+    "apache-airflow-providers-common-messaging>=1.0.3rc1"
 ]
 
 [dependency-groups]
@@ -206,8 +206,8 @@ apache-airflow-providers-common-sql = {workspace = true}
 apache-airflow-providers-standard = {workspace = true}
 
 [project.urls]
-"Documentation" = "https://airflow.staged.apache.org/docs/apache-airflow-providers-amazon/9.11.0"
-"Changelog" = "https://airflow.staged.apache.org/docs/apache-airflow-providers-amazon/9.11.0/changelog.html"
+"Documentation" = "https://airflow.staged.apache.org/docs/apache-airflow-providers-amazon/9.12.0"
+"Changelog" = "https://airflow.staged.apache.org/docs/apache-airflow-providers-amazon/9.12.0/changelog.html"
 "Bug Tracker" = "https://github.com/apache/airflow/issues"
 "Source Code" = "https://github.com/apache/airflow"
 "Slack Chat" = "https://s.apache.org/airflow-slack"

{apache_airflow_providers_amazon-9.11.0rc1 → apache_airflow_providers_amazon-9.12.0rc1}/src/airflow/providers/amazon/__init__.py

@@ -29,7 +29,7 @@ from airflow import __version__ as airflow_version
 
 __all__ = ["__version__"]
 
-__version__ = "9.11.0"
+__version__ = "9.12.0"
 
 if packaging.version.parse(packaging.version.parse(airflow_version).base_version) < packaging.version.parse(
     "2.10.0"

{apache_airflow_providers_amazon-9.11.0rc1 → apache_airflow_providers_amazon-9.12.0rc1}/src/airflow/providers/amazon/aws/bundles/s3.py

@@ -137,10 +137,23 @@ class S3DagBundle(BaseDagBundle):
             )
 
     def view_url(self, version: str | None = None) -> str | None:
+        """
+        Return a URL for viewing the DAGs in S3. Currently, versioning is not supported.
+
+        This method is deprecated and will be removed when the minimum supported Airflow version is 3.1.
+        Use `view_url_template` instead.
+        """
+        return self.view_url_template()
+
+    def view_url_template(self) -> str | None:
         """Return a URL for viewing the DAGs in S3. Currently, versioning is not supported."""
         if self.version:
             raise AirflowException("S3 url with version is not supported")
-
+        if hasattr(self, "_view_url_template") and self._view_url_template:
+            # Because we use this method in the view_url method, we need to handle
+            # backward compatibility for Airflow versions that doesn't have the
+            # _view_url_template attribute. Should be removed when we drop support for Airflow 3.0
+            return self._view_url_template
         # https://<bucket-name>.s3.<region>.amazonaws.com/<object-key>
         url = f"https://{self.bucket_name}.s3"
         if self.s3_hook.region_name:

{apache_airflow_providers_amazon-9.11.0rc1 → apache_airflow_providers_amazon-9.12.0rc1}/src/airflow/providers/amazon/aws/executors/aws_lambda/docker/app.py

@@ -38,6 +38,7 @@ S3_URI = os.environ.get("S3_URI", None)
 # Input and output keys
 TASK_KEY_KEY = "task_key"
 COMMAND_KEY = "command"
+EXECUTOR_CONFIG_KEY = "executor_config"
 RETURN_CODE_KEY = "return_code"
 
 
@@ -47,8 +48,9 @@ def lambda_handler(event, context):
 
     command = event.get(COMMAND_KEY)
     task_key = event.get(TASK_KEY_KEY)
+    executor_config = event.get(EXECUTOR_CONFIG_KEY, {})  # noqa: F841
 
-    # Any pre-processing or validation of the command or use of the context can be done here or above.
+    # Any pre-processing or validation of the command or use of the executor_config can be done here or above.
 
     # Sync dags from s3 to the local dags directory
     if S3_URI:

{apache_airflow_providers_amazon-9.11.0rc1 → apache_airflow_providers_amazon-9.12.0rc1}/src/airflow/providers/amazon/aws/executors/aws_lambda/lambda_executor.py

@@ -43,7 +43,11 @@ from airflow.providers.amazon.aws.executors.utils.exponential_backoff_retry impo
 from airflow.providers.amazon.aws.hooks.lambda_function import LambdaHook
 from airflow.providers.amazon.aws.hooks.sqs import SqsHook
 from airflow.stats import Stats
-from airflow.utils import timezone
+
+try:
+    from airflow.sdk import timezone
+except ImportError:
+    from airflow.utils import timezone  # type: ignore[attr-defined,no-redef]
 
 from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS
 
@@ -270,6 +274,7 @@ class AwsLambdaExecutor(BaseExecutor):
             payload = {
                 "task_key": ser_task_key,
                 "command": cmd,
+                "executor_config": task_to_run.executor_config,
             }
             if timezone.utcnow() < task_to_run.next_attempt_time:
                 self.pending_tasks.append(task_to_run)
@@ -421,17 +426,26 @@ class AwsLambdaExecutor(BaseExecutor):
                         "Successful Lambda invocation for task %s received from SQS queue.", task_key
                     )
                 else:
-                    # In this case the Lambda likely started but failed at run time since we got a non-zero
-                    # return code. We could consider retrying these tasks within the executor, because this _likely_
-                    # means the Airflow task did not run to completion, however we can't be sure (maybe the
-                    # lambda runtime code has a bug and is returning a non-zero when it actually passed?). So
-                    # perhaps not retrying is the safest option.
                     self.fail(task_key)
-                    self.log.error(
-                        "Lambda invocation for task: %s has failed to run with return code %s",
-                        task_key,
-                        return_code,
-                    )
+                    if queue_url == self.dlq_url and return_code is None:
+                        # DLQ failure: AWS Lambda service could not complete the invocation after retries.
+                        # This indicates a Lambda-level failure (timeout, memory limit, crash, etc.)
+                        # where the function was unable to successfully execute to return a result.
+                        self.log.error(
+                            "DLQ message received: Lambda invocation for task: %s was unable to successfully execute. This likely indicates a Lambda-level failure (timeout, memory limit, crash, etc.).",
+                            task_key,
+                        )
+                    else:
+                        # In this case the Lambda likely started but failed at run time since we got a non-zero
+                        # return code. We could consider retrying these tasks within the executor, because this _likely_
+                        # means the Airflow task did not run to completion, however we can't be sure (maybe the
+                        # lambda runtime code has a bug and is returning a non-zero when it actually passed?). So
+                        # perhaps not retrying is the safest option.
+                        self.log.debug(
+                            "Lambda invocation for task: %s completed but the underlying Airflow task has returned a non-zero exit code %s",
+                            task_key,
+                            return_code,
+                        )
                 # Remove the task from the tracking mapping.
                 self.running_tasks.pop(ser_task_key)
 

{apache_airflow_providers_amazon-9.11.0rc1 → apache_airflow_providers_amazon-9.12.0rc1}/src/airflow/providers/amazon/aws/executors/batch/batch_executor.py

@@ -38,7 +38,11 @@ from airflow.providers.amazon.aws.executors.utils.exponential_backoff_retry impo
 from airflow.providers.amazon.aws.hooks.batch_client import BatchClientHook
 from airflow.providers.amazon.version_compat import AIRFLOW_V_3_0_PLUS
 from airflow.stats import Stats
-from airflow.utils import timezone
+
+try:
+    from airflow.sdk import timezone
+except ImportError:
+    from airflow.utils import timezone  # type: ignore[attr-defined,no-redef]
 from airflow.utils.helpers import merge_dicts
 
 if TYPE_CHECKING:

{apache_airflow_providers_amazon-9.11.0rc1 → apache_airflow_providers_amazon-9.12.0rc1}/src/airflow/providers/amazon/aws/executors/ecs/ecs_executor.py

@@ -51,7 +51,11 @@ from airflow.providers.amazon.aws.executors.utils.exponential_backoff_retry impo
 from airflow.providers.amazon.aws.hooks.ecs import EcsHook
 from airflow.providers.amazon.version_compat import AIRFLOW_V_3_0_PLUS
 from airflow.stats import Stats
-from airflow.utils import timezone
+
+try:
+    from airflow.sdk import timezone
+except ImportError:
+    from airflow.utils import timezone  # type: ignore[attr-defined,no-redef]
 from airflow.utils.helpers import merge_dicts
 from airflow.utils.state import State
 

{apache_airflow_providers_amazon-9.11.0rc1 → apache_airflow_providers_amazon-9.12.0rc1}/src/airflow/providers/amazon/aws/executors/utils/exponential_backoff_retry.py

@@ -20,7 +20,10 @@ import logging
 from collections.abc import Callable
 from datetime import datetime, timedelta
 
-from airflow.utils import timezone
+try:
+    from airflow.sdk import timezone
+except ImportError:
+    from airflow.utils import timezone  # type: ignore[attr-defined,no-redef]
 
 log = logging.getLogger(__name__)
 

{apache_airflow_providers_amazon-9.11.0rc1 → apache_airflow_providers_amazon-9.12.0rc1}/src/airflow/providers/amazon/aws/hooks/sagemaker.py

@@ -36,7 +36,11 @@ from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
 from airflow.providers.amazon.aws.hooks.logs import AwsLogsHook
 from airflow.providers.amazon.aws.hooks.s3 import S3Hook
 from airflow.providers.amazon.aws.utils.tags import format_tags
-from airflow.utils import timezone
+
+try:
+    from airflow.sdk import timezone
+except ImportError:
+    from airflow.utils import timezone  # type: ignore[attr-defined,no-redef]
 
 
 class LogState:

{apache_airflow_providers_amazon-9.11.0rc1 → apache_airflow_providers_amazon-9.12.0rc1}/src/airflow/providers/amazon/aws/sensors/mwaa.py

@@ -24,10 +24,9 @@ from airflow.configuration import conf
 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.triggers.mwaa import MwaaDagRunCompletedTrigger
-from airflow.providers.amazon.aws.utils import validate_execute_complete_event
+from airflow.providers.amazon.aws.triggers.mwaa import MwaaDagRunCompletedTrigger, MwaaTaskCompletedTrigger
 from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
-from airflow.utils.state import DagRunState
+from airflow.utils.state import DagRunState, TaskInstanceState
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context
@@ -139,7 +138,7 @@ class MwaaDagRunSensor(AwsBaseSensor[MwaaHook]):
         return state in self.success_states
 
     def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> None:
-        validate_execute_complete_event(event)
+        return None
 
     def execute(self, context: Context):
         if self.deferrable:
@@ -150,10 +149,152 @@ class MwaaDagRunSensor(AwsBaseSensor[MwaaHook]):
                     external_dag_run_id=self.external_dag_run_id,
                     success_states=self.success_states,
                     failure_states=self.failure_states,
-                    # somehow the type of poke_interval is derived as float ??
-                    waiter_delay=self.poke_interval,  # type: ignore[arg-type]
+                    waiter_delay=int(self.poke_interval),
                     waiter_max_attempts=self.max_retries,
                     aws_conn_id=self.aws_conn_id,
+                    end_from_trigger=True,
+                ),
+                method_name="execute_complete",
+            )
+        else:
+            super().execute(context=context)
+
+
+class MwaaTaskSensor(AwsBaseSensor[MwaaHook]):
+    """
+    Waits for a task in an MWAA Environment to complete.
+
+    If the task fails, an AirflowException is thrown.
+
+    .. seealso::
+        For more information on how to use this sensor, take a look at the guide:
+        :ref:`howto/sensor:MwaaTaskSensor`
+
+    :param external_env_name: The external MWAA environment name that contains the Task Instance you want to wait for
+        (templated)
+    :param external_dag_id: The DAG ID in the external MWAA environment that contains the Task Instance 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 external_task_id: The Task ID in the external MWAA environment that you want to wait for (templated)
+    :param success_states: Collection of task instance states that would make this task marked as successful, default is
+        ``{airflow.utils.state.TaskInstanceState.SUCCESS}`` (templated)
+    :param failure_states: Collection of task instance states that would make this task marked as failed and raise an
+        AirflowException, default is ``{airflow.utils.state.TaskInstanceState.FAILED}`` (templated)
+    :param deferrable: If True, the sensor will operate in deferrable mode. This mode requires aiobotocore
+        module to be installed.
+        (default: False, but can be overridden in config file by setting default_deferrable to True)
+    :param poke_interval: Polling period in seconds to check for the status of the job. (default: 60)
+    :param max_retries: Number of times before returning the current state. (default: 720)
+    :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
+    :param botocore_config: Configuration dictionary (key-values) for botocore client. See:
+        https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
+    """
+
+    aws_hook_class = MwaaHook
+    template_fields: Sequence[str] = aws_template_fields(
+        "external_env_name",
+        "external_dag_id",
+        "external_dag_run_id",
+        "external_task_id",
+        "success_states",
+        "failure_states",
+        "deferrable",
+        "max_retries",
+        "poke_interval",
+    )
+
+    def __init__(
+        self,
+        *,
+        external_env_name: str,
+        external_dag_id: str,
+        external_dag_run_id: str | None = None,
+        external_task_id: str,
+        success_states: Collection[str] | None = None,
+        failure_states: Collection[str] | None = None,
+        deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+        poke_interval: int = 60,
+        max_retries: int = 720,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        self.success_states = set(success_states) if success_states else {TaskInstanceState.SUCCESS.value}
+        self.failure_states = set(failure_states) if failure_states else {TaskInstanceState.FAILED.value}
+
+        if len(self.success_states & self.failure_states):
+            raise ValueError("success_states and failure_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
+        self.external_task_id = external_task_id
+        self.deferrable = deferrable
+        self.poke_interval = poke_interval
+        self.max_retries = max_retries
+
+    def poke(self, context: Context) -> bool:
+        self.log.info(
+            "Poking for task %s of DAG run %s of DAG %s in MWAA environment %s",
+            self.external_task_id,
+            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}/taskInstances/{self.external_task_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 task
+
+        state = response["RestApiResponse"]["state"]
+
+        if state in self.failure_states:
+            raise AirflowException(
+                f"The task {self.external_task_id} of 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 state in self.success_states
+
+    def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> None:
+        return None
+
+    def execute(self, context: Context):
+        if self.external_dag_run_id is None:
+            response = self.hook.invoke_rest_api(
+                env_name=self.external_env_name,
+                path=f"/dags/{self.external_dag_id}/dagRuns",
+                method="GET",
+            )
+            self.external_dag_run_id = response["RestApiResponse"]["dag_runs"][-1]["dag_run_id"]
+
+        if self.deferrable:
+            self.defer(
+                trigger=MwaaTaskCompletedTrigger(
+                    external_env_name=self.external_env_name,
+                    external_dag_id=self.external_dag_id,
+                    external_dag_run_id=self.external_dag_run_id,
+                    external_task_id=self.external_task_id,
+                    success_states=self.success_states,
+                    failure_states=self.failure_states,
+                    waiter_delay=int(self.poke_interval),
+                    waiter_max_attempts=self.max_retries,
+                    aws_conn_id=self.aws_conn_id,
+                    end_from_trigger=True,
                 ),
                 method_name="execute_complete",
             )

{apache_airflow_providers_amazon-9.11.0rc1 → apache_airflow_providers_amazon-9.12.0rc1}/src/airflow/providers/amazon/aws/triggers/base.py

@@ -80,7 +80,7 @@ class AwsBaseWaiterTrigger(BaseTrigger):
         waiter_delay: int,
         waiter_max_attempts: int,
         waiter_config_overrides: dict[str, Any] | None = None,
-        aws_conn_id: str | None,
+        aws_conn_id: str | None = "aws_default",
         region_name: str | None = None,
         verify: bool | str | None = None,
         botocore_config: dict | None = None,