diff --git a/providers/amazon/docs/operators/mwaa_serverless.rst b/providers/amazon/docs/operators/mwaa_serverless.rst index 40790a2417dc2..93353ac30c440 100644 --- a/providers/amazon/docs/operators/mwaa_serverless.rst +++ b/providers/amazon/docs/operators/mwaa_serverless.rst @@ -22,6 +22,20 @@ Amazon MWAA Serverless (Managed Workflows) Amazon MWAA Serverless provides a serverless execution environment for Apache Airflow workflows. Use the operators below to manage MWAA Serverless workflow runs. +.. _howto/operator:MwaaServerlessCreateWorkflowOperator: + +Create a Workflow +----------------- + +To create an Amazon MWAA Serverless workflow, use +:class:`~airflow.providers.amazon.aws.operators.mwaa_serverless.MwaaServerlessCreateWorkflowOperator`. + +.. exampleinclude:: /../../amazon/tests/system/amazon/aws/example_mwaa_serverless.py + :language: python + :dedent: 4 + :start-after: [START howto_operator_mwaa_serverless_create_workflow] + :end-before: [END howto_operator_mwaa_serverless_create_workflow] + .. _howto/operator:MwaaServerlessStartWorkflowRunOperator: Start a Workflow Run diff --git a/providers/amazon/src/airflow/providers/amazon/aws/operators/mwaa_serverless.py b/providers/amazon/src/airflow/providers/amazon/aws/operators/mwaa_serverless.py index c283be89ed457..41633de37e155 100644 --- a/providers/amazon/src/airflow/providers/amazon/aws/operators/mwaa_serverless.py +++ b/providers/amazon/src/airflow/providers/amazon/aws/operators/mwaa_serverless.py @@ -18,7 +18,9 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, Literal + +from botocore.exceptions import ClientError from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook from airflow.providers.amazon.aws.operators.base_aws import AwsBaseOperator @@ -84,3 +86,74 @@ def execute(self, context: Context) -> str: run_id = response["RunId"] self.log.info("Started workflow run %s (status: %s)", run_id, response.get("Status")) return run_id + + +class MwaaServerlessCreateWorkflowOperator(AwsBaseOperator[AwsBaseHook]): + """ + Create an Amazon MWAA Serverless workflow. + + .. seealso:: + For more information on how to use this operator, take a look at the guide: + :ref:`howto/operator:MwaaServerlessCreateWorkflowOperator` + + :param workflow_name: The name of the workflow. (templated) + :param definition_s3_location: Dict with ``Bucket`` and ``ObjectKey`` for the YAML definition. (templated) + :param role_arn: The execution role ARN. (templated) + :param description: Optional description. (templated) + :param tags: Optional tags dict. + :param if_exists: Behavior when the workflow already exists. + ``"fail"`` raises an error, ``"skip"`` returns the existing ARN. + """ + + aws_hook_class = AwsBaseHook + template_fields: tuple[str, ...] = aws_template_fields( + "workflow_name", "definition_s3_location", "role_arn", "description" + ) + template_fields_renderers = {"definition_s3_location": "json"} + + def __init__( + self, + *, + workflow_name: str, + definition_s3_location: dict[str, str], + role_arn: str, + description: str | None = None, + tags: dict[str, str] | None = None, + if_exists: Literal["fail", "skip"] = "skip", + **kwargs, + ) -> None: + super().__init__(**kwargs) + self.workflow_name = workflow_name + self.definition_s3_location = definition_s3_location + self.role_arn = role_arn + self.description = description + self.tags = tags + self.if_exists = if_exists + + @property + def _hook_parameters(self) -> dict[str, Any]: + return {**super()._hook_parameters, "client_type": "mwaa-serverless"} + + def execute(self, context: Context) -> str: + self.log.info("Creating MWAA Serverless workflow %s", self.workflow_name) + kwargs: dict[str, Any] = prune_dict( + { + "Name": self.workflow_name, + "DefinitionS3Location": self.definition_s3_location, + "RoleArn": self.role_arn, + "Description": self.description, + "Tags": self.tags, + } + ) + try: + response = self.hook.conn.create_workflow(**kwargs) + workflow_arn = response["WorkflowArn"] + except ClientError as e: + if e.response["Error"]["Code"] == "ConflictException" and self.if_exists == "skip": + self.log.info("Workflow %s already exists, skipping.", self.workflow_name) + response = self.hook.conn.get_workflow(WorkflowArn=self.workflow_name) + workflow_arn = response["WorkflowArn"] + else: + raise + self.log.info("Workflow %s: %s", self.workflow_name, workflow_arn) + return workflow_arn 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 ebe01cd1ebe27..748f4df90ee7c 100644 --- a/providers/amazon/tests/system/amazon/aws/example_mwaa_serverless.py +++ b/providers/amazon/tests/system/amazon/aws/example_mwaa_serverless.py @@ -18,9 +18,8 @@ from datetime import datetime -from tenacity import retry, retry_if_exception_type, stop_after_attempt, wait_exponential - from airflow.providers.amazon.aws.operators.mwaa_serverless import ( + MwaaServerlessCreateWorkflowOperator, MwaaServerlessStartWorkflowRunOperator, ) from airflow.providers.amazon.aws.operators.s3 import ( @@ -60,25 +59,6 @@ sys_test_context_task = SystemTestContextBuilder().add_variable(ROLE_ARN_KEY).build() -@task -@retry( - retry=retry_if_exception_type(Exception), - stop=stop_after_attempt(5), - wait=wait_exponential(multiplier=2, min=4, max=30), - reraise=True, -) -def create_workflow(bucket: str, role_arn: str) -> str: - """Create the MWAA Serverless workflow with retry for IAM propagation.""" - import boto3 - - mwaa = boto3.client("mwaa-serverless") - return mwaa.create_workflow( - Name=bucket, - DefinitionS3Location={"Bucket": bucket, "ObjectKey": "workflow.yaml"}, - RoleArn=role_arn, - )["WorkflowArn"] - - @task(trigger_rule=TriggerRule.ALL_DONE) def stop_workflow_run(workflow_arn: str, run_id: str): """Stop the workflow run.""" @@ -88,12 +68,6 @@ def stop_workflow_run(workflow_arn: str, run_id: str): @task(trigger_rule=TriggerRule.ALL_DONE) -@retry( - retry=retry_if_exception_type(Exception), - stop=stop_after_attempt(3), - wait=wait_exponential(multiplier=2, min=5, max=30), - reraise=True, -) def delete_workflow(workflow_arn: str): """Delete the MWAA Serverless workflow.""" import boto3 @@ -121,7 +95,16 @@ def delete_workflow(workflow_arn: str): data=WORKFLOW_YAML.format(bucket=bucket_name), ) - workflow_arn = create_workflow(bucket=bucket_name, role_arn=role_arn) + # [START howto_operator_mwaa_serverless_create_workflow] + create_workflow = MwaaServerlessCreateWorkflowOperator( + task_id="create_workflow", + workflow_name=bucket_name, + definition_s3_location={"Bucket": bucket_name, "ObjectKey": "workflow.yaml"}, + role_arn=role_arn, + ) + # [END howto_operator_mwaa_serverless_create_workflow] + + workflow_arn = create_workflow.output # [START howto_operator_mwaa_serverless_start_workflow_run] start_workflow = MwaaServerlessStartWorkflowRunOperator( diff --git a/providers/amazon/tests/unit/amazon/aws/operators/test_mwaa_serverless.py b/providers/amazon/tests/unit/amazon/aws/operators/test_mwaa_serverless.py index b830146e4d4f1..b7535f825d3c6 100644 --- a/providers/amazon/tests/unit/amazon/aws/operators/test_mwaa_serverless.py +++ b/providers/amazon/tests/unit/amazon/aws/operators/test_mwaa_serverless.py @@ -18,8 +18,12 @@ from unittest import mock +import pytest +from botocore.exceptions import ClientError + from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook from airflow.providers.amazon.aws.operators.mwaa_serverless import ( + MwaaServerlessCreateWorkflowOperator, MwaaServerlessStartWorkflowRunOperator, ) @@ -76,3 +80,69 @@ def test_execute_with_overrides(self, mock_conn): def test_template_fields(self): validate_template_fields(self.operator) + + +WORKFLOW_NAME = "test-workflow" +WORKFLOW_ARN = "arn:aws:mwaa-serverless:us-east-1:123456789012:workflow/test-workflow" +S3_LOCATION = {"Bucket": "test-bucket", "ObjectKey": "workflow.yaml"} +ROLE_ARN = "arn:aws:iam::123456789012:role/test-role" + + +class TestMwaaServerlessCreateWorkflowOperator: + def setup_method(self): + self.operator = MwaaServerlessCreateWorkflowOperator( + task_id="create_workflow", + workflow_name=WORKFLOW_NAME, + definition_s3_location=S3_LOCATION, + role_arn=ROLE_ARN, + ) + + @mock.patch.object(AwsBaseHook, "conn", new_callable=mock.PropertyMock) + def test_execute(self, mock_conn): + mock_client = mock.MagicMock() + mock_client.create_workflow.return_value = {"WorkflowArn": WORKFLOW_ARN} + mock_conn.return_value = mock_client + + result = self.operator.execute({}) + + mock_client.create_workflow.assert_called_once_with( + Name=WORKFLOW_NAME, + DefinitionS3Location=S3_LOCATION, + RoleArn=ROLE_ARN, + ) + assert result == WORKFLOW_ARN + + @mock.patch.object(AwsBaseHook, "conn", new_callable=mock.PropertyMock) + def test_execute_skip_existing(self, mock_conn): + mock_client = mock.MagicMock() + mock_client.create_workflow.side_effect = ClientError( + {"Error": {"Code": "ConflictException", "Message": "Already exists"}}, + "CreateWorkflow", + ) + mock_client.get_workflow.return_value = {"WorkflowArn": WORKFLOW_ARN} + mock_conn.return_value = mock_client + + result = self.operator.execute({}) + assert result == WORKFLOW_ARN + + @mock.patch.object(AwsBaseHook, "conn", new_callable=mock.PropertyMock) + def test_execute_fail_on_conflict(self, mock_conn): + op = MwaaServerlessCreateWorkflowOperator( + task_id="create_workflow", + workflow_name=WORKFLOW_NAME, + definition_s3_location=S3_LOCATION, + role_arn=ROLE_ARN, + if_exists="fail", + ) + mock_client = mock.MagicMock() + mock_client.create_workflow.side_effect = ClientError( + {"Error": {"Code": "ConflictException", "Message": "Already exists"}}, + "CreateWorkflow", + ) + mock_conn.return_value = mock_client + + with pytest.raises(ClientError): + op.execute({}) + + def test_template_fields(self): + validate_template_fields(self.operator)