This is an automated email from the ASF dual-hosted git repository.
o-nikolas pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 0c8a2269b61 Support code packages in MWAA Serverless workflow
operators (#72622)
0c8a2269b61 is described below
commit 0c8a2269b616bf3714bc00950f30dda39e16bcad
Author: Kazushi Y <[email protected]>
AuthorDate: Tue Sep 15 15:39:50 2026 +0900
Support code packages in MWAA Serverless workflow operators (#72622)
MWAA Serverless added PythonOperator and BashOperator support, which
requires
shipping the Python modules and shell scripts those tasks run separately
from
the YAML workflow definition, through the Code request parameter.
Neither MwaaServerlessCreateWorkflowOperator nor
MwaaServerlessUpdateWorkflowOperator exposed Code, and neither accepts a
kwargs
passthrough, so the parameter was unreachable. Workflows using those
operators
could not be created, and their code could not be updated, from Airflow.
Add a `code` parameter to both operators, named after the API field and
passed
through untouched so a future member of the Code union needs no operator
change. prune_dict omits it when unset, so existing DAGs are unaffected.
---
.../amazon/aws/operators/mwaa_serverless.py | 20 +++++++--
.../amazon/aws/operators/test_mwaa_serverless.py | 49 ++++++++++++++++++++++
2 files changed, 65 insertions(+), 4 deletions(-)
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 b8a31f7f1a5..142affd82fe 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
@@ -98,6 +98,9 @@ class
MwaaServerlessCreateWorkflowOperator(AwsBaseOperator[AwsBaseHook]):
:param workflow_name: The name of the workflow. (templated)
:param definition_s3_location: Dict with ``Bucket`` and ``ObjectKey`` for
the YAML definition. (templated)
+ :param code: Optional code location for ``PythonOperator`` and
``BashOperator`` tasks,
+ as a dict with an ``S3Location`` key containing ``Bucket``,
``ObjectKey``, and
+ optionally ``VersionId``. (templated)
:param role_arn: The execution role ARN. (templated)
:param description: Optional description. (templated)
:param tags: Optional tags dict.
@@ -107,15 +110,16 @@ class
MwaaServerlessCreateWorkflowOperator(AwsBaseOperator[AwsBaseHook]):
aws_hook_class = AwsBaseHook
template_fields: tuple[str, ...] = aws_template_fields(
- "workflow_name", "definition_s3_location", "role_arn", "description"
+ "workflow_name", "definition_s3_location", "code", "role_arn",
"description"
)
- template_fields_renderers = {"definition_s3_location": "json"}
+ template_fields_renderers = {"definition_s3_location": "json", "code":
"json"}
def __init__(
self,
*,
workflow_name: str,
definition_s3_location: dict[str, str],
+ code: dict[str, Any] | None = None,
role_arn: str,
description: str | None = None,
tags: dict[str, str] | None = None,
@@ -125,6 +129,7 @@ class
MwaaServerlessCreateWorkflowOperator(AwsBaseOperator[AwsBaseHook]):
super().__init__(**kwargs)
self.workflow_name = workflow_name
self.definition_s3_location = definition_s3_location
+ self.code = code
self.role_arn = role_arn
self.description = description
self.tags = tags
@@ -140,6 +145,7 @@ class
MwaaServerlessCreateWorkflowOperator(AwsBaseOperator[AwsBaseHook]):
{
"Name": self.workflow_name,
"DefinitionS3Location": self.definition_s3_location,
+ "Code": self.code,
"RoleArn": self.role_arn,
"Description": self.description,
"Tags": self.tags,
@@ -174,21 +180,25 @@ class
MwaaServerlessUpdateWorkflowOperator(AwsBaseOperator[AwsBaseHook]):
:param workflow_arn: The ARN of the workflow to update. (templated)
:param definition_s3_location: Dict with ``Bucket``, ``ObjectKey``, and
optionally
``VersionId`` for the updated YAML definition. (templated)
+ :param code: Optional code location for ``PythonOperator`` and
``BashOperator`` tasks,
+ as a dict with an ``S3Location`` key containing ``Bucket``,
``ObjectKey``, and
+ optionally ``VersionId``. (templated)
:param role_arn: The execution role ARN. (templated)
:param description: Optional updated description. (templated)
"""
aws_hook_class = AwsBaseHook
template_fields: tuple[str, ...] = aws_template_fields(
- "workflow_arn", "definition_s3_location", "role_arn", "description"
+ "workflow_arn", "definition_s3_location", "code", "role_arn",
"description"
)
- template_fields_renderers = {"definition_s3_location": "json"}
+ template_fields_renderers = {"definition_s3_location": "json", "code":
"json"}
def __init__(
self,
*,
workflow_arn: str,
definition_s3_location: dict[str, str],
+ code: dict[str, Any] | None = None,
role_arn: str,
description: str | None = None,
**kwargs,
@@ -196,6 +206,7 @@ class
MwaaServerlessUpdateWorkflowOperator(AwsBaseOperator[AwsBaseHook]):
super().__init__(**kwargs)
self.workflow_arn = workflow_arn
self.definition_s3_location = definition_s3_location
+ self.code = code
self.role_arn = role_arn
self.description = description
@@ -209,6 +220,7 @@ class
MwaaServerlessUpdateWorkflowOperator(AwsBaseOperator[AwsBaseHook]):
{
"WorkflowArn": self.workflow_arn,
"DefinitionS3Location": self.definition_s3_location,
+ "Code": self.code,
"RoleArn": self.role_arn,
"Description": self.description,
}
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 9b2be23881b..3397120f70c 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
@@ -88,6 +88,7 @@ class TestMwaaServerlessStartWorkflowRunOperator:
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"}
+CODE = {"S3Location": {"Bucket": "test-bucket", "ObjectKey":
"code/my_package.zip"}}
ROLE_ARN = "arn:aws:iam::123456789012:role/test-role"
@@ -156,6 +157,27 @@ class TestMwaaServerlessCreateWorkflowOperator:
with pytest.raises(ClientError, match="ConflictException"):
op.execute({})
+ @mock.patch.object(AwsBaseHook, "conn", new_callable=mock.PropertyMock)
+ def test_execute_with_code_version_id(self, mock_conn):
+ code = {"S3Location": {**CODE["S3Location"], "VersionId": "abc123"}}
+ op = MwaaServerlessCreateWorkflowOperator(
+ task_id="create_workflow",
+ workflow_name=WORKFLOW_NAME,
+ definition_s3_location=S3_LOCATION,
+ code=code,
+ role_arn=ROLE_ARN,
+ )
+ mock_client = mock.MagicMock()
+ mock_client.create_workflow.return_value = {"WorkflowArn":
WORKFLOW_ARN}
+ mock_conn.return_value = mock_client
+
+ result = op.execute({})
+
+ mock_client.create_workflow.assert_called_once_with(
+ Name=WORKFLOW_NAME, DefinitionS3Location=S3_LOCATION, Code=code,
RoleArn=ROLE_ARN
+ )
+ assert result == WORKFLOW_ARN
+
def test_template_fields(self):
validate_template_fields(self.operator)
@@ -226,6 +248,33 @@ class TestMwaaServerlessUpdateWorkflowOperator:
with pytest.raises(ClientError):
self.operator.execute({})
+ @mock.patch.object(AwsBaseHook, "conn", new_callable=mock.PropertyMock)
+ def test_execute_with_code(self, mock_conn):
+ op = MwaaServerlessUpdateWorkflowOperator(
+ task_id="update_workflow",
+ workflow_arn=WORKFLOW_ARN,
+ definition_s3_location=S3_LOCATION,
+ code=CODE,
+ role_arn=ROLE_ARN,
+ )
+ mock_client = mock.MagicMock()
+ mock_client.update_workflow.return_value = {
+ "WorkflowArn": WORKFLOW_ARN,
+ "WorkflowVersion": "abc123",
+ "ModifiedAt": "2026-05-12T00:00:00Z",
+ }
+ mock_conn.return_value = mock_client
+
+ result = op.execute({})
+
+ mock_client.update_workflow.assert_called_once_with(
+ WorkflowArn=WORKFLOW_ARN,
+ DefinitionS3Location=S3_LOCATION,
+ Code=CODE,
+ RoleArn=ROLE_ARN,
+ )
+ assert result == WORKFLOW_ARN
+
def test_template_fields(self):
validate_template_fields(self.operator)
