From d6d0d71c5125871fe3a001882aa37629388021be Mon Sep 17 00:00:00 2001 From: Kacper Muda Date: Fri, 3 Jul 2026 12:26:38 +0200 Subject: [PATCH 1/2] Add providers E2E tests framework and OpenLineage tests --- .dockerignore | 1 + .../workflows/additional-prod-image-tests.yml | 37 +++ .github/workflows/ci-amd.yml | 6 +- .github/workflows/ci-arm.yml | 6 +- .github/workflows/providers-e2e-tests.yml | 157 +++++++++++ .rat-excludes | 1 + dev/breeze/doc/05_test_commands.rst | 21 ++ dev/breeze/doc/ci/04_selective_checks.md | 5 + ...utput_setup_check-all-params-in-groups.svg | 2 +- ...utput_setup_check-all-params-in-groups.txt | 2 +- ...output_setup_regenerate-command-images.svg | 8 +- ...output_setup_regenerate-command-images.txt | 2 +- dev/breeze/doc/images/output_testing.svg | 18 +- dev/breeze/doc/images/output_testing.txt | 2 +- .../output_testing_providers-e2e-tests.svg | 182 +++++++++++++ .../output_testing_providers-e2e-tests.txt | 1 + .../commands/testing_commands.py | 168 ++++++++++++ .../commands/testing_commands_config.py | 20 ++ .../src/airflow_breeze/utils/run_tests.py | 6 + .../airflow_breeze/utils/selective_checks.py | 13 + dev/breeze/tests/test_selective_checks.py | 36 ++- providers-e2e-tests/README.md | 111 ++++++++ providers-e2e-tests/openlineage/.gitignore | 7 + .../openlineage/.pre-commit-config.yaml | 33 +++ providers-e2e-tests/openlineage/Dockerfile | 28 ++ providers-e2e-tests/openlineage/README.md | 75 ++++++ .../dags_extra/e2e_versioned_bundle.py | 62 +++++ .../dags_extra/openlineage_warmup_dag.py | 40 +++ .../openlineage/docker-compose-local.yaml | 56 ++++ .../openlineage/docker-compose.yaml | 181 +++++++++++++ .../openlineage/prepare_dags.py | 121 +++++++++ .../openlineage/pyproject.toml | 85 ++++++ .../openlineage/tests/conftest.py | 194 ++++++++++++++ .../openlineage/tests/constants.py | 42 +++ .../openlineage/tests/harness.py | 246 ++++++++++++++++++ .../openlineage/tests/task_logs.py | 111 ++++++++ .../openlineage/tests/test_openlineage_e2e.py | 72 +++++ .../tests/system/openlineage/constants.py | 27 ++ .../example_openlineage_all_facets_dag.py | 2 + .../example_openlineage_base_complex_dag.py | 2 + .../example_openlineage_base_simple_dag.py | 2 + ...openlineage_custom_operator_failure_dag.py | 2 + ...nlineage_custom_operator_ol_methods_dag.py | 2 + .../example_openlineage_defer_simple_dag.py | 10 +- .../example_openlineage_docs_file_dag.py | 2 + .../example_openlineage_edge_labels_dag.py | 2 + .../example_openlineage_hitl_dag.py | 2 + .../example_openlineage_manual_lineage_dag.py | 2 + .../example_openlineage_mapped_simple_dag.py | 2 + ...mple_openlineage_policy_conf_locked_dag.py | 2 + ...openlineage_policy_conf_source_code_dag.py | 2 + ...e_openlineage_policy_dag_emit_false_dag.py | 2 + ...openlineage_policy_dag_events_false_dag.py | 2 + ...penlineage_policy_dag_override_task_dag.py | 2 + ...openlineage_policy_extract_metadata_dag.py | 2 + ...e_openlineage_policy_full_task_info_dag.py | 2 + ...ple_openlineage_policy_hook_lineage_dag.py | 2 + ...mple_openlineage_policy_source_code_dag.py | 2 + ..._openlineage_policy_task_emit_false_dag.py | 2 + ..._openlineage_schedule_asset_or_time_dag.py | 2 + .../example_openlineage_schedule_cron_dag.py | 2 + ...age_schedule_cron_trigger_timetable_dag.py | 2 + ...ge_schedule_delta_trigger_timetable_dag.py | 2 + ...ineage_schedule_list_complex_assets_dag.py | 2 + ...neage_schedule_list_multiple_assets_dag.py | 2 + ...nlineage_schedule_list_single_asset_dag.py | 2 + ...penlineage_schedule_multiple_assets_dag.py | 2 + ...ule_multiple_cron_trigger_timetable_dag.py | 2 + ...e_openlineage_schedule_single_asset_dag.py | 2 + ...mple_openlineage_schedule_timedelta_dag.py | 2 + ...mple_openlineage_schedule_timetable_dag.py | 2 + .../example_openlineage_setup_teardown_dag.py | 2 + .../example_openlineage_short_circuit_dag.py | 2 + .../example_openlineage_task_groups_dag.py | 2 + ...example_openlineage_taskflow_simple_dag.py | 9 +- .../example_openlineage_trigger_dag.py | 5 +- ...mple_openlineage_trigger_dag_deferrable.py | 4 + .../example_openlineage_trigger_failed_dag.py | 3 + .../example_openlineage_versioned_dag.py | 2 + ...enlineage_custom_operator_failure_dag.json | 2 +- pyproject.toml | 7 + ...py_full_dist_local_venv_or_breeze_in_ci.py | 4 + uv.lock | 29 +++ 83 files changed, 2312 insertions(+), 18 deletions(-) create mode 100644 .github/workflows/providers-e2e-tests.yml create mode 100644 dev/breeze/doc/images/output_testing_providers-e2e-tests.svg create mode 100644 dev/breeze/doc/images/output_testing_providers-e2e-tests.txt create mode 100644 providers-e2e-tests/README.md create mode 100644 providers-e2e-tests/openlineage/.gitignore create mode 100644 providers-e2e-tests/openlineage/.pre-commit-config.yaml create mode 100644 providers-e2e-tests/openlineage/Dockerfile create mode 100644 providers-e2e-tests/openlineage/README.md create mode 100644 providers-e2e-tests/openlineage/dags_extra/e2e_versioned_bundle.py create mode 100644 providers-e2e-tests/openlineage/dags_extra/openlineage_warmup_dag.py create mode 100644 providers-e2e-tests/openlineage/docker-compose-local.yaml create mode 100644 providers-e2e-tests/openlineage/docker-compose.yaml create mode 100644 providers-e2e-tests/openlineage/prepare_dags.py create mode 100644 providers-e2e-tests/openlineage/pyproject.toml create mode 100644 providers-e2e-tests/openlineage/tests/conftest.py create mode 100644 providers-e2e-tests/openlineage/tests/constants.py create mode 100644 providers-e2e-tests/openlineage/tests/harness.py create mode 100644 providers-e2e-tests/openlineage/tests/task_logs.py create mode 100644 providers-e2e-tests/openlineage/tests/test_openlineage_e2e.py create mode 100644 providers/openlineage/tests/system/openlineage/constants.py diff --git a/.dockerignore b/.dockerignore index 401e8fff9fb40..a7e55a8f4e78e 100644 --- a/.dockerignore +++ b/.dockerignore @@ -51,6 +51,7 @@ !airflow-ctl-tests !shared/ !airflow-e2e-tests +!providers-e2e-tests # Add scripts so that we can use them inside the container !scripts diff --git a/.github/workflows/additional-prod-image-tests.yml b/.github/workflows/additional-prod-image-tests.yml index bfb21345c4d20..27861ccce5165 100644 --- a/.github/workflows/additional-prod-image-tests.yml +++ b/.github/workflows/additional-prod-image-tests.yml @@ -60,6 +60,14 @@ on: # yamllint disable-line rule:truthy description: "Whether to run Go SDK e2e tests (true/false)" required: true type: string + run-providers-e2e-tests-openlineage: + description: "Whether to run OpenLineage e2e tests (true/false)" + required: true + type: string + providers-compatibility-tests-matrix: + description: "JSON matrix of released Airflow versions for provider e2e compat tests." + required: true + type: string constraints-branch: description: "Branch used to construct constraints URL from." required: true @@ -321,6 +329,35 @@ jobs: e2e_test_mode: "go_sdk" if: inputs.canary-run == 'true' || inputs.run-go-sdk-e2e-tests == 'true' + test-providers-e2e-tests-openlineage: + name: "Provider E2E tests openlineage" + uses: ./.github/workflows/providers-e2e-tests.yml + with: + provider: "openlineage" + provider-display-name: "OpenLineage" + runners: ${{ inputs.runners }} + platform: ${{ inputs.platform }} + default-python-version: "${{ inputs.default-python-version }}" + use-uv: ${{ inputs.use-uv }} + if: inputs.canary-run == 'true' || inputs.run-providers-e2e-tests-openlineage == 'true' + + test-providers-e2e-tests-openlineage-compat: + name: "Provider E2E tests openlineage compat" + strategy: + fail-fast: false + matrix: + compat: ${{ fromJSON(inputs.providers-compatibility-tests-matrix) }} + uses: ./.github/workflows/providers-e2e-tests.yml + with: + provider: "openlineage" + provider-display-name: "OpenLineage" + runners: ${{ inputs.runners }} + platform: ${{ inputs.platform }} + default-python-version: "${{ inputs.default-python-version }}" + use-uv: ${{ inputs.use-uv }} + airflow-version: ${{ matrix.compat.airflow-version }} + if: inputs.canary-run == 'true' || inputs.run-providers-e2e-tests-openlineage == 'true' + test-ui-e2e-chromium: name: "Chromium UI e2e tests with PROD image" uses: ./.github/workflows/ui-e2e-tests.yml diff --git a/.github/workflows/ci-amd.yml b/.github/workflows/ci-amd.yml index ab0a1a7ff606f..f404b67c35d43 100644 --- a/.github/workflows/ci-amd.yml +++ b/.github/workflows/ci-amd.yml @@ -141,6 +141,7 @@ jobs: run-event-driven-e2e-tests: ${{ steps.selective-checks.outputs.run-event-driven-e2e-tests }} run-java-sdk-e2e-tests: ${{ steps.selective-checks.outputs.run-java-sdk-e2e-tests }} run-go-sdk-e2e-tests: ${{ steps.selective-checks.outputs.run-go-sdk-e2e-tests }} + run-providers-e2e-tests-openlineage: ${{ steps.selective-checks.outputs.run-providers-e2e-tests-openlineage }} 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 }} @@ -435,7 +436,7 @@ jobs: DEFAULT_BRANCH: ${{ needs.build-info.outputs.default-branch }} providers: - name: "provider distributions tests" + name: "Provider distributions tests" uses: ./.github/workflows/test-providers.yml needs: [build-info, build-ci-images] permissions: @@ -913,6 +914,9 @@ jobs: run-event-driven-e2e-tests: ${{ needs.build-info.outputs.run-event-driven-e2e-tests }} run-java-sdk-e2e-tests: ${{ needs.build-info.outputs.run-java-sdk-e2e-tests }} run-go-sdk-e2e-tests: ${{ needs.build-info.outputs.run-go-sdk-e2e-tests }} + run-providers-e2e-tests-openlineage: ${{ needs.build-info.outputs.run-providers-e2e-tests-openlineage }} + providers-compatibility-tests-matrix: > + ${{ needs.build-info.outputs.providers-compatibility-tests-matrix }} 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 354155192a2d5..4f7dd7b54c4d0 100644 --- a/.github/workflows/ci-arm.yml +++ b/.github/workflows/ci-arm.yml @@ -134,6 +134,7 @@ jobs: run-event-driven-e2e-tests: ${{ steps.selective-checks.outputs.run-event-driven-e2e-tests }} run-java-sdk-e2e-tests: ${{ steps.selective-checks.outputs.run-java-sdk-e2e-tests }} run-go-sdk-e2e-tests: ${{ steps.selective-checks.outputs.run-go-sdk-e2e-tests }} + run-providers-e2e-tests-openlineage: ${{ steps.selective-checks.outputs.run-providers-e2e-tests-openlineage }} 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 }} @@ -428,7 +429,7 @@ jobs: DEFAULT_BRANCH: ${{ needs.build-info.outputs.default-branch }} providers: - name: "provider distributions tests" + name: "Provider distributions tests" uses: ./.github/workflows/test-providers.yml needs: [build-info, build-ci-images] permissions: @@ -906,6 +907,9 @@ jobs: run-event-driven-e2e-tests: ${{ needs.build-info.outputs.run-event-driven-e2e-tests }} run-java-sdk-e2e-tests: ${{ needs.build-info.outputs.run-java-sdk-e2e-tests }} run-go-sdk-e2e-tests: ${{ needs.build-info.outputs.run-go-sdk-e2e-tests }} + run-providers-e2e-tests-openlineage: ${{ needs.build-info.outputs.run-providers-e2e-tests-openlineage }} + providers-compatibility-tests-matrix: > + ${{ needs.build-info.outputs.providers-compatibility-tests-matrix }} 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/providers-e2e-tests.yml b/.github/workflows/providers-e2e-tests.yml new file mode 100644 index 0000000000000..1ed6d77eb754e --- /dev/null +++ b/.github/workflows/providers-e2e-tests.yml @@ -0,0 +1,157 @@ +# 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. +# +--- + +name: Providers E2E Tests + +permissions: + contents: read +on: # yamllint disable-line rule:truthy + workflow_dispatch: + inputs: + provider: + description: "The provider whose e2e suite to run (a directory under providers-e2e-tests/)." + type: string + default: 'openlineage' + provider-display-name: + description: "Display name for the provider, used only in the job name." + type: string + default: 'OpenLineage' + runners: + description: "The array of labels (in json form) determining runners." + type: string + default: '["ubuntu-24.04"]' + platform: + description: "Platform for the build - 'linux/amd64' or 'linux/arm64'" + type: string + default: 'linux/amd64' + default-python-version: + description: "Which version of python should be used by default" + type: string + default: '3.10' + use-uv: + description: "Whether to use uv to build the image (true/false)" + type: string + default: 'true' + docker-image-tag: + description: "Tag of the Docker image to test" + type: string + required: true + airflow-version: + description: "Released Airflow version to test with current providers (empty = default PROD image)" + type: string + default: "" + + workflow_call: + inputs: + provider: + description: "The provider whose e2e suite to run (a directory under providers-e2e-tests/)." + type: string + default: 'openlineage' + provider-display-name: + description: "Display name for the provider, used only in the job name." + type: string + default: 'OpenLineage' + runners: + description: "The array of labels (in json form) determining runners." + required: true + type: string + platform: + description: "Platform for the build - 'linux/amd64' or 'linux/arm64'" + required: true + type: string + default-python-version: + description: "Which version of python should be used by default" + required: true + type: string + use-uv: + description: "Whether to use uv to build the image (true/false)" + required: true + type: string + docker-image-tag: + description: "Tag of the Docker image to test" + type: string + default: "" + airflow-version: + description: "Released Airflow version to test with current providers (empty = default PROD image)" + type: string + default: "" + +jobs: + test-providers-e2e-tests: + timeout-minutes: 60 + name: "Provider E2E tests \ + ${{ inputs.provider-display-name }}\ + ${{ inputs.airflow-version != '' && format(' Compat {0}', inputs.airflow-version) || '' }}" + runs-on: ${{ fromJSON(inputs.runners) }} + env: + PYTHON_MAJOR_MINOR_VERSION: "${{ inputs.default-python-version }}" + PROVIDER: ${{ inputs.provider }} + GITHUB_REPOSITORY: ${{ github.repository }} + GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} + GITHUB_USERNAME: ${{ github.actor }} + VERBOSE: "true" + steps: + # Gating on job-level `if:` instead of per-step would skip the whole job — and GitHub Actions + # does not evaluate a dynamic `name:` (the one above, using `inputs.*`) for a skipped job, so it + # renders the raw unevaluated template instead. Gating each step keeps the job itself "running" + # (and its name properly resolved), ending in success with every real step skipped for AF2. + - name: "Skip notice for Airflow 2.x" + if: ${{ startsWith(inputs.airflow-version, '2.') }} + run: echo "Skipping $PROVIDER e2e tests — not supported against Airflow 2.x." + - name: "Cleanup repo" + if: ${{ !startsWith(inputs.airflow-version, '2.') }} + shell: bash + run: sudo rm -rf ${GITHUB_WORKSPACE}/* + - name: "Checkout ${{ github.ref }} ( ${{ github.sha }} )" + if: ${{ !startsWith(inputs.airflow-version, '2.') }} + uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6.0.3 + with: + fetch-depth: 2 + persist-credentials: false + - name: "Prepare breeze & PROD image: ${{ env.PYTHON_MAJOR_MINOR_VERSION }}" + if: ${{ !startsWith(inputs.airflow-version, '2.') }} + uses: ./.github/actions/prepare_breeze_and_image + with: + platform: ${{ inputs.platform }} + image-type: "prod" + python: ${{ env.PYTHON_MAJOR_MINOR_VERSION }} + use-uv: ${{ inputs.use-uv }} + make-mnt-writeable-and-cleanup: true + id: breeze + - name: "Run ${{ inputs.provider }} e2e tests" + if: ${{ !startsWith(inputs.airflow-version, '2.') }} + run: > + breeze testing providers-e2e-tests "$PROVIDER" + ${AIRFLOW_VERSION:+--airflow-version ${AIRFLOW_VERSION}} + --skip-mounting-local-volumes + env: + DOCKER_IMAGE: "${{ inputs.docker-image-tag }}" + AIRFLOW_VERSION: "${{ inputs.airflow-version }}" + - name: Zip logs + if: ${{ always() && !startsWith(inputs.airflow-version, '2.') }} + run: | + cd "./providers-e2e-tests/$PROVIDER" && zip -r logs.zip logs + - name: "Upload logs" + if: ${{ always() && !startsWith(inputs.airflow-version, '2.') }} + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + with: + name: "${{ inputs.provider }}-e2e-test-logs-${{ inputs.airflow-version || 'prod' }}" + path: './providers-e2e-tests/${{ inputs.provider }}/logs.zip' + retention-days: '7' + if-no-files-found: 'error' diff --git a/.rat-excludes b/.rat-excludes index f0cd2910f5b27..0510ce24783f9 100644 --- a/.rat-excludes +++ b/.rat-excludes @@ -45,6 +45,7 @@ venv clients/* files/* dags/* +providers-e2e-tests/openlineage/dags/dag_doc.md generated/* .gitmodules prod_image_installed_providers.txt diff --git a/dev/breeze/doc/05_test_commands.rst b/dev/breeze/doc/05_test_commands.rst index df3cfbb970f8f..52256c930fbfb 100644 --- a/dev/breeze/doc/05_test_commands.rst +++ b/dev/breeze/doc/05_test_commands.rst @@ -373,6 +373,27 @@ 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. +Running provider E2E tests +.......................... + +You can use Breeze to run provider end-to-end tests, run against a real built Airflow image +(``$DOCKER_IMAGE``) rather than the in-process ``dag.test()`` path used by provider system tests. +What "e2e" means, and how the stack under test gets stood up, is entirely up to the provider — the +OpenLineage suite, for example, deploys a real Airflow via docker-compose and runs its system-test +DAGs against it, but a different provider could drive it a completely different way. Those tests run +using the Production image by default; pass ``--airflow-version`` to run against an older released +Airflow version with the current providers installed from main. + +.. image:: ./images/output_testing_providers-e2e-tests.svg + :target: https://raw.githubusercontent.com/apache/airflow/main/dev/breeze/images/output_testing_providers-e2e-tests.svg + :width: 100% + :alt: Breeze testing providers-e2e-tests + +The provider E2E tests are in the ``providers-e2e-tests/`` folder in the main repo; see +``providers-e2e-tests/README.md`` for how to run any provider's suite, how it's wired into CI, and +what's involved in adding a new provider (``providers-e2e-tests/openlineage/`` is one +example implementation). + Running Airflow UI E2E tests ............................. diff --git a/dev/breeze/doc/ci/04_selective_checks.md b/dev/breeze/doc/ci/04_selective_checks.md index 13a7e39f66f62..bc131cdd066b7 100644 --- a/dev/breeze/doc/ci/04_selective_checks.md +++ b/dev/breeze/doc/ci/04_selective_checks.md @@ -427,6 +427,11 @@ together using `pytest-xdist` (pytest-xdist distributes the tests among parallel of affected providers (but not recursively - only direct dependencies are added) * if there are any changes to "common" provider code not belonging to any provider (usually system tests or tests), then tests for all Providers are run +* `OpenLineage E2E tests` (the deployed-stack tests under `providers-e2e-tests/openlineage`, exposed as + the `run-providers-e2e-tests-openlineage` output) run when the `openlineage` or `common` providers or the + `providers-e2e-tests` harness change — and always on `canary` runs (where `full tests needed` also + covers core/task-sdk changes). Like the other deployed e2e suites, enabling them forces + `PROD Image building`. * The specific unit test type is enabled only if changed files match the expected patterns for each type (`API`, `CLI`, `WWW`, `Providers` etc.). The `Always` test type is added always if any unit tests are run. `Providers` tests are removed if current branch is different than `main` diff --git a/dev/breeze/doc/images/output_setup_check-all-params-in-groups.svg b/dev/breeze/doc/images/output_setup_check-all-params-in-groups.svg index 202eba652be36..22be2e9030b39 100644 --- a/dev/breeze/doc/images/output_setup_check-all-params-in-groups.svg +++ b/dev/breeze/doc/images/output_setup_check-all-params-in-groups.svg @@ -240,7 +240,7 @@ │setup:config | setup:regenerate-command-images | setup:self-upgrade | setup:synchronize-local-mounts | │ │setup:version | shell | start-airflow | testing | testing:airflow-ctl-integration-tests | │ │testing:airflow-ctl-tests | testing:airflow-e2e-tests | testing:core-integration-tests | │ -│testing:core-tests | testing:docker-compose-tests | testing:helm-tests | │ +│testing:core-tests | testing:docker-compose-tests | testing:helm-tests | testing:providers-e2e-tests | │ │testing:providers-integration-tests | testing:providers-tests | testing:python-api-client-tests | │ │testing:system-tests | testing:task-sdk-integration-tests | testing:task-sdk-tests | testing:ui-e2e-tests│ │| ui | ui:check-translation-completeness | ui:compile-assets | workflow-run | workflow-run:publish-docs)│ diff --git a/dev/breeze/doc/images/output_setup_check-all-params-in-groups.txt b/dev/breeze/doc/images/output_setup_check-all-params-in-groups.txt index 92e75f17fd456..ef361b8db62cf 100644 --- a/dev/breeze/doc/images/output_setup_check-all-params-in-groups.txt +++ b/dev/breeze/doc/images/output_setup_check-all-params-in-groups.txt @@ -1 +1 @@ -fc5f65d76377357873c6b3ea200f2ec8 +ebfe38ba1a4b30b85d55382436ef1834 diff --git a/dev/breeze/doc/images/output_setup_regenerate-command-images.svg b/dev/breeze/doc/images/output_setup_regenerate-command-images.svg index de56522e8180b..88faae5eb3574 100644 --- a/dev/breeze/doc/images/output_setup_regenerate-command-images.svg +++ b/dev/breeze/doc/images/output_setup_regenerate-command-images.svg @@ -258,10 +258,10 @@ │| setup:synchronize-local-mounts | setup:version | shell | start-airflow | testing | │ │testing:airflow-ctl-integration-tests | testing:airflow-ctl-tests | testing:airflow-e2e-tests | │ │testing:core-integration-tests | testing:core-tests | testing:docker-compose-tests | │ -│testing:helm-tests | testing:providers-integration-tests | testing:providers-tests | │ -│testing:python-api-client-tests | testing:system-tests | testing:task-sdk-integration-tests | │ -│testing:task-sdk-tests | testing:ui-e2e-tests | ui | ui:check-translation-completeness | │ -│ui:compile-assets | workflow-run | workflow-run:publish-docs)│ +│testing:helm-tests | testing:providers-e2e-tests | testing:providers-integration-tests | │ +│testing:providers-tests | testing:python-api-client-tests | testing:system-tests | │ +│testing:task-sdk-integration-tests | testing:task-sdk-tests | testing:ui-e2e-tests | ui | │ +│ui:check-translation-completeness | ui:compile-assets | workflow-run | workflow-run:publish-docs)│ │--check-onlyOnly check if some images need to be regenerated. Return 0 if no need or 1 if needed. Cannot be used  │ │together with --command flag or --force.                                                              │ ╰──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯ diff --git a/dev/breeze/doc/images/output_setup_regenerate-command-images.txt b/dev/breeze/doc/images/output_setup_regenerate-command-images.txt index fe2e502097d32..81d6cfb046ac2 100644 --- a/dev/breeze/doc/images/output_setup_regenerate-command-images.txt +++ b/dev/breeze/doc/images/output_setup_regenerate-command-images.txt @@ -1 +1 @@ -141b4bc2fbf13eddd1301de1064bf9fe +021adb1e353dab626f2092d96b2e7b45 diff --git a/dev/breeze/doc/images/output_testing.svg b/dev/breeze/doc/images/output_testing.svg index 60192bbfea0d5..a05e3458ee1dd 100644 --- a/dev/breeze/doc/images/output_testing.svg +++ b/dev/breeze/doc/images/output_testing.svg @@ -1,4 +1,4 @@ - + ╭─ UI Tests ───────────────────────────────────────────────────────────────────────────────────────────────────────────╮ │ui-e2e-tests               Run UI End-to-End tests using Playwright.                                              │ ╰──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯ +╭─ Commands ───────────────────────────────────────────────────────────────────────────────────────────────────────────╮ +│providers-e2e-tests        Run provider end-to-end tests against a deployed Airflow stack.                        │ +╰──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯ diff --git a/dev/breeze/doc/images/output_testing.txt b/dev/breeze/doc/images/output_testing.txt index e80be8356c537..bee169a1b6781 100644 --- a/dev/breeze/doc/images/output_testing.txt +++ b/dev/breeze/doc/images/output_testing.txt @@ -1 +1 @@ -5c361e73595950fbdbe6d3aa8fd3855b +f7645849511cf80e1747ddb411a881f5 diff --git a/dev/breeze/doc/images/output_testing_providers-e2e-tests.svg b/dev/breeze/doc/images/output_testing_providers-e2e-tests.svg new file mode 100644 index 0000000000000..e727fbeb032e4 --- /dev/null +++ b/dev/breeze/doc/images/output_testing_providers-e2e-tests.svg @@ -0,0 +1,182 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + Command: testing providers-e2e-tests + + + + + + + + + + +Usage:breeze testing providers-e2e-tests[OPTIONS] {openlineage} [EXTRA_PYTEST_ARGS]... + +Run provider end-to-end tests against a deployed Airflow stack. +Defaults to the PROD image (current Airflow + providers from sources). Pass --airflow-version to run against an older  +released Airflow version with the current providers installed from main. + +╭─ Docker-compose tests flag ──────────────────────────────────────────────────────────────────────────────────────────╮ +│--image-name                  -nName of the image to verify (overrides --python). (TEXT)│ +│--python                      -pPython major/minor version used in Airflow image for images. [default: 3.10]│ +│(>3.10< | 3.11 | 3.12 | 3.13 | 3.14)│ +│--airflow-version             Run against a released Airflow version (e.g. 3.0.6) with current providers from │ +│main, instead of the default PROD image. Builds a lightweight image from        │ +│apache/airflow:<version>. (TEXT)│ +│--skip-docker-compose-deletionSkip deletion of docker-compose instance after the test│ +│--skip-mounting-local-volumes Skip mounting local volumes - useful when we do not want to iterate with        │ +│modified local files. For example, when the PROD image is built from packages   │ +│rather than from local sources.                                                 │ +│--down                        Shuts down the docker-compose setup without running any tests. Useful to make   │ +│sure to free resources.                                                         │ +╰──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯ +╭─ Common CI options ──────────────────────────────────────────────────────────────────────────────────────────────────╮ +│--include-success-outputsWhether to include outputs of successful runs (not shown by default).│ +│--github-repository      -gGitHub repository used to pull, push run images. [default: apache/airflow](TEXT)│ +╰──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯ +╭─ Common options ─────────────────────────────────────────────────────────────────────────────────────────────────────╮ +│--verbose-vPrint verbose information about performed steps.│ +│--dry-run-DIf dry-run is set, commands are only printed, not executed.│ +│--help   -hShow this message and exit.│ +╰──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯ + + + + diff --git a/dev/breeze/doc/images/output_testing_providers-e2e-tests.txt b/dev/breeze/doc/images/output_testing_providers-e2e-tests.txt new file mode 100644 index 0000000000000..806839eb74ae3 --- /dev/null +++ b/dev/breeze/doc/images/output_testing_providers-e2e-tests.txt @@ -0,0 +1 @@ +236c2c2140bb30b55a1c7a74f6044453 diff --git a/dev/breeze/src/airflow_breeze/commands/testing_commands.py b/dev/breeze/src/airflow_breeze/commands/testing_commands.py index c3c210d9e8c6b..fada767ddbcaa 100644 --- a/dev/breeze/src/airflow_breeze/commands/testing_commands.py +++ b/dev/breeze/src/airflow_breeze/commands/testing_commands.py @@ -115,6 +115,7 @@ ) from airflow_breeze.utils.path_utils import AIRFLOW_CTL_ROOT_PATH, FILES_PATH, cleanup_python_generated_files from airflow_breeze.utils.run_tests import ( + PROVIDERS_E2E_TESTS_ROOT_PATH, TASK_SDK_INTEGRATION_TESTS_ROOT_PATH, are_all_test_paths_excluded, file_name_from_test_type, @@ -922,6 +923,173 @@ def task_sdk_integration_tests( sys.exit(return_code) +def _available_e2e_providers() -> list[str]: + return sorted( + d.name + for d in PROVIDERS_E2E_TESTS_ROOT_PATH.iterdir() + if d.is_dir() and (d / "pyproject.toml").exists() + ) + + +def _build_providers_e2e_compat_image(provider: str, airflow_version: str, python: str) -> str: + """Build a lightweight image: released ``apache/airflow:`` + current providers from main. + + Replicates the provider-compatibility approach (current provider code on an older Airflow core) + without a full PROD image build — the released image is pulled and the providers are reinstalled + from wheels built from main. + + The list of providers to install is read from ``providers-e2e-tests//pyproject.toml`` + under ``[tool.e2e-tests] required-providers``. + """ + import shutil + from pathlib import Path + + try: + import tomllib + except ImportError: + import tomli as tomllib # type: ignore[no-redef] + + from airflow_breeze.utils.path_utils import AIRFLOW_ROOT_PATH + + provider_root = PROVIDERS_E2E_TESTS_ROOT_PATH / provider + pyproject_path = provider_root / "pyproject.toml" + required_providers: list[str] = tomllib.loads(pyproject_path.read_text())["tool"]["e2e-tests"][ + "required-providers" + ] + console_print(f"[info]Building provider wheels from main: {required_providers}[/]") + run_command( + [ + "breeze", + "release-management", + "prepare-provider-distributions", + *required_providers, + "--distribution-format", + "wheel", + "--skip-tag-check", + "--include-not-ready-providers", + ], + check=True, + cwd=AIRFLOW_ROOT_PATH, + ) + provider_dist = provider_root / "provider_dist" + if provider_dist.exists(): + shutil.rmtree(provider_dist) + provider_dist.mkdir(parents=True) + for provider_id in required_providers: + wheel_prefix = f"apache_airflow_providers_{provider_id.replace('.', '_')}-" + for wheel in (AIRFLOW_ROOT_PATH / "dist").glob(f"{wheel_prefix}*.whl"): + shutil.copy2(wheel, provider_dist / wheel.name) + + base_image = f"apache/airflow:{airflow_version}-python{python}" + image_name = f"{provider}-e2e/airflow:{airflow_version}-python{python}" + console_print(f"[info]Building {image_name} from {base_image}[/]") + run_command( + [ + "docker", + "build", + "--build-arg", + f"AIRFLOW_BASE_IMAGE={base_image}", + "-t", + image_name, + "-f", + (provider_root / "Dockerfile").as_posix(), + Path(provider_root).as_posix(), + ], + check=True, + ) + return image_name + + +@testing_group.command( + name="providers-e2e-tests", + context_settings=dict( + ignore_unknown_options=True, + allow_extra_args=True, + ), +) +@click.argument("provider", type=click.Choice(_available_e2e_providers())) +@option_python +@option_image_name +@click.option( + "--airflow-version", + help="Run against a released Airflow version (e.g. 3.0.6) with current providers from main, " + "instead of the default PROD image. Builds a lightweight image from apache/airflow:.", + default=None, +) +@option_skip_docker_compose_deletion +@option_skip_mounting_local_volumes +@click.option( + "--down", + help="Shuts down the docker-compose setup without running any tests. " + "Useful to make sure to free resources.", + is_flag=True, +) +@option_github_repository +@option_include_success_outputs +@option_verbose +@option_dry_run +@click.argument("extra_pytest_args", nargs=-1, type=click.Path(path_type=str)) +def providers_e2e_tests( + provider: str, + python: str, + image_name: str | None, + airflow_version: str | None, + down: bool, + skip_docker_compose_deletion: bool, + skip_mounting_local_volumes: bool, + github_repository: str, + include_success_outputs: bool, + extra_pytest_args: tuple, +): + """Run provider end-to-end tests against a deployed Airflow stack. + + Defaults to the PROD image (current Airflow + providers from sources). Pass --airflow-version to + run against an older released Airflow version with the current providers installed from main. + """ + perform_environment_checks() + + if airflow_version: + # prepare_dags.py reads this to drop DAGs that need a newer Airflow than the target. + os.environ["E2E_TARGET_AIRFLOW_VERSION"] = airflow_version + if image_name is None: + image_name = _build_providers_e2e_compat_image(provider, airflow_version, python) + else: + image_name = image_name or os.environ.get("DOCKER_IMAGE") + if not image_name: + build_params = BuildProdParams(python=python, github_repository=github_repository) + image_name = build_params.airflow_image_name + + if down: + env = { + **os.environ, + "DOCKER_IMAGE": image_name, + } + down_cmd = [ + "docker", + "compose", + "down", + "--remove-orphans", + "--volumes", + ] + console_print("[info]Running docker-compose down[/]") + run_command(down_cmd, output=None, check=False, env=env, cwd=PROVIDERS_E2E_TESTS_ROOT_PATH / provider) + sys.exit(0) + + console_print(f"[info]Running {provider} e2e tests with image: {image_name}[/]") + return_code, info = run_docker_compose_tests( + image_name=image_name, + python_version=python, + include_success_outputs=include_success_outputs, + extra_pytest_args=extra_pytest_args, + skip_docker_compose_deletion=skip_docker_compose_deletion, + skip_mounting_local_volumes=skip_mounting_local_volumes, + test_type="providers-e2e-tests", + provider=provider, + skip_image_check=bool(airflow_version), + ) + sys.exit(return_code) + + @testing_group.command( name="airflow-ctl-integration-tests", context_settings=dict( diff --git a/dev/breeze/src/airflow_breeze/commands/testing_commands_config.py b/dev/breeze/src/airflow_breeze/commands/testing_commands_config.py index 56a4b17450239..4e4c1601a7443 100644 --- a/dev/breeze/src/airflow_breeze/commands/testing_commands_config.py +++ b/dev/breeze/src/airflow_breeze/commands/testing_commands_config.py @@ -215,6 +215,26 @@ ], }, ], + "breeze testing providers-e2e-tests": [ + { + "name": "Docker-compose tests flag", + "options": [ + "--image-name", + "--python", + "--airflow-version", + "--skip-docker-compose-deletion", + "--skip-mounting-local-volumes", + "--down", + ], + }, + { + "name": "Common CI options", + "options": [ + "--include-success-outputs", + "--github-repository", + ], + }, + ], "breeze testing airflow-ctl-tests": [ { "name": "Test environment", diff --git a/dev/breeze/src/airflow_breeze/utils/run_tests.py b/dev/breeze/src/airflow_breeze/utils/run_tests.py index ef87dd1f6e3a7..d90b5f63dad93 100644 --- a/dev/breeze/src/airflow_breeze/utils/run_tests.py +++ b/dev/breeze/src/airflow_breeze/utils/run_tests.py @@ -49,6 +49,8 @@ AIRFLOW_E2E_TESTS_ROOT_PATH = AIRFLOW_ROOT_PATH / "airflow-e2e-tests" +PROVIDERS_E2E_TESTS_ROOT_PATH = AIRFLOW_ROOT_PATH / "providers-e2e-tests" + AIRFLOW_CTL_TESTS_ROOT_PATH = AIRFLOW_ROOT_PATH / "airflow-ctl-tests" IGNORE_DB_INIT_FOR_TEST_GROUPS = [ @@ -123,6 +125,7 @@ def run_docker_compose_tests( test_type: str = "docker-compose", skip_image_check: bool = False, test_mode: str = "basic", + provider: str = "", ) -> tuple[int, str]: if not skip_image_check: command_result = run_command( @@ -165,6 +168,9 @@ def run_docker_compose_tests( elif test_type == "airflow-e2e-tests": test_path = Path("tests") / "airflow_e2e_tests" / f"{test_mode}_tests" cwd = AIRFLOW_E2E_TESTS_ROOT_PATH.as_posix() + elif test_type == "providers-e2e-tests": + test_path = Path("tests") + cwd = (PROVIDERS_E2E_TESTS_ROOT_PATH / provider).as_posix() elif test_type == "airflow-ctl-integration": test_path = Path("tests") / "airflowctl_tests" cwd = AIRFLOW_CTL_TESTS_ROOT_PATH.as_posix() diff --git a/dev/breeze/src/airflow_breeze/utils/selective_checks.py b/dev/breeze/src/airflow_breeze/utils/selective_checks.py index 89baca951458e..1b70412b51402 100644 --- a/dev/breeze/src/airflow_breeze/utils/selective_checks.py +++ b/dev/breeze/src/airflow_breeze/utils/selective_checks.py @@ -132,6 +132,7 @@ class FileGroupForCi(Enum): EVENT_DRIVEN_E2E_FILES = auto() JAVA_SDK_E2E_FILES = auto() GO_SDK_E2E_FILES = auto() + PROVIDERS_E2E_OPENLINEAGE_FILES = auto() ALL_PYPROJECT_TOML_FILES = auto() ALL_PYTHON_FILES = auto() ALL_SOURCE_FILES = auto() @@ -250,6 +251,13 @@ def __hash__(self): r"^task-sdk/src/airflow/sdk/coordinators/_subprocess\.py$", r"^task-sdk/src/airflow/sdk/coordinators/executable/.*", ], + FileGroupForCi.PROVIDERS_E2E_OPENLINEAGE_FILES: [ + r"^providers-e2e-tests/openlineage/.*", + r"^providers/openlineage/.*", + r"^providers/common/compat/.*", + r"^providers/common/io/.*", + r"^providers/common/sql/.*", + ], FileGroupForCi.PYTHON_PRODUCTION_FILES: [ # Production Python source the runtime ships — excludes tests, docs, # dev tooling, and generated files within those trees. Used by @@ -1037,6 +1045,10 @@ def run_java_sdk_e2e_tests(self) -> bool: def run_go_sdk_e2e_tests(self) -> bool: return self._should_be_run(FileGroupForCi.GO_SDK_E2E_FILES) + @cached_property + def run_providers_e2e_tests_openlineage(self) -> bool: + return self._should_be_run(FileGroupForCi.PROVIDERS_E2E_OPENLINEAGE_FILES) + @cached_property def run_amazon_tests(self) -> bool: if self.providers_test_types_list_as_strings_in_json == "[]": @@ -1179,6 +1191,7 @@ def prod_image_build(self) -> bool: or self.run_event_driven_e2e_tests or self.run_java_sdk_e2e_tests or self.run_go_sdk_e2e_tests + or self.run_providers_e2e_tests_openlineage or self.run_ui_e2e_tests ) diff --git a/dev/breeze/tests/test_selective_checks.py b/dev/breeze/tests/test_selective_checks.py index 43a5e0081bccf..8a6dbd8f7bb25 100644 --- a/dev/breeze/tests/test_selective_checks.py +++ b/dev/breeze/tests/test_selective_checks.py @@ -1201,7 +1201,9 @@ def assert_outputs_are_printed(expected_outputs: dict[str, str], stderr: str): "python-versions": f"['{DEFAULT_PYTHON_MAJOR_MINOR_VERSION}']", "python-versions-list-as-string": DEFAULT_PYTHON_MAJOR_MINOR_VERSION, "ci-image-build": "true", - "prod-image-build": "false", + # common providers feed OpenLineage e2e tests, which need the PROD image. + "prod-image-build": "true", + "run-providers-e2e-tests-openlineage": "true", "run-helm-tests": "false", "run-unit-tests": "true", "run-amazon-tests": "false", @@ -1472,6 +1474,38 @@ def assert_outputs_are_printed(expected_outputs: dict[str, str], stderr: str): }, id="Run go e2e tests when ExecutableCoordinator changes", ), + pytest.param( + ("providers/openlineage/src/airflow/providers/openlineage/plugins/adapter.py",), + { + "run-providers-e2e-tests-openlineage": "true", + "prod-image-build": "true", + }, + id="Run OpenLineage e2e tests for openlineage provider change", + ), + pytest.param( + ("providers/common/compat/src/airflow/providers/common/compat/sdk.py",), + { + "run-providers-e2e-tests-openlineage": "true", + "prod-image-build": "true", + }, + id="Run OpenLineage e2e tests for common provider change", + ), + pytest.param( + ("providers-e2e-tests/openlineage/tests/openlineage_e2e_tests/test_openlineage_e2e.py",), + { + "run-providers-e2e-tests-openlineage": "true", + "prod-image-build": "true", + }, + id="Run OpenLineage e2e tests when the e2e harness changes", + ), + pytest.param( + ("providers/ftp/src/airflow/providers/ftp/hooks/ftp.py",), + { + "run-providers-e2e-tests-openlineage": "false", + "prod-image-build": "false", + }, + id="Do not run OpenLineage e2e tests for unrelated provider change", + ), ( pytest.param( ("devel-common/pyproject.toml",), diff --git a/providers-e2e-tests/README.md b/providers-e2e-tests/README.md new file mode 100644 index 0000000000000..588e13be45943 --- /dev/null +++ b/providers-e2e-tests/README.md @@ -0,0 +1,111 @@ + + +# Provider end-to-end tests + +Each subdirectory here is a self-contained e2e suite for one provider, that can be run against a real built +Airflow image (`$DOCKER_IMAGE`) rather than the in-process `dag.test()` path used by provider system tests. + +What "e2e" means, and how the stack under test gets stood up, is entirely up to the provider — see that provider's +own `README.md` for what its suite tests, how it works, and how to interpret its results. For an example implementation +(a docker-compose deployed stack), see [`openlineage`](openlineage) and its [README](openlineage/README.md). + +## Running locally + +Docker must be running. Every provider is driven by the same breeze command: + +```bash +# Default: PROD image (current Airflow from sources + providers from sources) +breeze testing providers-e2e-tests + +# An older released Airflow version, with the current providers installed from main +# (builds apache/airflow: + current provider wheels — no full PROD build) +breeze testing providers-e2e-tests --airflow-version 3.1.8 + +# Tear the stack down +breeze testing providers-e2e-tests --down +``` + +`` is any subdirectory of `providers-e2e-tests/` with a `pyproject.toml` — run +`breeze testing providers-e2e-tests --help` to see the current choices. + +Useful flags, passed to every provider's pytest run as env vars for its own harness to honor: + +- `--skip-docker-compose-deletion` (`SKIP_DOCKER_COMPOSE_DELETION`) — keep the stack up after the run to inspect it. +- `--skip-mounting-local-volumes` (`SKIP_MOUNTING_LOCAL_VOLUMES`) — run against the image only (what CI does). +Omit it during dev to mount and hot-reload local sources. +- Anything after the provider is passed through to pytest (e.g. `... -- -s`). + +The first run builds the PROD image; subsequent runs are much faster. +When running non-interactively without a pre-built image, prefix `ANSWER=yes` so the image-build prompt auto-confirms. + +## How this is wired into CI + +```text +ci-amd.yml / ci-arm.yml + └─ additional-prod-image-tests (job) ──uses──▶ additional-prod-image-tests.yml + └─ test-providers-e2e-tests-[-compat] (job) ──uses──▶ providers-e2e-tests.yml + └─ runs: breeze testing providers-e2e-tests [--airflow-version ] +``` + +`additional-prod-image-tests.yml` is a single reusable workflow called identically from both +`ci-amd.yml` and `ci-arm.yml`, so any job defined there — including each provider's e2e tests — +automatically gets both ARM and AMD coverage from one job definition. +[`scripts/ci/prek/check_ci_workflows_in_sync.py`](../scripts/ci/prek/check_ci_workflows_in_sync.py) +enforces that the two `ci-*.yml` files stay identical outside a small documented allowlist, so +don't add provider-e2e jobs directly to `ci-amd.yml`/`ci-arm.yml` — they belong in +`additional-prod-image-tests.yml`. + +Each provider's suite is gated by its own `run-providers-e2e-tests-` selective-checks output +(computed in `dev/breeze/src/airflow_breeze/utils/selective_checks.py` from the files changed in a +PR) — it always runs on canary/`main` runs, and on a PR only when that provider (or a file group it +depends on, e.g. `common` providers) changed. + +## Adding a new provider + +1. Create `providers-e2e-tests//`. The breeze harness (`breeze testing + providers-e2e-tests`, see `dev/breeze/src/airflow_breeze/commands/testing_commands.py` and + `run_docker_compose_tests` in `dev/breeze/src/airflow_breeze/utils/run_tests.py`) only requires: + - `pyproject.toml` — its presence is what makes `` a valid choice + (`_available_e2e_providers()`). It only needs a `[tool.e2e-tests] required-providers = [...]` + entry if the provider supports `--airflow-version` compat mode — that lists the provider + distributions to build from `main` for that mode (see `_build_providers_e2e_compat_image`). + - `tests/` — a pytest suite. Breeze runs `pytest tests/` from this directory with `$DOCKER_IMAGE` + (and the flags above) set as env vars; what it does with that image is entirely up to the + provider. + - `Dockerfile` — only needed if the provider supports `--airflow-version` compat mode: breeze + builds `apache/airflow:` + the provider wheels from `main` using it. +2. Wire the trigger into selective checks + (`dev/breeze/src/airflow_breeze/utils/selective_checks.py`): add a `FileGroupForCi` member and + its regex patterns, a `run_providers_e2e_tests_` cached property, and expose it as a + `run-providers-e2e-tests-` output — mirror the existing `PROVIDERS_E2E_OPENLINEAGE_FILES` / + `run_providers_e2e_tests_openlineage` pair. Update + [`dev/breeze/tests/test_selective_checks.py`](../dev/breeze/tests/test_selective_checks.py) and + [`dev/breeze/doc/ci/04_selective_checks.md`](../dev/breeze/doc/ci/04_selective_checks.md) in the + same PR. +3. Thread the new output through **both** `ci-amd.yml` and `ci-arm.yml` identically: add it to the + `build-info` job's `outputs:` and to the `additional-prod-image-tests` job's `with:` block. +4. Add the job(s) to `.github/workflows/additional-prod-image-tests.yml`, named + `test-providers-e2e-tests-` (and `test-providers-e2e-tests--compat` for the + compat matrix job, if any) — mirror the existing `test-providers-e2e-tests-openlineage` pair. + Call `providers-e2e-tests.yml` with `provider: ""` and + `provider-display-name: ""`, gated on + `inputs.canary-run == 'true' || inputs.run-providers-e2e-tests- == 'true'`. The compat + job additionally needs a `strategy.matrix` over `inputs.providers-compatibility-tests-matrix` and + `airflow-version: ${{ matrix.compat.airflow-version }}`. diff --git a/providers-e2e-tests/openlineage/.gitignore b/providers-e2e-tests/openlineage/.gitignore new file mode 100644 index 0000000000000..fdf48551bf27f --- /dev/null +++ b/providers-e2e-tests/openlineage/.gitignore @@ -0,0 +1,7 @@ +.env +*.iml +# DAGs are sourced from providers/openlineage/tests/system at runtime by prepare_dags.py +dags/ +logs/ +# Provider wheels built from main for older-version (compat) image builds +provider_dist/ diff --git a/providers-e2e-tests/openlineage/.pre-commit-config.yaml b/providers-e2e-tests/openlineage/.pre-commit-config.yaml new file mode 100644 index 0000000000000..85fd7040301b9 --- /dev/null +++ b/providers-e2e-tests/openlineage/.pre-commit-config.yaml @@ -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. +--- +default_stages: [pre-commit, pre-push] +minimum_prek_version: '0.3.4' +default_language_version: + python: python3 +repos: + - repo: local + hooks: + - id: mypy-openlineage-e2e-tests + name: Run mypy for providers-e2e-tests/openlineage + language: python + entry: >- + ../../scripts/ci/prek/run_mypy_full_dist_local_venv_or_breeze_in_ci.py + providers-e2e-tests/openlineage + pass_filenames: false + files: ^.*\.py$ + require_serial: true diff --git a/providers-e2e-tests/openlineage/Dockerfile b/providers-e2e-tests/openlineage/Dockerfile new file mode 100644 index 0000000000000..3c1d772cf4b06 --- /dev/null +++ b/providers-e2e-tests/openlineage/Dockerfile @@ -0,0 +1,28 @@ +# 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. + +# Lightweight image for running the OpenLineage e2e tests against an OLDER, released Airflow version. +# Instead of building a full PROD image, it starts from the published apache/airflow: image +# and reinstalls the providers under test (OpenLineage + the providers the system-test DAGs rely on) +# from wheels built from main — so the current provider code runs on an older Airflow core, the same +# idea as the provider compatibility tests. The default (prod) run does NOT use this Dockerfile; it +# uses the breeze PROD image with current sources. +ARG AIRFLOW_BASE_IMAGE +FROM ${AIRFLOW_BASE_IMAGE} + +COPY --chown=airflow:0 provider_dist/ /tmp/provider_dist/ +RUN pip install --no-cache-dir --upgrade /tmp/provider_dist/*.whl diff --git a/providers-e2e-tests/openlineage/README.md b/providers-e2e-tests/openlineage/README.md new file mode 100644 index 0000000000000..76c1bc5a56485 --- /dev/null +++ b/providers-e2e-tests/openlineage/README.md @@ -0,0 +1,75 @@ + + +# OpenLineage provider end-to-end tests + +These tests deploy a **real, multi-process Airflow** (via docker-compose) and run the OpenLineage +provider's system-test DAGs against it — exercising the integration the way it actually runs in a +deployment (separate scheduler / dag-processor / triggerer / api-server / workers), rather than the +in-process `dag.test()` path. + +See the [top-level README](../README.md) for how to run any provider's e2e suite with breeze, how +this is wired into CI, and what's involved in adding a new provider. + +## How it works + +- The DAGs, `OpenLineageTestOperator`, `VariableTransport`, and expected-event JSON are the ones in + [`providers/openlineage/tests/system/openlineage`](../../providers/openlineage/tests/system/openlineage) — + that directory stays the single source of truth. `prepare_dags.py` copies them into a (gitignored) + `dags/` folder at runtime and strips the pytest-only `get_test_run` footer so they parse in a + deployment. +- OpenLineage is configured to emit through the test `VariableTransport` (events land in Airflow + Variables); each DAG's terminal `OpenLineageTestOperator` task validates the emitted events against + the expected templates. **A DAG run that ends `success` means its lineage matched.** +- `harness.py` triggers every DAG through the REST API, waits for completion, retries a failed DAG + once, and asserts all expected DAGs succeeded. +- A small `MockVersionedLocalDagBundle` provides a versioned bundle (so `dag_bundle_version` is + populated) while still serving DAGs from the local dags folder. + +## Checking results + +- **pytest summary:** `1 passed` means every expected DAG is green. On failure it prints the + offenders, e.g. `... not successful: {'openlineage_docs_file_dag': 'failed'}`, and logs + `⚠ DAGs that passed only on retry (flaky first run): [...]`. +- **On failure, the relevant task logs are printed right in the pytest output** (`check_events` plus + any task that didn't succeed, filtered to warning/error lines — see `tests/task_logs.py`) — no + need to open the log artifact for the common case, e.g. the exact mismatch + (``Path `job > facets > ...`: expected X but got Y``). +- **Per-DAG logs** are still written in full to `logs/dag_id=/.../task_id=/attempt=1.log` + for deeper digging beyond what gets printed. +- With `--skip-docker-compose-deletion`, the stack stays up: get a token from + `curl -s http://localhost:8080/auth/token` and check `GET /api/v2/importErrors` for DAG parse + errors, or browse the runs in the UI at (login `airflow` / `airflow`). + +## Version-specific DAGs + +Some DAGs require a newer Airflow than the compat targets and are dropped for older versions in +`prepare_dags.py` (`MIN_AIRFLOW_VERSION_FOR_DAG`) — e.g. `example_openlineage_hitl_dag` needs 3.1+ +(its operators import-raise on older cores). Add an entry there when a new DAG is version-gated. + +## Layout + +```text +docker-compose.yaml # the deployed stack (postgres + apiserver + scheduler + dag-processor + triggerer) +docker-compose-local.yaml # overlay that mounts local sources (used unless --skip-mounting-local-volumes) +Dockerfile # lightweight image for --airflow-version: apache/airflow: + current providers +prepare_dags.py # sources DAGs from the provider system tests into dags/ at runtime +dags_extra/ # harness-only DAGs/modules (warmup DAG, versioned bundle) copied into dags/ +tests/ # conftest (compose lifecycle), harness, constants, and the pytest entrypoint +``` diff --git a/providers-e2e-tests/openlineage/dags_extra/e2e_versioned_bundle.py b/providers-e2e-tests/openlineage/dags_extra/e2e_versioned_bundle.py new file mode 100644 index 0000000000000..20d1c022e4582 --- /dev/null +++ b/providers-e2e-tests/openlineage/dags_extra/e2e_versioned_bundle.py @@ -0,0 +1,62 @@ +# 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 minimal *versioned* DAG bundle for the OpenLineage e2e deployment. + +example_openlineage_versioned_dag asserts that OpenLineage emits a non-null ``dag_bundle_version``, +which Airflow only populates when the bundle reports ``supports_versioning = True``. The real +versioned bundle (the git provider's GitDagBundle) serves DAGs from a separate bundle storage path, +which would break the ``PYTHONPATH=/opt/airflow/dags`` that the test ``VariableTransport`` class path +relies on. + +This bundle instead keeps serving DAGs straight from the local dags folder (so PYTHONPATH stays +valid) and reports a deterministic, content-derived version string — enough to exercise the +version-related OpenLineage attributes without standing up a real VCS-backed bundle. The module is +copied into the dags folder by prepare_dags.py so it is importable via PYTHONPATH. +""" + +from __future__ import annotations + +import hashlib +from typing import TYPE_CHECKING + +from airflow.dag_processing.bundles.local import LocalDagBundle + +if TYPE_CHECKING: + from airflow.dag_processing.bundles.base import BundleVersion + + +class MockVersionedLocalDagBundle(LocalDagBundle): + """LocalDagBundle that reports a content-derived version so ``dag_bundle_version`` is populated.""" + + supports_versioning = True + + # Intentionally broadens LocalDagBundle.get_current_version (which returns None, as the base + # bundle is non-versioned) to report a real version. + def get_current_version(self) -> str | BundleVersion: # type: ignore[override] + digest = hashlib.sha1(usedforsecurity=False) + for dag_file in sorted(self.path.rglob("*.py")): + digest.update(dag_file.relative_to(self.path).as_posix().encode()) + digest.update(str(dag_file.stat().st_size).encode()) + version = digest.hexdigest()[:12] + try: + # Preferred on newer cores; bare str is deprecated for versioned bundles. + from airflow.dag_processing.bundles.base import BundleVersion + + return BundleVersion(version=version) + except ImportError: + # Older Airflow cores (compat runs) predate BundleVersion — a bare string still works. + return version diff --git a/providers-e2e-tests/openlineage/dags_extra/openlineage_warmup_dag.py b/providers-e2e-tests/openlineage/dags_extra/openlineage_warmup_dag.py new file mode 100644 index 0000000000000..22ce6bb2a92fe --- /dev/null +++ b/providers-e2e-tests/openlineage/dags_extra/openlineage_warmup_dag.py @@ -0,0 +1,40 @@ +# 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. +""" +Warmup DAG triggered before all others to verify the worker is ready to pick up tasks. + +This DAG is specific to the deployed e2e harness (it is not part of the provider system +tests). It is copied into the generated dags folder by prepare_dags.py. +""" + +from __future__ import annotations + +from datetime import datetime + +from airflow import DAG +from airflow.providers.standard.operators.python import PythonOperator + +DAG_ID = "openlineage_warmup_dag" + +with DAG( + dag_id=DAG_ID, + start_date=datetime(2021, 1, 1), + schedule=None, + catchup=False, + default_args={"retries": 0}, +) as dag: + PythonOperator(task_id="noop", python_callable=lambda: None) diff --git a/providers-e2e-tests/openlineage/docker-compose-local.yaml b/providers-e2e-tests/openlineage/docker-compose-local.yaml new file mode 100644 index 0000000000000..70d543f809a96 --- /dev/null +++ b/providers-e2e-tests/openlineage/docker-compose-local.yaml @@ -0,0 +1,56 @@ +# 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. +# +--- +# If the image is locally built from the repo we can mount sources of Airflow and the +# OpenLineage provider so we do not have to rebuild the image on every change. +services: + airflow-init: + volumes: + - ../../airflow-core/src:/opt/airflow/airflow-core/src + - ../../task-sdk/src:/opt/airflow/task-sdk/src + - ../../providers/openlineage/src:/opt/airflow/providers/openlineage/src + environment: + - DEV_MODE=true + airflow-apiserver: + volumes: + - ../../airflow-core/src:/opt/airflow/airflow-core/src + - ../../task-sdk/src:/opt/airflow/task-sdk/src + - ../../providers/openlineage/src:/opt/airflow/providers/openlineage/src + environment: + - DEV_MODE=true + airflow-scheduler: + volumes: + - ../../airflow-core/src:/opt/airflow/airflow-core/src + - ../../task-sdk/src:/opt/airflow/task-sdk/src + - ../../providers/openlineage/src:/opt/airflow/providers/openlineage/src + environment: + - DEV_MODE=true + airflow-dag-processor: + volumes: + - ../../airflow-core/src:/opt/airflow/airflow-core/src + - ../../task-sdk/src:/opt/airflow/task-sdk/src + - ../../providers/openlineage/src:/opt/airflow/providers/openlineage/src + environment: + - DEV_MODE=true + airflow-triggerer: + volumes: + - ../../airflow-core/src:/opt/airflow/airflow-core/src + - ../../task-sdk/src:/opt/airflow/task-sdk/src + - ../../providers/openlineage/src:/opt/airflow/providers/openlineage/src + environment: + - DEV_MODE=true diff --git a/providers-e2e-tests/openlineage/docker-compose.yaml b/providers-e2e-tests/openlineage/docker-compose.yaml new file mode 100644 index 0000000000000..8adcc95e80884 --- /dev/null +++ b/providers-e2e-tests/openlineage/docker-compose.yaml @@ -0,0 +1,181 @@ +# 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. +# +--- +x-airflow-common: + &airflow-common + image: ${AIRFLOW_IMAGE_NAME} + env_file: + - ${ENV_FILE_PATH:-.env} + environment: + &airflow-common-env + AIRFLOW__CORE__EXECUTOR: LocalExecutor + # yamllint disable rule:line-length + AIRFLOW__CORE__AUTH_MANAGER: 'airflow.api_fastapi.auth.managers.simple.simple_auth_manager.SimpleAuthManager' + AIRFLOW__CORE__SIMPLE_AUTH_MANAGER_ALL_ADMINS: 'true' + AIRFLOW__DATABASE__SQL_ALCHEMY_CONN: postgresql+psycopg2://airflow:airflow@postgres/airflow + AIRFLOW__CORE__FERNET_KEY: ${FERNET_KEY} + AIRFLOW__CORE__LOAD_EXAMPLES: 'false' + AIRFLOW__CORE__DAGS_FOLDER: '/opt/airflow/dags' + AIRFLOW__CORE__EXECUTION_API_SERVER_URL: 'http://airflow-apiserver:8080/execution/' + AIRFLOW__API__BASE_URL: 'http://airflow-apiserver:8080/' + # Do not reuse this secret in any deployment — it exists only for this ephemeral CI/test stack. + AIRFLOW__API_AUTH__JWT_SECRET: 'test-secret-key-for-testing' + # The OpenLineage system-test DAGs (and the custom VariableTransport / OpenLineageTestOperator + # they import) live under /system/openlineage; the dags folder is added to PYTHONPATH so + # both DAG parsing and the dotted-path transport class resolve in every component. + PYTHONPATH: '/opt/airflow/dags' + # Route OpenLineage events to the test VariableTransport (stores events in Airflow Variables, + # read back and validated by the in-DAG OpenLineageTestOperator). The transport type is the + # fully-qualified class path; the OpenLineage client imports it directly. + AIRFLOW__OPENLINEAGE__TRANSPORT: '{"type": "system.openlineage.transport.variable.VariableTransport"}' + AIRFLOW__OPENLINEAGE__DISABLED: 'false' + # Deployment-wide policy required by the two conf-driven policy DAGs. Each rule is scoped by + # dag_id, so both rules coexist and do not affect any other DAG. + AIRFLOW__OPENLINEAGE__EMISSION_POLICY: '[{"scope": {"dag_id": "openlineage_policy_conf_locked_dag"}, "locked": true, "controls": {"include_source_code": false}}, {"scope": {"dag_id": "openlineage_policy_conf_source_code_dag"}, "controls": {"include_source_code": false}}]' + AIRFLOW__OPENLINEAGE__EXECUTION_TIMEOUT: '60' + # Both control how quickly the dag-processor picks up dags: REFRESH_INTERVAL is how often it + # re-scans the folder for new/removed files, MIN_FILE_PROCESS_INTERVAL is how often it re-parses + # a file it already knows about. Tightened from the 300s/30s defaults so the e2e run doesn't wait + # on the dag-processor's own cadence on top of everything else. + AIRFLOW__DAG_PROCESSOR__REFRESH_INTERVAL: '3' + AIRFLOW__DAG_PROCESSOR__MIN_FILE_PROCESS_INTERVAL: '3' + # openlineage_defer_simple_dag's TimeDeltaSensor defaults to a 180s delta so it reliably defers at least + # once even before a worker picks up the task; tighten it here since we do not need to wait for worker. + SYSTEM_TESTS_OL_TIMEDELTA: '60' + # Every DAG's dagrun_timeout (system.openlineage.constants.DEFAULT_DAGRUN_TIMEOUT) + # defaults to 10 minutes; tighten it here since we do not need to wait for worker. + SYSTEM_TESTS_OL_DAGRUN_TIMEOUT_MINUTES: '5' + AIRFLOW__CORE__TASK_SUCCESS_OVERTIME: '70' + # Serve the dags folder through a versioned bundle so example_openlineage_versioned_dag gets a + # non-null dag_bundle_version. MockVersionedLocalDagBundle still serves from /opt/airflow/dags + # (keeping PYTHONPATH valid); it is copied into the dags folder by prepare_dags.py. + AIRFLOW__DAG_PROCESSOR__DAG_BUNDLE_CONFIG_LIST: '[{"name": "dags-folder", "classpath": "e2e_versioned_bundle.MockVersionedLocalDagBundle", "kwargs": {}}]' + # The breeze PROD image installs every provider from source, including informatica, whose + # listener fires on every task and errors without its connection — suppressing OL events. + # A curated deployment (e.g. Astro) does not ship it; disable its listener to match. + AIRFLOW__INFORMATICA__LISTENER_DISABLED: 'true' + HOST_OS: ${HOST_OS:-linux} + user: "${AIRFLOW_UID:-50000}:0" + volumes: + - ./dags:/opt/airflow/dags + - ./logs:/opt/airflow/logs + # example_openlineage_docs_file_dag sets doc_md="dag_doc.md", which Airflow resolves with a + # CWD-relative open(); components run with CWD=/opt/airflow. prepare_dags.py generates a clean, + # license-free dag_doc.md (matching the expected event) in the dags folder; expose it here. + - ./dags/dag_doc.md:/opt/airflow/dag_doc.md + depends_on: + &airflow-common-depends-on + postgres: + condition: service_healthy + +services: + postgres: + image: postgres:13 + environment: + POSTGRES_USER: airflow + POSTGRES_PASSWORD: airflow + POSTGRES_DB: airflow + healthcheck: + test: ["CMD", "pg_isready", "-U", "airflow"] + interval: 10s + retries: 5 + start_period: 5s + restart: always + + airflow-init: + <<: *airflow-common + entrypoint: /bin/bash + command: + - -c + - | + echo "Creating missing opt dirs if missing:" + mkdir -v -p /opt/airflow/{logs,dags,plugins,config} + echo "Airflow version:" + /entrypoint airflow version + echo "Running airflow config list to create default config file if missing." + /entrypoint airflow config list >/dev/null + if [ "${HOST_OS}" == "linux" ]; then + echo "Change ownership of files in /opt/airflow to ${AIRFLOW_UID}:0" + chown -R "${AIRFLOW_UID}:0" /opt/airflow/ + else + echo "Skipping changing ownership as platform is not linux (${HOST_OS})" + fi + environment: + <<: *airflow-common-env + _AIRFLOW_DB_MIGRATE: 'true' + _AIRFLOW_WWW_USER_CREATE: 'true' + _AIRFLOW_WWW_USER_USERNAME: ${_AIRFLOW_WWW_USER_USERNAME:-airflow} + _AIRFLOW_WWW_USER_PASSWORD: ${_AIRFLOW_WWW_USER_PASSWORD:-airflow} + _PIP_ADDITIONAL_REQUIREMENTS: '' + user: "0:0" + depends_on: + <<: *airflow-common-depends-on + + airflow-apiserver: + <<: *airflow-common + command: api-server + ports: + - "8080:8080" + healthcheck: + test: ["CMD", "curl", "--fail", "http://localhost:8080/api/v2/monitor/health"] + interval: 30s + timeout: 10s + retries: 5 + start_period: 30s + restart: always + depends_on: + <<: *airflow-common-depends-on + airflow-init: + condition: service_completed_successfully + + airflow-scheduler: + <<: *airflow-common + command: scheduler + healthcheck: + test: ["CMD", "airflow", "jobs", "check", "--job-type", "SchedulerJob"] + interval: 30s + timeout: 10s + retries: 5 + restart: always + depends_on: + postgres: + condition: service_healthy + airflow-init: + condition: service_completed_successfully + + airflow-dag-processor: + <<: *airflow-common + command: dag-processor + depends_on: + <<: *airflow-common-depends-on + airflow-init: + condition: service_completed_successfully + + airflow-triggerer: + <<: *airflow-common + command: triggerer + healthcheck: + test: ["CMD", "airflow", "jobs", "check", "--job-type", "TriggererJob"] + interval: 30s + timeout: 10s + retries: 5 + restart: always + depends_on: + <<: *airflow-common-depends-on + airflow-init: + condition: service_completed_successfully diff --git a/providers-e2e-tests/openlineage/prepare_dags.py b/providers-e2e-tests/openlineage/prepare_dags.py new file mode 100644 index 0000000000000..228cd16943d72 --- /dev/null +++ b/providers-e2e-tests/openlineage/prepare_dags.py @@ -0,0 +1,121 @@ +# 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. +""" +Generate the dags folder for the OpenLineage e2e deployment from the provider system tests. + +The OpenLineage system-test DAGs in ``providers/openlineage/tests/system/openlineage`` are the +single source of truth. They are *designed* to run in-process via pytest, so each one ends with a +``get_test_run`` block that imports ``tests_common`` — which is not present in the PROD image. This +script copies the whole ``system`` package into the e2e dags folder and strips that pytest-only +footer so the DAGs parse and run in a real deployment. + +The package layout is preserved as ``/system/openlineage/...`` and the dags folder is added to +``PYTHONPATH`` in docker-compose, so the DAGs' ``from system.openlineage... import ...`` and the +dotted-path ``VariableTransport`` both resolve unchanged — no import rewriting needed. +""" + +from __future__ import annotations + +import os +import shutil +from pathlib import Path + +HERE = Path(__file__).resolve().parent +AIRFLOW_ROOT_PATH = HERE.parents[1] + +# Example DAGs that need a minimum Airflow version. When running against an older version, the DAG +# file is removed entirely — a module-level import (e.g. the HITL operators) would otherwise raise at +# parse time and take the whole dag-processor down, not just skip that one DAG. +MIN_AIRFLOW_VERSION_FOR_DAG: dict[str, tuple[int, int]] = { + # HITL operators import-raise on Airflow < 3.1 (taking the dag-processor down at parse time). + "example_openlineage_hitl_dag.py": (3, 1), +} + +SYSTEM_TESTS_SOURCE = AIRFLOW_ROOT_PATH / "providers" / "openlineage" / "tests" / "system" +DAGS_DEST = HERE / "dags" +# Harness-only modules (warmup DAG, versioned bundle) that are not part of the provider system tests. +DAGS_EXTRA_SOURCE = HERE / "dags_extra" + +# The marker after which everything is Airflow's in-process test harness (pytest-only). +PYTEST_FOOTER_MARKER = "from tests_common.test_utils.system_tests import get_test_run" + + +def _target_airflow_version() -> tuple[int, int] | None: + """Parse the (major, minor) of the targeted Airflow version, or None for the default/prod run.""" + raw = os.environ.get("E2E_TARGET_AIRFLOW_VERSION", "").strip() + if not raw: + return None + parts = raw.split(".") + if len(parts) < 2 or not (parts[0].isdigit() and parts[1].isdigit()): + return None + return int(parts[0]), int(parts[1]) + + +def _strip_pytest_footer(dag_file: Path) -> None: + lines = dag_file.read_text().splitlines(keepends=True) + kept: list[str] = [] + for line in lines: + if PYTEST_FOOTER_MARKER in line: + break + kept.append(line) + dag_file.write_text("".join(kept)) + + +def prepare_dags() -> Path: + """Populate the e2e dags folder from the provider system tests. Returns the dags folder path.""" + if not SYSTEM_TESTS_SOURCE.is_dir(): + raise FileNotFoundError(f"OpenLineage system tests not found at {SYSTEM_TESTS_SOURCE}") + + if DAGS_DEST.exists(): + shutil.rmtree(DAGS_DEST) + DAGS_DEST.mkdir(parents=True) + + # Copy the whole `system` package so `system.openlineage.{operator,transport,expected_events}` + # imports resolve from the dags folder. + shutil.copytree(SYSTEM_TESTS_SOURCE, DAGS_DEST / "system") + + openlineage_dir = DAGS_DEST / "system" / "openlineage" + + # The pytest conftest imports `pytest`, which is absent from the PROD image; drop it so the + # dag-processor does not choke on it. + (openlineage_dir / "conftest.py").unlink(missing_ok=True) + + target_version = _target_airflow_version() + if target_version is not None: + for dag_file_name, min_version in MIN_AIRFLOW_VERSION_FOR_DAG.items(): + if target_version < min_version: + (openlineage_dir / dag_file_name).unlink(missing_ok=True) + + for dag_file in openlineage_dir.glob("example_openlineage_*.py"): + _strip_pytest_footer(dag_file) + + for extra_file in DAGS_EXTRA_SOURCE.glob("*.py"): + shutil.copy2(extra_file, DAGS_DEST / extra_file.name) + + # example_openlineage_docs_file_dag uses doc_md="dag_doc.md" and the expected event's + # documentation.description is exactly "# MD doc file". The provider's own dag_doc.md carries an + # Apache license header (required for repo files), so OpenLineage would emit the header too. + # Generate a clean, license-free copy here (it lives only in the gitignored dags folder, never in + # the repo); docker-compose mounts it into the parsing CWD so doc_md resolves to this content. + (DAGS_DEST / "dag_doc.md").write_text("# MD doc file") + + return DAGS_DEST + + +if __name__ == "__main__": + dest = prepare_dags() + print(f"Prepared OpenLineage e2e dags in {dest}") diff --git a/providers-e2e-tests/openlineage/pyproject.toml b/providers-e2e-tests/openlineage/pyproject.toml new file mode 100644 index 0000000000000..981a424dfab15 --- /dev/null +++ b/providers-e2e-tests/openlineage/pyproject.toml @@ -0,0 +1,85 @@ + +# 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. + +[build-system] +requires = [ + "hatchling==1.30.1", + "packaging==26.2", + "pathspec==1.1.1", + "pluggy==1.6.0", + "tomli==2.4.1; python_version < '3.11'", + "trove-classifiers==2026.6.1.19", +] +build-backend = "hatchling.build" + +[project] +name = "apache-airflow-openlineage-e2e-tests" +description = "OpenLineage provider end-to-end tests for Apache Airflow" +classifiers = [ + "Private :: Do Not Upload", +] +requires-python = ">=3.10,!=3.15" +authors = [ + { name = "Apache Software Foundation", email = "dev@airflow.apache.org" }, +] +maintainers = [ + { name = "Apache Software Foundation", email="dev@airflow.apache.org" }, +] +version = "0.0.1" + +dependencies = [ + "apache-airflow-core", + "apache-airflow-devel-common", +] + +[tool.pytest] +addopts = [ + "-rasl", + "--verbosity=2", + "-p", "no:flaky", + "-p", "no:nose", + "-p", "no:legacypath", +] +norecursedirs = [ + ".eggs", +] +log_level = "INFO" +filterwarnings = [ + "error::pytest.PytestCollectionWarning", +] +python_files = [ + "*.py", +] + +# Keep temporary directories (created by `tmp_path`) for 2 recent runs only failed tests. +tmp_path_retention_count = "2" +tmp_path_retention_policy = "failed" + +[tool.e2e-tests] +required-providers = ["openlineage", "standard", "common.compat", "common.sql", "common.io"] + +[tool.hatch.build.targets.sdist] +exclude = ["*"] + +[tool.hatch.build.targets.wheel] +bypass-selection = true + +[dependency-groups] +mypy = [ + "apache-airflow-devel-common[mypy]", +] diff --git a/providers-e2e-tests/openlineage/tests/conftest.py b/providers-e2e-tests/openlineage/tests/conftest.py new file mode 100644 index 0000000000000..917ac4a677283 --- /dev/null +++ b/providers-e2e-tests/openlineage/tests/conftest.py @@ -0,0 +1,194 @@ +# 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 + +import os +import platform +import re +import subprocess +import sys +import time + +import pytest +import requests +from constants import ( + API_BASE_URL, + DOCKER_COMPOSE_FILE_PATH, + DOCKER_IMAGE, + ENV_FILE, + LOCAL_DOCKER_COMPOSE_FILE_PATH, + MIN_DOCKER_COMPOSE_VERSION, + MIN_DOCKER_VERSION, + OPENLINEAGE_E2E_TESTS_ROOT, + console, +) +from packaging import version + +from tests_common.test_utils.fernet import generate_fernet_key_string + +COMPOSE_PROJECT_NAME = "breeze-openlineage-e2e" +FALSEY_VALUES = ("false", "0") + + +def _env_flag(name: str) -> bool: + """True when the environment variable is set to a non-falsey value.""" + value = os.environ.get(name, "") + return bool(value) and value.lower() not in FALSEY_VALUES + + +def _compare_gte(a: str, b: str) -> bool: + try: + return version.parse(a) >= version.parse(b) + except Exception: + return False + + +def _docker_is_running() -> bool: + try: + return ( + subprocess.run(["docker", "info"], capture_output=True, check=False, timeout=10).returncode == 0 + ) + except Exception: + return False + + +def _tool_version(command: list[str]) -> str: + try: + result = subprocess.run(command, capture_output=True, text=True, check=False, timeout=10) + if result.returncode != 0: + return "unknown" + match = re.search(r"(\d+)\.(\d+)\.(\d+)", result.stdout.strip()) + return ".".join(match.groups()) if match else "unknown" + except Exception: + return "unknown" + + +def check_docker_requirements() -> None: + """Verify Docker is running and Docker / Docker Compose meet the minimum versions.""" + console.print("[yellow]Checking Docker requirements...[/]") + if not _docker_is_running(): + console.print("[red]Docker is not running. Please start Docker and try again.[/]") + sys.exit(1) + + docker_version = _tool_version(["docker", "version", "--format", "{{.Client.Version}}"]) + if not _compare_gte(docker_version, MIN_DOCKER_VERSION): + console.print(f"[red]Docker version {docker_version} is too old (min {MIN_DOCKER_VERSION}).[/]") + sys.exit(1) + console.print(f"[green]Docker version {docker_version} meets requirements[/]") + + compose_version = _tool_version(["docker", "compose", "version"]) + if not _compare_gte(compose_version, MIN_DOCKER_COMPOSE_VERSION): + console.print( + f"[red]Docker Compose version {compose_version} is too old (min {MIN_DOCKER_COMPOSE_VERSION}).[/]" + ) + sys.exit(1) + console.print(f"[green]Docker Compose version {compose_version} meets requirements[/]") + + +def _print_diagnostics(compose) -> None: + console.print("[red]=== DIAGNOSTIC INFORMATION ===[/]") + try: + for container in compose.compose.ps(): + console.print(f" {container.name}: {container.state}") + console.print("\n[yellow]Container Logs:[/]") + console.print(compose.compose.logs()) + except Exception as exc: + console.print(f" Error collecting diagnostics: {exc}") + + +def pytest_sessionstart(session): + """Check Docker and source the OpenLineage DAGs into the dags folder before the stack starts.""" + check_docker_requirements() + + # prepare_dags lives at the project root (next to docker-compose.yaml), not in the test package. + sys.path.insert(0, OPENLINEAGE_E2E_TESTS_ROOT.as_posix()) + from prepare_dags import prepare_dags + + console.print("[yellow]Preparing OpenLineage DAGs from the provider system tests...[/]") + dags_folder = prepare_dags() + console.print(f"[green]Prepared dags in {dags_folder}[/]") + + +@pytest.fixture(scope="session") +def docker_compose_setup(): + """Start docker-compose once per session and tear it down at the end.""" + from python_on_whales import DockerClient + + debugging_on = _env_flag("VERBOSE") + mount_volumes = not _env_flag("SKIP_MOUNTING_LOCAL_VOLUMES") + delete_compose = not _env_flag("SKIP_DOCKER_COMPOSE_DELETION") + if mount_volumes: + delete_compose = False + + with open(ENV_FILE, "w") as f: + print(f"AIRFLOW_IMAGE_NAME={DOCKER_IMAGE}", file=f) + print(f"AIRFLOW_UID={os.getuid()}", file=f) + print(f"HOST_OS={platform.system().lower()}", file=f) + # Please do not reuse this Fernet key in any deployment — it is generated only for these tests. + print(f"FERNET_KEY={generate_fernet_key_string()}", file=f) + + docker_compose_files = [DOCKER_COMPOSE_FILE_PATH.as_posix()] + if mount_volumes: + docker_compose_files.append(LOCAL_DOCKER_COMPOSE_FILE_PATH.as_posix()) + log_level = "debug" if debugging_on else "info" + + compose = DockerClient( + compose_files=docker_compose_files, + debug=debugging_on, + log_level=log_level, + compose_project_name=COMPOSE_PROJECT_NAME, + ) + + start_new_compose = True + processes = compose.compose.ps(["airflow-apiserver"]) + if processes and processes[0].state.status == "running": + if mount_volumes: + console.print( + "[yellow]Docker compose already running. Reusing it. Run `docker compose down` to reset.[/]" + ) + start_new_compose = False + else: + compose.compose.down(remove_orphans=True, volumes=True, quiet=True) + + try: + if start_new_compose: + console.print("[yellow]Starting docker-compose for session...[/]") + files = " ".join(f'-f "{file}"' for file in docker_compose_files) + console.print(f"[info]Equivalent manual command: docker compose {files} up --detach[/]") + compose.compose.up(detach=True, wait=True) + console.print("[green]Docker compose started successfully![/]") + yield compose + except Exception as exc: + console.print(f"[red]Docker compose failed to start: {exc}[/]") + _print_diagnostics(compose) + raise + finally: + if delete_compose: + console.print("[yellow]Cleaning up docker-compose...[/]") + compose.compose.down(remove_orphans=True, volumes=True, quiet=True) + console.print("[green]Docker compose cleaned up[/]") + + +@pytest.fixture(scope="session") +def auth_headers(docker_compose_setup) -> dict[str, str]: + """Wait for Airflow and return REST API auth headers (SimpleAuthManager grants an admin token).""" + time.sleep(15) + auth_response = requests.get(f"{API_BASE_URL}/auth/token", timeout=30) + auth_response.raise_for_status() + token = auth_response.json()["access_token"] + console.print("[green]✅ Got Airflow auth token[/]") + return {"Authorization": f"Bearer {token}", "Content-Type": "application/json"} diff --git a/providers-e2e-tests/openlineage/tests/constants.py b/providers-e2e-tests/openlineage/tests/constants.py new file mode 100644 index 0000000000000..cce0aa92e91bd --- /dev/null +++ b/providers-e2e-tests/openlineage/tests/constants.py @@ -0,0 +1,42 @@ +# 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 + +import os +from pathlib import Path + +from rich.console import Console + +console = Console(width=400, color_system="standard") + +AIRFLOW_ROOT_PATH = Path(__file__).resolve().parents[3] +OPENLINEAGE_E2E_TESTS_ROOT = AIRFLOW_ROOT_PATH / "providers-e2e-tests" / "openlineage" + +DEFAULT_PYTHON_MAJOR_MINOR_VERSION = "3.10" +MIN_DOCKER_VERSION = "25.0.0" +MIN_DOCKER_COMPOSE_VERSION = "2.20.2" +DEFAULT_DOCKER_IMAGE = f"ghcr.io/apache/airflow/main/prod/python{DEFAULT_PYTHON_MAJOR_MINOR_VERSION}:latest" +DOCKER_IMAGE = os.environ.get("DOCKER_IMAGE") or DEFAULT_DOCKER_IMAGE + +API_HOST_PORT = os.environ.get("HOST_PORT", "localhost:8080") +API_BASE_URL = f"http://{API_HOST_PORT}" + +DOCKER_COMPOSE_FILE_PATH = OPENLINEAGE_E2E_TESTS_ROOT / "docker-compose.yaml" +LOCAL_DOCKER_COMPOSE_FILE_PATH = OPENLINEAGE_E2E_TESTS_ROOT / "docker-compose-local.yaml" +DAGS_FOLDER = OPENLINEAGE_E2E_TESTS_ROOT / "dags" +LOGS_FOLDER = OPENLINEAGE_E2E_TESTS_ROOT / "logs" +ENV_FILE = OPENLINEAGE_E2E_TESTS_ROOT / ".env" diff --git a/providers-e2e-tests/openlineage/tests/harness.py b/providers-e2e-tests/openlineage/tests/harness.py new file mode 100644 index 0000000000000..b45a7c67ad218 --- /dev/null +++ b/providers-e2e-tests/openlineage/tests/harness.py @@ -0,0 +1,246 @@ +# 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. +""" +Drive all OpenLineage DAGs in the deployed stack and collect their final run state. + +Adapted from the OpenLineage dags-dashboard ``trigger_dag.py`` integration harness: it talks to a +local Airflow REST API v2 (rather than an Astro deployment), authenticates with a SimpleAuthManager +bearer token, and reports through logs only. The actual OpenLineage event validation happens inside +each DAG's terminal ``OpenLineageTestOperator`` task, so a run that ends ``success`` means its +emitted events matched the expected templates. +""" + +from __future__ import annotations + +import datetime as dt +import re +import time +from pathlib import Path + +import requests +from constants import console +from requests.adapters import HTTPAdapter +from urllib3.util.retry import Retry + +WARMUP_DAG_ID = "openlineage_warmup_dag" +TRIGGER_DELAY_SECONDS = 2 +# Shared wall-clock budget for one _trigger_and_wait() call, covering every DAG triggered in that +# batch — bounds total polling time regardless of how many DAGs are involved +BATCH_TIMEOUT_SECONDS = 360 # 6 min, each dagrun timeout is set to 5 minutes + +# DAGs that take longer to complete — trigger these first so they run while the rest are triggered. +LONG_RUNNING_DAG_IDS = ("openlineage_defer_simple_dag",) + +# Child DAGs spawned by TriggerDagRunOperator carry this marker; they must not be triggered or +# state-checked directly (their parent triggers them). +NO_TRIGGER_MARKER = "__notrigger" + + +def discover_expected_dag_ids(dags_folder: Path) -> set[str]: + """ + Top-level ``DAG_ID = "..."`` declarations across the prepared example DAGs. + + Used as a coverage check: every expected DAG must actually load in the deployment (a missing one + signals an import/parse failure rather than a test failure). + """ + pattern = re.compile(r'^DAG_ID\s*=\s*["\'](?P[^"\']+)["\']') + dag_ids: set[str] = set() + for dag_file in (dags_folder / "system" / "openlineage").glob("example_openlineage_*.py"): + for line in dag_file.read_text().splitlines(): + match = pattern.match(line) + if match: + dag_ids.add(match.group("dag_id")) + return dag_ids + + +def _session_with_retries() -> requests.Session: + retry = Retry( + total=10, + connect=5, + read=10, + backoff_factor=2, + raise_on_status=False, + allowed_methods=["GET", "POST", "PATCH", "DELETE"], + status_forcelist=[429, 500, 502, 503, 504], + ) + session = requests.Session() + session.mount("http://", HTTPAdapter(max_retries=retry)) + session.mount("https://", HTTPAdapter(max_retries=retry)) + return session + + +class OpenLineageE2ERunner: + """Triggers all OpenLineage DAGs against a running deployment and collects their final states.""" + + def __init__(self, api_base_url: str, headers: dict[str, str]): + self.api_url = f"{api_base_url}/api/v2" + self.headers = headers + self.session = _session_with_retries() + now = dt.datetime.now(tz=dt.timezone.utc) + self.run_id = f"ci_triggered_{now.isoformat()}" + self.retry_run_id = f"{self.run_id}_retry1" + # dag_id -> whether its final state came from retry_run_id rather than run_id; populated by run(). + self.retried_dag_ids: set[str] = set() + + def wait_for_airflow_api(self, max_attempts: int = 30, poll_interval: int = 10) -> None: + for attempt in range(1, max_attempts + 1): + try: + response = requests.get(f"{self.api_url}/dags?limit=1", headers=self.headers, timeout=15) + if response.status_code < 500: + console.print(f"[green]Airflow API ready after {attempt} attempt(s)") + return + except requests.exceptions.RequestException as exc: + console.print(f"[yellow]Airflow API not ready (attempt {attempt}/{max_attempts}): {exc}") + time.sleep(poll_interval) + raise RuntimeError(f"Airflow API did not become ready after {max_attempts} attempts.") + + def list_dags(self) -> list[str]: + response = self.session.get(f"{self.api_url}/dags?limit=500", headers=self.headers) + response.raise_for_status() + return [dag["dag_id"] for dag in response.json()["dags"]] + + def get_task_states(self, dag_id: str, run_id: str) -> dict[str, str]: + """task_id -> state for every task instance in a dag run.""" + response = self.session.get( + f"{self.api_url}/dags/{dag_id}/dagRuns/{run_id}/taskInstances", headers=self.headers + ) + response.raise_for_status() + return {ti["task_id"]: ti["state"] for ti in response.json()["task_instances"]} + + def wait_for_dags_loaded(self, timeout: int = 60, poll_interval: int = 3) -> list[str]: + """Poll until the dag-processor has parsed the DAGs (the warmup DAG is the readiness marker).""" + deadline = time.monotonic() + timeout + dag_ids: list[str] = [] + while time.monotonic() < deadline: + dag_ids = self.list_dags() + if WARMUP_DAG_ID in dag_ids: + console.print(f"[green]DAGs loaded ({len(dag_ids)} found)") + return dag_ids + console.print(f"[yellow]Waiting for DAGs to load (have {len(dag_ids)})...") + time.sleep(poll_interval) + raise RuntimeError(f"DAGs did not load within {timeout}s (warmup DAG missing; have {dag_ids}).") + + def unpause_dag(self, dag_id: str) -> None: + response = self.session.patch( + f"{self.api_url}/dags/{dag_id}", headers=self.headers, json={"is_paused": False} + ) + if response.status_code != 200: + console.print(f"[red]Failed to unpause DAG `{dag_id}`: {response.text}") + + def trigger_dag_run(self, dag_id: str, run_id: str) -> bool: + now = dt.datetime.now(tz=dt.timezone.utc).isoformat() + payload = {"dag_run_id": run_id, "logical_date": now, "conf": {}} + response = self.session.post( + f"{self.api_url}/dags/{dag_id}/dagRuns", headers=self.headers, json=payload + ) + if response.status_code not in (200, 201): + console.print(f"[red]Failed to trigger DAG `{dag_id}`: {response.text}") + return False + return True + + def wait_for_dag_run_to_complete(self, dag_id: str, run_id: str, timeout: int = 300) -> str: + url = f"{self.api_url}/dags/{dag_id}/dagRuns/{run_id}" + deadline = time.monotonic() + timeout + state = "unknown" + while time.monotonic() < deadline: + response = self.session.get(url, headers=self.headers) + response.raise_for_status() + state = response.json()["state"] + if state not in ("running", "queued"): + break + time.sleep(5) + console.print(f"[blue]DAG `{dag_id}` finished in state: {state}") + return state + + def clear_airflow_variables(self) -> None: + response = self.session.get(f"{self.api_url}/variables?limit=500", headers=self.headers) + response.raise_for_status() + keys = [variable["key"] for variable in response.json()["variables"]] + for key in keys: + self.session.delete(f"{self.api_url}/variables/{key}", headers=self.headers) + + def warmup(self, dag_ids: list[str]) -> None: + """Unpause all DAGs and run the warmup DAG so the worker is confirmed ready.""" + if WARMUP_DAG_ID not in dag_ids: + raise KeyError(f"Warmup DAG `{WARMUP_DAG_ID}` not found in deployment.") + for dag_id in dag_ids: + self.unpause_dag(dag_id) + time.sleep(TRIGGER_DELAY_SECONDS) + if not self.trigger_dag_run(WARMUP_DAG_ID, self.run_id): + raise RuntimeError(f"Failed to trigger warmup DAG `{WARMUP_DAG_ID}`") + self.wait_for_dag_run_to_complete(WARMUP_DAG_ID, self.run_id) + + def _trigger_and_wait( + self, dag_ids: list[str], run_id: str, batch_timeout: int = BATCH_TIMEOUT_SECONDS + ) -> dict[str, str]: + statuses: dict[str, str] = {} + ordered = [d for d in LONG_RUNNING_DAG_IDS if d in dag_ids] + ordered += [d for d in dag_ids if d not in LONG_RUNNING_DAG_IDS] + for dag_id in ordered: + if not self.trigger_dag_run(dag_id, run_id): + statuses[dag_id] = "trigger_error" + time.sleep(TRIGGER_DELAY_SECONDS) + deadline = time.monotonic() + batch_timeout + for dag_id in ordered: + if dag_id in statuses: + continue + remaining = deadline - time.monotonic() + if remaining <= 0: + console.print( + f"[red]Batch deadline exceeded before checking `{dag_id}`; marking as timed_out." + ) + statuses[dag_id] = "timed_out" + continue + statuses[dag_id] = self.wait_for_dag_run_to_complete(dag_id, run_id, timeout=int(remaining)) + return statuses + + def run(self, expected_dag_ids: set[str]) -> dict[str, str]: + """Run the full cycle and return ``{dag_id: final_state}`` for every triggered DAG.""" + self.wait_for_airflow_api() + dag_ids = self.wait_for_dags_loaded() + if not dag_ids: + raise ValueError("No DAGs found in the deployment.") + + self.warmup(dag_ids) + # Auto-runs of cron/timetable DAGs may have started while unpaused — clear their events so + # the test runs start from a clean slate. + time.sleep(10) + self.clear_airflow_variables() + + triggerable = [ + dag_id for dag_id in dag_ids if dag_id != WARMUP_DAG_ID and NO_TRIGGER_MARKER not in dag_id + ] + statuses = self._trigger_and_wait(triggerable, self.run_id) + self.clear_airflow_variables() + + failed = [dag_id for dag_id, state in statuses.items() if state != "success"] + if failed: + self.retried_dag_ids = set(failed) + console.print(f"[yellow]Retrying {len(failed)} failed DAG(s) once: {failed}") + retry_statuses = self._trigger_and_wait(failed, self.retry_run_id) + self.clear_airflow_variables() + recovered = sorted(dag_id for dag_id, state in retry_statuses.items() if state == "success") + if recovered: + # The retry counts as a pass, but surface it so first-run flakiness is not hidden. + console.print(f"[yellow]⚠ DAGs that passed only on retry (flaky first run): {recovered}") + statuses.update(retry_statuses) + + for dag_id in expected_dag_ids: + if dag_id not in statuses: + statuses[dag_id] = "missing" + + return statuses diff --git a/providers-e2e-tests/openlineage/tests/task_logs.py b/providers-e2e-tests/openlineage/tests/task_logs.py new file mode 100644 index 0000000000000..f5d3c19170b9f --- /dev/null +++ b/providers-e2e-tests/openlineage/tests/task_logs.py @@ -0,0 +1,111 @@ +# 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. +""" +Read and format Airflow task logs for display in pytest output on failure. + +Task logs are structlog JSON-lines written by the deployed Airflow stack; they land on the host via +the ``./logs:/opt/airflow/logs`` bind mount in ``docker-compose.yaml``. Rather than requiring someone +to unzip the CI log artifact and open a specific ``attempt=N.log`` file, the warning/error highlights +are printed straight into the failing test's captured output. +""" + +from __future__ import annotations + +import json +import re +from collections.abc import Iterable +from pathlib import Path +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from rich.console import Console + +LOG_LEVELS_OF_INTEREST = frozenset({"warning", "error", "critical"}) +MAX_HIGHLIGHT_LINES = 100 + +_ATTEMPT_NUMBER_PATTERN = re.compile(r"attempt=(\d+)\.log$") + + +def format_log_line(record: dict) -> str: + """Render one parsed structlog record as a compact, human-readable line.""" + level = str(record.get("level", "info")).upper() + event = record.get("event", "") + lines = [f"[{level}] {event}"] + for exc in record.get("error_detail") or []: + exc_type = exc.get("exc_type", "Exception") + exc_value = exc.get("exc_value", "") + lines.append(f" {exc_type}: {exc_value}") + frames = exc.get("frames") or [] + if frames: + frame = frames[-1] + lines.append(f" at {Path(frame['filename']).name}:{frame['lineno']} in {frame['name']}") + return "\n".join(lines) + + +def filter_log_highlights(raw_lines: Iterable[str], max_lines: int = MAX_HIGHLIGHT_LINES) -> list[str]: + """Parse structlog JSON lines and keep only warning/error/critical ones, formatted for display.""" + highlights = [] + for raw_line in raw_lines: + stripped_line = raw_line.strip() + if not stripped_line: + continue + try: + record = json.loads(stripped_line) + except json.JSONDecodeError: + # Task logs aren't guaranteed to be 100% JSON (e.g. stray print output) — skip those lines. + continue + if record.get("level") in LOG_LEVELS_OF_INTEREST: + highlights.append(format_log_line(record)) + if not highlights: + return ["(no warning/error lines in this task's log)"] + if len(highlights) > max_lines: + dropped = len(highlights) - max_lines + highlights = highlights[:max_lines] + [f"... ({dropped} more lines suppressed)"] + return highlights + + +def _attempt_number(path: Path) -> int: + match = _ATTEMPT_NUMBER_PATTERN.search(path.name) + return int(match.group(1)) if match else -1 + + +def find_log_paths(logs_folder: Path, dag_id: str, run_id: str, task_id: str) -> list[Path]: + """Latest-attempt log file(s) for a task — one per map index if the task is dynamically mapped.""" + task_root = logs_folder / f"dag_id={dag_id}" / f"run_id={run_id}" / f"task_id={task_id}" + latest_by_dir: dict[Path, Path] = {} + for attempt_path in task_root.glob("**/attempt=*.log"): + current = latest_by_dir.get(attempt_path.parent) + if current is None or _attempt_number(attempt_path) > _attempt_number(current): + latest_by_dir[attempt_path.parent] = attempt_path + return sorted(latest_by_dir.values()) + + +def print_task_log(console: Console, logs_folder: Path, dag_id: str, run_id: str, task_id: str) -> None: + """Print the warning/error highlights of a task's log(s), or a note if none are found.""" + console.print(f"[bold cyan]--- {dag_id} / {task_id} / START---[/]") + log_paths = find_log_paths(logs_folder, dag_id, run_id, task_id) + if not log_paths: + console.print( + f"[yellow](no log file found under dag_id={dag_id}/run_id={run_id}/task_id={task_id})[/]" + ) + return + for log_path in log_paths: + if len(log_paths) > 1: + console.print(f"[cyan]{log_path.relative_to(logs_folder)}[/]") + for line in filter_log_highlights(log_path.read_text().splitlines()): + console.print(line) + console.print(f"[bold cyan]--- {dag_id} / {task_id} / END---[/]") diff --git a/providers-e2e-tests/openlineage/tests/test_openlineage_e2e.py b/providers-e2e-tests/openlineage/tests/test_openlineage_e2e.py new file mode 100644 index 0000000000000..277ce2a942214 --- /dev/null +++ b/providers-e2e-tests/openlineage/tests/test_openlineage_e2e.py @@ -0,0 +1,72 @@ +# 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 + +import pytest +import requests +from constants import API_BASE_URL, DAGS_FOLDER, LOGS_FOLDER, console +from harness import OpenLineageE2ERunner, discover_expected_dag_ids +from task_logs import find_log_paths, print_task_log + + +def _print_failure_logs(runner: OpenLineageE2ERunner, failed: dict[str, str]) -> None: + """Print the relevant task logs for each failed DAG, so CI output doesn't require the log artifact.""" + for dag_id, state in failed.items(): + if state == "missing": + console.print(f"[yellow]{dag_id} never loaded — check the dag-processor logs, not task logs.[/]") + continue + run_id = runner.retry_run_id if dag_id in runner.retried_dag_ids else runner.run_id + try: + task_states = runner.get_task_states(dag_id, run_id) + except requests.exceptions.RequestException as exc: + console.print(f"[yellow]Could not fetch task states for {dag_id}: {exc}[/]") + task_states = {} + # check_events always included: it shows what was actually compared, even if a different + # task is what actually failed the DAG. Every other non-success task is shown too, EXCEPT a + # routine SKIPPED task (e.g. a branch not taken) never ran and has no log — showing it would + # look like a false-positive failure. A task SKIPPED by dagrun_timeout cutting off a + # still-running task, though, usually does have a partial log, so still show those. + tasks_to_show = {"check_events"} + for task_id, ti_state in task_states.items(): + if ti_state == "success" or task_id in tasks_to_show: + continue + if ti_state == "skipped" and not find_log_paths(LOGS_FOLDER, dag_id, run_id, task_id): + continue + tasks_to_show.add(task_id) + for task_id in sorted(tasks_to_show): + print_task_log(console, LOGS_FOLDER, dag_id, run_id, task_id) + + +@pytest.mark.execution_timeout(900) # 15 min +def test_all_openlineage_dags_succeed(auth_headers): + """ + Trigger every OpenLineage system-test DAG in the deployment and require each run to succeed. + + The terminal ``OpenLineageTestOperator`` task inside each DAG validates the emitted OpenLineage + events against the expected templates, so a ``success`` run state means the lineage matched. + """ + expected_dag_ids = discover_expected_dag_ids(DAGS_FOLDER) + assert expected_dag_ids, "No expected OpenLineage DAG ids were discovered from the prepared dags" + + runner = OpenLineageE2ERunner(API_BASE_URL, auth_headers) + statuses = runner.run(expected_dag_ids) + + console.print(f"[blue]OpenLineage e2e results ({len(statuses)} DAGs): {statuses}") + failed = {dag_id: state for dag_id, state in sorted(statuses.items()) if state != "success"} + if failed: + _print_failure_logs(runner, failed) + assert not failed, f"OpenLineage e2e DAG runs were not successful: {failed}" diff --git a/providers/openlineage/tests/system/openlineage/constants.py b/providers/openlineage/tests/system/openlineage/constants.py new file mode 100644 index 0000000000000..15a650d94266f --- /dev/null +++ b/providers/openlineage/tests/system/openlineage/constants.py @@ -0,0 +1,27 @@ +# 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. +"""Shared configuration constants for the OpenLineage system-test DAGs.""" + +from __future__ import annotations + +import os +from datetime import timedelta + +# Defensive per-DagRun timeout so a hung/misbehaving DAG fails fast instead of running forever. +# Overridable via env var — the e2e harness runs every DAG in the same deployment and wants a +# tighter value than a single real system-test run needs. +DEFAULT_DAGRUN_TIMEOUT = timedelta(minutes=int(os.getenv("SYSTEM_TESTS_OL_DAGRUN_TIMEOUT_MINUTES", "10"))) diff --git a/providers/openlineage/tests/system/openlineage/example_openlineage_all_facets_dag.py b/providers/openlineage/tests/system/openlineage/example_openlineage_all_facets_dag.py index dce37a1d91448..ed6560fc99cea 100644 --- a/providers/openlineage/tests/system/openlineage/example_openlineage_all_facets_dag.py +++ b/providers/openlineage/tests/system/openlineage/example_openlineage_all_facets_dag.py @@ -71,6 +71,7 @@ from airflow.providers.common.compat.sdk import BaseOperator from airflow.providers.openlineage.extractors.base import OperatorLineage +from system.openlineage.constants import DEFAULT_DAGRUN_TIMEOUT from system.openlineage.expected_events import get_expected_event_file_path from system.openlineage.operator import OpenLineageTestOperator @@ -487,6 +488,7 @@ def get_openlineage_facets_on_start(self) -> OperatorLineage: DAG_ID = "openlineage_all_facets_dag" with DAG( + dagrun_timeout=DEFAULT_DAGRUN_TIMEOUT, dag_id=DAG_ID, start_date=datetime(2021, 1, 1), schedule=None, diff --git a/providers/openlineage/tests/system/openlineage/example_openlineage_base_complex_dag.py b/providers/openlineage/tests/system/openlineage/example_openlineage_base_complex_dag.py index b31bd57bd3610..d843fb6c7c454 100644 --- a/providers/openlineage/tests/system/openlineage/example_openlineage_base_complex_dag.py +++ b/providers/openlineage/tests/system/openlineage/example_openlineage_base_complex_dag.py @@ -45,6 +45,7 @@ except ImportError: from airflow.utils.task_group import TaskGroup # type: ignore[no-redef] +from system.openlineage.constants import DEFAULT_DAGRUN_TIMEOUT from system.openlineage.expected_events import AIRFLOW_VERSION, get_expected_event_file_path from system.openlineage.operator import OpenLineageTestOperator @@ -74,6 +75,7 @@ def execute(self, context): DAG_ID = "openlineage_base_complex_dag" with DAG( + dagrun_timeout=DEFAULT_DAGRUN_TIMEOUT, dag_id=DAG_ID, start_date=datetime(2021, 1, 1), schedule=None, diff --git a/providers/openlineage/tests/system/openlineage/example_openlineage_base_simple_dag.py b/providers/openlineage/tests/system/openlineage/example_openlineage_base_simple_dag.py index f58f500c1d159..71be41773ca0f 100644 --- a/providers/openlineage/tests/system/openlineage/example_openlineage_base_simple_dag.py +++ b/providers/openlineage/tests/system/openlineage/example_openlineage_base_simple_dag.py @@ -30,6 +30,7 @@ from airflow import DAG from airflow.providers.standard.operators.python import PythonOperator +from system.openlineage.constants import DEFAULT_DAGRUN_TIMEOUT from system.openlineage.expected_events import get_expected_event_file_path from system.openlineage.operator import OpenLineageTestOperator @@ -46,6 +47,7 @@ def do_nothing(): DAG_ID = "openlineage_base_simple_dag" with DAG( + dagrun_timeout=DEFAULT_DAGRUN_TIMEOUT, dag_id=DAG_ID, start_date=datetime(2021, 1, 1), schedule=None, diff --git a/providers/openlineage/tests/system/openlineage/example_openlineage_custom_operator_failure_dag.py b/providers/openlineage/tests/system/openlineage/example_openlineage_custom_operator_failure_dag.py index 6767985c47aff..1829b058c87b2 100644 --- a/providers/openlineage/tests/system/openlineage/example_openlineage_custom_operator_failure_dag.py +++ b/providers/openlineage/tests/system/openlineage/example_openlineage_custom_operator_failure_dag.py @@ -34,6 +34,7 @@ from airflow.providers.standard.operators.empty import EmptyOperator from airflow.utils.trigger_rule import TriggerRule +from system.openlineage.constants import DEFAULT_DAGRUN_TIMEOUT from system.openlineage.expected_events import get_expected_event_file_path from system.openlineage.operator import OpenLineageTestOperator @@ -62,6 +63,7 @@ def get_openlineage_facets_on_failure(self, task_instance) -> OperatorLineage: DAG_ID = "openlineage_custom_operator_failure_dag" with DAG( + dagrun_timeout=DEFAULT_DAGRUN_TIMEOUT, dag_id=DAG_ID, start_date=datetime(2021, 1, 1), schedule=None, diff --git a/providers/openlineage/tests/system/openlineage/example_openlineage_custom_operator_ol_methods_dag.py b/providers/openlineage/tests/system/openlineage/example_openlineage_custom_operator_ol_methods_dag.py index 846c8a8e7b98a..c8a4377de3eb6 100644 --- a/providers/openlineage/tests/system/openlineage/example_openlineage_custom_operator_ol_methods_dag.py +++ b/providers/openlineage/tests/system/openlineage/example_openlineage_custom_operator_ol_methods_dag.py @@ -39,6 +39,7 @@ from airflow.providers.common.compat.sdk import BaseOperator from airflow.providers.openlineage.extractors.base import OperatorLineage +from system.openlineage.constants import DEFAULT_DAGRUN_TIMEOUT from system.openlineage.expected_events import get_expected_event_file_path from system.openlineage.operator import OpenLineageTestOperator @@ -321,6 +322,7 @@ def get_openlineage_facets_on_failure(self, task_instance) -> OperatorLineage: DAG_ID = "openlineage_custom_operator_ol_methods_dag" with DAG( + dagrun_timeout=DEFAULT_DAGRUN_TIMEOUT, dag_id=DAG_ID, start_date=datetime(2021, 1, 1), schedule=None, diff --git a/providers/openlineage/tests/system/openlineage/example_openlineage_defer_simple_dag.py b/providers/openlineage/tests/system/openlineage/example_openlineage_defer_simple_dag.py index 1925d0209373e..d0937e0964ab3 100644 --- a/providers/openlineage/tests/system/openlineage/example_openlineage_defer_simple_dag.py +++ b/providers/openlineage/tests/system/openlineage/example_openlineage_defer_simple_dag.py @@ -24,17 +24,20 @@ from __future__ import annotations +import os from datetime import datetime, timedelta from airflow import DAG from airflow.providers.standard.sensors.time_delta import TimeDeltaSensor +from system.openlineage.constants import DEFAULT_DAGRUN_TIMEOUT from system.openlineage.expected_events import get_expected_event_file_path from system.openlineage.operator import OpenLineageTestOperator DAG_ID = "openlineage_defer_simple_dag" with DAG( + dagrun_timeout=DEFAULT_DAGRUN_TIMEOUT, dag_id=DAG_ID, start_date=datetime(2021, 1, 1), schedule=None, @@ -43,7 +46,12 @@ ) as dag: # Timedelta is compared to the DAGRun start timestamp, which can occur long before a worker picks up the # task. We need to ensure the sensor gets deferred at least once, so setting 180s. - wait = TimeDeltaSensor(task_id="wait", delta=timedelta(seconds=180), poke_interval=10, deferrable=True) + wait = TimeDeltaSensor( + task_id="wait", + delta=timedelta(seconds=int(os.getenv("SYSTEM_TESTS_OL_TIMEDELTA", "180"))), + poke_interval=10, + deferrable=True, + ) check_events = OpenLineageTestOperator( task_id="check_events", diff --git a/providers/openlineage/tests/system/openlineage/example_openlineage_docs_file_dag.py b/providers/openlineage/tests/system/openlineage/example_openlineage_docs_file_dag.py index a30bf400c34a7..70acb73ed5b73 100644 --- a/providers/openlineage/tests/system/openlineage/example_openlineage_docs_file_dag.py +++ b/providers/openlineage/tests/system/openlineage/example_openlineage_docs_file_dag.py @@ -28,12 +28,14 @@ from airflow import DAG from airflow.providers.standard.operators.bash import BashOperator +from system.openlineage.constants import DEFAULT_DAGRUN_TIMEOUT from system.openlineage.expected_events import get_expected_event_file_path from system.openlineage.operator import OpenLineageTestOperator DAG_ID = "openlineage_docs_file_dag" with DAG( + dagrun_timeout=DEFAULT_DAGRUN_TIMEOUT, dag_id=DAG_ID, start_date=datetime(2021, 1, 1), schedule=None, diff --git a/providers/openlineage/tests/system/openlineage/example_openlineage_edge_labels_dag.py b/providers/openlineage/tests/system/openlineage/example_openlineage_edge_labels_dag.py index d869005bbb7ec..8822e0bf7f786 100644 --- a/providers/openlineage/tests/system/openlineage/example_openlineage_edge_labels_dag.py +++ b/providers/openlineage/tests/system/openlineage/example_openlineage_edge_labels_dag.py @@ -34,12 +34,14 @@ from airflow.providers.standard.operators.bash import BashOperator from airflow.utils.edgemodifier import Label +from system.openlineage.constants import DEFAULT_DAGRUN_TIMEOUT from system.openlineage.expected_events import get_expected_event_file_path from system.openlineage.operator import OpenLineageTestOperator DAG_ID = "openlineage_edge_labels_dag" with DAG( + dagrun_timeout=DEFAULT_DAGRUN_TIMEOUT, dag_id=DAG_ID, start_date=datetime(2021, 1, 1), schedule=None, diff --git a/providers/openlineage/tests/system/openlineage/example_openlineage_hitl_dag.py b/providers/openlineage/tests/system/openlineage/example_openlineage_hitl_dag.py index 6227c12fed861..b4f0644e13fb3 100644 --- a/providers/openlineage/tests/system/openlineage/example_openlineage_hitl_dag.py +++ b/providers/openlineage/tests/system/openlineage/example_openlineage_hitl_dag.py @@ -41,12 +41,14 @@ from airflow.sdk.definitions.param import Param from airflow.utils.trigger_rule import TriggerRule +from system.openlineage.constants import DEFAULT_DAGRUN_TIMEOUT from system.openlineage.expected_events import get_expected_event_file_path from system.openlineage.operator import OpenLineageTestOperator DAG_ID = "openlineage_hitl_dag" with DAG( + dagrun_timeout=DEFAULT_DAGRUN_TIMEOUT, dag_id=DAG_ID, start_date=datetime(2021, 1, 1), schedule=None, diff --git a/providers/openlineage/tests/system/openlineage/example_openlineage_manual_lineage_dag.py b/providers/openlineage/tests/system/openlineage/example_openlineage_manual_lineage_dag.py index a829184c989d9..f754dcdfb87d4 100644 --- a/providers/openlineage/tests/system/openlineage/example_openlineage_manual_lineage_dag.py +++ b/providers/openlineage/tests/system/openlineage/example_openlineage_manual_lineage_dag.py @@ -46,6 +46,7 @@ from airflow.decorators import dag, task # type: ignore[no-redef, attr-defined] from airflow.operators.python import get_current_context # type: ignore[no-redef] +from system.openlineage.constants import DEFAULT_DAGRUN_TIMEOUT from system.openlineage.expected_events import get_expected_event_file_path from system.openlineage.operator import OpenLineageTestOperator @@ -71,6 +72,7 @@ def _dataset_with_assertions(namespace: str, name: str) -> Dataset: @dag( + dagrun_timeout=DEFAULT_DAGRUN_TIMEOUT, dag_id=DAG_ID, start_date=dt.datetime(2024, 1, 1), schedule=None, diff --git a/providers/openlineage/tests/system/openlineage/example_openlineage_mapped_simple_dag.py b/providers/openlineage/tests/system/openlineage/example_openlineage_mapped_simple_dag.py index c58651974c828..dcc7d4bab80a5 100644 --- a/providers/openlineage/tests/system/openlineage/example_openlineage_mapped_simple_dag.py +++ b/providers/openlineage/tests/system/openlineage/example_openlineage_mapped_simple_dag.py @@ -35,12 +35,14 @@ # Airflow 2 path from airflow.decorators import task # type: ignore[attr-defined,no-redef] +from system.openlineage.constants import DEFAULT_DAGRUN_TIMEOUT from system.openlineage.expected_events import get_expected_event_file_path from system.openlineage.operator import OpenLineageTestOperator DAG_ID = "openlineage_mapped_simple_dag" with DAG( + dagrun_timeout=DEFAULT_DAGRUN_TIMEOUT, dag_id=DAG_ID, start_date=datetime(2021, 1, 1), schedule=None, diff --git a/providers/openlineage/tests/system/openlineage/example_openlineage_policy_conf_locked_dag.py b/providers/openlineage/tests/system/openlineage/example_openlineage_policy_conf_locked_dag.py index 42420c199afc0..9ce8bd6c178dc 100644 --- a/providers/openlineage/tests/system/openlineage/example_openlineage_policy_conf_locked_dag.py +++ b/providers/openlineage/tests/system/openlineage/example_openlineage_policy_conf_locked_dag.py @@ -30,12 +30,14 @@ from airflow.providers.openlineage.api.emission_policy import extend_global_openlineage_emission_policy from airflow.providers.standard.operators.bash import BashOperator +from system.openlineage.constants import DEFAULT_DAGRUN_TIMEOUT from system.openlineage.expected_events import get_expected_event_file_path from system.openlineage.operator import OpenLineageTestOperator DAG_ID = "openlineage_policy_conf_locked_dag" with DAG( + dagrun_timeout=DEFAULT_DAGRUN_TIMEOUT, dag_id=DAG_ID, start_date=datetime(2021, 1, 1), schedule=None, diff --git a/providers/openlineage/tests/system/openlineage/example_openlineage_policy_conf_source_code_dag.py b/providers/openlineage/tests/system/openlineage/example_openlineage_policy_conf_source_code_dag.py index 8900f11fcf2e8..7740267ee69e3 100644 --- a/providers/openlineage/tests/system/openlineage/example_openlineage_policy_conf_source_code_dag.py +++ b/providers/openlineage/tests/system/openlineage/example_openlineage_policy_conf_source_code_dag.py @@ -30,12 +30,14 @@ from airflow.providers.openlineage.api.emission_policy import extend_global_openlineage_emission_policy from airflow.providers.standard.operators.bash import BashOperator +from system.openlineage.constants import DEFAULT_DAGRUN_TIMEOUT from system.openlineage.expected_events import get_expected_event_file_path from system.openlineage.operator import OpenLineageTestOperator DAG_ID = "openlineage_policy_conf_source_code_dag" with DAG( + dagrun_timeout=DEFAULT_DAGRUN_TIMEOUT, dag_id=DAG_ID, start_date=datetime(2021, 1, 1), schedule=None, diff --git a/providers/openlineage/tests/system/openlineage/example_openlineage_policy_dag_emit_false_dag.py b/providers/openlineage/tests/system/openlineage/example_openlineage_policy_dag_emit_false_dag.py index 11f8b9d64254d..0f1ec517150ad 100644 --- a/providers/openlineage/tests/system/openlineage/example_openlineage_policy_dag_emit_false_dag.py +++ b/providers/openlineage/tests/system/openlineage/example_openlineage_policy_dag_emit_false_dag.py @@ -28,6 +28,7 @@ from airflow.providers.openlineage.api.emission_policy import extend_global_openlineage_emission_policy from airflow.providers.standard.operators.python import PythonOperator +from system.openlineage.constants import DEFAULT_DAGRUN_TIMEOUT from system.openlineage.expected_events import get_expected_event_file_path from system.openlineage.operator import OpenLineageTestOperator @@ -39,6 +40,7 @@ def _say_hello(): DAG_ID = "openlineage_policy_dag_emit_false_dag" with DAG( + dagrun_timeout=DEFAULT_DAGRUN_TIMEOUT, dag_id=DAG_ID, start_date=datetime(2021, 1, 1), schedule=None, diff --git a/providers/openlineage/tests/system/openlineage/example_openlineage_policy_dag_events_false_dag.py b/providers/openlineage/tests/system/openlineage/example_openlineage_policy_dag_events_false_dag.py index 52d754b3132e2..f550d162d6fc9 100644 --- a/providers/openlineage/tests/system/openlineage/example_openlineage_policy_dag_events_false_dag.py +++ b/providers/openlineage/tests/system/openlineage/example_openlineage_policy_dag_events_false_dag.py @@ -30,6 +30,7 @@ from airflow.providers.openlineage.api.emission_policy import extend_global_openlineage_emission_policy from airflow.providers.standard.operators.python import PythonOperator +from system.openlineage.constants import DEFAULT_DAGRUN_TIMEOUT from system.openlineage.expected_events import get_expected_event_file_path from system.openlineage.operator import OpenLineageTestOperator @@ -47,6 +48,7 @@ def _register_hook_lineage(): DAG_ID = "openlineage_policy_dag_events_false_dag" with DAG( + dagrun_timeout=DEFAULT_DAGRUN_TIMEOUT, dag_id=DAG_ID, start_date=datetime(2021, 1, 1), schedule=None, diff --git a/providers/openlineage/tests/system/openlineage/example_openlineage_policy_dag_override_task_dag.py b/providers/openlineage/tests/system/openlineage/example_openlineage_policy_dag_override_task_dag.py index b326e49330143..1ae0d8f8c569e 100644 --- a/providers/openlineage/tests/system/openlineage/example_openlineage_policy_dag_override_task_dag.py +++ b/providers/openlineage/tests/system/openlineage/example_openlineage_policy_dag_override_task_dag.py @@ -28,12 +28,14 @@ from airflow.providers.openlineage.api.emission_policy import extend_global_openlineage_emission_policy from airflow.providers.standard.operators.bash import BashOperator +from system.openlineage.constants import DEFAULT_DAGRUN_TIMEOUT from system.openlineage.expected_events import get_expected_event_file_path from system.openlineage.operator import OpenLineageTestOperator DAG_ID = "openlineage_policy_dag_override_task_dag" with DAG( + dagrun_timeout=DEFAULT_DAGRUN_TIMEOUT, dag_id=DAG_ID, start_date=datetime(2021, 1, 1), schedule=None, diff --git a/providers/openlineage/tests/system/openlineage/example_openlineage_policy_extract_metadata_dag.py b/providers/openlineage/tests/system/openlineage/example_openlineage_policy_extract_metadata_dag.py index 112638a587ba8..f30cc112b1d74 100644 --- a/providers/openlineage/tests/system/openlineage/example_openlineage_policy_extract_metadata_dag.py +++ b/providers/openlineage/tests/system/openlineage/example_openlineage_policy_extract_metadata_dag.py @@ -31,6 +31,7 @@ from airflow.providers.openlineage.api.emission_policy import extend_global_openlineage_emission_policy from airflow.providers.openlineage.extractors.base import OperatorLineage +from system.openlineage.constants import DEFAULT_DAGRUN_TIMEOUT from system.openlineage.expected_events import get_expected_event_file_path from system.openlineage.operator import OpenLineageTestOperator @@ -51,6 +52,7 @@ def get_openlineage_facets_on_complete(self, task_instance): DAG_ID = "openlineage_policy_extract_metadata_dag" with DAG( + dagrun_timeout=DEFAULT_DAGRUN_TIMEOUT, dag_id=DAG_ID, start_date=datetime(2021, 1, 1), schedule=None, diff --git a/providers/openlineage/tests/system/openlineage/example_openlineage_policy_full_task_info_dag.py b/providers/openlineage/tests/system/openlineage/example_openlineage_policy_full_task_info_dag.py index 1bf264af68b93..38d48a55137bc 100644 --- a/providers/openlineage/tests/system/openlineage/example_openlineage_policy_full_task_info_dag.py +++ b/providers/openlineage/tests/system/openlineage/example_openlineage_policy_full_task_info_dag.py @@ -28,12 +28,14 @@ from airflow.providers.openlineage.api.emission_policy import extend_global_openlineage_emission_policy from airflow.providers.standard.operators.bash import BashOperator +from system.openlineage.constants import DEFAULT_DAGRUN_TIMEOUT from system.openlineage.expected_events import get_expected_event_file_path from system.openlineage.operator import OpenLineageTestOperator DAG_ID = "openlineage_policy_full_task_info_dag" with DAG( + dagrun_timeout=DEFAULT_DAGRUN_TIMEOUT, dag_id=DAG_ID, start_date=datetime(2021, 1, 1), schedule=None, diff --git a/providers/openlineage/tests/system/openlineage/example_openlineage_policy_hook_lineage_dag.py b/providers/openlineage/tests/system/openlineage/example_openlineage_policy_hook_lineage_dag.py index 6f1e8a76f972f..d04641530ea47 100644 --- a/providers/openlineage/tests/system/openlineage/example_openlineage_policy_hook_lineage_dag.py +++ b/providers/openlineage/tests/system/openlineage/example_openlineage_policy_hook_lineage_dag.py @@ -32,6 +32,7 @@ from airflow.providers.openlineage.api.emission_policy import extend_global_openlineage_emission_policy from airflow.providers.standard.operators.python import PythonOperator +from system.openlineage.constants import DEFAULT_DAGRUN_TIMEOUT from system.openlineage.expected_events import get_expected_event_file_path from system.openlineage.operator import OpenLineageTestOperator @@ -61,6 +62,7 @@ def _register_sql_hook_lineage(): DAG_ID = "openlineage_policy_hook_lineage_dag" with DAG( + dagrun_timeout=DEFAULT_DAGRUN_TIMEOUT, dag_id=DAG_ID, start_date=datetime(2021, 1, 1), schedule=None, diff --git a/providers/openlineage/tests/system/openlineage/example_openlineage_policy_source_code_dag.py b/providers/openlineage/tests/system/openlineage/example_openlineage_policy_source_code_dag.py index 7ccff61e7a44c..0f194a222a616 100644 --- a/providers/openlineage/tests/system/openlineage/example_openlineage_policy_source_code_dag.py +++ b/providers/openlineage/tests/system/openlineage/example_openlineage_policy_source_code_dag.py @@ -32,12 +32,14 @@ from airflow.providers.openlineage.api.emission_policy import extend_global_openlineage_emission_policy from airflow.providers.standard.operators.bash import BashOperator +from system.openlineage.constants import DEFAULT_DAGRUN_TIMEOUT from system.openlineage.expected_events import get_expected_event_file_path from system.openlineage.operator import OpenLineageTestOperator DAG_ID = "openlineage_policy_source_code_dag" with DAG( + dagrun_timeout=DEFAULT_DAGRUN_TIMEOUT, dag_id=DAG_ID, start_date=datetime(2021, 1, 1), schedule=None, diff --git a/providers/openlineage/tests/system/openlineage/example_openlineage_policy_task_emit_false_dag.py b/providers/openlineage/tests/system/openlineage/example_openlineage_policy_task_emit_false_dag.py index 19569c0f82b7f..e964424cf7873 100644 --- a/providers/openlineage/tests/system/openlineage/example_openlineage_policy_task_emit_false_dag.py +++ b/providers/openlineage/tests/system/openlineage/example_openlineage_policy_task_emit_false_dag.py @@ -28,6 +28,7 @@ from airflow.providers.openlineage.api.emission_policy import extend_global_openlineage_emission_policy from airflow.providers.standard.operators.python import PythonOperator +from system.openlineage.constants import DEFAULT_DAGRUN_TIMEOUT from system.openlineage.expected_events import get_expected_event_file_path from system.openlineage.operator import OpenLineageTestOperator @@ -39,6 +40,7 @@ def _say_hello(): DAG_ID = "openlineage_policy_task_emit_false_dag" with DAG( + dagrun_timeout=DEFAULT_DAGRUN_TIMEOUT, dag_id=DAG_ID, start_date=datetime(2021, 1, 1), schedule=None, diff --git a/providers/openlineage/tests/system/openlineage/example_openlineage_schedule_asset_or_time_dag.py b/providers/openlineage/tests/system/openlineage/example_openlineage_schedule_asset_or_time_dag.py index f72c35d61d251..9eacc660b5ff5 100644 --- a/providers/openlineage/tests/system/openlineage/example_openlineage_schedule_asset_or_time_dag.py +++ b/providers/openlineage/tests/system/openlineage/example_openlineage_schedule_asset_or_time_dag.py @@ -30,6 +30,7 @@ from airflow.providers.standard.operators.bash import BashOperator from airflow.timetables.trigger import CronTriggerTimetable +from system.openlineage.constants import DEFAULT_DAGRUN_TIMEOUT from system.openlineage.expected_events import AIRFLOW_VERSION, get_expected_event_file_path from system.openlineage.operator import OpenLineageTestOperator @@ -57,6 +58,7 @@ ) with DAG( + dagrun_timeout=DEFAULT_DAGRUN_TIMEOUT, dag_id=DAG_ID, start_date=datetime(2025, 5, 1), schedule=schedule, diff --git a/providers/openlineage/tests/system/openlineage/example_openlineage_schedule_cron_dag.py b/providers/openlineage/tests/system/openlineage/example_openlineage_schedule_cron_dag.py index d44c3e01719c6..7d29ac4d33d6b 100644 --- a/providers/openlineage/tests/system/openlineage/example_openlineage_schedule_cron_dag.py +++ b/providers/openlineage/tests/system/openlineage/example_openlineage_schedule_cron_dag.py @@ -28,12 +28,14 @@ from airflow import DAG from airflow.providers.standard.operators.bash import BashOperator +from system.openlineage.constants import DEFAULT_DAGRUN_TIMEOUT from system.openlineage.expected_events import get_expected_event_file_path from system.openlineage.operator import OpenLineageTestOperator DAG_ID = "openlineage_schedule_cron_dag" with DAG( + dagrun_timeout=DEFAULT_DAGRUN_TIMEOUT, dag_id=DAG_ID, start_date=datetime(2021, 1, 1), schedule="21 13 29 2 4", # Unlikely to ever get triggered by itself, February 29th and Thursday diff --git a/providers/openlineage/tests/system/openlineage/example_openlineage_schedule_cron_trigger_timetable_dag.py b/providers/openlineage/tests/system/openlineage/example_openlineage_schedule_cron_trigger_timetable_dag.py index 2a0979c2d8e7a..f8dee74e93ac0 100644 --- a/providers/openlineage/tests/system/openlineage/example_openlineage_schedule_cron_trigger_timetable_dag.py +++ b/providers/openlineage/tests/system/openlineage/example_openlineage_schedule_cron_trigger_timetable_dag.py @@ -31,12 +31,14 @@ from airflow.providers.standard.operators.bash import BashOperator from airflow.timetables.trigger import CronTriggerTimetable +from system.openlineage.constants import DEFAULT_DAGRUN_TIMEOUT from system.openlineage.expected_events import get_expected_event_file_path from system.openlineage.operator import OpenLineageTestOperator DAG_ID = "openlineage_schedule_cron_trigger_timetable_dag" with DAG( + dagrun_timeout=DEFAULT_DAGRUN_TIMEOUT, dag_id=DAG_ID, start_date=datetime(2021, 1, 1), schedule=CronTriggerTimetable("21 13 29 2 4", timezone="UTC"), diff --git a/providers/openlineage/tests/system/openlineage/example_openlineage_schedule_delta_trigger_timetable_dag.py b/providers/openlineage/tests/system/openlineage/example_openlineage_schedule_delta_trigger_timetable_dag.py index 7ed88d38e6811..ead0ddf95ec14 100644 --- a/providers/openlineage/tests/system/openlineage/example_openlineage_schedule_delta_trigger_timetable_dag.py +++ b/providers/openlineage/tests/system/openlineage/example_openlineage_schedule_delta_trigger_timetable_dag.py @@ -31,12 +31,14 @@ from airflow.providers.standard.operators.bash import BashOperator from airflow.timetables.trigger import DeltaTriggerTimetable +from system.openlineage.constants import DEFAULT_DAGRUN_TIMEOUT from system.openlineage.expected_events import get_expected_event_file_path from system.openlineage.operator import OpenLineageTestOperator DAG_ID = "openlineage_schedule_delta_trigger_timetable_dag" with DAG( + dagrun_timeout=DEFAULT_DAGRUN_TIMEOUT, dag_id=DAG_ID, start_date=datetime(2021, 1, 1), schedule=DeltaTriggerTimetable(timedelta(days=987)), diff --git a/providers/openlineage/tests/system/openlineage/example_openlineage_schedule_list_complex_assets_dag.py b/providers/openlineage/tests/system/openlineage/example_openlineage_schedule_list_complex_assets_dag.py index 88e5399fd619d..558dcb47e792d 100644 --- a/providers/openlineage/tests/system/openlineage/example_openlineage_schedule_list_complex_assets_dag.py +++ b/providers/openlineage/tests/system/openlineage/example_openlineage_schedule_list_complex_assets_dag.py @@ -30,6 +30,7 @@ from airflow.providers.common.compat.assets import Asset from airflow.providers.standard.operators.bash import BashOperator +from system.openlineage.constants import DEFAULT_DAGRUN_TIMEOUT from system.openlineage.expected_events import AIRFLOW_VERSION, get_expected_event_file_path from system.openlineage.operator import OpenLineageTestOperator @@ -46,6 +47,7 @@ DAG_ID = "openlineage_schedule_list_complex_assets_dag" with DAG( + dagrun_timeout=DEFAULT_DAGRUN_TIMEOUT, dag_id=DAG_ID, start_date=datetime(2021, 1, 1), schedule=schedule, diff --git a/providers/openlineage/tests/system/openlineage/example_openlineage_schedule_list_multiple_assets_dag.py b/providers/openlineage/tests/system/openlineage/example_openlineage_schedule_list_multiple_assets_dag.py index e51f18619ac55..649c377d9b064 100644 --- a/providers/openlineage/tests/system/openlineage/example_openlineage_schedule_list_multiple_assets_dag.py +++ b/providers/openlineage/tests/system/openlineage/example_openlineage_schedule_list_multiple_assets_dag.py @@ -29,12 +29,14 @@ from airflow.providers.common.compat.assets import Asset from airflow.providers.standard.operators.bash import BashOperator +from system.openlineage.constants import DEFAULT_DAGRUN_TIMEOUT from system.openlineage.expected_events import get_expected_event_file_path from system.openlineage.operator import OpenLineageTestOperator DAG_ID = "openlineage_schedule_list_multiple_assets_dag" with DAG( + dagrun_timeout=DEFAULT_DAGRUN_TIMEOUT, dag_id=DAG_ID, start_date=datetime(2021, 1, 1), schedule=[Asset(uri="s3://bucket/file.txt", extra={"a": 1}), Asset(uri="s3://bucket2/file.txt")], diff --git a/providers/openlineage/tests/system/openlineage/example_openlineage_schedule_list_single_asset_dag.py b/providers/openlineage/tests/system/openlineage/example_openlineage_schedule_list_single_asset_dag.py index 1af5a6112b8e6..e1f0832114d23 100644 --- a/providers/openlineage/tests/system/openlineage/example_openlineage_schedule_list_single_asset_dag.py +++ b/providers/openlineage/tests/system/openlineage/example_openlineage_schedule_list_single_asset_dag.py @@ -29,12 +29,14 @@ from airflow.providers.common.compat.assets import Asset from airflow.providers.standard.operators.bash import BashOperator +from system.openlineage.constants import DEFAULT_DAGRUN_TIMEOUT from system.openlineage.expected_events import get_expected_event_file_path from system.openlineage.operator import OpenLineageTestOperator DAG_ID = "openlineage_schedule_list_single_asset_dag" with DAG( + dagrun_timeout=DEFAULT_DAGRUN_TIMEOUT, dag_id=DAG_ID, start_date=datetime(2021, 1, 1), schedule=[Asset(uri="s3://bucket/file.txt", extra={"some_extra": 123})], diff --git a/providers/openlineage/tests/system/openlineage/example_openlineage_schedule_multiple_assets_dag.py b/providers/openlineage/tests/system/openlineage/example_openlineage_schedule_multiple_assets_dag.py index a0e1646430088..3625133c38b8e 100644 --- a/providers/openlineage/tests/system/openlineage/example_openlineage_schedule_multiple_assets_dag.py +++ b/providers/openlineage/tests/system/openlineage/example_openlineage_schedule_multiple_assets_dag.py @@ -29,12 +29,14 @@ from airflow.providers.common.compat.assets import Asset from airflow.providers.standard.operators.bash import BashOperator +from system.openlineage.constants import DEFAULT_DAGRUN_TIMEOUT from system.openlineage.expected_events import get_expected_event_file_path from system.openlineage.operator import OpenLineageTestOperator DAG_ID = "openlineage_schedule_multiple_assets_dag" with DAG( + dagrun_timeout=DEFAULT_DAGRUN_TIMEOUT, dag_id=DAG_ID, start_date=datetime(2021, 1, 1), schedule=( diff --git a/providers/openlineage/tests/system/openlineage/example_openlineage_schedule_multiple_cron_trigger_timetable_dag.py b/providers/openlineage/tests/system/openlineage/example_openlineage_schedule_multiple_cron_trigger_timetable_dag.py index 7fe5e06d48b34..f296fad0a6fd9 100644 --- a/providers/openlineage/tests/system/openlineage/example_openlineage_schedule_multiple_cron_trigger_timetable_dag.py +++ b/providers/openlineage/tests/system/openlineage/example_openlineage_schedule_multiple_cron_trigger_timetable_dag.py @@ -31,12 +31,14 @@ from airflow.providers.standard.operators.bash import BashOperator from airflow.timetables.trigger import MultipleCronTriggerTimetable +from system.openlineage.constants import DEFAULT_DAGRUN_TIMEOUT from system.openlineage.expected_events import get_expected_event_file_path from system.openlineage.operator import OpenLineageTestOperator DAG_ID = "openlineage_schedule_multiple_cron_trigger_timetable_dag" with DAG( + dagrun_timeout=DEFAULT_DAGRUN_TIMEOUT, dag_id=DAG_ID, start_date=datetime(2021, 1, 1), schedule=MultipleCronTriggerTimetable("21 13 29 2 4", "9 27 29 2 4", timezone="UTC"), diff --git a/providers/openlineage/tests/system/openlineage/example_openlineage_schedule_single_asset_dag.py b/providers/openlineage/tests/system/openlineage/example_openlineage_schedule_single_asset_dag.py index 64cd6ed34ecb9..277dab5adb6d9 100644 --- a/providers/openlineage/tests/system/openlineage/example_openlineage_schedule_single_asset_dag.py +++ b/providers/openlineage/tests/system/openlineage/example_openlineage_schedule_single_asset_dag.py @@ -29,12 +29,14 @@ from airflow.providers.common.compat.assets import Asset from airflow.providers.standard.operators.bash import BashOperator +from system.openlineage.constants import DEFAULT_DAGRUN_TIMEOUT from system.openlineage.expected_events import get_expected_event_file_path from system.openlineage.operator import OpenLineageTestOperator DAG_ID = "openlineage_schedule_single_asset_dag" with DAG( + dagrun_timeout=DEFAULT_DAGRUN_TIMEOUT, dag_id=DAG_ID, start_date=datetime(2021, 1, 1), schedule=Asset(uri="s3://bucket/file.txt", extra={"some_extra": 123}), diff --git a/providers/openlineage/tests/system/openlineage/example_openlineage_schedule_timedelta_dag.py b/providers/openlineage/tests/system/openlineage/example_openlineage_schedule_timedelta_dag.py index db2370e855108..f28f4eac955ba 100644 --- a/providers/openlineage/tests/system/openlineage/example_openlineage_schedule_timedelta_dag.py +++ b/providers/openlineage/tests/system/openlineage/example_openlineage_schedule_timedelta_dag.py @@ -30,12 +30,14 @@ from airflow import DAG from airflow.providers.standard.operators.bash import BashOperator +from system.openlineage.constants import DEFAULT_DAGRUN_TIMEOUT from system.openlineage.expected_events import get_expected_event_file_path from system.openlineage.operator import OpenLineageTestOperator DAG_ID = "openlineage_schedule_timedelta_dag" with DAG( + dagrun_timeout=DEFAULT_DAGRUN_TIMEOUT, dag_id=DAG_ID, start_date=datetime(2021, 1, 1), schedule=timedelta(days=987), diff --git a/providers/openlineage/tests/system/openlineage/example_openlineage_schedule_timetable_dag.py b/providers/openlineage/tests/system/openlineage/example_openlineage_schedule_timetable_dag.py index c90c0b16a915f..bf7d3ec8e3ab2 100644 --- a/providers/openlineage/tests/system/openlineage/example_openlineage_schedule_timetable_dag.py +++ b/providers/openlineage/tests/system/openlineage/example_openlineage_schedule_timetable_dag.py @@ -31,12 +31,14 @@ from airflow.providers.standard.operators.bash import BashOperator from airflow.timetables.events import EventsTimetable +from system.openlineage.constants import DEFAULT_DAGRUN_TIMEOUT from system.openlineage.expected_events import get_expected_event_file_path from system.openlineage.operator import OpenLineageTestOperator DAG_ID = "openlineage_schedule_timetable_dag" with DAG( + dagrun_timeout=DEFAULT_DAGRUN_TIMEOUT, dag_id=DAG_ID, start_date=datetime(2021, 1, 1), schedule=EventsTimetable( # Unlikely to ever get triggered by itself diff --git a/providers/openlineage/tests/system/openlineage/example_openlineage_setup_teardown_dag.py b/providers/openlineage/tests/system/openlineage/example_openlineage_setup_teardown_dag.py index 46c7c980481a8..8e61a55c952c8 100644 --- a/providers/openlineage/tests/system/openlineage/example_openlineage_setup_teardown_dag.py +++ b/providers/openlineage/tests/system/openlineage/example_openlineage_setup_teardown_dag.py @@ -29,6 +29,7 @@ from airflow import DAG from airflow.providers.standard.operators.python import PythonOperator +from system.openlineage.constants import DEFAULT_DAGRUN_TIMEOUT from system.openlineage.expected_events import get_expected_event_file_path from system.openlineage.operator import OpenLineageTestOperator @@ -40,6 +41,7 @@ def do_nothing(): DAG_ID = "openlineage_setup_teardown_dag" with DAG( + dagrun_timeout=DEFAULT_DAGRUN_TIMEOUT, dag_id=DAG_ID, start_date=datetime(2021, 1, 1), schedule=None, diff --git a/providers/openlineage/tests/system/openlineage/example_openlineage_short_circuit_dag.py b/providers/openlineage/tests/system/openlineage/example_openlineage_short_circuit_dag.py index 5ce68293ed43e..f404977fd4d6f 100644 --- a/providers/openlineage/tests/system/openlineage/example_openlineage_short_circuit_dag.py +++ b/providers/openlineage/tests/system/openlineage/example_openlineage_short_circuit_dag.py @@ -29,6 +29,7 @@ from airflow import DAG from airflow.providers.standard.operators.python import PythonOperator, ShortCircuitOperator +from system.openlineage.constants import DEFAULT_DAGRUN_TIMEOUT from system.openlineage.expected_events import get_expected_event_file_path from system.openlineage.operator import OpenLineageTestOperator @@ -40,6 +41,7 @@ def do_nothing(): DAG_ID = "openlineage_short_circuit_dag" with DAG( + dagrun_timeout=DEFAULT_DAGRUN_TIMEOUT, dag_id=DAG_ID, start_date=datetime(2021, 1, 1), schedule=None, diff --git a/providers/openlineage/tests/system/openlineage/example_openlineage_task_groups_dag.py b/providers/openlineage/tests/system/openlineage/example_openlineage_task_groups_dag.py index 9c081991d1dda..01f3c9a029f2e 100644 --- a/providers/openlineage/tests/system/openlineage/example_openlineage_task_groups_dag.py +++ b/providers/openlineage/tests/system/openlineage/example_openlineage_task_groups_dag.py @@ -33,12 +33,14 @@ except ImportError: from airflow.utils.task_group import TaskGroup # type: ignore[no-redef] +from system.openlineage.constants import DEFAULT_DAGRUN_TIMEOUT from system.openlineage.expected_events import get_expected_event_file_path from system.openlineage.operator import OpenLineageTestOperator DAG_ID = "openlineage_task_groups_dag" with DAG( + dagrun_timeout=DEFAULT_DAGRUN_TIMEOUT, dag_id=DAG_ID, start_date=datetime(2021, 1, 1), schedule=None, diff --git a/providers/openlineage/tests/system/openlineage/example_openlineage_taskflow_simple_dag.py b/providers/openlineage/tests/system/openlineage/example_openlineage_taskflow_simple_dag.py index 7008331ceb47a..ab0c26d3cb42c 100644 --- a/providers/openlineage/tests/system/openlineage/example_openlineage_taskflow_simple_dag.py +++ b/providers/openlineage/tests/system/openlineage/example_openlineage_taskflow_simple_dag.py @@ -27,6 +27,7 @@ from datetime import datetime +from system.openlineage.constants import DEFAULT_DAGRUN_TIMEOUT from system.openlineage.expected_events import get_expected_event_file_path from system.openlineage.operator import OpenLineageTestOperator @@ -39,7 +40,13 @@ DAG_ID = "openlineage_taskflow_simple_dag" -@dag(schedule=None, start_date=datetime(2021, 1, 1), catchup=False, default_args={"retries": 0}) +@dag( + dagrun_timeout=DEFAULT_DAGRUN_TIMEOUT, + schedule=None, + start_date=datetime(2021, 1, 1), + catchup=False, + default_args={"retries": 0}, +) def openlineage_taskflow_simple_dag(): @task def do_nothing_task(**context): diff --git a/providers/openlineage/tests/system/openlineage/example_openlineage_trigger_dag.py b/providers/openlineage/tests/system/openlineage/example_openlineage_trigger_dag.py index 1fa1474a99280..779b13ce9b500 100644 --- a/providers/openlineage/tests/system/openlineage/example_openlineage_trigger_dag.py +++ b/providers/openlineage/tests/system/openlineage/example_openlineage_trigger_dag.py @@ -31,12 +31,14 @@ from airflow.providers.standard.operators.bash import BashOperator from airflow.providers.standard.operators.trigger_dagrun import TriggerDagRunOperator +from system.openlineage.constants import DEFAULT_DAGRUN_TIMEOUT from system.openlineage.expected_events import get_expected_event_file_path from system.openlineage.operator import OpenLineageTestOperator DAG_ID = "openlineage_trigger_dag" with DAG( + dagrun_timeout=DEFAULT_DAGRUN_TIMEOUT, dag_id=DAG_ID, start_date=datetime(2021, 1, 1), schedule=None, @@ -59,7 +61,7 @@ "rootParentJobName": "generate_report_sales_e2e", }, }, - poke_interval=10, + poke_interval=5, ) check_events = OpenLineageTestOperator( @@ -70,6 +72,7 @@ with DAG( + dagrun_timeout=DEFAULT_DAGRUN_TIMEOUT, dag_id="openlineage_trigger_dag_child__notrigger", start_date=datetime(2021, 1, 1), schedule=None, diff --git a/providers/openlineage/tests/system/openlineage/example_openlineage_trigger_dag_deferrable.py b/providers/openlineage/tests/system/openlineage/example_openlineage_trigger_dag_deferrable.py index ebe79dab4c04f..5546baca4defd 100644 --- a/providers/openlineage/tests/system/openlineage/example_openlineage_trigger_dag_deferrable.py +++ b/providers/openlineage/tests/system/openlineage/example_openlineage_trigger_dag_deferrable.py @@ -32,12 +32,14 @@ from airflow.providers.standard.operators.bash import BashOperator from airflow.providers.standard.operators.trigger_dagrun import TriggerDagRunOperator +from system.openlineage.constants import DEFAULT_DAGRUN_TIMEOUT from system.openlineage.expected_events import get_expected_event_file_path from system.openlineage.operator import OpenLineageTestOperator DAG_ID = "openlineage_trigger_dag_deferrable" with DAG( + dagrun_timeout=DEFAULT_DAGRUN_TIMEOUT, dag_id=DAG_ID, start_date=datetime(2021, 1, 1), schedule=None, @@ -63,6 +65,7 @@ with DAG( + dagrun_timeout=DEFAULT_DAGRUN_TIMEOUT, dag_id="openlineage_trigger_dag_deferrable_child__notrigger", start_date=datetime(2021, 1, 1), schedule=None, @@ -78,6 +81,7 @@ with DAG( + dagrun_timeout=DEFAULT_DAGRUN_TIMEOUT, dag_id="openlineage_trigger_dag_deferrable_child2__notrigger", start_date=datetime(2021, 1, 1), schedule=None, diff --git a/providers/openlineage/tests/system/openlineage/example_openlineage_trigger_failed_dag.py b/providers/openlineage/tests/system/openlineage/example_openlineage_trigger_failed_dag.py index 6f25749a87f26..b62bcd4708c1f 100644 --- a/providers/openlineage/tests/system/openlineage/example_openlineage_trigger_failed_dag.py +++ b/providers/openlineage/tests/system/openlineage/example_openlineage_trigger_failed_dag.py @@ -35,12 +35,14 @@ from airflow.providers.standard.operators.trigger_dagrun import TriggerDagRunOperator from airflow.utils.trigger_rule import TriggerRule +from system.openlineage.constants import DEFAULT_DAGRUN_TIMEOUT from system.openlineage.expected_events import get_expected_event_file_path from system.openlineage.operator import OpenLineageTestOperator DAG_ID = "openlineage_trigger_failed_dag" with DAG( + dagrun_timeout=DEFAULT_DAGRUN_TIMEOUT, dag_id=DAG_ID, start_date=datetime(2021, 1, 1), schedule=None, @@ -80,6 +82,7 @@ with DAG( + dagrun_timeout=DEFAULT_DAGRUN_TIMEOUT, dag_id="openlineage_trigger_failed_dag_child__notrigger", start_date=datetime(2021, 1, 1), schedule=None, diff --git a/providers/openlineage/tests/system/openlineage/example_openlineage_versioned_dag.py b/providers/openlineage/tests/system/openlineage/example_openlineage_versioned_dag.py index d26d8d7cdea35..df6b6dc2db288 100644 --- a/providers/openlineage/tests/system/openlineage/example_openlineage_versioned_dag.py +++ b/providers/openlineage/tests/system/openlineage/example_openlineage_versioned_dag.py @@ -30,6 +30,7 @@ from airflow import DAG from airflow.providers.standard.operators.python import PythonOperator +from system.openlineage.constants import DEFAULT_DAGRUN_TIMEOUT from system.openlineage.expected_events import get_expected_event_file_path from system.openlineage.operator import OpenLineageTestOperator @@ -41,6 +42,7 @@ def do_nothing(): DAG_ID = "openlineage_versioned_dag" with DAG( + dagrun_timeout=DEFAULT_DAGRUN_TIMEOUT, dag_id=DAG_ID, start_date=datetime(2021, 1, 1), schedule=None, diff --git a/providers/openlineage/tests/system/openlineage/expected_events/openlineage_custom_operator_failure_dag.json b/providers/openlineage/tests/system/openlineage/expected_events/openlineage_custom_operator_failure_dag.json index 448e53e09b9a9..b66a4acae336e 100644 --- a/providers/openlineage/tests/system/openlineage/expected_events/openlineage_custom_operator_failure_dag.json +++ b/providers/openlineage/tests/system/openlineage/expected_events/openlineage_custom_operator_failure_dag.json @@ -35,7 +35,7 @@ "errorMessage": { "message": "Intentional failure for OpenLineage on_failure testing", "programmingLanguage": "python", - "stackTrace": "{{ result.startswith('Traceback (most recent call last):\n') and result.endswith('example_openlineage_custom_operator_failure_dag.py\", line 48, in execute\n raise ValueError(\"Intentional failure for OpenLineage on_failure testing\")\nValueError: Intentional failure for OpenLineage on_failure testing\n') }}" + "stackTrace": "{{ result.startswith('Traceback (most recent call last):\n') and result.endswith('example_openlineage_custom_operator_failure_dag.py\", line 49, in execute\n raise ValueError(\"Intentional failure for OpenLineage on_failure testing\")\nValueError: Intentional failure for OpenLineage on_failure testing\n') }}" }, "parent": { "job": { diff --git a/pyproject.toml b/pyproject.toml index 633a208e4dcfe..d197d7a7d7bdc 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -819,6 +819,7 @@ testing = ["dev", "providers.tests", "tests_common", "tests", "system", "unit", "airflow-core/tests/*" = ["D", "TID253", "S101", "TRY002"] "airflow-ctl-tests/tests/*" = ["D", "TID253", "S101", "TRY002"] "airflow-e2e-tests/tests/*" = ["D", "TID253", "S101", "TRY002"] +"providers-e2e-tests/*" = ["D", "TID253", "S101", "TRY002"] "docker-tests/*" = ["D", "TID253", "S101", "TRY002"] "task-sdk-integration-tests/*" = ["D", "TID253", "S101", "TRY002"] "kubernetes-tests/*" = ["D", "TID253", "S101", "TRY002"] @@ -1085,6 +1086,7 @@ mypy_path = [ "$MYPY_CONFIG_FILE_DIR/kubernetes-tests/tests", "$MYPY_CONFIG_FILE_DIR/docker-tests/tests", "$MYPY_CONFIG_FILE_DIR/task-sdk-integration-tests/tests", + "$MYPY_CONFIG_FILE_DIR/providers-e2e-tests/openlineage/tests", # Automatically generated mypy paths (update_airflow_pyproject_toml.py) "$MYPY_CONFIG_FILE_DIR/providers/airbyte/src", "$MYPY_CONFIG_FILE_DIR/providers/airbyte/tests", @@ -1344,6 +1346,7 @@ dev = [ "apache-airflow-devel-common[no-doc]", "apache-airflow-docker-tests", "apache-airflow-task-sdk-integration-tests", + "apache-airflow-openlineage-e2e-tests", "apache-airflow-helm-chart", "apache-airflow-kubernetes-tests", "apache-airflow-task-sdk[all]", @@ -1428,6 +1431,7 @@ apache-airflow-e2e-tests = false apache-airflow-helm-chart = false apache-airflow-kubernetes-tests = false apache-airflow-mypy = false +apache-airflow-openlineage-e2e-tests = false apache-airflow-providers = false apache-airflow-providers-airbyte = false apache-airflow-providers-akeyless = false @@ -1573,6 +1577,7 @@ apache-airflow-e2e-tests = false apache-airflow-helm-chart = false apache-airflow-kubernetes-tests = false apache-airflow-mypy = false +apache-airflow-openlineage-e2e-tests = false apache-airflow-providers = false apache-airflow-providers-airbyte = false apache-airflow-providers-akeyless = false @@ -1715,6 +1720,7 @@ apache-airflow-devel-common = { workspace = true } apache-airflow-docker-tests = { workspace = true } apache-airflow-e2e-tests = { workspace = true } apache-airflow-task-sdk-integration-tests = { workspace = true } +apache-airflow-openlineage-e2e-tests = { workspace = true } apache-airflow-helm-chart = { workspace = true } apache-airflow-kubernetes-tests = { workspace = true } apache-airflow-providers = { workspace = true } @@ -1854,6 +1860,7 @@ members = [ "scripts", "docker-tests", "task-sdk-integration-tests", + "providers-e2e-tests/openlineage", "chart", "kubernetes-tests", "task-sdk", diff --git a/scripts/ci/prek/run_mypy_full_dist_local_venv_or_breeze_in_ci.py b/scripts/ci/prek/run_mypy_full_dist_local_venv_or_breeze_in_ci.py index ca9f09adae6c7..ce1f79af0ac4e 100755 --- a/scripts/ci/prek/run_mypy_full_dist_local_venv_or_breeze_in_ci.py +++ b/scripts/ci/prek/run_mypy_full_dist_local_venv_or_breeze_in_ci.py @@ -68,6 +68,7 @@ "chart/tests", "airflow-e2e-tests", "task-sdk-integration-tests", + "providers-e2e-tests/openlineage", "docker-tests", "kubernetes-tests", ] @@ -127,6 +128,9 @@ r"^.*/node_modules/.*", r"^.*\\..*", r"^.*/src/airflow/__init__.py$", + # Generated-at-runtime artifacts of the provider e2e tests (gitignored; not source to check). + r"^.*/providers-e2e-tests/[^/]+/dags/.*", + r"^.*/providers-e2e-tests/[^/]+/provider_dist/.*", ] ] diff --git a/uv.lock b/uv.lock index 1ff0a8117e71f..351a0abdde7f7 100644 --- a/uv.lock +++ b/uv.lock @@ -21,6 +21,7 @@ exclude-newer = "0001-01-01T00:00:00Z" # This has no effect and is included for exclude-newer-span = "P4D" [options.exclude-newer-package] +apache-airflow-openlineage-e2e-tests = false apache-airflow-providers-informatica = false apache-airflow-providers-clickhousedb = false apache-airflow-providers-amazon = false @@ -171,6 +172,7 @@ members = [ "apache-airflow-helm-chart", "apache-airflow-kubernetes-tests", "apache-airflow-mypy", + "apache-airflow-openlineage-e2e-tests", "apache-airflow-providers", "apache-airflow-providers-airbyte", "apache-airflow-providers-akeyless", @@ -1504,6 +1506,7 @@ ci-image = [ { name = "apache-airflow-docker-tests" }, { name = "apache-airflow-helm-chart" }, { name = "apache-airflow-kubernetes-tests" }, + { name = "apache-airflow-openlineage-e2e-tests" }, { name = "apache-airflow-scripts" }, { name = "apache-airflow-shared-configuration" }, { name = "apache-airflow-shared-dagnode" }, @@ -1533,6 +1536,7 @@ dev = [ { name = "apache-airflow-docker-tests" }, { name = "apache-airflow-helm-chart" }, { name = "apache-airflow-kubernetes-tests" }, + { name = "apache-airflow-openlineage-e2e-tests" }, { name = "apache-airflow-scripts" }, { name = "apache-airflow-shared-configuration" }, { name = "apache-airflow-shared-dagnode" }, @@ -1814,6 +1818,7 @@ ci-image = [ { name = "apache-airflow-docker-tests", editable = "docker-tests" }, { name = "apache-airflow-helm-chart", editable = "chart" }, { name = "apache-airflow-kubernetes-tests", editable = "kubernetes-tests" }, + { name = "apache-airflow-openlineage-e2e-tests", editable = "providers-e2e-tests/openlineage" }, { name = "apache-airflow-scripts", editable = "scripts" }, { name = "apache-airflow-shared-configuration", editable = "shared/configuration" }, { name = "apache-airflow-shared-dagnode", editable = "shared/dagnode" }, @@ -1843,6 +1848,7 @@ dev = [ { name = "apache-airflow-docker-tests", editable = "docker-tests" }, { name = "apache-airflow-helm-chart", editable = "chart" }, { name = "apache-airflow-kubernetes-tests", editable = "kubernetes-tests" }, + { name = "apache-airflow-openlineage-e2e-tests", editable = "providers-e2e-tests/openlineage" }, { name = "apache-airflow-scripts", editable = "scripts" }, { name = "apache-airflow-shared-configuration", editable = "shared/configuration" }, { name = "apache-airflow-shared-dagnode", editable = "shared/dagnode" }, @@ -2885,6 +2891,29 @@ dependencies = [ [package.metadata] requires-dist = [{ name = "mypy", specifier = ">=1.0.0" }] +[[package]] +name = "apache-airflow-openlineage-e2e-tests" +version = "0.0.1" +source = { editable = "providers-e2e-tests/openlineage" } +dependencies = [ + { name = "apache-airflow-core" }, + { name = "apache-airflow-devel-common" }, +] + +[package.dev-dependencies] +mypy = [ + { name = "apache-airflow-devel-common", extra = ["mypy"] }, +] + +[package.metadata] +requires-dist = [ + { name = "apache-airflow-core", editable = "airflow-core" }, + { name = "apache-airflow-devel-common", editable = "devel-common" }, +] + +[package.metadata.requires-dev] +mypy = [{ name = "apache-airflow-devel-common", extras = ["mypy"], editable = "devel-common" }] + [[package]] name = "apache-airflow-providers" version = "0.0.1" From c1cccfdb157ad22b2c464c3a327645fa630c8780 Mon Sep 17 00:00:00 2001 From: Kacper Muda Date: Wed, 8 Jul 2026 11:29:17 +0200 Subject: [PATCH 2/2] Address review comments --- .../commands/testing_commands.py | 9 ++-- dev/breeze/tests/test_selective_checks.py | 2 +- .../openlineage/prepare_dags.py | 9 +++- .../openlineage/tests/harness.py | 48 +++++++++---------- 4 files changed, 37 insertions(+), 31 deletions(-) diff --git a/dev/breeze/src/airflow_breeze/commands/testing_commands.py b/dev/breeze/src/airflow_breeze/commands/testing_commands.py index fada767ddbcaa..504ac0f8a5757 100644 --- a/dev/breeze/src/airflow_breeze/commands/testing_commands.py +++ b/dev/breeze/src/airflow_breeze/commands/testing_commands.py @@ -924,6 +924,8 @@ def task_sdk_integration_tests( def _available_e2e_providers() -> list[str]: + if not PROVIDERS_E2E_TESTS_ROOT_PATH.exists(): + return [] return sorted( d.name for d in PROVIDERS_E2E_TESTS_ROOT_PATH.iterdir() @@ -967,6 +969,7 @@ def _build_providers_e2e_compat_image(provider: str, airflow_version: str, pytho "wheel", "--skip-tag-check", "--include-not-ready-providers", + "--clean-dist", ], check=True, cwd=AIRFLOW_ROOT_PATH, @@ -975,10 +978,8 @@ def _build_providers_e2e_compat_image(provider: str, airflow_version: str, pytho if provider_dist.exists(): shutil.rmtree(provider_dist) provider_dist.mkdir(parents=True) - for provider_id in required_providers: - wheel_prefix = f"apache_airflow_providers_{provider_id.replace('.', '_')}-" - for wheel in (AIRFLOW_ROOT_PATH / "dist").glob(f"{wheel_prefix}*.whl"): - shutil.copy2(wheel, provider_dist / wheel.name) + for wheel in (AIRFLOW_ROOT_PATH / "dist").glob("*.whl"): # With --clean-dist above we can copy all + shutil.copy2(wheel, provider_dist / wheel.name) base_image = f"apache/airflow:{airflow_version}-python{python}" image_name = f"{provider}-e2e/airflow:{airflow_version}-python{python}" diff --git a/dev/breeze/tests/test_selective_checks.py b/dev/breeze/tests/test_selective_checks.py index 8a6dbd8f7bb25..1e0e2644ffd5d 100644 --- a/dev/breeze/tests/test_selective_checks.py +++ b/dev/breeze/tests/test_selective_checks.py @@ -1491,7 +1491,7 @@ def assert_outputs_are_printed(expected_outputs: dict[str, str], stderr: str): id="Run OpenLineage e2e tests for common provider change", ), pytest.param( - ("providers-e2e-tests/openlineage/tests/openlineage_e2e_tests/test_openlineage_e2e.py",), + ("providers-e2e-tests/openlineage/tests/harness.py",), { "run-providers-e2e-tests-openlineage": "true", "prod-image-build": "true", diff --git a/providers-e2e-tests/openlineage/prepare_dags.py b/providers-e2e-tests/openlineage/prepare_dags.py index 228cd16943d72..7d6b51f45f855 100644 --- a/providers-e2e-tests/openlineage/prepare_dags.py +++ b/providers-e2e-tests/openlineage/prepare_dags.py @@ -81,8 +81,13 @@ def prepare_dags() -> Path: raise FileNotFoundError(f"OpenLineage system tests not found at {SYSTEM_TESTS_SOURCE}") if DAGS_DEST.exists(): - shutil.rmtree(DAGS_DEST) - DAGS_DEST.mkdir(parents=True) + for child in DAGS_DEST.iterdir(): + if child.is_dir(): + shutil.rmtree(child) + else: + child.unlink() + else: + DAGS_DEST.mkdir(parents=True) # Copy the whole `system` package so `system.openlineage.{operator,transport,expected_events}` # imports resolve from the dags folder. diff --git a/providers-e2e-tests/openlineage/tests/harness.py b/providers-e2e-tests/openlineage/tests/harness.py index b45a7c67ad218..05c2b23674273 100644 --- a/providers-e2e-tests/openlineage/tests/harness.py +++ b/providers-e2e-tests/openlineage/tests/harness.py @@ -15,12 +15,12 @@ # specific language governing permissions and limitations # under the License. """ -Drive all OpenLineage DAGs in the deployed stack and collect their final run state. +Drive all OpenLineage Dags in the deployed stack and collect their final run state. Adapted from the OpenLineage dags-dashboard ``trigger_dag.py`` integration harness: it talks to a local Airflow REST API v2 (rather than an Astro deployment), authenticates with a SimpleAuthManager bearer token, and reports through logs only. The actual OpenLineage event validation happens inside -each DAG's terminal ``OpenLineageTestOperator`` task, so a run that ends ``success`` means its +each Dag's terminal ``OpenLineageTestOperator`` task, so a run that ends ``success`` means its emitted events matched the expected templates. """ @@ -38,23 +38,23 @@ WARMUP_DAG_ID = "openlineage_warmup_dag" TRIGGER_DELAY_SECONDS = 2 -# Shared wall-clock budget for one _trigger_and_wait() call, covering every DAG triggered in that -# batch — bounds total polling time regardless of how many DAGs are involved +# Shared wall-clock budget for one _trigger_and_wait() call, covering every Dag triggered in that +# batch — bounds total polling time regardless of how many Dags are involved BATCH_TIMEOUT_SECONDS = 360 # 6 min, each dagrun timeout is set to 5 minutes -# DAGs that take longer to complete — trigger these first so they run while the rest are triggered. +# Dags that take longer to complete — trigger these first so they run while the rest are triggered. LONG_RUNNING_DAG_IDS = ("openlineage_defer_simple_dag",) -# Child DAGs spawned by TriggerDagRunOperator carry this marker; they must not be triggered or +# Child Dags spawned by TriggerDagRunOperator carry this marker; they must not be triggered or # state-checked directly (their parent triggers them). NO_TRIGGER_MARKER = "__notrigger" def discover_expected_dag_ids(dags_folder: Path) -> set[str]: """ - Top-level ``DAG_ID = "..."`` declarations across the prepared example DAGs. + Top-level ``DAG_ID = "..."`` declarations across the prepared example Dags. - Used as a coverage check: every expected DAG must actually load in the deployment (a missing one + Used as a coverage check: every expected Dag must actually load in the deployment (a missing one signals an import/parse failure rather than a test failure). """ pattern = re.compile(r'^DAG_ID\s*=\s*["\'](?P[^"\']+)["\']') @@ -84,7 +84,7 @@ def _session_with_retries() -> requests.Session: class OpenLineageE2ERunner: - """Triggers all OpenLineage DAGs against a running deployment and collects their final states.""" + """Triggers all OpenLineage Dags against a running deployment and collects their final states.""" def __init__(self, api_base_url: str, headers: dict[str, str]): self.api_url = f"{api_base_url}/api/v2" @@ -122,24 +122,24 @@ def get_task_states(self, dag_id: str, run_id: str) -> dict[str, str]: return {ti["task_id"]: ti["state"] for ti in response.json()["task_instances"]} def wait_for_dags_loaded(self, timeout: int = 60, poll_interval: int = 3) -> list[str]: - """Poll until the dag-processor has parsed the DAGs (the warmup DAG is the readiness marker).""" + """Poll until the dag-processor has parsed the Dags (the warmup Dag is the readiness marker).""" deadline = time.monotonic() + timeout dag_ids: list[str] = [] while time.monotonic() < deadline: dag_ids = self.list_dags() if WARMUP_DAG_ID in dag_ids: - console.print(f"[green]DAGs loaded ({len(dag_ids)} found)") + console.print(f"[green]Dags loaded ({len(dag_ids)} found)") return dag_ids - console.print(f"[yellow]Waiting for DAGs to load (have {len(dag_ids)})...") + console.print(f"[yellow]Waiting for Dags to load (have {len(dag_ids)})...") time.sleep(poll_interval) - raise RuntimeError(f"DAGs did not load within {timeout}s (warmup DAG missing; have {dag_ids}).") + raise RuntimeError(f"Dags did not load within {timeout}s (warmup Dag missing; have {dag_ids}).") def unpause_dag(self, dag_id: str) -> None: response = self.session.patch( f"{self.api_url}/dags/{dag_id}", headers=self.headers, json={"is_paused": False} ) if response.status_code != 200: - console.print(f"[red]Failed to unpause DAG `{dag_id}`: {response.text}") + console.print(f"[red]Failed to unpause Dag `{dag_id}`: {response.text}") def trigger_dag_run(self, dag_id: str, run_id: str) -> bool: now = dt.datetime.now(tz=dt.timezone.utc).isoformat() @@ -148,7 +148,7 @@ def trigger_dag_run(self, dag_id: str, run_id: str) -> bool: f"{self.api_url}/dags/{dag_id}/dagRuns", headers=self.headers, json=payload ) if response.status_code not in (200, 201): - console.print(f"[red]Failed to trigger DAG `{dag_id}`: {response.text}") + console.print(f"[red]Failed to trigger Dag `{dag_id}`: {response.text}") return False return True @@ -163,7 +163,7 @@ def wait_for_dag_run_to_complete(self, dag_id: str, run_id: str, timeout: int = if state not in ("running", "queued"): break time.sleep(5) - console.print(f"[blue]DAG `{dag_id}` finished in state: {state}") + console.print(f"[blue]Dag `{dag_id}` finished in state: {state}") return state def clear_airflow_variables(self) -> None: @@ -174,14 +174,14 @@ def clear_airflow_variables(self) -> None: self.session.delete(f"{self.api_url}/variables/{key}", headers=self.headers) def warmup(self, dag_ids: list[str]) -> None: - """Unpause all DAGs and run the warmup DAG so the worker is confirmed ready.""" + """Unpause all Dags and run the warmup Dag so the worker is confirmed ready.""" if WARMUP_DAG_ID not in dag_ids: - raise KeyError(f"Warmup DAG `{WARMUP_DAG_ID}` not found in deployment.") + raise KeyError(f"Warmup Dag `{WARMUP_DAG_ID}` not found in deployment.") for dag_id in dag_ids: self.unpause_dag(dag_id) time.sleep(TRIGGER_DELAY_SECONDS) if not self.trigger_dag_run(WARMUP_DAG_ID, self.run_id): - raise RuntimeError(f"Failed to trigger warmup DAG `{WARMUP_DAG_ID}`") + raise RuntimeError(f"Failed to trigger warmup Dag `{WARMUP_DAG_ID}`") self.wait_for_dag_run_to_complete(WARMUP_DAG_ID, self.run_id) def _trigger_and_wait( @@ -209,14 +209,14 @@ def _trigger_and_wait( return statuses def run(self, expected_dag_ids: set[str]) -> dict[str, str]: - """Run the full cycle and return ``{dag_id: final_state}`` for every triggered DAG.""" + """Run the full cycle and return ``{dag_id: final_state}`` for every triggered Dag.""" self.wait_for_airflow_api() dag_ids = self.wait_for_dags_loaded() if not dag_ids: - raise ValueError("No DAGs found in the deployment.") + raise ValueError("No Dags found in the deployment.") self.warmup(dag_ids) - # Auto-runs of cron/timetable DAGs may have started while unpaused — clear their events so + # Auto-runs of cron/timetable Dags may have started while unpaused — clear their events so # the test runs start from a clean slate. time.sleep(10) self.clear_airflow_variables() @@ -230,13 +230,13 @@ def run(self, expected_dag_ids: set[str]) -> dict[str, str]: failed = [dag_id for dag_id, state in statuses.items() if state != "success"] if failed: self.retried_dag_ids = set(failed) - console.print(f"[yellow]Retrying {len(failed)} failed DAG(s) once: {failed}") + console.print(f"[yellow]Retrying {len(failed)} failed Dag(s) once: {failed}") retry_statuses = self._trigger_and_wait(failed, self.retry_run_id) self.clear_airflow_variables() recovered = sorted(dag_id for dag_id, state in retry_statuses.items() if state == "success") if recovered: # The retry counts as a pass, but surface it so first-run flakiness is not hidden. - console.print(f"[yellow]⚠ DAGs that passed only on retry (flaky first run): {recovered}") + console.print(f"[yellow]⚠ Dags that passed only on retry (flaky first run): {recovered}") statuses.update(retry_statuses) for dag_id in expected_dag_ids: