diff --git a/providers/amazon/docs/operators/mwaa_serverless.rst b/providers/amazon/docs/operators/mwaa_serverless.rst index 93353ac30c440..86906990c591e 100644 --- a/providers/amazon/docs/operators/mwaa_serverless.rst +++ b/providers/amazon/docs/operators/mwaa_serverless.rst @@ -54,3 +54,17 @@ Reference ~~~~~~~~~ * `AWS boto3 Library Documentation for MWAA Serverless `__ + +.. _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] diff --git a/providers/amazon/provider.yaml b/providers/amazon/provider.yaml index d4475e0a4b686..e85fd1a378e8c 100644 --- a/providers/amazon/provider.yaml +++ b/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 diff --git a/providers/amazon/src/airflow/providers/amazon/aws/sensors/mwaa_serverless.py b/providers/amazon/src/airflow/providers/amazon/aws/sensors/mwaa_serverless.py new file mode 100644 index 0000000000000..d76179f33994b --- /dev/null +++ b/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 diff --git a/providers/amazon/src/airflow/providers/amazon/get_provider_info.py b/providers/amazon/src/airflow/providers/amazon/get_provider_info.py index 1e50f1d32960c..1363bd84d2ab2 100644 --- a/providers/amazon/src/airflow/providers/amazon/get_provider_info.py +++ b/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"], diff --git a/providers/amazon/tests/system/amazon/aws/example_mwaa_serverless.py b/providers/amazon/tests/system/amazon/aws/example_mwaa_serverless.py index 748f4df90ee7c..18ea3e8ddf9ac 100644 --- a/providers/amazon/tests/system/amazon/aws/example_mwaa_serverless.py +++ b/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, diff --git a/providers/amazon/tests/unit/amazon/aws/sensors/test_mwaa_serverless.py b/providers/amazon/tests/unit/amazon/aws/sensors/test_mwaa_serverless.py new file mode 100644 index 0000000000000..6dab7fc2add2d --- /dev/null +++ b/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)