This is an automated email from the ASF dual-hosted git repository. vatsrahul1001 pushed a commit to branch v3-4-test in repository https://gitbox.apache.org/repos/asf/airflow.git
commit 18a51dfadbabd49aba236e55db8f3cf6b03ea8ca Author: Jason(Zhe-You) Liu <[email protected]> AuthorDate: Fri Oct 30 23:54:02 2026 +0800 [v3-4-test] Add an end-to-end test for a Dag written in YAML (#75928) (#76026) The YAML Dag importer was covered by unit tests only. Nothing ran a YAML Dag through a real Dag processor, and the `run:` plus queue-routing path had no coverage at all. Add a `yaml_dag` e2e mode whose Dag exercises every task body the importer supports: provider operators and a sensor driving S3 on LocalStack, a `run:` task routed by its queue to the Go coordinator, and a bodiless task that becomes an EmptyOperator. The Go example bundle gains the handler behind that one code task. (cherry picked from commit 406474723a2cdb4de9dcc7547565769a8845a4d5) Co-authored-by: Claude <[email protected]> --- .github/workflows/additional-prod-image-tests.yml | 18 ++ .github/workflows/airflow-e2e-tests.yml | 12 +- .github/workflows/ci-amd.yml | 2 + .github/workflows/ci-arm.yml | 2 + .../scheduled-legacy-prod-image-canary.yml | 2 + airflow-e2e-tests/docker/yaml-dag.yml | 33 ++++ .../tests/airflow_e2e_tests/conftest.py | 74 +++++++- .../tests/airflow_e2e_tests/constants.py | 6 + .../airflow_e2e_tests/yaml_dag_tests/__init__.py | 16 ++ .../yaml_dag_tests/test_yaml_dag.py | 204 +++++++++++++++++++++ airflow-e2e-tests/yaml-dags/yaml_s3_roundtrip.yaml | 88 +++++++++ dev/breeze/doc/05_test_commands.rst | 2 +- dev/breeze/doc/ci/04_selective_checks.md | 11 ++ .../images/output_testing_airflow-e2e-tests.svg | 2 +- .../images/output_testing_airflow-e2e-tests.txt | 2 +- .../airflow_breeze/commands/testing_commands.py | 1 + .../src/airflow_breeze/utils/selective_checks.py | 19 ++ .../src/airflow_breeze/utils/verification_plan.py | 1 + dev/breeze/tests/test_selective_checks.py | 40 ++++ .../cmd/airflow-go-pack/pack_integration_test.go | 3 + go-sdk/example/bundle/main.go | 5 + go-sdk/example/bundle/yamldag/handlers.go | 43 +++++ 22 files changed, 573 insertions(+), 13 deletions(-) diff --git a/.github/workflows/additional-prod-image-tests.yml b/.github/workflows/additional-prod-image-tests.yml index 47657cbe678..5df6fd20f00 100644 --- a/.github/workflows/additional-prod-image-tests.yml +++ b/.github/workflows/additional-prod-image-tests.yml @@ -76,6 +76,10 @@ on: # yamllint disable-line rule:truthy description: "Whether to run TypeScript SDK e2e tests (true/false)" required: true type: string + run-yaml-dag-e2e-tests: + description: "Whether to run YAML Dag e2e tests (true/false)" + required: true + type: string providers-compatibility-tests-matrix: description: "JSON matrix of released Airflow versions for the OpenLineage e2e compat tests." required: true @@ -382,6 +386,20 @@ jobs: (inputs.canary-run == 'true' && inputs.run-generic-prod-image-tests == 'true') || inputs.run-ts-sdk-e2e-tests == 'true' + test-e2e-integration-tests-yaml-dag: + name: "YAML Dag e2e tests with PROD image" + uses: ./.github/workflows/airflow-e2e-tests.yml + with: + workflow-name: "YAML Dag e2e test" + runners: ${{ format(inputs.runners, '8+16', '32+64') }} + platform: ${{ inputs.platform }} + default-python-version: "${{ inputs.default-python-version }}" + use-uv: ${{ inputs.use-uv }} + e2e_test_mode: "yaml_dag" + if: >- + (inputs.canary-run == 'true' && inputs.run-generic-prod-image-tests == 'true') || + inputs.run-yaml-dag-e2e-tests == 'true' + test-e2e-integration-tests-openlineage: name: "OpenLineage e2e tests with PROD image" uses: ./.github/workflows/airflow-e2e-tests.yml diff --git a/.github/workflows/airflow-e2e-tests.yml b/.github/workflows/airflow-e2e-tests.yml index 2b33bdf8948..184eeb6ac73 100644 --- a/.github/workflows/airflow-e2e-tests.yml +++ b/.github/workflows/airflow-e2e-tests.yml @@ -49,7 +49,7 @@ on: # yamllint disable-line rule:truthy type: string required: true e2e_test_mode: - description: "Test mode - basic, remote_log, remote_log_elasticsearch, remote_log_opensearch, xcom_object_storage, event_driven, java_sdk, go_sdk, or ts_sdk" # yamllint disable-line rule:line-length + description: "Test mode - basic, remote_log, remote_log_elasticsearch, remote_log_opensearch, xcom_object_storage, event_driven, java_sdk, go_sdk, ts_sdk, or yaml_dag" # yamllint disable-line rule:line-length type: string default: "basic" java-sdk-version: @@ -84,7 +84,7 @@ on: # yamllint disable-line rule:truthy type: string default: "" e2e_test_mode: - description: "Test mode - basic, remote_log, remote_log_elasticsearch, remote_log_opensearch, xcom_object_storage, event_driven, java_sdk, go_sdk, or ts_sdk" # yamllint disable-line rule:line-length + description: "Test mode - basic, remote_log, remote_log_elasticsearch, remote_log_opensearch, xcom_object_storage, event_driven, java_sdk, go_sdk, ts_sdk, or yaml_dag" # yamllint disable-line rule:line-length type: string default: "basic" java-sdk-version: @@ -157,13 +157,13 @@ jobs: # separate version pin to keep in sync; cache: false because the explicit # restore/save pair below persists the module + build caches even on red runs. - name: "Setup Go for the Go SDK build" - if: inputs.e2e_test_mode == 'go_sdk' + if: inputs.e2e_test_mode == 'go_sdk' || inputs.e2e_test_mode == 'yaml_dag' uses: actions/setup-go@b7ad1dad31e06c5925ef5d2fc7ad053ef454303e # v7.0.0 with: go-version-file: go-sdk/go.mod cache: false - name: "Restore Go SDK module + build cache" - if: inputs.e2e_test_mode == 'go_sdk' + if: inputs.e2e_test_mode == 'go_sdk' || inputs.e2e_test_mode == 'yaml_dag' id: go-cache uses: actions/cache/restore@55cc8345863c7cc4c66a329aec7e433d2d1c52a9 # v6.1.0 with: @@ -191,7 +191,7 @@ jobs: env: DOCKER_IMAGE: "${{ inputs.docker-image-tag }}" E2E_TEST_MODE: "${{ inputs.e2e_test_mode }}" - LANG_SDK_NATIVE_TOOLCHAIN: "${{ (inputs.e2e_test_mode == 'java_sdk' || inputs.e2e_test_mode == 'go_sdk' || inputs.e2e_test_mode == 'ts_sdk') && 'true' || 'false' }}" # yamllint disable-line rule:line-length + LANG_SDK_NATIVE_TOOLCHAIN: "${{ (inputs.e2e_test_mode == 'java_sdk' || inputs.e2e_test_mode == 'go_sdk' || inputs.e2e_test_mode == 'ts_sdk' || inputs.e2e_test_mode == 'yaml_dag') && 'true' || 'false' }}" # yamllint disable-line rule:line-length - name: "Save Java SDK Gradle dependency cache" # Saved even when the e2e tests fail: the Gradle warm-up is independent of # test outcome, and actions/cache's post step would drop it on every red run. @@ -207,7 +207,7 @@ jobs: key: ${{ steps.gradle-cache.outputs.cache-primary-key }} - name: "Save Go SDK module + build cache" # Saved even when the e2e tests fail, for the same reason as the Gradle cache above. - if: always() && inputs.e2e_test_mode == 'go_sdk' && steps.go-cache.outputs.cache-hit != 'true' + if: always() && (inputs.e2e_test_mode == 'go_sdk' || inputs.e2e_test_mode == 'yaml_dag') && steps.go-cache.outputs.cache-hit != 'true' # yamllint disable-line rule:line-length uses: actions/cache/save@55cc8345863c7cc4c66a329aec7e433d2d1c52a9 # v6.1.0 with: path: | diff --git a/.github/workflows/ci-amd.yml b/.github/workflows/ci-amd.yml index 20f2309f5fd..8af50f522b8 100644 --- a/.github/workflows/ci-amd.yml +++ b/.github/workflows/ci-amd.yml @@ -353,6 +353,7 @@ jobs: run-openlineage-e2e-tests: ${{ steps.selective-checks.outputs.run-openlineage-e2e-tests }} run-openlineage-e2e-compat-tests: ${{ steps.selective-checks.outputs.run-openlineage-e2e-compat-tests }} run-ts-sdk-e2e-tests: ${{ steps.selective-checks.outputs.run-ts-sdk-e2e-tests }} + run-yaml-dag-e2e-tests: ${{ steps.selective-checks.outputs.run-yaml-dag-e2e-tests }} run-system-tests: ${{ steps.selective-checks.outputs.run-system-tests }} run-task-sdk-tests: ${{ steps.selective-checks.outputs.run-task-sdk-tests }} run-task-sdk-integration-tests: ${{ steps.selective-checks.outputs.run-task-sdk-integration-tests }} @@ -1176,6 +1177,7 @@ jobs: providers-compatibility-tests-matrix: > ${{ needs.build-info.outputs.providers-compatibility-tests-matrix }} run-ts-sdk-e2e-tests: ${{ needs.build-info.outputs.run-ts-sdk-e2e-tests }} + run-yaml-dag-e2e-tests: ${{ needs.build-info.outputs.run-yaml-dag-e2e-tests }} use-uv: ${{ needs.build-info.outputs.use-uv }} run-ui-e2e-tests: ${{ needs.build-info.outputs.run-ui-e2e-tests }} run-airflow-ctl-integration-tests: ${{ needs.build-info.outputs.run-airflow-ctl-integration-tests }} diff --git a/.github/workflows/ci-arm.yml b/.github/workflows/ci-arm.yml index 24d903088e0..09b0967b490 100644 --- a/.github/workflows/ci-arm.yml +++ b/.github/workflows/ci-arm.yml @@ -234,6 +234,7 @@ jobs: run-openlineage-e2e-tests: ${{ steps.selective-checks.outputs.run-openlineage-e2e-tests }} run-openlineage-e2e-compat-tests: ${{ steps.selective-checks.outputs.run-openlineage-e2e-compat-tests }} run-ts-sdk-e2e-tests: ${{ steps.selective-checks.outputs.run-ts-sdk-e2e-tests }} + run-yaml-dag-e2e-tests: ${{ steps.selective-checks.outputs.run-yaml-dag-e2e-tests }} run-system-tests: ${{ steps.selective-checks.outputs.run-system-tests }} run-task-sdk-tests: ${{ steps.selective-checks.outputs.run-task-sdk-tests }} run-task-sdk-integration-tests: ${{ steps.selective-checks.outputs.run-task-sdk-integration-tests }} @@ -1057,6 +1058,7 @@ jobs: providers-compatibility-tests-matrix: > ${{ needs.build-info.outputs.providers-compatibility-tests-matrix }} run-ts-sdk-e2e-tests: ${{ needs.build-info.outputs.run-ts-sdk-e2e-tests }} + run-yaml-dag-e2e-tests: ${{ needs.build-info.outputs.run-yaml-dag-e2e-tests }} use-uv: ${{ needs.build-info.outputs.use-uv }} run-ui-e2e-tests: ${{ needs.build-info.outputs.run-ui-e2e-tests }} run-airflow-ctl-integration-tests: ${{ needs.build-info.outputs.run-airflow-ctl-integration-tests }} diff --git a/.github/workflows/scheduled-legacy-prod-image-canary.yml b/.github/workflows/scheduled-legacy-prod-image-canary.yml index 646d69d8e8d..db8b8a154eb 100644 --- a/.github/workflows/scheduled-legacy-prod-image-canary.yml +++ b/.github/workflows/scheduled-legacy-prod-image-canary.yml @@ -68,6 +68,7 @@ jobs: run-remote-logging-s3-e2e-tests: ${{ steps.selective-checks.outputs.run-remote-logging-s3-e2e-tests }} run-task-sdk-integration-tests: ${{ steps.selective-checks.outputs.run-task-sdk-integration-tests }} run-ts-sdk-e2e-tests: ${{ steps.selective-checks.outputs.run-ts-sdk-e2e-tests }} + run-yaml-dag-e2e-tests: ${{ steps.selective-checks.outputs.run-yaml-dag-e2e-tests }} run-ui-e2e-tests: ${{ steps.selective-checks.outputs.run-ui-e2e-tests }} use-uv: ${{ steps.selective-checks.outputs.force-pip == 'true' && 'false' || 'true' }} steps: @@ -235,6 +236,7 @@ jobs: providers-compatibility-tests-matrix: > ${{ needs.build-info.outputs.providers-compatibility-tests-matrix }} run-ts-sdk-e2e-tests: ${{ needs.build-info.outputs.run-ts-sdk-e2e-tests }} + run-yaml-dag-e2e-tests: ${{ needs.build-info.outputs.run-yaml-dag-e2e-tests }} use-uv: ${{ needs.build-info.outputs.use-uv }} run-ui-e2e-tests: ${{ needs.build-info.outputs.run-ui-e2e-tests }} run-airflow-ctl-integration-tests: ${{ needs.build-info.outputs.run-airflow-ctl-integration-tests }} diff --git a/airflow-e2e-tests/docker/yaml-dag.yml b/airflow-e2e-tests/docker/yaml-dag.yml new file mode 100644 index 00000000000..eb9c8101aca --- /dev/null +++ b/airflow-e2e-tests/docker/yaml-dag.yml @@ -0,0 +1,33 @@ +# 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. + +# Docker Compose override for the yaml_dag E2E test mode. +# +# The Dag is a YAML file in the Dags folder, so nothing extra is needed to parse +# it beyond the importer registered in the .env file. Its "run:" task is a stub +# routed to the Go coordinator, which is why the worker listens on the "golang" +# queue and the packed Go bundle is mounted where the coordinator scans for it +# (the same arrangement go.yml makes; see its comments for the USER / HOME ones). +--- +services: + airflow-worker: + environment: + USER: airflow + HOME: /home/airflow + volumes: + - ./go-bundles:/opt/airflow/go-bundles:ro + command: celery worker -q golang,default diff --git a/airflow-e2e-tests/tests/airflow_e2e_tests/conftest.py b/airflow-e2e-tests/tests/airflow_e2e_tests/conftest.py index ecb6bbe3ed2..ddeb41756b1 100644 --- a/airflow-e2e-tests/tests/airflow_e2e_tests/conftest.py +++ b/airflow-e2e-tests/tests/airflow_e2e_tests/conftest.py @@ -80,6 +80,8 @@ from airflow_e2e_tests.constants import ( TS_SDK_EXAMPLE_PATH, TS_SDK_ROOT_PATH, XCOM_BUCKET, + YAML_DAG_COMPOSE_PATH, + YAML_DAGS_PATH, ) from tests_common.test_utils.fernet import generate_fernet_key_string @@ -100,6 +102,11 @@ def _is_image_present_locally(image: str) -> bool: ) +# Every mode that composes localstack.yml points the AWS SDK at it this way; the +# connection each mode's Dags use is set next to the rest of that mode's config. +LOCALSTACK_AWS_ENV = "AWS_DEFAULT_REGION=us-east-1\nAWS_ENDPOINT_URL_S3=http://localstack:4566\n"; + + def _copy_localstack_files(tmp_dir): """Copy localstack compose file and init script into the temp directory.""" copyfile(LOCALSTACK_PATH, tmp_dir / "localstack.yml") @@ -124,8 +131,7 @@ def _setup_s3_integration(dot_env_file, tmp_dir): dot_env_file.write_text( f"AIRFLOW_UID={os.getuid()}\n" - "AWS_DEFAULT_REGION=us-east-1\n" - "AWS_ENDPOINT_URL_S3=http://localstack:4566\n"; + f"{LOCALSTACK_AWS_ENV}" "AIRFLOW__LOGGING__REMOTE_LOGGING=true\n" "AIRFLOW_CONN_AWS_S3_LOGS=aws://test:test@\n" "AIRFLOW__LOGGING__REMOTE_LOG_CONN_ID=aws_s3_logs\n" @@ -268,8 +274,7 @@ def _setup_xcom_object_storage_integration(dot_env_file, tmp_dir): # for authentication rather than parsing credentials from the connection URI "AWS_ACCESS_KEY_ID=test\n" "AWS_SECRET_ACCESS_KEY=test\n" - "AWS_DEFAULT_REGION=us-east-1\n" - "AWS_ENDPOINT_URL_S3=http://localstack:4566\n"; + f"{LOCALSTACK_AWS_ENV}" "AIRFLOW_CONN_AWS_DEFAULT=aws://test:test@\n" "AIRFLOW__CORE__XCOM_BACKEND=airflow.providers.common.io.xcom.backend.XComObjectStorageBackend\n" f"AIRFLOW__COMMON_IO__XCOM_OBJECTSTORAGE_PATH=s3://aws_default@{XCOM_BUCKET}/xcom\n" @@ -783,6 +788,64 @@ def _setup_go_sdk_integration(dot_env_file, tmp_dir): _write_state_store_retention_config(tmp_dir) +def _setup_yaml_dag_integration(dot_env_file, tmp_dir): + """Set up the yaml_dag E2E test mode. + + The Dag under test is written in YAML, so there is nothing to build for the + Dag itself: the file is copied into the stack's Dags folder and the opt-in + ``YamlDagImporter`` is registered so the Dag processor reads it. + + Its ``uses:`` tasks drive S3 through LocalStack, and its one ``run:`` task + carries no code, so it is routed by its ``golang`` queue to the Go + coordinator, which finds the handler in the packed Go example bundle (the + same bundle the go_sdk mode uses -- see ``_setup_go_sdk_integration``). + """ + _pack_go_sdk_bundle(GO_SDK_EXAMPLE_BUNDLE_PKG, GO_SDK_BUNDLE_NAME, native=LANG_SDK_NATIVE_TOOLCHAIN) + + _copy_localstack_files(tmp_dir) + copyfile(YAML_DAG_COMPOSE_PATH, tmp_dir / "yaml-dag.yml") + + # Where the compose bind-mount (./go-bundles) exposes the bundle to the worker. The + # bundle scanner requires the file to be executable, so preserve the exec bit. + go_bundles_dir = tmp_dir / "go-bundles" + go_bundles_dir.mkdir() + packed_bundle = go_bundles_dir / GO_SDK_BUNDLE_NAME + copyfile(GO_SDK_BIN_PATH / GO_SDK_BUNDLE_NAME, packed_bundle) + os.chmod(packed_bundle, 0o755) + # This bundle only registers task handlers, so the Dag processor must not parse it if the + # directory is ever mounted there. Tasks do not read .airflowignore. + (go_bundles_dir / ".airflowignore").write_text("*\n") + + copytree(YAML_DAGS_PATH, tmp_dir / "dags" / "yaml") + + dag_bundle_config = _build_dag_bundle_config({"go-task-handlers": "/opt/airflow/go-bundles"}) + # The YAML importer is opt-in, and is what makes the .yaml file a Dag file rather than a + # file the Dag processor ignores. + dag_importer_configs = json.dumps([{"classpath": "airflow.sdk.importers.yaml_importer.YamlDagImporter"}]) + coordinator_config = json.dumps( + { + "go-sdk": { + "classpath": "airflow.sdk.coordinators.executable.ExecutableCoordinator", + "kwargs": {"task_handler_bundle_name": "go-task-handlers"}, + } + } + ) + queue_to_coordinator = json.dumps({"golang": "go-sdk"}) + + dot_env_file.write_text( + f"AIRFLOW_UID={os.getuid()}\n" + # Single-quote the JSON values so Docker Compose reads them literally. + f"AIRFLOW__DAG_PROCESSOR__DAG_IMPORTER_CONFIGS='{dag_importer_configs}'\n" + f"AIRFLOW__DAG_PROCESSOR__DAG_BUNDLE_CONFIG_LIST='{dag_bundle_config}'\n" + f"AIRFLOW__SDK__COORDINATORS='{coordinator_config}'\n" + f"AIRFLOW__SDK__QUEUE_TO_COORDINATOR='{queue_to_coordinator}'\n" + # The connection the Dag's S3 tasks name. + "AIRFLOW_CONN_AWS_LOCALSTACK=aws://test:test@\n" + f"{LOCALSTACK_AWS_ENV}" + ) + os.environ["ENV_FILE_PATH"] = str(dot_env_file) + + def _setup_openlineage_integration(dot_env_file, tmp_dir, compose_file_names): """Set up the openlineage E2E test mode. @@ -960,6 +1023,9 @@ def spin_up_airflow_environment(tmp_path_factory: pytest.TempPathFactory): elif E2E_TEST_MODE == "ts_sdk": compose_file_names.append("ts.yml") _setup_ts_sdk_integration(dot_env_file, tmp_dir) + elif E2E_TEST_MODE == "yaml_dag": + compose_file_names.extend(["localstack.yml", "yaml-dag.yml"]) + _setup_yaml_dag_integration(dot_env_file, tmp_dir) # # Please Do not use this Fernet key in any deployments! Please generate your own key. diff --git a/airflow-e2e-tests/tests/airflow_e2e_tests/constants.py b/airflow-e2e-tests/tests/airflow_e2e_tests/constants.py index 40de0c37ffc..2863d95ad1d 100644 --- a/airflow-e2e-tests/tests/airflow_e2e_tests/constants.py +++ b/airflow-e2e-tests/tests/airflow_e2e_tests/constants.py @@ -127,6 +127,12 @@ NODE_IMAGE = os.environ.get("NODE_IMAGE", "node:22-slim") # Writable HOME for the containerized pnpm build; gitignored, caches persist. TS_SDK_BUILD_HOME_PATH = AIRFLOW_ROOT_PATH / "files" / "pnpm-home" +# YAML Dag E2E test paths +YAML_DAGS_PATH = AIRFLOW_ROOT_PATH / "airflow-e2e-tests" / "yaml-dags" +YAML_DAG_COMPOSE_PATH = AIRFLOW_ROOT_PATH / "airflow-e2e-tests" / "docker" / "yaml-dag.yml" +# Must match "params.bucket" in yaml-dags/yaml_s3_roundtrip.yaml. +YAML_DAG_BUCKET = "airflow-yaml-dag" + # Local provider sources are mounted into the airflow containers under this directory so # ``_PIP_ADDITIONAL_REQUIREMENTS`` can install the in-tree (latest, possibly unreleased) # provider rather than the published one from PyPI. diff --git a/airflow-e2e-tests/tests/airflow_e2e_tests/yaml_dag_tests/__init__.py b/airflow-e2e-tests/tests/airflow_e2e_tests/yaml_dag_tests/__init__.py new file mode 100644 index 00000000000..13a83393a91 --- /dev/null +++ b/airflow-e2e-tests/tests/airflow_e2e_tests/yaml_dag_tests/__init__.py @@ -0,0 +1,16 @@ +# 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. diff --git a/airflow-e2e-tests/tests/airflow_e2e_tests/yaml_dag_tests/test_yaml_dag.py b/airflow-e2e-tests/tests/airflow_e2e_tests/yaml_dag_tests/test_yaml_dag.py new file mode 100644 index 00000000000..4daa9a29f1f --- /dev/null +++ b/airflow-e2e-tests/tests/airflow_e2e_tests/yaml_dag_tests/test_yaml_dag.py @@ -0,0 +1,204 @@ +# 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. +""" +End-to-end test of a Dag written in YAML. + +Run with:: + + E2E_TEST_MODE=yaml_dag uv run --project airflow-e2e-tests pytest \\ + tests/airflow_e2e_tests/yaml_dag_tests/ -xvs + +No Python declares this Dag: ``airflow-e2e-tests/yaml-dags/yaml_s3_roundtrip.yaml`` is dropped +into the Dags folder and the opt-in ``YamlDagImporter`` builds it (see +``conftest._setup_yaml_dag_integration``). The run covers every task body the importer supports: + +* ``create_bucket`` / ``write_object`` / ``list_objects`` are provider operators named by ``uses``, + and ``wait_for_object`` is a sensor. They drive S3 on LocalStack, so the templated ``with:`` + arguments, the ``$t`` markers and the ``extends`` merge are all proven against a real service + rather than against the serialized Dag alone. +* ``summarize`` is a ``run:`` task. It carries no code at all: the ``golang`` queue routes it to the + Go coordinator, which runs the handler registered for this dag_id in the packed Go bundle. The + upstream listing reaches it through ``{$x: list_objects}``. +* ``done`` has no body, so the importer turns it into an ``EmptyOperator`` join. +""" + +from __future__ import annotations + +import json +import time +from dataclasses import dataclass +from datetime import UTC, datetime + +import pytest + +from airflow_e2e_tests.constants import YAML_DAG_BUCKET +from airflow_e2e_tests.e2e_test_utils.clients import AirflowClient, get_s3_client + +# The sensor alone may poke for up to 60 s, and the Go coordinator starts before the +# run: task's first line; allow the same headroom the other lang-SDK suites do. +_RUN_TIMEOUT = 600 +# Task logs are written when the task finishes; allow a little slack for them to +# become retrievable through the API after the run reaches a terminal state. +_LOG_FETCH_TIMEOUT = 120 + +_DAG_ID = "yaml_s3_roundtrip" + + +@dataclass +class _CompletedRun: + """The single ``yaml_s3_roundtrip`` run shared across the assertions in this module.""" + + client: AirflowClient + run_id: str + logical_date: datetime + state: str + ti_states: dict[str, str] + + @property + def object_key(self) -> str: + """The key the Dag templates from ``{{ ds }}`` and ``{{ ts_nodash }}``; see its ``s3_key``. + + The timestamp keeps every run in its own prefix, so the scheduled runs the + ``@daily`` schedule produces never land on this run's object. + """ + ds = self.logical_date.strftime("%Y-%m-%d") + ts_nodash = self.logical_date.strftime("%Y%m%dT%H%M%S") + return f"runs/{ds}/{ts_nodash}/payload.json" + + def xcom(self, task_id: str, key: str = "return_value"): + return self.client.get_xcom_value(dag_id=_DAG_ID, task_id=task_id, run_id=self.run_id, key=key).get( + "value" + ) + + def logs(self, task_id: str, try_number: int = 1) -> str: + """Return the concatenated task-log records for *task_id*, retrying until present.""" + deadline = time.monotonic() + _LOG_FETCH_TIMEOUT + while True: + resp = self.client.get_task_logs( + dag_id=_DAG_ID, run_id=self.run_id, task_id=task_id, try_number=try_number + ) + text = "\n".join(str(entry) for entry in resp.get("content", [])) + if text.strip() or time.monotonic() > deadline: + return text + time.sleep(3) + + [email protected](scope="module") +def completed_run() -> _CompletedRun: + """Trigger ``yaml_s3_roundtrip`` once and wait for it to finish. + + The logical date is passed explicitly because the Dag templates its S3 key from + ``{{ ds }}``; a run without one renders nothing and the write fails. + """ + client = AirflowClient() + logical_date = datetime.now(UTC) + resp = client.trigger_dag(_DAG_ID, json={"logical_date": logical_date.isoformat()}) + run_id = resp["dag_run_id"] + state = client.wait_for_dag_run(dag_id=_DAG_ID, run_id=run_id, timeout=_RUN_TIMEOUT) + ti_resp = client.get_task_instances(dag_id=_DAG_ID, run_id=run_id) + ti_states = {ti["task_id"]: ti.get("state") for ti in ti_resp.get("task_instances", [])} + return _CompletedRun( + client=client, + run_id=run_id, + logical_date=logical_date, + state=state, + ti_states=ti_states, + ) + + +def test_every_task_succeeded(completed_run: _CompletedRun): + """The importer built all six tasks and the run drove each of them to success.""" + expected = { + "create_bucket": "success", + "write_object": "success", + "wait_for_object": "success", + "list_objects": "success", + "summarize": "success", + "done": "success", + } + assert completed_run.ti_states == expected + + +def test_dag_run_succeeded(completed_run: _CompletedRun): + assert completed_run.state == "success", ( + f"expected the run to succeed; got {completed_run.state!r}. task states: {completed_run.ti_states}" + ) + + +def test_each_task_body_became_the_operator_it_names(completed_run: _CompletedRun): + """``uses`` resolves the import path, ``run`` becomes a stub, and no body at all is an EmptyOperator.""" + tasks = {task["task_id"]: task for task in completed_run.client.get_tasks(_DAG_ID).get("tasks", [])} + operators = {task_id: task["operator_name"] for task_id, task in tasks.items()} + + assert operators == { + "create_bucket": "S3CreateBucketOperator", + "write_object": "S3CreateObjectOperator", + "wait_for_object": "S3KeySensor", + "list_objects": "S3ListOperator", + "summarize": "@task.stub", + "done": "EmptyOperator", + } + # The run: task is the only one that leaves the Python runtime, and only its + # task-level queue sends it there. + assert tasks["summarize"]["queue"] == "golang" + assert {task_id for task_id, task in tasks.items() if task["queue"] != "default"} == {"summarize"} + + +def test_extends_merged_both_task_level_and_operator_arguments(completed_run: _CompletedRun): + """``extends: [retryable, s3]`` reached the task model, and the tasks that skip it keep the defaults.""" + tasks = {task["task_id"]: task for task in completed_run.client.get_tasks(_DAG_ID).get("tasks", [])} + + assert tasks["create_bucket"]["retries"] == 2 + assert tasks["write_object"]["retries"] == 2 + assert tasks["wait_for_object"]["retries"] == 0 + # The operator argument half of the merge is proven by the run itself: without the + # aws_conn_id the s3 template contributes, every S3 task would fail. + + +def test_the_code_view_shows_the_yaml_file(completed_run: _CompletedRun): + """The Code tab serves the YAML document that declares the Dag.""" + content = completed_run.client.get_dag_source(_DAG_ID)["content"] + + assert "dag_id: yaml_s3_roundtrip" in content + assert "$schema: https://airflow.apache.org/schemas/dag/"; in content + + +def test_the_object_round_tripped_through_s3(completed_run: _CompletedRun): + """The templated write really happened: the key and its body carry this run's values.""" + body = ( + get_s3_client() + .get_object(Bucket=YAML_DAG_BUCKET, Key=completed_run.object_key)["Body"] + .read() + .decode() + ) + + assert json.loads(body) == {"run_id": completed_run.run_id} + + +def test_the_listing_reached_the_code_task(completed_run: _CompletedRun): + """``{$x: list_objects}`` bound the upstream XCom as the Go handler's argument.""" + assert completed_run.xcom("list_objects") == [completed_run.object_key] + assert completed_run.xcom("summarize") == { + "label": "yaml", + "count": 1, + "first_key": completed_run.object_key, + } + + +def test_the_code_task_ran_on_the_go_runtime(completed_run: _CompletedRun): + """The handler's own log line reaches Airflow, so the task body ran outside Python.""" + assert "Summarizing the YAML Dag's listing" in completed_run.logs("summarize") diff --git a/airflow-e2e-tests/yaml-dags/yaml_s3_roundtrip.yaml b/airflow-e2e-tests/yaml-dags/yaml_s3_roundtrip.yaml new file mode 100644 index 00000000000..63a082aea71 --- /dev/null +++ b/airflow-e2e-tests/yaml-dags/yaml_s3_roundtrip.yaml @@ -0,0 +1,88 @@ +# 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. + +# A Dag defined entirely in YAML, driven by the yaml_dag e2e test mode. +# +# It writes an object to S3 (LocalStack), waits for it, lists it back, hands the +# listing to a Go task handler over the "golang" queue, and joins on an +# EmptyOperator. Between them the four task bodies the importer supports are all +# exercised: a provider operator (uses), a sensor, a code task (run), and a task +# with no body at all. +--- +$schema: https://airflow.apache.org/schemas/dag/2026-10-30.json +dag_id: yaml_s3_roundtrip +schedule: "@daily" +start_date: "2026-01-01" +catchup: false +tags: [yaml, aws] +params: + bucket: airflow-yaml-dag +templates: + # Proves `extends` reaches task-level arguments: every task that extends this + # is serialized with retries 2. + retryable: + retries: 2 + retry_delay: "5s" + # ... and operator arguments, which merge into `with` one key at a time. + s3: + with: + aws_conn_id: aws_localstack +tasks: + - id: create_bucket + extends: [retryable, s3] + uses: airflow.providers.amazon.aws.operators.s3.S3CreateBucketOperator + with: + bucket_name: {$t: "{{ params.bucket }}"} + + - id: write_object + extends: [retryable, s3] + needs: [create_bucket] + uses: airflow.providers.amazon.aws.operators.s3.S3CreateObjectOperator + with: + s3_bucket: {$t: "{{ params.bucket }}"} + s3_key: {$t: "runs/{{ ds }}/{{ ts_nodash }}/payload.json"} + data: {$t: '{"run_id": "{{ run_id }}"}'} + replace: true + + - id: wait_for_object + extends: [s3] + needs: [write_object] + uses: airflow.providers.amazon.aws.sensors.s3.S3KeySensor + with: + bucket_name: {$t: "{{ params.bucket }}"} + bucket_key: {$t: "runs/{{ ds }}/{{ ts_nodash }}/payload.json"} + poke_interval: 5 + timeout: 60 + + - id: list_objects + extends: [s3] + needs: [wait_for_object] + uses: airflow.providers.amazon.aws.operators.s3.S3ListOperator + with: + bucket: {$t: "{{ params.bucket }}"} + prefix: {$t: "runs/{{ ds }}/{{ ts_nodash }}/"} + + # A code task: its queue routes it to the Go TaskHandler. + - id: summarize + queue: golang + run: + keys: {$x: list_objects} + label: yaml + + # No body at all, so the importer makes this an EmptyOperator join. + - id: done + needs: [summarize] diff --git a/dev/breeze/doc/05_test_commands.rst b/dev/breeze/doc/05_test_commands.rst index 84b1edcb590..78dd6c6fe8f 100644 --- a/dev/breeze/doc/05_test_commands.rst +++ b/dev/breeze/doc/05_test_commands.rst @@ -387,7 +387,7 @@ You can override the ``DOCKER_IMAGE`` environment variable to point to the image The Airflow E2E tests are in ``airflow-e2e-tests/`` folder in the main repo. Each suite is a ``--e2e-test-mode`` (``basic``, ``remote_log``, ``event_driven``, ``java_sdk``, ``go_sdk``, -``openlineage``, ...); the mode selects a docker-compose overlay and the matching test package under +``openlineage``, ``yaml_dag``, ...); the mode selects a docker-compose overlay and the matching test package under ``airflow-e2e-tests/tests/airflow_e2e_tests/<mode>_tests``. For example, the ``openlineage`` mode deploys a real Airflow and runs the OpenLineage provider's system-test DAGs against it, asserting the emitted lineage events match — run it with ``breeze testing airflow-e2e-tests --e2e-test-mode openlineage``. diff --git a/dev/breeze/doc/ci/04_selective_checks.md b/dev/breeze/doc/ci/04_selective_checks.md index 409e2cba52e..8b750aa6280 100644 --- a/dev/breeze/doc/ci/04_selective_checks.md +++ b/dev/breeze/doc/ci/04_selective_checks.md @@ -548,6 +548,16 @@ together using `pytest-xdist` (pytest-xdist distributes the tests among parallel they trigger the TypeScript SDK E2E tests too. Shared Dag processing files such as `manager.py` are left to `canary` runs. Like the other deployed e2e suites, enabling them forces `PROD Image building`. +* `YAML Dag E2E tests` (the `yaml_dag` mode of the deployed-stack tests, exposed as the + `run-yaml-dag-e2e-tests` output) run when the YAML Dag fixture or its suite change + (`airflow-e2e-tests/yaml-dags/`, `airflow-e2e-tests/tests/airflow_e2e_tests/yaml_dag_tests/`, + `airflow-e2e-tests/docker/yaml-dag.yml`), when the Go example bundle's handler for that Dag changes + (`go-sdk/example/bundle/yamldag/`, `go-sdk/example/bundle/main.go` — the mode runs the Dag's one + code task on that bundle, so the rest of the Go SDK is left to the Go SDK E2E tests), when the + executable coordinator changes (`task-sdk/src/airflow/sdk/coordinators/executable/`, + `_subprocess.py`), or when the native Lang-SDK Dag parsing sources listed above change — the YAML + importer itself lives in `task-sdk/src/airflow/sdk/importers/`. Like the other deployed e2e + suites, enabling them forces `PROD Image building`. * `OpenLineage E2E tests` (the `openlineage` mode of the deployed-stack tests under `airflow-e2e-tests/tests/airflow_e2e_tests/openlineage_tests`, exposed as the `run-openlineage-e2e-tests` output) run when the `openlineage` or `common` providers or the @@ -735,6 +745,7 @@ GitHub Actions to pass the list of parameters to a command to execute | run-ui-tests | Whether UI tests should be run ("true"/"false") | true | | | run-unit-tests | Whether unit tests should be run ("true"/"false") | true | | | run-www-tests | Whether Legacy WWW tests should be run ("true"/"false") | true | | +| run-yaml-dag-e2e-tests | Whether YAML Dag e2e tests should be run — on YAML Dag fixture/suite, its Go handler, executable coordinator, or native Lang-SDK Dag parsing changes ("true"/"false") | true | | | amd-runners | List of labels assigned for runners for that build for public AMD runners | \["ubuntu-22.04"\] | | | arm-runners | List of labels assigned for runners for that build for public ARM runners | \["ubuntu-22.04-arm"\] | | | runner-type | Runner labels for the CI jobs on this platform: public, or a RunsOn template jobs fill with cpu/ram | \["ubuntu-22.04"\] | | diff --git a/dev/breeze/doc/images/output_testing_airflow-e2e-tests.svg b/dev/breeze/doc/images/output_testing_airflow-e2e-tests.svg index a0cd6236df1..f135e187251 100644 --- a/dev/breeze/doc/images/output_testing_airflow-e2e-tests.svg +++ b/dev/breeze/doc/images/output_testing_airflow-e2e-tests.svg @@ -148,7 +148,7 @@ </text><text class="breeze-testing-airflow-e2e-tests-r5" x="0" y="312.8" textLength="12.2" clip-path="url(#breeze-testing-airflow-e2e-tests-line-12)">│</text><text class="breeze-testing-airflow-e2e-tests-r7" x="463.6" y="312.8" textLength="73.2" clip-path="url(#breeze-testing-airflow-e2e-tests-line-12)">(TEXT)</text><text class="breeze-testing-airflow-e2e-tests-r5" x="1451.8" y="312.8" textLength="12.2" clip-path="url(#breeze-testing-airflow-e2e-tests-line-12)">│</text><text class="breez [...] </text><text class="breeze-testing-airflow-e2e-tests-r5" x="0" y="337.2" textLength="12.2" clip-path="url(#breeze-testing-airflow-e2e-tests-line-13)">│</text><text class="breeze-testing-airflow-e2e-tests-r4" x="24.4" y="337.2" textLength="366" clip-path="url(#breeze-testing-airflow-e2e-tests-line-13)">--e2e-test-mode               </text><text class="breeze-testing-airflow-e2e-tests-r1" x="463.6" y="337.2" textLen [...] </text><text class="breeze-testing-airflow-e2e-tests-r5" x="0" y="361.6" textLength="12.2" clip-path="url(#breeze-testing-airflow-e2e-tests-line-14)">│</text><text class="breeze-testing-airflow-e2e-tests-r7" x="463.6" y="361.6" textLength="976" clip-path="url(#breeze-testing-airflow-e2e-tests-line-14)">(basic|remote_log|remote_log_elasticsearch|remote_log_opensearch|xcom_object_sto</text><text class="breeze-testing-airflow-e2e-tests-r5" x="1451.8" y="361.6" textLength="12.2" clip-path="u [...] -</text><text class="breeze-testing-airflow-e2e-tests-r5" x="0" y="386" textLength="12.2" clip-path="url(#breeze-testing-airflow-e2e-tests-line-15)">│</text><text class="breeze-testing-airflow-e2e-tests-r7" x="463.6" y="386" textLength="646.6" clip-path="url(#breeze-testing-airflow-e2e-tests-line-15)">rage|event_driven|java_sdk|go_sdk|openlineage|ts_sdk)</text><text class="breeze-testing-airflow-e2e-tests-r5" x="1451.8" y="386" textLength="12.2" clip-path="url(#breeze-testing-airflow-e2e- [...] +</text><text class="breeze-testing-airflow-e2e-tests-r5" x="0" y="386" textLength="12.2" clip-path="url(#breeze-testing-airflow-e2e-tests-line-15)">│</text><text class="breeze-testing-airflow-e2e-tests-r7" x="463.6" y="386" textLength="756.4" clip-path="url(#breeze-testing-airflow-e2e-tests-line-15)">rage|event_driven|java_sdk|go_sdk|openlineage|ts_sdk|yaml_dag)</text><text class="breeze-testing-airflow-e2e-tests-r5" x="1451.8" y="386" textLength="12.2" clip-path="url(#breeze-testing-air [...] </text><text class="breeze-testing-airflow-e2e-tests-r5" x="0" y="410.4" textLength="12.2" clip-path="url(#breeze-testing-airflow-e2e-tests-line-16)">│</text><text class="breeze-testing-airflow-e2e-tests-r4" x="24.4" y="410.4" textLength="366" clip-path="url(#breeze-testing-airflow-e2e-tests-line-16)">--airflow-version             </text><text class="breeze-testing-airflow-e2e-tests-r1" x="463.6" y="410.4" textLength="976" [...] </text><text class="breeze-testing-airflow-e2e-tests-r5" x="0" y="434.8" textLength="12.2" clip-path="url(#breeze-testing-airflow-e2e-tests-line-17)">│</text><text class="breeze-testing-airflow-e2e-tests-r1" x="463.6" y="434.8" textLength="976" clip-path="url(#breeze-testing-airflow-e2e-tests-line-17)">current providers from main, instead of the default PROD image. Builds a        </text><text [...] </text><text class="breeze-testing-airflow-e2e-tests-r5" x="0" y="459.2" textLength="12.2" clip-path="url(#breeze-testing-airflow-e2e-tests-line-18)">│</text><text class="breeze-testing-airflow-e2e-tests-r1" x="463.6" y="459.2" textLength="463.6" clip-path="url(#breeze-testing-airflow-e2e-tests-line-18)">lightweight image from apache/airflow:</text><text class="breeze-testing-airflow-e2e-tests-r8" x="927.2" y="459.2" textLength="109.8" clip-path="url(#breeze-testing-airflo [...] diff --git a/dev/breeze/doc/images/output_testing_airflow-e2e-tests.txt b/dev/breeze/doc/images/output_testing_airflow-e2e-tests.txt index e5a1ee13134..295865b6347 100644 --- a/dev/breeze/doc/images/output_testing_airflow-e2e-tests.txt +++ b/dev/breeze/doc/images/output_testing_airflow-e2e-tests.txt @@ -1 +1 @@ -bf145c62e5a2d5f6edceca06cce9fd8b +f4905f43679e01c3b60e203a787acee2 diff --git a/dev/breeze/src/airflow_breeze/commands/testing_commands.py b/dev/breeze/src/airflow_breeze/commands/testing_commands.py index 8934b40cbf6..94db1ad18a4 100644 --- a/dev/breeze/src/airflow_breeze/commands/testing_commands.py +++ b/dev/breeze/src/airflow_breeze/commands/testing_commands.py @@ -1552,6 +1552,7 @@ option_e2e_test_mode = click.option( "go_sdk", "openlineage", "ts_sdk", + "yaml_dag", ], case_sensitive=False, ), diff --git a/dev/breeze/src/airflow_breeze/utils/selective_checks.py b/dev/breeze/src/airflow_breeze/utils/selective_checks.py index 7bde4b0237b..fca38367b45 100644 --- a/dev/breeze/src/airflow_breeze/utils/selective_checks.py +++ b/dev/breeze/src/airflow_breeze/utils/selective_checks.py @@ -172,6 +172,7 @@ class FileGroupForCi(Enum): OPENLINEAGE_E2E_FILES = auto() OPENLINEAGE_E2E_COMPAT_FILES = auto() TS_SDK_E2E_FILES = auto() + YAML_DAG_E2E_FILES = auto() ALL_PYPROJECT_TOML_FILES = auto() ALL_PYTHON_FILES = auto() ALL_SOURCE_FILES = auto() @@ -342,6 +343,18 @@ CI_FILE_GROUP_MATCHES: HashableDict[FileGroupForCi] = HashableDict( r"^task-sdk/src/airflow/sdk/coordinators/node/.*", *LANG_SDK_NATIVE_DAG_PARSING_FILES, ], + FileGroupForCi.YAML_DAG_E2E_FILES: [ + r"^airflow-e2e-tests/tests/airflow_e2e_tests/yaml_dag_tests/.*", + r"^airflow-e2e-tests/yaml-dags/.*", + r"^airflow-e2e-tests/docker/yaml-dag\.yml$", + # The mode runs the Dag's one code task on the Go example bundle, so the handler + # it registers counts, but the rest of the Go SDK is the go_sdk mode's business. + r"^go-sdk/example/bundle/yamldag/.*", + r"^go-sdk/example/bundle/main\.go$", + r"^task-sdk/src/airflow/sdk/coordinators/_subprocess\.py$", + r"^task-sdk/src/airflow/sdk/coordinators/executable/.*", + *LANG_SDK_NATIVE_DAG_PARSING_FILES, + ], FileGroupForCi.PYTHON_PRODUCTION_FILES: [ # Production Python source the runtime ships — excludes tests, docs, # dev tooling, and generated files within those trees. Used by @@ -1475,6 +1488,11 @@ class SelectiveChecks: def run_go_sdk_e2e_tests(self) -> bool: return self._should_be_run(FileGroupForCi.GO_SDK_E2E_FILES) + @cached_property + @_scoped_to_component_branch + def run_yaml_dag_e2e_tests(self) -> bool: + return self._should_be_run(FileGroupForCi.YAML_DAG_E2E_FILES) + @cached_property @_scoped_to_component_branch def run_openlineage_e2e_tests(self) -> bool: @@ -1707,6 +1725,7 @@ class SelectiveChecks: or self.run_openlineage_e2e_tests or self.run_openlineage_e2e_compat_tests or self.run_ts_sdk_e2e_tests + or self.run_yaml_dag_e2e_tests or self.run_ui_e2e_tests ) diff --git a/dev/breeze/src/airflow_breeze/utils/verification_plan.py b/dev/breeze/src/airflow_breeze/utils/verification_plan.py index 808bed1f7ff..1ecd9f11c3d 100644 --- a/dev/breeze/src/airflow_breeze/utils/verification_plan.py +++ b/dev/breeze/src/airflow_breeze/utils/verification_plan.py @@ -93,6 +93,7 @@ NOT_RUNNABLE_LOCALLY: frozenset[str] = frozenset( "run_openlineage_e2e_tests", "run_openlineage_e2e_compat_tests", "run_ts_sdk_e2e_tests", + "run_yaml_dag_e2e_tests", "run_generic_prod_image_tests", "run_generic_basic_tests", } diff --git a/dev/breeze/tests/test_selective_checks.py b/dev/breeze/tests/test_selective_checks.py index bf4e2370a39..fd36a7c4975 100644 --- a/dev/breeze/tests/test_selective_checks.py +++ b/dev/breeze/tests/test_selective_checks.py @@ -1795,6 +1795,46 @@ def assert_outputs_are_printed(expected_outputs: dict[str, str], stderr: str): }, id="Run go e2e tests when ExecutableCoordinator changes", ), + pytest.param( + ("airflow-e2e-tests/yaml-dags/yaml_s3_roundtrip.yaml",), + { + "run-yaml-dag-e2e-tests": "true", + "run-go-sdk-e2e-tests": "false", + "prod-image-build": "true", + }, + id="Run only the YAML Dag e2e tests when the YAML Dag fixture changes", + ), + pytest.param( + ("airflow-e2e-tests/tests/airflow_e2e_tests/yaml_dag_tests/test_yaml_dag.py",), + { + "run-yaml-dag-e2e-tests": "true", + "run-go-sdk-e2e-tests": "false", + }, + id="Run only the YAML Dag e2e tests when its suite changes", + ), + pytest.param( + ("go-sdk/example/bundle/yamldag/handlers.go",), + { + "run-yaml-dag-e2e-tests": "true", + "run-go-sdk-e2e-tests": "true", + }, + id="Run the YAML Dag e2e tests alongside go's when the Dag's Go handler changes", + ), + pytest.param( + ("task-sdk/src/airflow/sdk/importers/yaml_importer/builder.py",), + { + "run-yaml-dag-e2e-tests": "true", + }, + id="Run the YAML Dag e2e tests when the YAML importer changes", + ), + pytest.param( + ("go-sdk/sdk/variable.go",), + { + "run-yaml-dag-e2e-tests": "false", + "run-go-sdk-e2e-tests": "true", + }, + id="Skip the YAML Dag e2e tests for a go-sdk change outside the Dag's handler", + ), pytest.param( ("providers/openlineage/src/airflow/providers/openlineage/plugins/adapter.py",), { diff --git a/go-sdk/cmd/airflow-go-pack/pack_integration_test.go b/go-sdk/cmd/airflow-go-pack/pack_integration_test.go index 019fc6c3fd2..80c7e39ab27 100644 --- a/go-sdk/cmd/airflow-go-pack/pack_integration_test.go +++ b/go-sdk/cmd/airflow-go-pack/pack_integration_test.go @@ -181,6 +181,9 @@ dags: "variable_write_dag": tasks: - "write_and_delete_variable" + "yaml_s3_roundtrip": + tasks: + - "summarize" ` assert.Equal(t, expectedManifest, string(metadata)) diff --git a/go-sdk/example/bundle/main.go b/go-sdk/example/bundle/main.go index 25e74681cd2..5bca037f7f1 100644 --- a/go-sdk/example/bundle/main.go +++ b/go-sdk/example/bundle/main.go @@ -29,6 +29,7 @@ import ( "github.com/apache/airflow/go-sdk/example/bundle/taskflowbinding" "github.com/apache/airflow/go-sdk/example/bundle/taskstate" "github.com/apache/airflow/go-sdk/example/bundle/variablewrite" + "github.com/apache/airflow/go-sdk/example/bundle/yamldag" ) func main() { @@ -86,6 +87,10 @@ func main() { ), airflow.TaskHandler("task_state_dag", "roundtrip_task_state", taskstate.RoundtripTaskState), + + // The Dag this one belongs to is written in YAML, not in Python; a handler + // is registered for it the same way. + airflow.TaskHandler("yaml_s3_roundtrip", "summarize", yamldag.Summarize), ) if err := bundle.Serve(); err != nil { diff --git a/go-sdk/example/bundle/yamldag/handlers.go b/go-sdk/example/bundle/yamldag/handlers.go new file mode 100644 index 00000000000..50dc2c096a6 --- /dev/null +++ b/go-sdk/example/bundle/yamldag/handlers.go @@ -0,0 +1,43 @@ +// 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. + +// Package yamldag holds the handler behind a "run:" task of a Dag written in +// YAML. The Dag declares the task, its arguments and the queue; this package +// supplies the code, exactly as it would for a @task.stub written in Python. +package yamldag + +import ( + "fmt" + + "github.com/apache/airflow/go-sdk/airflow" +) + +// Summarize is the body of the "summarize" task of the yaml_s3_roundtrip Dag. +// +// Both arguments come from the Dag's "run:" mapping: label is a literal, and +// keys is the return value of the upstream S3 listing, bound through {$x: ...}. +func Summarize(actx airflow.Context, keys []string, label string) (any, error) { + if len(keys) == 0 { + return nil, fmt.Errorf("the upstream listing bound no keys (label %q)", label) + } + actx.Logger().InfoContext(actx, "Summarizing the YAML Dag's listing", + "label", label, + "count", len(keys), + "first_key", keys[0], + ) + return map[string]any{"label": label, "count": len(keys), "first_key": keys[0]}, nil +}
