Add MwaaServerlessWorkflowRunSensor

providers/amazon/docs/operators/mwaa_serverless.rst

@@ -54,3 +54,17 @@ Reference
 ~~~~~~~~~
 
 * `AWS boto3 Library Documentation for MWAA Serverless <https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/mwaa-serverless.html>`__
+
+.. _howto/sensor:MwaaServerlessWorkflowRunSensor:
+
+Wait for a Workflow Run
+-----------------------
+
+To wait for an Amazon MWAA Serverless workflow run to complete, use
+:class:`~airflow.providers.amazon.aws.sensors.mwaa_serverless.MwaaServerlessWorkflowRunSensor`.
+
+.. exampleinclude:: /../../amazon/tests/system/amazon/aws/example_mwaa_serverless.py
+    :language: python
+    :dedent: 4
+    :start-after: [START howto_sensor_mwaa_serverless_workflow_run]
+    :end-before: [END howto_sensor_mwaa_serverless_workflow_run]

providers/amazon/provider.yaml

@@ -579,6 +579,9 @@ sensors:
   - integration-name: Amazon Managed Workflows for Apache Airflow (MWAA)
     python-modules:
       - airflow.providers.amazon.aws.sensors.mwaa
+  - integration-name: Amazon MWAA Serverless
+    python-modules:
+      - airflow.providers.amazon.aws.sensors.mwaa_serverless
   - integration-name: Amazon OpenSearch Serverless
     python-modules:
       - airflow.providers.amazon.aws.sensors.opensearch_serverless

providers/amazon/src/airflow/providers/amazon/aws/sensors/mwaa_serverless.py

@@ -0,0 +1,77 @@
+# 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.
+"""Amazon MWAA Serverless sensors."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
+from airflow.providers.amazon.aws.sensors.base_aws import AwsBaseSensor
+from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
+
+if TYPE_CHECKING:
+    from airflow.sdk import Context
+
+
+class MwaaServerlessWorkflowRunSensor(AwsBaseSensor[AwsBaseHook]):
+    """
+    Wait for an Amazon MWAA Serverless workflow run to reach a terminal state.
+
+    .. seealso::
+        For more information on how to use this sensor, take a look at the guide:
+        :ref:`howto/sensor:MwaaServerlessWorkflowRunSensor`
+
+    :param workflow_arn: The ARN of the workflow. (templated)
+    :param run_id: The ID of the workflow run to monitor. (templated)
+    :param success_states: Set of states considered successful. Default: ``{"SUCCESS"}``.
+    :param failure_states: Set of states that raise an exception.
+        Default: ``{"FAILED", "TIMEOUT", "STOPPED"}``.
+    """
+
+    aws_hook_class = AwsBaseHook
+    template_fields: tuple[str, ...] = aws_template_fields("workflow_arn", "run_id")
+
+    def __init__(
+        self,
+        *,
+        workflow_arn: str,
+        run_id: str,
+        success_states: set[str] | None = None,
+        failure_states: set[str] | None = None,
+        **kwargs,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.workflow_arn = workflow_arn
+        self.run_id = run_id
+        self.success_states = success_states or {"SUCCESS"}
+        self.failure_states = failure_states or {"FAILED", "TIMEOUT", "STOPPED"}
+
+    @property
+    def _hook_parameters(self) -> dict[str, Any]:
+        return {**super()._hook_parameters, "client_type": "mwaa-serverless"}
+
+    def poke(self, context: Context) -> bool:
+        response = self.hook.conn.get_workflow_run(WorkflowArn=self.workflow_arn, RunId=self.run_id)
+        state = response["RunDetail"]["RunState"]
+        self.log.info("Workflow run %s state: %s", self.run_id, state)
+
+        if state in self.failure_states:
+            error_msg = response["RunDetail"].get("ErrorMessage", "")
+            raise RuntimeError(f"Workflow run {self.run_id} failed with state {state}: {error_msg}")
+
+        return state in self.success_states

providers/amazon/src/airflow/providers/amazon/get_provider_info.py

@@ -613,6 +613,10 @@ def get_provider_info():
                 "integration-name": "Amazon Managed Workflows for Apache Airflow (MWAA)",
                 "python-modules": ["airflow.providers.amazon.aws.sensors.mwaa"],
             },
+            {
+                "integration-name": "Amazon MWAA Serverless",
+                "python-modules": ["airflow.providers.amazon.aws.sensors.mwaa_serverless"],
+            },
             {
                 "integration-name": "Amazon OpenSearch Serverless",
                 "python-modules": ["airflow.providers.amazon.aws.sensors.opensearch_serverless"],

providers/amazon/tests/system/amazon/aws/example_mwaa_serverless.py

@@ -27,6 +27,7 @@
     S3CreateObjectOperator,
     S3DeleteBucketOperator,
 )
+from airflow.providers.amazon.aws.sensors.mwaa_serverless import MwaaServerlessWorkflowRunSensor
 from airflow.providers.common.compat.sdk import DAG, chain
 
 from system.amazon.aws.utils import ENV_ID_KEY, SystemTestContextBuilder
@@ -113,6 +114,16 @@ def delete_workflow(workflow_arn: str):
     )
     # [END howto_operator_mwaa_serverless_start_workflow_run]
 
+    # [START howto_sensor_mwaa_serverless_workflow_run]
+    wait_for_run = MwaaServerlessWorkflowRunSensor(
+        task_id="wait_for_run",
+        workflow_arn=workflow_arn,
+        run_id=start_workflow.output,
+        poke_interval=30,
+        timeout=600,
+    )
+    # [END howto_sensor_mwaa_serverless_workflow_run]
+
     delete_bucket = S3DeleteBucketOperator(
         task_id="delete_bucket",
         bucket_name=bucket_name,
@@ -126,6 +137,7 @@ def delete_workflow(workflow_arn: str):
         upload_workflow_yaml,
         workflow_arn,
         start_workflow,
+        wait_for_run,
         stop_workflow_run(workflow_arn=workflow_arn, run_id=start_workflow.output),
         delete_workflow(workflow_arn=workflow_arn),
         delete_bucket,

providers/amazon/tests/unit/amazon/aws/sensors/test_mwaa_serverless.py

@@ -0,0 +1,84 @@
+# 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 unittest import mock
+
+import pytest
+
+from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
+from airflow.providers.amazon.aws.sensors.mwaa_serverless import MwaaServerlessWorkflowRunSensor
+
+from unit.amazon.aws.utils.test_template_fields import validate_template_fields
+
+WORKFLOW_ARN = "arn:aws:mwaa-serverless:us-east-1:123456789012:workflow/test"
+RUN_ID = "run-abc123"
+
+
+class TestMwaaServerlessWorkflowRunSensor:
+    def setup_method(self):
+        self.sensor = MwaaServerlessWorkflowRunSensor(
+            task_id="wait_for_run",
+            workflow_arn=WORKFLOW_ARN,
+            run_id=RUN_ID,
+            poke_interval=5,
+        )
+
+    @mock.patch.object(AwsBaseHook, "conn", new_callable=mock.PropertyMock)
+    def test_poke_success(self, mock_conn):
+        mock_client = mock.MagicMock()
+        mock_client.get_workflow_run.return_value = {"RunDetail": {"RunState": "SUCCESS", "ErrorMessage": ""}}
+        mock_conn.return_value = mock_client
+
+        assert self.sensor.poke({}) is True
+
+    @mock.patch.object(AwsBaseHook, "conn", new_callable=mock.PropertyMock)
+    def test_poke_running(self, mock_conn):
+        mock_client = mock.MagicMock()
+        mock_client.get_workflow_run.return_value = {"RunDetail": {"RunState": "RUNNING", "ErrorMessage": ""}}
+        mock_conn.return_value = mock_client
+
+        assert self.sensor.poke({}) is False
+
+    @mock.patch.object(AwsBaseHook, "conn", new_callable=mock.PropertyMock)
+    def test_poke_failed(self, mock_conn):
+        mock_client = mock.MagicMock()
+        mock_client.get_workflow_run.return_value = {
+            "RunDetail": {"RunState": "FAILED", "ErrorMessage": "Task failed"}
+        }
+        mock_conn.return_value = mock_client
+
+        with pytest.raises(RuntimeError, match="Task failed"):
+            self.sensor.poke({})
+
+    @mock.patch.object(AwsBaseHook, "conn", new_callable=mock.PropertyMock)
+    def test_poke_custom_states(self, mock_conn):
+        sensor = MwaaServerlessWorkflowRunSensor(
+            task_id="wait_for_run",
+            workflow_arn=WORKFLOW_ARN,
+            run_id=RUN_ID,
+            success_states={"SUCCESS", "STOPPED"},
+            failure_states={"FAILED"},
+        )
+        mock_client = mock.MagicMock()
+        mock_client.get_workflow_run.return_value = {"RunDetail": {"RunState": "STOPPED", "ErrorMessage": ""}}
+        mock_conn.return_value = mock_client
+
+        assert sensor.poke({}) is True
+
+    def test_template_fields(self):
+        validate_template_fields(self.sensor)