Add MwaaServerlessCreateWorkflowOperator

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

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

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(

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)