diff --git a/.github/boring-cyborg.yml b/.github/boring-cyborg.yml index 779060fcd70ae..40470aa6ae598 100644 --- a/.github/boring-cyborg.yml +++ b/.github/boring-cyborg.yml @@ -401,7 +401,6 @@ labelPRBasedOnFilePath: - airflow/operators/**/* - airflow/hooks/**/* - airflow/sensors/**/* - - tests/operators/**/* - tests/hooks/**/* - tests/sensors/**/* - docs/apache-airflow/operators-and-hooks-ref.rst diff --git a/tests/operators/__init__.py b/airflow/api_fastapi/execution_api/datamodels/dagrun.py similarity index 58% rename from tests/operators/__init__.py rename to airflow/api_fastapi/execution_api/datamodels/dagrun.py index 217e5db960782..1f73156cebf5a 100644 --- a/tests/operators/__init__.py +++ b/airflow/api_fastapi/execution_api/datamodels/dagrun.py @@ -1,4 +1,3 @@ -# # 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 @@ -15,3 +14,25 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. + +from __future__ import annotations + +from pydantic import Field + +from airflow.api_fastapi.common.types import UtcDateTime +from airflow.api_fastapi.core_api.base import BaseModel, StrictBaseModel +from airflow.utils.state import DagRunState + + +class TriggerDAGRunPayload(StrictBaseModel): + """Schema for Trigger DAG Run API request.""" + + logical_date: UtcDateTime | None = None + conf: dict = Field(default_factory=dict) + reset_dag_run: bool = False + + +class DagRunStateResponse(BaseModel): + """Schema for DAG Run State response.""" + + state: DagRunState diff --git a/airflow/api_fastapi/execution_api/routes/__init__.py b/airflow/api_fastapi/execution_api/routes/__init__.py index 2f4dac922351f..0c21b9c053599 100644 --- a/airflow/api_fastapi/execution_api/routes/__init__.py +++ b/airflow/api_fastapi/execution_api/routes/__init__.py @@ -23,6 +23,7 @@ asset_events, assets, connections, + dag_runs, health, task_instances, variables, @@ -38,6 +39,7 @@ authenticated_router.include_router(assets.router, prefix="/assets", tags=["Assets"]) authenticated_router.include_router(asset_events.router, prefix="/asset-events", tags=["Asset Events"]) authenticated_router.include_router(connections.router, prefix="/connections", tags=["Connections"]) +authenticated_router.include_router(dag_runs.router, prefix="/dag-runs", tags=["Dag Runs"]) authenticated_router.include_router(task_instances.router, prefix="/task-instances", tags=["Task Instances"]) authenticated_router.include_router(variables.router, prefix="/variables", tags=["Variables"]) authenticated_router.include_router(xcoms.router, prefix="/xcoms", tags=["XComs"]) diff --git a/airflow/api_fastapi/execution_api/routes/dag_runs.py b/airflow/api_fastapi/execution_api/routes/dag_runs.py new file mode 100644 index 0000000000000..3a680c1ef8c69 --- /dev/null +++ b/airflow/api_fastapi/execution_api/routes/dag_runs.py @@ -0,0 +1,152 @@ +# 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 logging + +from fastapi import HTTPException, status +from sqlalchemy import select + +from airflow.api.common.trigger_dag import trigger_dag +from airflow.api_fastapi.common.db.common import SessionDep +from airflow.api_fastapi.common.router import AirflowRouter +from airflow.api_fastapi.execution_api.datamodels.dagrun import DagRunStateResponse, TriggerDAGRunPayload +from airflow.exceptions import DagRunAlreadyExists +from airflow.models.dag import DagModel +from airflow.models.dagbag import DagBag +from airflow.models.dagrun import DagRun +from airflow.utils.types import DagRunTriggeredByType + +router = AirflowRouter() + + +log = logging.getLogger(__name__) + + +@router.post( + "/{dag_id}/{run_id}", + status_code=status.HTTP_204_NO_CONTENT, + responses={ + status.HTTP_400_BAD_REQUEST: {"description": "DAG has import errors and cannot be triggered"}, + status.HTTP_404_NOT_FOUND: {"description": "DAG not found for the given dag_id"}, + status.HTTP_409_CONFLICT: {"description": "DAG Run already exists for the given dag_id"}, + status.HTTP_422_UNPROCESSABLE_ENTITY: {"description": "Invalid payload"}, + }, +) +def trigger_dag_run( + dag_id: str, + run_id: str, + payload: TriggerDAGRunPayload, + session: SessionDep, +): + """Trigger a DAG Run.""" + dm = session.scalar(select(DagModel).where(DagModel.is_active, DagModel.dag_id == dag_id).limit(1)) + if not dm: + raise HTTPException( + status.HTTP_404_NOT_FOUND, + detail={"reason": "not_found", "message": f"DAG with dag_id: '{dag_id}' not found"}, + ) + + if dm.has_import_errors: + raise HTTPException( + status.HTTP_400_BAD_REQUEST, + detail={ + "reason": "import_errors", + "message": f"DAG with dag_id: '{dag_id}' has import errors and cannot be triggered", + }, + ) + + try: + trigger_dag( + dag_id=dag_id, + run_id=run_id, + conf=payload.conf, + logical_date=payload.logical_date, + triggered_by=DagRunTriggeredByType.OPERATOR, + replace_microseconds=False, + session=session, + ) + except DagRunAlreadyExists: + raise HTTPException( + status.HTTP_409_CONFLICT, + detail={ + "reason": "already_exists", + "message": f"A DAG Run already exists for DAG {dag_id} with run id {run_id}", + }, + ) + + +@router.post( + "/{dag_id}/{run_id}/clear", + status_code=status.HTTP_204_NO_CONTENT, + responses={ + status.HTTP_400_BAD_REQUEST: {"description": "DAG has import errors and cannot be triggered"}, + status.HTTP_404_NOT_FOUND: {"description": "DAG not found for the given dag_id"}, + status.HTTP_422_UNPROCESSABLE_ENTITY: {"description": "Invalid payload"}, + }, +) +def clear_dag_run( + dag_id: str, + run_id: str, + session: SessionDep, +): + """Clear a DAG Run.""" + dm = session.scalar(select(DagModel).where(DagModel.is_active, DagModel.dag_id == dag_id).limit(1)) + if not dm: + raise HTTPException( + status.HTTP_404_NOT_FOUND, + detail={"reason": "not_found", "message": f"DAG with dag_id: '{dag_id}' not found"}, + ) + + if dm.has_import_errors: + raise HTTPException( + status.HTTP_400_BAD_REQUEST, + detail={ + "reason": "import_errors", + "message": f"DAG with dag_id: '{dag_id}' has import errors and cannot be triggered", + }, + ) + + dag_bag = DagBag(dag_folder=dm.fileloc, read_dags_from_db=True) + dag = dag_bag.get_dag(dag_id) + dag.clear(run_id=run_id) + + +@router.get( + "/{dag_id}/{run_id}/state", + responses={ + status.HTTP_404_NOT_FOUND: {"description": "DAG not found for the given dag_id"}, + }, +) +def get_dagrun_state( + dag_id: str, + run_id: str, + session: SessionDep, +) -> DagRunStateResponse: + """Get a DAG Run State.""" + dag_run = session.scalar(select(DagRun).where(DagRun.dag_id == dag_id, DagRun.run_id == run_id)) + if dag_run is None: + raise HTTPException( + status.HTTP_404_NOT_FOUND, + detail={ + "reason": "not_found", + "message": f"The DagRun with dag_id: `{dag_id}` and run_id: `{run_id}` was not found", + }, + ) + + return DagRunStateResponse(state=dag_run.state) diff --git a/airflow/exceptions.py b/airflow/exceptions.py index 4d3fba6101ea2..ecbd871d76da3 100644 --- a/airflow/exceptions.py +++ b/airflow/exceptions.py @@ -23,7 +23,7 @@ import warnings from collections.abc import Collection, Sequence -from datetime import timedelta +from datetime import datetime, timedelta from http import HTTPStatus from typing import TYPE_CHECKING, Any, NamedTuple @@ -34,6 +34,7 @@ from airflow.models import DagRun from airflow.sdk.definitions.asset import AssetNameRef, AssetUniqueKey, AssetUriRef + from airflow.utils.state import DagRunState class AirflowException(Exception): @@ -418,6 +419,41 @@ def __init__(self, *, tasks: Sequence[str | tuple[str, int]]): self.tasks = tasks +class DagRunTriggerException(AirflowException): + """ + Signal by an operator to trigger a specific Dag Run of a dag. + + Special exception raised to signal that the operator it was raised from wishes to trigger + a specific Dag Run of a dag. This is used in the ``TriggerDagRunOperator``. + """ + + def __init__( + self, + *, + trigger_dag_id: str, + dag_run_id: str, + conf: dict | None, + logical_date: datetime | None, + reset_dag_run: bool, + skip_when_already_exists: bool, + wait_for_completion: bool, + allowed_states: list[str | DagRunState], + failed_states: list[str | DagRunState], + poke_interval: int, + ): + super().__init__() + self.trigger_dag_id = trigger_dag_id + self.dag_run_id = dag_run_id + self.conf = conf + self.logical_date = logical_date + self.reset_dag_run = reset_dag_run + self.skip_when_already_exists = skip_when_already_exists + self.wait_for_completion = wait_for_completion + self.allowed_states = allowed_states + self.failed_states = failed_states + self.poke_interval = poke_interval + + class TaskDeferred(BaseException): """ Signal an operator moving to deferred state. diff --git a/dev/breeze/doc/ci/04_selective_checks.md b/dev/breeze/doc/ci/04_selective_checks.md index 4fb073aa040d2..40331b0be5995 100644 --- a/dev/breeze/doc/ci/04_selective_checks.md +++ b/dev/breeze/doc/ci/04_selective_checks.md @@ -107,7 +107,7 @@ together using `pytest-xdist` (pytest-xdist distributes the tests among parallel * 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 * The specific unit test type is enabled only if changed files match the expected patterns for each type - (`API`, `CLI`, `WWW`, `Providers`, `Operators` etc.). The `Always` test type is added always if any unit + (`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` * If there are no files left in sources after matching the test types and Kubernetes files, then apparently some Core/Other files have been changed. This automatically adds all test diff --git a/dev/breeze/doc/images/output-commands.svg b/dev/breeze/doc/images/output-commands.svg index dcb17ac367dd0..2d9668f5a21af 100644 --- a/dev/breeze/doc/images/output-commands.svg +++ b/dev/breeze/doc/images/output-commands.svg @@ -326,8 +326,8 @@ │kafka | kerberos | keycloak | mongo | mssql |          │ │openlineage | otel | pinot | qdrant | redis | statsd | │ │trino | ydb)                                           │ -│--standalone-dag-processor/--no-standalone-dag-process…Run standalone dag processor for start-airflow         │ -│(required for Airflow 3).                              │ +│--standalone-dag-processor/--no-standalone-dag-processoRun standalone dag processor for start-airflow         │ +│r(required for Airflow 3).                              │ │[default: standalone-dag-processor]                    │ ╰──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯ ╭─ Docker Compose selection and cleanup ───────────────────────────────────────────────────────────────────────────────╮ diff --git a/dev/breeze/doc/images/output_testing_core-tests.svg b/dev/breeze/doc/images/output_testing_core-tests.svg index 6bf94a4627c33..50b7ae2e654e3 100644 --- a/dev/breeze/doc/images/output_testing_core-tests.svg +++ b/dev/breeze/doc/images/output_testing_core-tests.svg @@ -380,14 +380,14 @@ ╭─ Select test types to run (tests can also be selected by command args individually) ─────────────────────────────────╮ │--test-typeType of tests to run for core test group                                          │ │(All | All-Long | All-Quarantined | All-Postgres | All-MySQL | API | Always | CLI │ -│| Core | Operators | Other | Serialization)                                       │ +│| Core | Other | Serialization)                                                   │ │[default: All]                                                                    │ │--parallel-test-typesSpace separated list of core test types used for testing in parallel.│ -│(API | Always | CLI | Core | Operators | Other | Serialization)      │ -│[default: API Always CLI Core Operators Other Serialization]         │ +│(API | Always | CLI | Core | Other | Serialization)                  │ +│[default: API Always CLI Core Other Serialization]                   │ │--excluded-parallel-test-typesSpace separated list of core test types that will be excluded from parallel tes   │ │runs.                                                                             │ -│(API | Always | CLI | Core | Operators | Other | Serialization)                   │ +│(API | Always | CLI | Core | Other | Serialization)                               │ │[default: ""]                                                                     │ ╰──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯ ╭─ Test options ───────────────────────────────────────────────────────────────────────────────────────────────────────╮ @@ -458,8 +458,8 @@ │--force-lowest-dependenciesRun tests for the lowest direct dependencies of Airflow │ │or selected provider if `Provider[PROVIDER_ID]` is used │ │as test type.                                           │ -│--install-airflow-with-constraints/--no-install-airflow…Install airflow in a separate step, with constraints    │ -│determined from package or airflow version.             │ +│--install-airflow-with-constraints/--no-install-airflow-Install airflow in a separate step, with constraints    │ +│with-constraintsdetermined from package or airflow version.             │ │[default: no-install-airflow-with-constraints]          │ │--package-formatFormat of packages.(wheel | sdist | both)│ │[default: wheel]   │ diff --git a/dev/breeze/doc/images/output_testing_core-tests.txt b/dev/breeze/doc/images/output_testing_core-tests.txt index fbdb26945452a..7595fe2535ca0 100644 --- a/dev/breeze/doc/images/output_testing_core-tests.txt +++ b/dev/breeze/doc/images/output_testing_core-tests.txt @@ -1 +1 @@ -912ce9d271cc63cc41f0256bc6f7516e +f3560899cf0f64593214c507936359cb diff --git a/dev/breeze/src/airflow_breeze/global_constants.py b/dev/breeze/src/airflow_breeze/global_constants.py index 4fed68c9d5e03..e1a7be2385abb 100644 --- a/dev/breeze/src/airflow_breeze/global_constants.py +++ b/dev/breeze/src/airflow_breeze/global_constants.py @@ -228,7 +228,6 @@ class SelectiveCoreTestType(SelectiveTestType): CORE = "Core" SERIALIZATION = "Serialization" OTHER = "Other" - OPERATORS = "Operators" class SelectiveProvidersTestType(SelectiveTestType): diff --git a/dev/breeze/src/airflow_breeze/utils/run_tests.py b/dev/breeze/src/airflow_breeze/utils/run_tests.py index 814e1c91eacde..3927ed414c2ad 100644 --- a/dev/breeze/src/airflow_breeze/utils/run_tests.py +++ b/dev/breeze/src/airflow_breeze/utils/run_tests.py @@ -164,7 +164,6 @@ def get_excluded_provider_args(python_version: str) -> list[str]: "tests/utils", ], "Integration": ["tests/integration"], - "Operators": ["tests/operators"], "Serialization": [ "tests/serialization", ], diff --git a/dev/breeze/src/airflow_breeze/utils/selective_checks.py b/dev/breeze/src/airflow_breeze/utils/selective_checks.py index 6a596a90f3dea..dba5e49553b6f 100644 --- a/dev/breeze/src/airflow_breeze/utils/selective_checks.py +++ b/dev/breeze/src/airflow_breeze/utils/selective_checks.py @@ -86,7 +86,7 @@ USE_PUBLIC_RUNNERS_LABEL = "use public runners" USE_SELF_HOSTED_RUNNERS_LABEL = "use self-hosted runners" -ALL_CI_SELECTIVE_TEST_TYPES = "API Always CLI Core Operators Other Serialization" +ALL_CI_SELECTIVE_TEST_TYPES = "API Always CLI Core Other Serialization" ALL_PROVIDERS_SELECTIVE_TEST_TYPES = ( "Providers[-amazon,google,standard] Providers[amazon] Providers[google] Providers[standard]" @@ -290,10 +290,6 @@ def __hash__(self): r"^airflow/cli/", r"^tests/cli/", ], - SelectiveCoreTestType.OPERATORS: [ - r"^airflow/operators/", - r"^tests/operators/", - ], SelectiveProvidersTestType.PROVIDERS: [ r"^providers/.*/src/airflow/providers/", r"^providers/.*/tests/", diff --git a/dev/breeze/tests/test_pytest_args_for_test_types.py b/dev/breeze/tests/test_pytest_args_for_test_types.py index 741e4ab0fb789..7e77718a330d6 100644 --- a/dev/breeze/tests/test_pytest_args_for_test_types.py +++ b/dev/breeze/tests/test_pytest_args_for_test_types.py @@ -88,11 +88,6 @@ def _find_all_integration_folders() -> list[str]: "Serialization", ["tests/serialization"], ), - ( - GroupOfTests.CORE, - "Operators", - ["tests/operators"], - ), ( GroupOfTests.PROVIDERS, "Providers", diff --git a/dev/breeze/tests/test_selective_checks.py b/dev/breeze/tests/test_selective_checks.py index 419e3394982d1..17e6614232584 100644 --- a/dev/breeze/tests/test_selective_checks.py +++ b/dev/breeze/tests/test_selective_checks.py @@ -220,33 +220,6 @@ def assert_outputs_are_printed(expected_outputs: dict[str, str], stderr: str): id="All tests should run when API test files change", ) ), - ( - pytest.param( - ("airflow/operators/file.py",), - { - "selected-providers-list-as-string": None, - "all-python-versions": "['3.9']", - "all-python-versions-list-as-string": "3.9", - "python-versions": "['3.9']", - "python-versions-list-as-string": "3.9", - "ci-image-build": "true", - "prod-image-build": "false", - "needs-helm-tests": "false", - "run-tests": "true", - "run-amazon-tests": "false", - "docs-build": "true", - "skip-pre-commits": "check-provider-yaml-valid,identity,lint-helm-chart,mypy-airflow,mypy-dev," - "mypy-docs,mypy-providers,mypy-task-sdk,ts-compile-format-lint-ui", - "upgrade-to-newer-dependencies": "false", - "core-test-types-list-as-string": "Always Operators", - "providers-test-types-list-as-string": "", - "individual-providers-test-types-list-as-string": "", - "needs-mypy": "true", - "mypy-checks": "['mypy-airflow']", - }, - id="Only Operator tests and DOCS should run", - ) - ), ( pytest.param( ("providers/standard/src/airflow/providers/standard/operators/python.py",), @@ -510,7 +483,7 @@ def assert_outputs_are_printed(expected_outputs: dict[str, str], stderr: str): ), "skip-providers-tests": "false", "upgrade-to-newer-dependencies": "false", - "core-test-types-list-as-string": ("API Always CLI Core Operators Other Serialization"), + "core-test-types-list-as-string": ("API Always CLI Core Other Serialization"), "providers-test-types-list-as-string": "Providers[-amazon,google,standard] Providers[amazon] Providers[google] Providers[standard]", "needs-mypy": "true", "mypy-checks": "['mypy-providers', 'mypy-task-sdk']", @@ -1320,7 +1293,7 @@ def test_full_test_needed_when_scripts_changes(files: tuple[str, ...], expected_ "full-tests-needed": "true", "skip-pre-commits": "check-airflow-provider-compatibility,check-extra-packages-references,check-provider-yaml-valid,identity,lint-helm-chart,mypy-airflow,mypy-dev,mypy-docs,mypy-providers,mypy-task-sdk,validate-operators-init", "upgrade-to-newer-dependencies": "false", - "core-test-types-list-as-string": "API Always CLI Core Operators Other Serialization", + "core-test-types-list-as-string": "API Always CLI Core Other Serialization", "needs-mypy": "true", "mypy-checks": "['mypy-airflow', 'mypy-docs', 'mypy-dev', 'mypy-task-sdk']", }, @@ -1441,7 +1414,7 @@ def test_expected_output_full_tests_needed( "full-tests-needed": "false", "run-kubernetes-tests": "false", "upgrade-to-newer-dependencies": "false", - "core-test-types-list-as-string": "API Always CLI Core Operators Other Serialization", + "core-test-types-list-as-string": "API Always CLI Core Other Serialization", "needs-mypy": "true", "mypy-checks": "['mypy-airflow']", }, @@ -1504,7 +1477,7 @@ def test_expected_output_pull_request_v2_7( "skip-pre-commits": "check-airflow-provider-compatibility,check-extra-packages-references,check-provider-yaml-valid,identity,lint-helm-chart,mypy-airflow,mypy-dev,mypy-docs,mypy-providers,mypy-task-sdk,validate-operators-init", "docs-list-as-string": "apache-airflow docker-stack", "upgrade-to-newer-dependencies": "true", - "core-test-types-list-as-string": "API Always CLI Core Operators Other Serialization", + "core-test-types-list-as-string": "API Always CLI Core Other Serialization", "needs-mypy": "true", "mypy-checks": "['mypy-airflow', 'mypy-docs', 'mypy-dev', 'mypy-task-sdk']", }, @@ -1674,7 +1647,7 @@ def test_expected_output_push( "docs-list-as-string": "", "upgrade-to-newer-dependencies": "false", "skip-pre-commits": "identity,mypy-airflow,mypy-dev,mypy-docs,mypy-providers,mypy-task-sdk", - "core-test-types-list-as-string": "API Always CLI Core Operators Other Serialization", + "core-test-types-list-as-string": "API Always CLI Core Other Serialization", "needs-mypy": "true", "mypy-checks": "['mypy-airflow', 'mypy-providers', 'mypy-docs', 'mypy-dev', 'mypy-task-sdk']", }, @@ -1703,7 +1676,7 @@ def test_expected_output_push( "ts-compile-format-lint-ui", "run-kubernetes-tests": "false", "upgrade-to-newer-dependencies": "false", - "core-test-types-list-as-string": "API Always CLI Core Operators Other Serialization", + "core-test-types-list-as-string": "API Always CLI Core Other Serialization", "providers-test-types-list-as-string": "Providers[amazon] Providers[common.compat,common.io,common.sql,dbt.cloud,ftp,microsoft.mssql,mysql,openlineage,postgres,sftp,snowflake,trino] Providers[google]", "needs-mypy": "false", "mypy-checks": "[]", diff --git a/providers/standard/src/airflow/providers/standard/operators/trigger_dagrun.py b/providers/standard/src/airflow/providers/standard/operators/trigger_dagrun.py index e00468d1bc42b..11977a081d3ff 100644 --- a/providers/standard/src/airflow/providers/standard/operators/trigger_dagrun.py +++ b/providers/standard/src/airflow/providers/standard/operators/trigger_dagrun.py @@ -43,7 +43,7 @@ from airflow.utils import timezone from airflow.utils.session import provide_session from airflow.utils.state import DagRunState -from airflow.utils.types import DagRunTriggeredByType, DagRunType +from airflow.utils.types import DagRunType XCOM_LOGICAL_DATE_ISO = "trigger_logical_date_iso" XCOM_RUN_ID = "trigger_run_id" @@ -78,15 +78,15 @@ class TriggerDagRunLink(BaseOperatorLink): name = "Triggered DAG" def get_link(self, operator: BaseOperator, *, ti_key: TaskInstanceKey) -> str: - from airflow.models.renderedtifields import RenderedTaskInstanceFields - if TYPE_CHECKING: assert isinstance(operator, TriggerDagRunOperator) - if template_fields := RenderedTaskInstanceFields.get_templated_fields(ti_key): - trigger_dag_id: str = template_fields.get("trigger_dag_id", operator.trigger_dag_id) - else: - trigger_dag_id = operator.trigger_dag_id + trigger_dag_id = operator.trigger_dag_id + if not AIRFLOW_V_3_0_PLUS: + from airflow.models.renderedtifields import RenderedTaskInstanceFields + + if template_fields := RenderedTaskInstanceFields.get_templated_fields(ti_key): + trigger_dag_id: str = template_fields.get("trigger_dag_id", operator.trigger_dag_id) # type: ignore[no-redef] # Fetch the correct dag_run_id for the triggerED dag which is # stored in xcom during execution of the triggerING task. @@ -174,7 +174,7 @@ def __init__( self.allowed_states = [DagRunState(s) for s in allowed_states] else: self.allowed_states = [DagRunState.SUCCESS] - if failed_states or failed_states == []: + if failed_states is not None: self.failed_states = [DagRunState(s) for s in failed_states] else: self.failed_states = [DagRunState.FAILED] @@ -198,25 +198,51 @@ def execute(self, context: Context): try: json.dumps(self.conf) except TypeError: - raise AirflowException("conf parameter should be JSON Serializable") + raise ValueError("conf parameter should be JSON Serializable") if self.trigger_run_id: run_id = str(self.trigger_run_id) else: - run_id = DagRun.generate_run_id( - run_type=DagRunType.MANUAL, - logical_date=parsed_logical_date, - run_after=parsed_logical_date or timezone.utcnow(), - ) + if AIRFLOW_V_3_0_PLUS: + run_id = DagRun.generate_run_id( + run_type=DagRunType.MANUAL, + logical_date=parsed_logical_date, + run_after=parsed_logical_date or timezone.utcnow(), + ) + else: + run_id = DagRun.generate_run_id(DagRunType.MANUAL, parsed_logical_date or timezone.utcnow()) # type: ignore[misc,call-arg] + if AIRFLOW_V_3_0_PLUS: + self._trigger_dag_af_3(context=context, run_id=run_id, parsed_logical_date=parsed_logical_date) + else: + self._trigger_dag_af_2(context=context, run_id=run_id, parsed_logical_date=parsed_logical_date) + + def _trigger_dag_af_3(self, context, run_id, parsed_logical_date): + from airflow.exceptions import DagRunTriggerException + + raise DagRunTriggerException( + trigger_dag_id=self.trigger_dag_id, + dag_run_id=run_id, + conf=self.conf, + logical_date=parsed_logical_date, + reset_dag_run=self.reset_dag_run, + skip_when_already_exists=self.skip_when_already_exists, + wait_for_completion=self.wait_for_completion, + allowed_states=self.allowed_states, + failed_states=self.failed_states, + poke_interval=self.poke_interval, + ) + + # TODO: Support deferral + + def _trigger_dag_af_2(self, context, run_id, parsed_logical_date): try: dag_run = trigger_dag( dag_id=self.trigger_dag_id, run_id=run_id, conf=self.conf, - logical_date=parsed_logical_date, + execution_date=parsed_logical_date, replace_microseconds=False, - triggered_by=DagRunTriggeredByType.OPERATOR, ) except DagRunAlreadyExists as e: @@ -232,7 +258,7 @@ def execute(self, context: Context): # Note: here execution fails on database isolation mode. Needs structural changes for AIP-72 dag_bag = DagBag(dag_folder=dag_model.fileloc, read_dags_from_db=True) dag = dag_bag.get_dag(self.trigger_dag_id) - dag.clear(run_id=dag_run.run_id) + dag.clear(start_date=dag_run.logical_date, end_date=dag_run.logical_date) else: if self.skip_when_already_exists: raise AirflowSkipException( @@ -253,6 +279,7 @@ def execute(self, context: Context): trigger=DagStateTrigger( dag_id=self.trigger_dag_id, states=self.allowed_states + self.failed_states, + execution_dates=[dag_run.logical_date], run_ids=[run_id], poll_interval=self.poke_interval, ), @@ -279,16 +306,18 @@ def execute(self, context: Context): @provide_session def execute_complete(self, context: Context, session: Session, event: tuple[str, dict[str, Any]]): - # This run_ids is parsed from the return trigger event - provided_run_id = event[1]["run_ids"][0] + # This logical_date is parsed from the return trigger event + provided_logical_date = event[1]["execution_dates"][0] try: # Note: here execution fails on database isolation mode. Needs structural changes for AIP-72 dag_run = session.execute( - select(DagRun).where(DagRun.dag_id == self.trigger_dag_id, DagRun.run_id == provided_run_id) + select(DagRun).where( + DagRun.dag_id == self.trigger_dag_id, DagRun.execution_date == provided_logical_date + ) ).scalar_one() except NoResultFound: raise AirflowException( - f"No DAG run found for DAG {self.trigger_dag_id} and run ID {provided_run_id}" + f"No DAG run found for DAG {self.trigger_dag_id} and logical date {self.logical_date}" ) state = dag_run.state diff --git a/tests/operators/test_trigger_dagrun.py b/providers/standard/tests/unit/standard/operators/test_trigger_dagrun.py similarity index 66% rename from tests/operators/test_trigger_dagrun.py rename to providers/standard/tests/unit/standard/operators/test_trigger_dagrun.py index e9a8a8a4f0ef9..ff1574cf1b8a9 100644 --- a/tests/operators/test_trigger_dagrun.py +++ b/providers/standard/tests/unit/standard/operators/test_trigger_dagrun.py @@ -32,15 +32,17 @@ from airflow.models.taskinstance import TaskInstance from airflow.providers.standard.operators.trigger_dagrun import TriggerDagRunOperator from airflow.providers.standard.triggers.external_task import DagStateTrigger -from airflow.sdk.execution_time.comms import XComResult from airflow.utils import timezone from airflow.utils.session import create_session -from airflow.utils.state import DagRunState, State, TaskInstanceState +from airflow.utils.state import DagRunState, TaskInstanceState from airflow.utils.types import DagRunType from tests_common.test_utils.db import parse_and_sync_to_db from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS +if AIRFLOW_V_3_0_PLUS: + from airflow.exceptions import DagRunTriggerException + pytestmark = pytest.mark.db_test DEFAULT_DATE = datetime(2019, 1, 1, tzinfo=timezone.utc) @@ -85,50 +87,149 @@ def teardown_method(self): synchronize_session=False ) - # pathlib.Path(self._tmpfile).unlink() + @pytest.mark.skipif(not AIRFLOW_V_3_0_PLUS, reason="Implementation is different for Airflow 2 & 3") + def test_trigger_dagrun(self): + """ + Test TriggerDagRunOperator. - def assert_extra_link(self, triggered_dag_run, triggering_task, session): - """Asserts whether the correct extra links url will be created.""" - triggering_ti = ( - session.query(TaskInstance) - .filter_by( - task_id=triggering_task.task_id, - dag_id=triggering_task.dag_id, + We only verify that the operator runs and raises correct exception. The actual execution logic + after the exception is in Task SDK code. + """ + with time_machine.travel("2025-02-18T08:04:46Z", tick=False): + task = TriggerDagRunOperator( + task_id="test_task", + trigger_dag_id=TRIGGERED_DAG_ID, + conf={"foo": "bar"}, ) - .one() + + # Ensure correct exception is raised + with pytest.raises(DagRunTriggerException) as exc_info: + task.execute(context={}) + + assert exc_info.value.trigger_dag_id == TRIGGERED_DAG_ID + assert exc_info.value.conf == {"foo": "bar"} + assert exc_info.value.logical_date is None + assert exc_info.value.reset_dag_run is False + assert exc_info.value.skip_when_already_exists is False + assert exc_info.value.wait_for_completion is False + assert exc_info.value.allowed_states == [DagRunState.SUCCESS] + assert exc_info.value.failed_states == [DagRunState.FAILED] + + expected_run_id = DagRun.generate_run_id( + run_type=DagRunType.MANUAL, run_after=timezone.utcnow() + ).rsplit("_", 1)[0] + # rsplit because last few characters are random. + assert exc_info.value.dag_run_id.rsplit("_", 1)[0] == expected_run_id + + @pytest.mark.skipif(not AIRFLOW_V_3_0_PLUS, reason="Implementation is different for Airflow 2 & 3") + @mock.patch("airflow.providers.standard.operators.trigger_dagrun.XCom.get_one") + def test_extra_operator_link(self, mock_xcom_get_one, dag_maker): + with dag_maker(TEST_DAG_ID, default_args={"start_date": DEFAULT_DATE}, serialized=True): + task = TriggerDagRunOperator( + task_id="test_task", + trigger_dag_id=TRIGGERED_DAG_ID, + trigger_run_id="test_run_id", + conf={"foo": "bar"}, + ) + + dr = dag_maker.create_dagrun(run_id="test_run_id") + ti = dr.get_task_instance(task_id=task.task_id) + + mock_xcom_get_one.return_value = ti.run_id + + link = task.operator_extra_links[0].get_link(operator=task, ti_key=ti.key) + + base_url = conf.get_mandatory_value("api", "base_url").lower() + expected_url = f"{base_url}dags/{TRIGGERED_DAG_ID}/runs/test_run_id" + assert link == expected_url, f"Expected {expected_url}, but got {link}" + + @pytest.mark.skipif(not AIRFLOW_V_3_0_PLUS, reason="Implementation is different for Airflow 2 & 3") + def test_trigger_dagrun_custom_run_id(self): + task = TriggerDagRunOperator( + task_id="test_task", + trigger_dag_id=TRIGGERED_DAG_ID, + trigger_run_id="custom_run_id", ) - if AIRFLOW_V_3_0_PLUS: - base_url = conf.get_mandatory_value("api", "base_url").lower() - expected_url = f"{base_url}dags/{triggered_dag_run.dag_id}/runs/{triggered_dag_run.run_id}" + with pytest.raises(DagRunTriggerException) as exc_info: + task.execute(context={}) - link = triggering_task.operator_extra_links[0].get_link( - operator=triggering_task, ti_key=triggering_ti.key + assert exc_info.value.dag_run_id == "custom_run_id" + + @pytest.mark.skipif(not AIRFLOW_V_3_0_PLUS, reason="Implementation is different for Airflow 2 & 3") + def test_trigger_dagrun_with_logical_date(self): + """Test TriggerDagRunOperator with custom logical_date.""" + task = TriggerDagRunOperator( + task_id="test_trigger_dagrun_with_logical_date", + trigger_dag_id=TRIGGERED_DAG_ID, + logical_date=timezone.datetime(2021, 1, 2, 3, 4, 5), + ) + + with pytest.raises(DagRunTriggerException) as exc_info: + task.execute(context={}) + + assert exc_info.value.logical_date == timezone.datetime(2021, 1, 2, 3, 4, 5) + + def test_trigger_dagrun_operator_templated_invalid_conf(self, dag_maker): + """Test passing a conf that is not JSON Serializable raise error.""" + with dag_maker( + TEST_DAG_ID, default_args={"owner": "airflow", "start_date": DEFAULT_DATE}, serialized=True + ): + task = TriggerDagRunOperator( + task_id="test_trigger_dagrun_with_invalid_conf", + trigger_dag_id=TRIGGERED_DAG_ID, + conf={"foo": "{{ dag.dag_id }}", "datetime": timezone.utcnow()}, ) + dag_maker.sync_dagbag_to_db() + parse_and_sync_to_db(self.f_name) + dag_maker.create_dagrun() + with pytest.raises(ValueError, match="^conf parameter should be JSON Serializable$"): + task.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE) + + def test_trigger_dagrun_with_no_failed_state(self, dag_maker): + task = TriggerDagRunOperator( + task_id="test_task", + trigger_dag_id=TRIGGERED_DAG_ID, + logical_date=DEFAULT_DATE, + wait_for_completion=True, + poke_interval=10, + failed_states=[], + ) + + assert task.failed_states == [] + + +# TODO: To be removed once the provider drops support for Airflow 2 +@pytest.mark.skipif(AIRFLOW_V_3_0_PLUS, reason="Test only for Airflow 2") +class TestDagRunOperatorAF2: + """Test TriggerDagRunOperator for Airflow 2.""" + + def setup_method(self): + # Airflow relies on reading the DAG from disk when triggering it. + # Therefore write a temp file holding the DAG to trigger. + with tempfile.NamedTemporaryFile(mode="w", delete=False) as f: + self._tmpfile = f.name + f.write(DAG_SCRIPT) + f.flush() + self.f_name = f.name + + with create_session() as session: + session.add(DagModel(dag_id=TRIGGERED_DAG_ID, fileloc=self._tmpfile)) + session.commit() - assert link == expected_url, f"Expected {expected_url}, but got {link}" - else: - with mock.patch( - "airflow.providers.standard.operators.trigger_dagrun.build_airflow_url_with_query" - ) as mock_build_url: - # This is equivalent of a task run calling this and pushing to xcom - triggering_task.operator_extra_links[0].get_link( - operator=triggering_task, ti_key=triggering_ti.key + def teardown_method(self): + """Cleanup state after testing in DB.""" + with create_session() as session: + session.query(Log).filter(Log.dag_id == TEST_DAG_ID).delete(synchronize_session=False) + for dbmodel in [DagModel, DagRun, TaskInstance]: + session.query(dbmodel).filter(dbmodel.dag_id.in_([TRIGGERED_DAG_ID, TEST_DAG_ID])).delete( + synchronize_session=False ) - assert mock_build_url.called - args, _ = mock_build_url.call_args - expected_args = { - "dag_id": triggered_dag_run.dag_id, - "dag_run_id": triggered_dag_run.run_id, - } - assert expected_args in args def test_trigger_dagrun(self, dag_maker, mock_supervisor_comms): """Test TriggerDagRunOperator.""" with time_machine.travel("2025-02-18T08:04:46Z", tick=False): - with dag_maker( - TEST_DAG_ID, default_args={"owner": "airflow", "start_date": DEFAULT_DATE}, serialized=True - ): + with dag_maker(TEST_DAG_ID, default_args={"start_date": DEFAULT_DATE}, serialized=True): task = TriggerDagRunOperator(task_id="test_task", trigger_dag_id=TRIGGERED_DAG_ID) dag_maker.sync_dagbag_to_db() parse_and_sync_to_db(self.f_name) @@ -137,50 +238,39 @@ def test_trigger_dagrun(self, dag_maker, mock_supervisor_comms): dagrun = dag_maker.session.query(DagRun).filter(DagRun.dag_id == TRIGGERED_DAG_ID).one() assert dagrun.run_type == DagRunType.MANUAL - actual_run_id = dagrun.run_id.rsplit("_", 1)[0] - - expected_run_id = DagRun.generate_run_id( - run_type=DagRunType.MANUAL, run_after=timezone.utcnow() - ).rsplit("_", 1)[0] - - assert actual_run_id == expected_run_id - - mock_supervisor_comms.get_message.return_value = XComResult(key="xcom_key", value=dagrun.run_id) - - self.assert_extra_link(dagrun, task, dag_maker.session) + assert dagrun.run_id == DagRun.generate_run_id(DagRunType.MANUAL, dagrun.logical_date) - def test_trigger_dagrun_custom_run_id(self, dag_maker): - with dag_maker( - TEST_DAG_ID, default_args={"owner": "airflow", "start_date": DEFAULT_DATE}, serialized=True - ): + def test_extra_operator_link(self, dag_maker, session): + """Asserts whether the correct extra links url will be created.""" + with dag_maker(TEST_DAG_ID, default_args={"start_date": DEFAULT_DATE}, serialized=True): task = TriggerDagRunOperator( - task_id="test_task", - trigger_dag_id=TRIGGERED_DAG_ID, - trigger_run_id="custom_run_id", + task_id="test_task", trigger_dag_id=TRIGGERED_DAG_ID, trigger_run_id="test_run_id" ) - dag_maker.sync_dagbag_to_db() - parse_and_sync_to_db(self.f_name) dag_maker.create_dagrun() - task.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE) + task.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE, ignore_ti_state=True) - with create_session() as session: - dagruns = session.query(DagRun).filter(DagRun.dag_id == TRIGGERED_DAG_ID).all() - assert len(dagruns) == 1 - assert dagruns[0].run_id == "custom_run_id" + triggering_ti = session.query(TaskInstance).filter_by(task_id=task.task_id, dag_id=task.dag_id).one() + + with mock.patch("airflow.utils.helpers.build_airflow_url_with_query") as mock_build_url: + # This is equivalent of a task run calling this and pushing to xcom + task.operator_extra_links[0].get_link(operator=task, ti_key=triggering_ti.key) + assert mock_build_url.called + args, _ = mock_build_url.call_args + expected_args = { + "dag_id": TRIGGERED_DAG_ID, + "dag_run_id": "test_run_id", + } + assert expected_args in args - def test_trigger_dagrun_with_logical_date(self, dag_maker, mock_supervisor_comms): + def test_trigger_dagrun_with_logical_date(self, dag_maker): """Test TriggerDagRunOperator with custom logical_date.""" custom_logical_date = timezone.datetime(2021, 1, 2, 3, 4, 5) - with dag_maker( - TEST_DAG_ID, default_args={"owner": "airflow", "start_date": DEFAULT_DATE}, serialized=True - ): + with dag_maker(TEST_DAG_ID, default_args={"start_date": DEFAULT_DATE}, serialized=True): task = TriggerDagRunOperator( task_id="test_trigger_dagrun_with_logical_date", trigger_dag_id=TRIGGERED_DAG_ID, logical_date=custom_logical_date, ) - dag_maker.sync_dagbag_to_db() - parse_and_sync_to_db(self.f_name) dag_maker.create_dagrun() task.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE, ignore_ti_state=True) @@ -188,13 +278,9 @@ def test_trigger_dagrun_with_logical_date(self, dag_maker, mock_supervisor_comms dagrun = session.query(DagRun).filter(DagRun.dag_id == TRIGGERED_DAG_ID).one() assert dagrun.run_type == DagRunType.MANUAL assert dagrun.logical_date == custom_logical_date - assert dagrun.run_id == DagRun.generate_run_id( - run_type=DagRunType.MANUAL, logical_date=custom_logical_date, run_after=custom_logical_date - ) - mock_supervisor_comms.get_message.return_value = XComResult(key="xcom_key", value=dagrun.run_id) - self.assert_extra_link(dagrun, task, session) + assert dagrun.run_id == DagRun.generate_run_id(DagRunType.MANUAL, custom_logical_date) - def test_trigger_dagrun_twice(self, dag_maker, mock_supervisor_comms): + def test_trigger_dagrun_twice(self, dag_maker): """Test TriggerDagRunOperator with custom logical_date.""" utc_now = timezone.utcnow() run_id = f"manual__{utc_now.isoformat()}" @@ -215,10 +301,8 @@ def test_trigger_dagrun_twice(self, dag_maker, mock_supervisor_comms): dag_maker.create_dagrun() dag_run = DagRun( dag_id=TRIGGERED_DAG_ID, - logical_date=utc_now, - data_interval=(utc_now, utc_now), - run_after=utc_now, - state=State.SUCCESS, + execution_date=utc_now, + state=DagRunState.SUCCESS, run_type="manual", run_id=run_id, ) @@ -231,10 +315,6 @@ def test_trigger_dagrun_twice(self, dag_maker, mock_supervisor_comms): triggered_dag_run = dagruns[0] assert triggered_dag_run.run_type == DagRunType.MANUAL assert triggered_dag_run.logical_date == utc_now - mock_supervisor_comms.get_message.return_value = XComResult( - key="xcom_key", value=triggered_dag_run.run_id - ) - self.assert_extra_link(triggered_dag_run, task, dag_maker.session) def test_trigger_dagrun_with_scheduled_dag_run(self, dag_maker, mock_supervisor_comms): """Test TriggerDagRunOperator with custom logical_date and scheduled dag_run.""" @@ -252,16 +332,12 @@ def test_trigger_dagrun_with_scheduled_dag_run(self, dag_maker, mock_supervisor_ reset_dag_run=True, wait_for_completion=True, ) - dag_maker.sync_dagbag_to_db() - parse_and_sync_to_db(self.f_name) dag_maker.create_dagrun() run_id = f"scheduled__{utc_now.isoformat()}" dag_run = DagRun( dag_id=TRIGGERED_DAG_ID, - logical_date=utc_now, - data_interval=(utc_now, utc_now), - run_after=utc_now, - state=State.SUCCESS, + execution_date=utc_now, + state=DagRunState.SUCCESS, run_type="scheduled", run_id=run_id, ) @@ -273,117 +349,6 @@ def test_trigger_dagrun_with_scheduled_dag_run(self, dag_maker, mock_supervisor_ assert len(dagruns) == 1 triggered_dag_run = dagruns[0] assert triggered_dag_run.logical_date == utc_now - mock_supervisor_comms.get_message.return_value = XComResult( - key="xcom_key", value=triggered_dag_run.run_id - ) - self.assert_extra_link(triggered_dag_run, task, dag_maker.session) - - def test_trigger_dagrun_with_templated_logical_date(self, dag_maker, mock_supervisor_comms): - """Test TriggerDagRunOperator with templated logical_date.""" - with dag_maker( - TEST_DAG_ID, default_args={"owner": "airflow", "start_date": DEFAULT_DATE}, serialized=True - ): - task = TriggerDagRunOperator( - task_id="test_trigger_dagrun_with_str_logical_date", - trigger_dag_id=TRIGGERED_DAG_ID, - logical_date="{{ logical_date }}", - ) - dag_maker.sync_dagbag_to_db() - parse_and_sync_to_db(self.f_name) - dag_maker.create_dagrun() - task.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE, ignore_ti_state=True) - - with create_session() as session: - dagruns = session.query(DagRun).filter(DagRun.dag_id == TRIGGERED_DAG_ID).all() - assert len(dagruns) == 1 - triggered_dag_run = dagruns[0] - assert triggered_dag_run.run_type == DagRunType.MANUAL - assert triggered_dag_run.logical_date == DEFAULT_DATE - mock_supervisor_comms.get_message.return_value = XComResult( - key="xcom_key", value=triggered_dag_run.run_id - ) - self.assert_extra_link(triggered_dag_run, task, session) - - def test_trigger_dagrun_with_templated_trigger_dag_id(self, dag_maker, mock_supervisor_comms): - """Test TriggerDagRunOperator with templated trigger dag id.""" - with dag_maker( - TEST_DAG_ID, default_args={"owner": "airflow", "start_date": DEFAULT_DATE}, serialized=True - ): - task = TriggerDagRunOperator( - task_id="__".join(["test_trigger_dagrun_with_templated_trigger_dag_id", TRIGGERED_DAG_ID]), - trigger_dag_id="{{ ti.task_id.rsplit('.', 1)[-1].split('__')[-1] }}", - ) - dag_maker.sync_dagbag_to_db() - parse_and_sync_to_db(self.f_name) - dag_maker.create_dagrun() - task.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE, ignore_ti_state=True) - - with create_session() as session: - dagruns = session.query(DagRun).filter(DagRun.dag_id == TRIGGERED_DAG_ID).all() - assert len(dagruns) == 1 - triggered_dag_run = dagruns[0] - assert triggered_dag_run.run_type == DagRunType.MANUAL - assert triggered_dag_run.dag_id == TRIGGERED_DAG_ID - mock_supervisor_comms.get_message.return_value = XComResult( - key="xcom_key", value=triggered_dag_run.run_id - ) - self.assert_extra_link(triggered_dag_run, task, session) - - def test_trigger_dagrun_operator_conf(self, dag_maker): - """Test passing conf to the triggered DagRun.""" - with dag_maker( - TEST_DAG_ID, default_args={"owner": "airflow", "start_date": DEFAULT_DATE}, serialized=True - ): - task = TriggerDagRunOperator( - task_id="test_trigger_dagrun_with_str_logical_date", - trigger_dag_id=TRIGGERED_DAG_ID, - conf={"foo": "bar"}, - ) - dag_maker.sync_dagbag_to_db() - parse_and_sync_to_db(self.f_name) - dag_maker.create_dagrun() - task.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE, ignore_ti_state=True) - - with create_session() as session: - dagruns = session.query(DagRun).filter(DagRun.dag_id == TRIGGERED_DAG_ID).all() - assert len(dagruns) == 1 - assert dagruns[0].conf == {"foo": "bar"} - - def test_trigger_dagrun_operator_templated_invalid_conf(self, dag_maker): - """Test passing a conf that is not JSON Serializable raise error.""" - with dag_maker( - TEST_DAG_ID, default_args={"owner": "airflow", "start_date": DEFAULT_DATE}, serialized=True - ): - task = TriggerDagRunOperator( - task_id="test_trigger_dagrun_with_invalid_conf", - trigger_dag_id=TRIGGERED_DAG_ID, - conf={"foo": "{{ dag.dag_id }}", "datetime": timezone.utcnow()}, - ) - dag_maker.sync_dagbag_to_db() - parse_and_sync_to_db(self.f_name) - dag_maker.create_dagrun() - with pytest.raises(AirflowException, match="^conf parameter should be JSON Serializable$"): - task.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE) - - def test_trigger_dagrun_operator_templated_conf(self, dag_maker): - """Test passing a templated conf to the triggered DagRun.""" - with dag_maker( - TEST_DAG_ID, default_args={"owner": "airflow", "start_date": DEFAULT_DATE}, serialized=True - ): - task = TriggerDagRunOperator( - task_id="test_trigger_dagrun_with_str_logical_date", - trigger_dag_id=TRIGGERED_DAG_ID, - conf={"foo": "{{ dag.dag_id }}"}, - ) - dag_maker.sync_dagbag_to_db() - parse_and_sync_to_db(self.f_name) - dag_maker.create_dagrun() - task.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE, ignore_ti_state=True) - - with create_session() as session: - dagruns = session.query(DagRun).filter(DagRun.dag_id == TRIGGERED_DAG_ID).all() - assert len(dagruns) == 1 - assert dagruns[0].conf == {"foo": TEST_DAG_ID} def test_trigger_dagrun_with_reset_dag_run_false(self, dag_maker): """Test TriggerDagRunOperator without reset_dag_run.""" @@ -452,8 +417,6 @@ def test_trigger_dagrun_with_skip_when_already_exists(self, dag_maker): reset_dag_run=False, skip_when_already_exists=True, ) - dag_maker.sync_dagbag_to_db() - parse_and_sync_to_db(self.f_name) dr: DagRun = dag_maker.create_dagrun() task.run(start_date=logical_date, end_date=logical_date, ignore_ti_state=True) assert dr.get_task_instance("test_task").state == TaskInstanceState.SUCCESS @@ -484,8 +447,6 @@ def test_trigger_dagrun_with_reset_dag_run_true( logical_date=trigger_logical_date, reset_dag_run=True, ) - dag_maker.sync_dagbag_to_db() - parse_and_sync_to_db(self.f_name) dag_maker.create_dagrun() task.run(start_date=logical_date, end_date=logical_date, ignore_ti_state=True) task.run(start_date=logical_date, end_date=logical_date, ignore_ti_state=True) @@ -493,21 +454,19 @@ def test_trigger_dagrun_with_reset_dag_run_true( with create_session() as session: dag_runs = session.query(DagRun).filter(DagRun.dag_id == TRIGGERED_DAG_ID).all() assert len(dag_runs) == expected_dagruns_count - assert dag_runs[0].run_type == DagRunType.MANUAL + assert dag_runs[0].external_trigger def test_trigger_dagrun_with_wait_for_completion_true(self, dag_maker): """Test TriggerDagRunOperator with wait_for_completion.""" logical_date = DEFAULT_DATE - with dag_maker( - TEST_DAG_ID, default_args={"owner": "airflow", "start_date": DEFAULT_DATE}, serialized=True - ): + with dag_maker(TEST_DAG_ID, default_args={"start_date": DEFAULT_DATE}, serialized=True): task = TriggerDagRunOperator( task_id="test_task", trigger_dag_id=TRIGGERED_DAG_ID, logical_date=logical_date, wait_for_completion=True, poke_interval=10, - allowed_states=[State.QUEUED], + allowed_states=[DagRunState.QUEUED], ) dag_maker.sync_dagbag_to_db() parse_and_sync_to_db(self.f_name) @@ -521,66 +480,32 @@ def test_trigger_dagrun_with_wait_for_completion_true(self, dag_maker): def test_trigger_dagrun_with_wait_for_completion_true_fail(self, dag_maker): """Test TriggerDagRunOperator with wait_for_completion but triggered dag fails.""" logical_date = DEFAULT_DATE - with dag_maker( - TEST_DAG_ID, default_args={"owner": "airflow", "start_date": DEFAULT_DATE}, serialized=True - ): + with dag_maker(TEST_DAG_ID, default_args={"start_date": DEFAULT_DATE}, serialized=True): task = TriggerDagRunOperator( task_id="test_task", trigger_dag_id=TRIGGERED_DAG_ID, logical_date=logical_date, wait_for_completion=True, poke_interval=10, - failed_states=[State.QUEUED], + failed_states=[DagRunState.QUEUED], ) - dag_maker.sync_dagbag_to_db() - parse_and_sync_to_db(self.f_name) dag_maker.create_dagrun() with pytest.raises(AirflowException): task.run(start_date=logical_date, end_date=logical_date) - def test_trigger_dagrun_triggering_itself(self, dag_maker): - """Test TriggerDagRunOperator that triggers itself""" - logical_date = DEFAULT_DATE - with dag_maker( - TEST_DAG_ID, default_args={"owner": "airflow", "start_date": DEFAULT_DATE}, serialized=True - ): - task = TriggerDagRunOperator( - task_id="test_task", - trigger_dag_id=TEST_DAG_ID, - logical_date=timezone.utcnow(), - ) - dag_maker.sync_dagbag_to_db() - parse_and_sync_to_db(self.f_name) - dag_maker.create_dagrun() - task.run(start_date=logical_date, end_date=logical_date) - - dagruns = ( - dag_maker.session.query(DagRun) - .filter(DagRun.dag_id == TEST_DAG_ID) - .order_by(DagRun.logical_date) - .all() - ) - assert len(dagruns) == 2 - triggered_dag_run = dagruns[1] - assert triggered_dag_run.state == State.QUEUED - def test_trigger_dagrun_with_wait_for_completion_true_defer_false(self, dag_maker): """Test TriggerDagRunOperator with wait_for_completion.""" logical_date = DEFAULT_DATE - with dag_maker( - TEST_DAG_ID, default_args={"owner": "airflow", "start_date": DEFAULT_DATE}, serialized=True - ): + with dag_maker(TEST_DAG_ID, default_args={"start_date": DEFAULT_DATE}, serialized=True): task = TriggerDagRunOperator( task_id="test_task", trigger_dag_id=TRIGGERED_DAG_ID, logical_date=logical_date, wait_for_completion=True, poke_interval=10, - allowed_states=[State.QUEUED], + allowed_states=[DagRunState.QUEUED], deferrable=False, ) - dag_maker.sync_dagbag_to_db() - parse_and_sync_to_db(self.f_name) dag_maker.create_dagrun() task.run(start_date=logical_date, end_date=logical_date) @@ -591,21 +516,17 @@ def test_trigger_dagrun_with_wait_for_completion_true_defer_false(self, dag_make def test_trigger_dagrun_with_wait_for_completion_true_defer_true(self, dag_maker): """Test TriggerDagRunOperator with wait_for_completion.""" logical_date = DEFAULT_DATE - with dag_maker( - TEST_DAG_ID, default_args={"owner": "airflow", "start_date": DEFAULT_DATE}, serialized=True - ): + with dag_maker(TEST_DAG_ID, default_args={"start_date": DEFAULT_DATE}, serialized=True): task = TriggerDagRunOperator( task_id="test_task", trigger_dag_id=TRIGGERED_DAG_ID, logical_date=logical_date, wait_for_completion=True, poke_interval=10, - allowed_states=[State.QUEUED], + allowed_states=[DagRunState.QUEUED], deferrable=True, trigger_run_id=DEFAULT_RUN_ID, ) - dag_maker.sync_dagbag_to_db() - parse_and_sync_to_db(self.f_name) dag_maker.create_dagrun() task.run(start_date=logical_date, end_date=logical_date) @@ -615,6 +536,7 @@ def test_trigger_dagrun_with_wait_for_completion_true_defer_true(self, dag_maker assert len(dagruns) == 1 trigger = DagStateTrigger( dag_id="down_stream", + execution_dates=[DEFAULT_DATE], run_ids=[DEFAULT_RUN_ID], poll_interval=20, states=["success", "failed"], @@ -634,12 +556,10 @@ def test_trigger_dagrun_with_wait_for_completion_true_defer_true_failure(self, d logical_date=logical_date, wait_for_completion=True, poke_interval=10, - allowed_states=[State.SUCCESS], + allowed_states=[DagRunState.SUCCESS], deferrable=True, trigger_run_id=DEFAULT_RUN_ID, ) - dag_maker.sync_dagbag_to_db() - parse_and_sync_to_db(self.f_name) dag_maker.create_dagrun() task.run(start_date=logical_date, end_date=logical_date) @@ -650,6 +570,7 @@ def test_trigger_dagrun_with_wait_for_completion_true_defer_true_failure(self, d trigger = DagStateTrigger( dag_id="down_stream", + execution_dates=[DEFAULT_DATE], run_ids=[DEFAULT_RUN_ID], poll_interval=20, states=["success", "failed"], @@ -672,13 +593,11 @@ def test_trigger_dagrun_with_wait_for_completion_true_defer_true_failure_2(self, logical_date=logical_date, wait_for_completion=True, poke_interval=10, - allowed_states=[State.SUCCESS], - failed_states=[State.QUEUED], + allowed_states=[DagRunState.SUCCESS], + failed_states=[DagRunState.QUEUED], deferrable=True, trigger_run_id=DEFAULT_RUN_ID, ) - dag_maker.sync_dagbag_to_db() - parse_and_sync_to_db(self.f_name) dag_maker.create_dagrun() task.run(start_date=logical_date, end_date=logical_date) @@ -689,6 +608,7 @@ def test_trigger_dagrun_with_wait_for_completion_true_defer_true_failure_2(self, trigger = DagStateTrigger( dag_id="down_stream", + execution_dates=[DEFAULT_DATE], run_ids=[DEFAULT_RUN_ID], poll_interval=20, states=["success", "failed"], @@ -718,8 +638,6 @@ def test_dagstatetrigger_run_id(self, trigger_logical_date, dag_maker): allowed_states=[DagRunState.QUEUED], deferrable=True, ) - dag_maker.sync_dagbag_to_db() - parse_and_sync_to_db(self.f_name) dag_maker.create_dagrun() mock_task_defer = mock.MagicMock(side_effect=task.defer) @@ -747,8 +665,6 @@ def test_dagstatetrigger_run_id_with_clear_and_reset(self, dag_maker): deferrable=True, reset_dag_run=True, ) - dag_maker.sync_dagbag_to_db() - parse_and_sync_to_db(self.f_name) dag_maker.create_dagrun() mock_task_defer = mock.MagicMock(side_effect=task.defer) @@ -778,22 +694,3 @@ def test_dagstatetrigger_run_id_with_clear_and_reset(self, dag_maker): # The second DagStateTrigger call should still use the original `logical_date` value. assert mock_task_defer.call_args_list[1].kwargs["trigger"].run_ids == [run_id] - - def test_trigger_dagrun_with_no_failed_state(self, dag_maker): - logical_date = DEFAULT_DATE - with dag_maker( - TEST_DAG_ID, default_args={"owner": "airflow", "start_date": DEFAULT_DATE}, serialized=True - ): - task = TriggerDagRunOperator( - task_id="test_task", - trigger_dag_id=TRIGGERED_DAG_ID, - logical_date=logical_date, - wait_for_completion=True, - poke_interval=10, - failed_states=[], - ) - dag_maker.sync_dagbag_to_db() - parse_and_sync_to_db(self.f_name) - dag_maker.create_dagrun() - - assert task.failed_states == [] diff --git a/task-sdk/pyproject.toml b/task-sdk/pyproject.toml index 46543faffe129..d8423b9fad1ce 100644 --- a/task-sdk/pyproject.toml +++ b/task-sdk/pyproject.toml @@ -102,6 +102,7 @@ codegen = [ "apache-airflow", "rich", "ruff", + "svcs>=25.1.0", ] dev = [ "apache-airflow", diff --git a/task-sdk/src/airflow/sdk/api/client.py b/task-sdk/src/airflow/sdk/api/client.py index 5bcf353dcf00f..334d283b688e1 100644 --- a/task-sdk/src/airflow/sdk/api/client.py +++ b/task-sdk/src/airflow/sdk/api/client.py @@ -38,6 +38,7 @@ AssetEventsResponse, AssetResponse, ConnectionResponse, + DagRunStateResponse, DagRunType, PrevSuccessfulDagRunResponse, TerminalStateNonSuccess, @@ -50,6 +51,7 @@ TISkippedDownstreamTasksStatePayload, TISuccessStatePayload, TITerminalStatePayload, + TriggerDAGRunPayload, ValidationError as RemoteValidationError, VariablePostBody, VariableResponse, @@ -405,6 +407,52 @@ def get( return AssetEventsResponse.model_validate_json(resp.read()) +class DagRunOperations: + __slots__ = ("client",) + + def __init__(self, client: Client): + self.client = client + + def trigger( + self, + dag_id: str, + run_id: str, + conf: dict | None = None, + logical_date: datetime | None = None, + reset_dag_run: bool = False, + ): + """Trigger a DAG run via the API server.""" + body = TriggerDAGRunPayload(logical_date=logical_date, conf=conf or {}, reset_dag_run=reset_dag_run) + + try: + self.client.post( + f"dag-runs/{dag_id}/{run_id}", content=body.model_dump_json(exclude_defaults=True) + ) + except ServerResponseError as e: + if e.response.status_code == HTTPStatus.CONFLICT: + if reset_dag_run: + log.info("DAG Run already exists; Resetting DAG Run.", dag_id=dag_id, run_id=run_id) + return self.clear(run_id=run_id, dag_id=dag_id) + + log.info("DAG Run already exists!", detail=e.detail, dag_id=dag_id, run_id=run_id) + return ErrorResponse(error=ErrorType.DAGRUN_ALREADY_EXISTS) + else: + raise + + return OKResponse(ok=True) + + def clear(self, dag_id: str, run_id: str): + """Clear a DAG run via the API server.""" + self.client.post(f"dag-runs/{dag_id}/{run_id}/clear") + # TODO: Error handling + return OKResponse(ok=True) + + def get_state(self, dag_id: str, run_id: str) -> DagRunStateResponse: + """Get the state of a DAG run via the API server.""" + resp = self.client.get(f"dag-runs/{dag_id}/{run_id}/state") + return DagRunStateResponse.model_validate_json(resp.read()) + + class BearerAuth(httpx.Auth): def __init__(self, token: str): self.token: str = token @@ -495,6 +543,12 @@ def task_instances(self) -> TaskInstanceOperations: """Operations related to TaskInstances.""" return TaskInstanceOperations(self) + @lru_cache() # type: ignore[misc] + @property + def dag_runs(self) -> DagRunOperations: + """Operations related to DagRuns.""" + return DagRunOperations(self) + @lru_cache() # type: ignore[misc] @property def connections(self) -> ConnectionOperations: diff --git a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py index 622f3d0a74e72..8725cdc04cfa0 100644 --- a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py +++ b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py @@ -94,6 +94,29 @@ class DagRunAssetReference(BaseModel): data_interval_end: Annotated[datetime | None, Field(title="Data Interval End")] = None +class DagRunState(str, Enum): + """ + All possible states that a DagRun can be in. + + These are "shared" with TaskInstanceState in some parts of the code, + so please ensure that their values always match the ones with the + same name in TaskInstanceState. + """ + + QUEUED = "queued" + RUNNING = "running" + SUCCESS = "success" + FAILED = "failed" + + +class DagRunStateResponse(BaseModel): + """ + Schema for DAG Run State response. + """ + + state: DagRunState + + class DagRunType(str, Enum): """ Class with DagRun types. @@ -245,6 +268,19 @@ class TerminalStateNonSuccess(str, Enum): FAIL_WITHOUT_RETRY = "fail_without_retry" +class TriggerDAGRunPayload(BaseModel): + """ + Schema for Trigger DAG Run API request. + """ + + model_config = ConfigDict( + extra="forbid", + ) + logical_date: Annotated[datetime | None, Field(title="Logical Date")] = None + conf: Annotated[dict[str, Any] | None, Field(title="Conf")] = None + reset_dag_run: Annotated[bool | None, Field(title="Reset Dag Run")] = False + + class ValidationError(BaseModel): loc: Annotated[list[str | int], Field(title="Location")] msg: Annotated[str, Field(title="Message")] diff --git a/task-sdk/src/airflow/sdk/exceptions.py b/task-sdk/src/airflow/sdk/exceptions.py index 4a007be8a874a..60c620cf0ce84 100644 --- a/task-sdk/src/airflow/sdk/exceptions.py +++ b/task-sdk/src/airflow/sdk/exceptions.py @@ -39,6 +39,7 @@ class ErrorType(enum.Enum): VARIABLE_NOT_FOUND = "VARIABLE_NOT_FOUND" XCOM_NOT_FOUND = "XCOM_NOT_FOUND" ASSET_NOT_FOUND = "ASSET_NOT_FOUND" + DAGRUN_ALREADY_EXISTS = "DAGRUN_ALREADY_EXISTS" GENERIC_ERROR = "GENERIC_ERROR" diff --git a/task-sdk/src/airflow/sdk/execution_time/comms.py b/task-sdk/src/airflow/sdk/execution_time/comms.py index bd081b534abfd..ae05956235fb4 100644 --- a/task-sdk/src/airflow/sdk/execution_time/comms.py +++ b/task-sdk/src/airflow/sdk/execution_time/comms.py @@ -55,6 +55,7 @@ AssetResponse, BundleInfo, ConnectionResponse, + DagRunStateResponse, PrevSuccessfulDagRunResponse, TaskInstance, TerminalTIState, @@ -64,6 +65,7 @@ TIRuntimeCheckPayload, TISkippedDownstreamTasksStatePayload, TISuccessStatePayload, + TriggerDAGRunPayload, VariableResponse, XComResponse, ) @@ -182,6 +184,23 @@ def from_variable_response(cls, variable_response: VariableResponse) -> Variable return cls(**variable_response.model_dump(exclude_defaults=True), type="VariableResult") +class DagRunStateResult(DagRunStateResponse): + type: Literal["DagRunStateResult"] = "DagRunStateResult" + + # TODO: Create a convert api_response to result classes so we don't need to do this + # for all the classes above + @classmethod + def from_api_response(cls, dr_state_response: DagRunStateResponse) -> DagRunStateResult: + """ + Create result class from API Response. + + API Response is autogenerated from the API schema, so we need to convert it to Result + for communication between the Supervisor and the task process since it needs a + discriminator field. + """ + return cls(**dr_state_response.model_dump(exclude_defaults=True), type="DagRunStateResult") + + class PrevSuccessfulDagRunResult(PrevSuccessfulDagRunResponse): type: Literal["PrevSuccessfulDagRunResult"] = "PrevSuccessfulDagRunResult" @@ -212,6 +231,7 @@ class OKResponse(BaseModel): AssetResult, AssetEventsResult, ConnectionResult, + DagRunStateResult, ErrorResponse, PrevSuccessfulDagRunResult, StartupDetails, @@ -371,6 +391,18 @@ class SetRenderedFields(BaseModel): type: Literal["SetRenderedFields"] = "SetRenderedFields" +class TriggerDagRun(TriggerDAGRunPayload): + dag_id: str + run_id: Annotated[str, Field(title="Dag Run Id")] + type: Literal["TriggerDagRun"] = "TriggerDagRun" + + +class GetDagRunState(BaseModel): + dag_id: str + run_id: str + type: Literal["GetDagRunState"] = "GetDagRunState" + + class GetAssetByName(BaseModel): name: str type: Literal["GetAssetByName"] = "GetAssetByName" @@ -406,6 +438,7 @@ class GetPrevSuccessfulDagRun(BaseModel): GetAssetEventByAsset, GetAssetEventByAssetAlias, GetConnection, + GetDagRunState, GetPrevSuccessfulDagRun, GetVariable, GetXCom, @@ -416,6 +449,7 @@ class GetPrevSuccessfulDagRun(BaseModel): SetRenderedFields, SetXCom, TaskState, + TriggerDagRun, RuntimeCheckOnTask, DeleteXCom, ], diff --git a/task-sdk/src/airflow/sdk/execution_time/supervisor.py b/task-sdk/src/airflow/sdk/execution_time/supervisor.py index e27e631419001..b73db7d568f04 100644 --- a/task-sdk/src/airflow/sdk/execution_time/supervisor.py +++ b/task-sdk/src/airflow/sdk/execution_time/supervisor.py @@ -65,6 +65,7 @@ AssetEventsResult, AssetResult, ConnectionResult, + DagRunStateResult, DeferTask, DeleteXCom, GetAssetByName, @@ -72,6 +73,7 @@ GetAssetEventByAsset, GetAssetEventByAssetAlias, GetConnection, + GetDagRunState, GetPrevSuccessfulDagRun, GetVariable, GetXCom, @@ -87,6 +89,7 @@ SucceedTask, TaskState, ToSupervisor, + TriggerDagRun, VariableResult, XComCountResponse, XComResult, @@ -934,6 +937,18 @@ def _handle_request(self, msg: ToSupervisor, log: FilteringBoundLogger): dagrun_resp = self.client.task_instances.get_previous_successful_dagrun(self.id) dagrun_result = PrevSuccessfulDagRunResult.from_dagrun_response(dagrun_resp) resp = dagrun_result.model_dump_json(exclude_unset=True).encode() + elif isinstance(msg, TriggerDagRun): + dr_resp = self.client.dag_runs.trigger( + msg.dag_id, + msg.run_id, + msg.conf, + msg.logical_date, + msg.reset_dag_run, + ) + resp = dr_resp.model_dump_json().encode() + elif isinstance(msg, GetDagRunState): + dr_resp = self.client.dag_runs.get_state(msg.dag_id, msg.run_id) + resp = DagRunStateResult.from_api_response(dr_resp).model_dump_json().encode() else: log.error("Unhandled request", msg=msg) return diff --git a/task-sdk/src/airflow/sdk/execution_time/task_runner.py b/task-sdk/src/airflow/sdk/execution_time/task_runner.py index 15cc8de3a2d47..3242a6eb9c66b 100644 --- a/task-sdk/src/airflow/sdk/execution_time/task_runner.py +++ b/task-sdk/src/airflow/sdk/execution_time/task_runner.py @@ -23,6 +23,7 @@ import functools import os import sys +import time from collections.abc import Iterable, Iterator, Mapping from datetime import datetime, timezone from io import FileIO @@ -50,8 +51,12 @@ from airflow.sdk.definitions.baseoperator import BaseOperator, ExecutorSafeguard from airflow.sdk.definitions.mappedoperator import MappedOperator from airflow.sdk.definitions.param import process_params +from airflow.sdk.exceptions import ErrorType from airflow.sdk.execution_time.comms import ( + DagRunStateResult, DeferTask, + ErrorResponse, + GetDagRunState, OKResponse, RescheduleTask, RuntimeCheckOnTask, @@ -62,6 +67,7 @@ TaskState, ToSupervisor, ToTask, + TriggerDagRun, ) from airflow.sdk.execution_time.context import ( ConnectionAccessor, @@ -80,6 +86,7 @@ import jinja2 from structlog.typing import FilteringBoundLogger as Logger + from airflow.exceptions import DagRunTriggerException from airflow.sdk.definitions._internal.abstractoperator import AbstractOperator from airflow.sdk.definitions.context import Context from airflow.sdk.types import OutletEventAccessorsProtocol @@ -547,6 +554,7 @@ def run( AirflowSkipException, AirflowTaskTerminated, AirflowTaskTimeout, + DagRunTriggerException, DownstreamTasksSkipped, TaskDeferred, ) @@ -558,6 +566,8 @@ def run( msg: ToSupervisor | None = None state: IntermediateTIState | TerminalTIState error: BaseException | None = None + + context = ti.get_template_context() try: # First, clear the xcom data sent from server if ti._ti_context_from_server and (keys_to_delete := ti._ti_context_from_server.xcom_keys_to_clear): @@ -570,7 +580,6 @@ def run( run_id=ti.run_id, ) - context = ti.get_template_context() with set_current_context(context): # This is the earliest that we can render templates -- as if it excepts for any reason we need to # catch it and handle it like a normal task failure @@ -585,11 +594,12 @@ def run( msg, state = _handle_current_task_success(context, ti) except DownstreamTasksSkipped as skip: - context = ti.get_template_context() log.info("Skipping downstream tasks.") tasks_to_skip = skip.tasks if isinstance(skip.tasks, list) else [skip.tasks] SUPERVISOR_COMMS.send_request(log=log, msg=SkipDownstreamTasks(tasks=tasks_to_skip)) msg, state = _handle_current_task_success(context, ti) + except DagRunTriggerException as drte: + msg, state = _handle_trigger_dag_run(drte, context, ti, log) except TaskDeferred as defer: # TODO: Should we use structlog.bind_contextvars here for dag_id, task_id & run_id? log.info("Pausing task as DEFERRED. ", dag_id=ti.dag_id, task_id=ti.task_id, run_id=ti.run_id) @@ -670,7 +680,7 @@ def run( return state, msg, error -def _handle_current_task_success(context, ti) -> tuple[SucceedTask, TerminalTIState]: +def _handle_current_task_success(context: Context, ti: RuntimeTaskInstance): task_outlets = list(_build_asset_profiles(ti.task.outlets)) outlet_events = list(_serialize_outlet_events(context["outlet_events"])) msg = SucceedTask( @@ -681,6 +691,83 @@ def _handle_current_task_success(context, ti) -> tuple[SucceedTask, TerminalTISt return msg, TerminalTIState.SUCCESS +def _handle_trigger_dag_run( + drte: DagRunTriggerException, context: Context, ti: RuntimeTaskInstance, log: Logger +): + """Handle exception from TriggerDagRunOperator.""" + log.info("Triggering Dag Run.", trigger_dag_id=drte.trigger_dag_id) + SUPERVISOR_COMMS.send_request( + log=log, + msg=TriggerDagRun( + dag_id=drte.trigger_dag_id, + run_id=drte.dag_run_id, + logical_date=drte.logical_date, + conf=drte.conf, + reset_dag_run=drte.reset_dag_run, + ), + ) + + comms_msg = SUPERVISOR_COMMS.get_message() + if isinstance(comms_msg, ErrorResponse) and comms_msg.error == ErrorType.DAGRUN_ALREADY_EXISTS: + if drte.skip_when_already_exists: + log.info( + "Dag Run already exists, skipping task as skip_when_already_exists is set to True.", + dag_id=drte.trigger_dag_id, + ) + msg = TaskState(state=TerminalTIState.SKIPPED, end_date=datetime.now(tz=timezone.utc)) + state = TerminalTIState.SKIPPED + else: + log.error("Dag Run already exists, marking task as failed.", dag_id=drte.trigger_dag_id) + msg = TaskState(state=TerminalTIState.FAILED, end_date=datetime.now(tz=timezone.utc)) + state = TerminalTIState.FAILED + + return msg, state + + log.info("Dag Run triggered successfully.", trigger_dag_id=drte.trigger_dag_id) + + # Store the run id from the dag run (either created or found above) to + # be used when creating the extra link on the webserver. + ti.xcom_push(key="trigger_run_id", value=drte.dag_run_id) + + if drte.wait_for_completion: + while True: + log.info( + "Waiting for dag run to complete execution in allowed state.", + dag_id=drte.trigger_dag_id, + run_id=drte.dag_run_id, + allowed_state=drte.allowed_states, + ) + time.sleep(drte.poke_interval) + + SUPERVISOR_COMMS.send_request( + log=log, msg=GetDagRunState(dag_id=drte.trigger_dag_id, run_id=drte.dag_run_id) + ) + comms_msg = SUPERVISOR_COMMS.get_message() + if TYPE_CHECKING: + assert isinstance(comms_msg, DagRunStateResult) + if comms_msg.state in drte.failed_states: + log.error( + "DagRun finished with failed state.", dag_id=drte.trigger_dag_id, state=comms_msg.state + ) + msg = TaskState(state=TerminalTIState.FAILED, end_date=datetime.now(tz=timezone.utc)) + state = TerminalTIState.FAILED + return msg, state + if comms_msg.state in drte.allowed_states: + log.info( + "DagRun finished with allowed state.", dag_id=drte.trigger_dag_id, state=comms_msg.state + ) + break + log.debug( + "DagRun not yet in allowed or failed state.", + dag_id=drte.trigger_dag_id, + state=comms_msg.state, + ) + + msg, state = _handle_current_task_success(context, ti) + + return msg, state + + def _execute_task(context: Context, ti: RuntimeTaskInstance): """Execute Task (optionally with a Timeout) and push Xcom results.""" from airflow.exceptions import AirflowTaskTimeout diff --git a/task-sdk/tests/task_sdk/api/test_client.py b/task-sdk/tests/task_sdk/api/test_client.py index 1731baf86defc..ff86b9c6184f9 100644 --- a/task-sdk/tests/task_sdk/api/test_client.py +++ b/task-sdk/tests/task_sdk/api/test_client.py @@ -29,11 +29,13 @@ from airflow.sdk.api.datamodels._generated import ( AssetResponse, ConnectionResponse, + DagRunState, + DagRunStateResponse, VariableResponse, XComResponse, ) from airflow.sdk.exceptions import ErrorType -from airflow.sdk.execution_time.comms import DeferTask, ErrorResponse, RescheduleTask +from airflow.sdk.execution_time.comms import DeferTask, ErrorResponse, OKResponse, RescheduleTask from airflow.utils import timezone from airflow.utils.state import TerminalTIState @@ -711,3 +713,102 @@ def handle_request(request: httpx.Request) -> httpx.Response: assert isinstance(result, ErrorResponse) assert result.error == ErrorType.ASSET_NOT_FOUND + + +class TestDagRunOperations: + def test_trigger(self): + # Simulate a successful response from the server when triggering a dag run + def handle_request(request: httpx.Request) -> httpx.Response: + if request.url.path == "/dag-runs/test_trigger/test_run_id": + actual_body = json.loads(request.read()) + assert actual_body["logical_date"] == "2025-01-01T00:00:00Z" + assert actual_body["conf"] == {} + + # Since the value for `reset_dag_run` is default, it should not be present in the request body + assert "reset_dag_run" not in actual_body + return httpx.Response(status_code=204) + return httpx.Response(status_code=400, json={"detail": "Bad Request"}) + + client = make_client(transport=httpx.MockTransport(handle_request)) + result = client.dag_runs.trigger( + dag_id="test_trigger", run_id="test_run_id", logical_date=timezone.datetime(2025, 1, 1) + ) + + assert result == OKResponse(ok=True) + + def test_trigger_conflict(self): + """Test that if the dag run already exists, the client returns an error when default reset_dag_run=False""" + + def handle_request(request: httpx.Request) -> httpx.Response: + if request.url.path == "/dag-runs/test_trigger_conflict/test_run_id": + return httpx.Response( + status_code=409, + json={ + "detail": { + "reason": "already_exists", + "message": "A DAG Run already exists for DAG test_trigger_conflict with run id test_run_id", + } + }, + ) + return httpx.Response(status_code=422) + + client = make_client(transport=httpx.MockTransport(handle_request)) + result = client.dag_runs.trigger(dag_id="test_trigger_conflict", run_id="test_run_id") + + assert result == ErrorResponse(error=ErrorType.DAGRUN_ALREADY_EXISTS) + + def test_trigger_conflict_reset_dag_run(self): + """Test that if dag run already exists and reset_dag_run=True, the client clears the dag run""" + + def handle_request(request: httpx.Request) -> httpx.Response: + if request.url.path == "/dag-runs/test_trigger_conflict_reset/test_run_id": + return httpx.Response( + status_code=409, + json={ + "detail": { + "reason": "already_exists", + "message": "A DAG Run already exists for DAG test_trigger_conflict with run id test_run_id", + } + }, + ) + elif request.url.path == "/dag-runs/test_trigger_conflict_reset/test_run_id/clear": + return httpx.Response(status_code=204) + return httpx.Response(status_code=422) + + client = make_client(transport=httpx.MockTransport(handle_request)) + result = client.dag_runs.trigger( + dag_id="test_trigger_conflict_reset", + run_id="test_run_id", + reset_dag_run=True, + ) + + assert result == OKResponse(ok=True) + + def test_clear(self): + """Test that the client can clear a dag run""" + + def handle_request(request: httpx.Request) -> httpx.Response: + if request.url.path == "/dag-runs/test_clear/test_run_id/clear": + return httpx.Response(status_code=204) + return httpx.Response(status_code=422) + + client = make_client(transport=httpx.MockTransport(handle_request)) + result = client.dag_runs.clear(dag_id="test_clear", run_id="test_run_id") + + assert result == OKResponse(ok=True) + + def test_get_state(self): + """Test that the client can get the state of a dag run""" + + def handle_request(request: httpx.Request) -> httpx.Response: + if request.url.path == "/dag-runs/test_state/test_run_id/state": + return httpx.Response( + status_code=200, + json={"state": "running"}, + ) + return httpx.Response(status_code=422) + + client = make_client(transport=httpx.MockTransport(handle_request)) + result = client.dag_runs.get_state(dag_id="test_state", run_id="test_run_id") + + assert result == DagRunStateResponse(state=DagRunState.RUNNING) diff --git a/task-sdk/tests/task_sdk/execution_time/test_supervisor.py b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py index 993d4520753a3..848b7be44a663 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_supervisor.py +++ b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py @@ -49,20 +49,25 @@ AssetEventResponse, AssetProfile, AssetResponse, + DagRunState, TaskInstance, TerminalTIState, ) +from airflow.sdk.exceptions import ErrorType from airflow.sdk.execution_time.comms import ( AssetEventsResult, AssetResult, ConnectionResult, + DagRunStateResult, DeferTask, DeleteXCom, + ErrorResponse, GetAssetByName, GetAssetByUri, GetAssetEventByAsset, GetAssetEventByAssetAlias, GetConnection, + GetDagRunState, GetPrevSuccessfulDagRun, GetVariable, GetXCom, @@ -75,6 +80,7 @@ SetXCom, SucceedTask, TaskState, + TriggerDagRun, VariableResult, XComResult, ) @@ -1312,6 +1318,39 @@ def watched_subprocess(self, mocker): OKResponse(ok=True), id="runtime_check_on_task", ), + pytest.param( + TriggerDagRun( + dag_id="test_dag", + run_id="test_run", + conf={"key": "value"}, + logical_date=timezone.datetime(2025, 1, 1), + reset_dag_run=True, + ), + b'{"ok":true,"type":"OKResponse"}\n', + "dag_runs.trigger", + ("test_dag", "test_run", {"key": "value"}, timezone.datetime(2025, 1, 1), True), + {}, + OKResponse(ok=True), + id="dag_run_trigger", + ), + pytest.param( + TriggerDagRun(dag_id="test_dag", run_id="test_run"), + b'{"error":"DAGRUN_ALREADY_EXISTS","detail":null,"type":"ErrorResponse"}\n', + "dag_runs.trigger", + ("test_dag", "test_run", None, None, False), + {}, + ErrorResponse(error=ErrorType.DAGRUN_ALREADY_EXISTS), + id="dag_run_trigger_already_exists", + ), + pytest.param( + GetDagRunState(dag_id="test_dag", run_id="test_run"), + b'{"state":"running","type":"DagRunStateResult"}\n', + "dag_runs.get_state", + ("test_dag", "test_run"), + {}, + DagRunStateResult(state=DagRunState.RUNNING), + id="get_dag_run_state", + ), ], ) def test_handle_requests( diff --git a/task-sdk/tests/task_sdk/execution_time/test_task_runner.py b/task-sdk/tests/task_sdk/execution_time/test_task_runner.py index 851a9c29b7ff1..414d6b6b2329f 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_task_runner.py +++ b/task-sdk/tests/task_sdk/execution_time/test_task_runner.py @@ -50,18 +50,23 @@ AssetEventResponse, AssetProfile, AssetResponse, + DagRunState, TaskInstance, TerminalTIState, ) from airflow.sdk.definitions.asset import Asset, AssetAlias, Dataset, Model from airflow.sdk.definitions.param import DagParam from airflow.sdk.definitions.variable import Variable +from airflow.sdk.exceptions import ErrorType from airflow.sdk.execution_time.comms import ( AssetEventsResult, BundleInfo, ConnectionResult, + DagRunStateResult, DeferTask, + ErrorResponse, GetConnection, + GetDagRunState, GetVariable, GetXCom, OKResponse, @@ -73,6 +78,7 @@ StartupDetails, SucceedTask, TaskState, + TriggerDagRun, VariableResult, XComResult, ) @@ -1862,3 +1868,182 @@ def execute(self, context): assert listener.state == [TaskInstanceState.RUNNING, TaskInstanceState.FAILED] assert listener.error == error + + +class TestTriggerDagRunOperator: + """Tests to verify various aspects of TriggerDagRunOperator""" + + def test_handle_trigger_dag_run(self, create_runtime_ti, mock_supervisor_comms): + """Test that TriggerDagRunOperator (with default args) sends the correct message to the Supervisor""" + from airflow.providers.standard.operators.trigger_dagrun import TriggerDagRunOperator + + task = TriggerDagRunOperator( + task_id="test_task", + trigger_dag_id="test_dag", + trigger_run_id="test_run_id", + ) + ti = create_runtime_ti(dag_id="test_handle_trigger_dag_run", run_id="test_run", task=task) + + log = mock.MagicMock() + mock_supervisor_comms.get_message.return_value = OKResponse(ok=True) + state, msg, _ = run(ti, log=log) + + assert state == TaskInstanceState.SUCCESS + assert msg.state == TaskInstanceState.SUCCESS + + expected_calls = [ + mock.call.send_request( + msg=TriggerDagRun( + dag_id="test_dag", + run_id="test_run_id", + reset_dag_run=False, + ), + log=mock.ANY, + ), + mock.call.get_message(), + mock.call.send_request( + msg=SetXCom( + key="trigger_run_id", + value="test_run_id", + dag_id="test_handle_trigger_dag_run", + task_id="test_task", + run_id="test_run", + map_index=-1, + ), + log=mock.ANY, + ), + ] + mock_supervisor_comms.assert_has_calls(expected_calls) + + @pytest.mark.parametrize( + ["skip_when_already_exists", "expected_state"], + [ + (True, TaskInstanceState.SKIPPED), + (False, TaskInstanceState.FAILED), + ], + ) + def test_handle_trigger_dag_run_conflict( + self, skip_when_already_exists, expected_state, create_runtime_ti, mock_supervisor_comms + ): + """Test that TriggerDagRunOperator (when dagrun already exists) sends the correct message to the Supervisor""" + from airflow.providers.standard.operators.trigger_dagrun import TriggerDagRunOperator + + task = TriggerDagRunOperator( + task_id="test_task", + trigger_dag_id="test_dag", + trigger_run_id="test_run_id", + skip_when_already_exists=skip_when_already_exists, + ) + ti = create_runtime_ti(dag_id="test_handle_trigger_dag_run_conflict", run_id="test_run", task=task) + + log = mock.MagicMock() + mock_supervisor_comms.get_message.return_value = ErrorResponse(error=ErrorType.DAGRUN_ALREADY_EXISTS) + state, msg, _ = run(ti, log=log) + + assert state == expected_state + assert msg.state == expected_state + + expected_calls = [ + mock.call.send_request( + msg=TriggerDagRun( + dag_id="test_dag", + run_id="test_run_id", + reset_dag_run=False, + ), + log=mock.ANY, + ), + mock.call.get_message(), + ] + mock_supervisor_comms.assert_has_calls(expected_calls) + + @pytest.mark.parametrize( + ["allowed_states", "failed_states", "target_dr_state", "expected_task_state"], + [ + (None, None, DagRunState.FAILED, TaskInstanceState.FAILED), + (None, None, DagRunState.SUCCESS, TaskInstanceState.SUCCESS), + ([DagRunState.FAILED], [], DagRunState.FAILED, DagRunState.SUCCESS), + ([DagRunState.FAILED], None, DagRunState.FAILED, DagRunState.FAILED), + ([DagRunState.SUCCESS], None, DagRunState.FAILED, DagRunState.FAILED), + ], + ) + def test_handle_trigger_dag_run_wait_for_completion( + self, + allowed_states, + failed_states, + target_dr_state, + expected_task_state, + create_runtime_ti, + mock_supervisor_comms, + ): + """ + Test that TriggerDagRunOperator (with wait_for_completion) sends the correct message to the Supervisor + + It also polls the Supervisor for the DagRun state until it completes execution. + """ + from airflow.providers.standard.operators.trigger_dagrun import TriggerDagRunOperator + + task = TriggerDagRunOperator( + task_id="test_task", + trigger_dag_id="test_dag", + trigger_run_id="test_run_id", + poke_interval=5, + wait_for_completion=True, + allowed_states=allowed_states, + failed_states=failed_states, + ) + ti = create_runtime_ti( + dag_id="test_handle_trigger_dag_run_wait_for_completion", run_id="test_run", task=task + ) + + log = mock.MagicMock() + mock_supervisor_comms.get_message.side_effect = [ + # Successful Dag Run trigger + OKResponse(ok=True), + # Dag Run is still running + DagRunStateResult(state=DagRunState.RUNNING), + # Dag Run completes execution on the next poll + DagRunStateResult(state=target_dr_state), + ] + with mock.patch("time.sleep", return_value=None): + state, msg, _ = run(ti, log=log) + + assert state == expected_task_state + assert msg.state == expected_task_state + + expected_calls = [ + mock.call.send_request( + msg=TriggerDagRun( + dag_id="test_dag", + run_id="test_run_id", + ), + log=mock.ANY, + ), + mock.call.get_message(), + mock.call.send_request( + msg=SetXCom( + key="trigger_run_id", + value="test_run_id", + dag_id="test_handle_trigger_dag_run_wait_for_completion", + task_id="test_task", + run_id="test_run", + map_index=-1, + ), + log=mock.ANY, + ), + mock.call.send_request( + msg=GetDagRunState( + dag_id="test_dag", + run_id="test_run_id", + ), + log=mock.ANY, + ), + mock.call.get_message(), + mock.call.send_request( + msg=GetDagRunState( + dag_id="test_dag", + run_id="test_run_id", + ), + log=mock.ANY, + ), + ] + mock_supervisor_comms.assert_has_calls(expected_calls) diff --git a/tests/api_fastapi/execution_api/routes/test_dag_runs.py b/tests/api_fastapi/execution_api/routes/test_dag_runs.py new file mode 100644 index 0000000000000..c5624f188b51a --- /dev/null +++ b/tests/api_fastapi/execution_api/routes/test_dag_runs.py @@ -0,0 +1,220 @@ +# 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 + +from airflow.models import DagModel +from airflow.models.dagrun import DagRun +from airflow.providers.standard.operators.empty import EmptyOperator +from airflow.utils import timezone +from airflow.utils.state import DagRunState + +from tests_common.test_utils.db import clear_db_runs + +pytestmark = pytest.mark.db_test + + +class TestDagRunTrigger: + def setup_method(self): + clear_db_runs() + + def teardown_method(self): + clear_db_runs() + + def test_trigger_dag_run(self, client, session, dag_maker): + dag_id = "test_trigger_dag_run" + run_id = "test_run_id" + logical_date = timezone.datetime(2025, 2, 20) + + with dag_maker(dag_id=dag_id, session=session, serialized=True): + EmptyOperator(task_id="test_task") + + session.commit() + + response = client.post( + f"/execution/dag-runs/{dag_id}/{run_id}", + json={"logical_date": logical_date.isoformat(), "conf": {"key1": "value1"}}, + ) + + assert response.status_code == 204 + + dag_run = session.query(DagRun).filter(DagRun.run_id == run_id).one() + assert dag_run.conf == {"key1": "value1"} + assert dag_run.logical_date == logical_date + + def test_trigger_dag_run_dag_not_found(self, client): + """Test that a DAG that does not exist cannot be triggered.""" + dag_id = "dag_not_found" + logical_date = timezone.datetime(2025, 2, 20) + + response = client.post( + f"/execution/dag-runs/{dag_id}/test_run_id", json={"logical_date": logical_date.isoformat()} + ) + + assert response.status_code == 404 + + def test_trigger_dag_run_import_error(self, client, session, dag_maker): + """Test that a DAG with import errors cannot be triggered.""" + + dag_id = "test_trigger_dag_run_import_error" + run_id = "test_run_id" + logical_date = timezone.datetime(2025, 2, 20) + + with dag_maker(dag_id=dag_id, session=session, serialized=True): + EmptyOperator(task_id="test_task") + + session.query(DagModel).filter(DagModel.dag_id == dag_id).update({"has_import_errors": True}) + + session.commit() + + response = client.post( + f"/execution/dag-runs/{dag_id}/{run_id}", + json={"logical_date": logical_date.isoformat()}, + ) + + assert response.status_code == 400 + assert response.json() == { + "detail": { + "message": "DAG with dag_id: 'test_trigger_dag_run_import_error' has import errors and cannot be triggered", + "reason": "import_errors", + } + } + + def test_trigger_dag_run_already_exists(self, client, session, dag_maker): + """Test that error is raised when a DAG Run already exists.""" + + dag_id = "test_trigger_dag_run_already_exists" + run_id = "test_run_id" + logical_date = timezone.datetime(2025, 2, 20) + + with dag_maker(dag_id=dag_id, session=session, serialized=True): + EmptyOperator(task_id="test_task") + + session.commit() + + response = client.post( + f"/execution/dag-runs/{dag_id}/{run_id}", + json={"logical_date": logical_date.isoformat()}, + ) + + assert response.status_code == 204 + + response = client.post( + f"/execution/dag-runs/{dag_id}/{run_id}", + json={"logical_date": logical_date.isoformat()}, + ) + + assert response.status_code == 409 + assert response.json() == { + "detail": { + "message": "A DAG Run already exists for DAG test_trigger_dag_run_already_exists with run id test_run_id", + "reason": "already_exists", + } + } + + +class TestDagRunClear: + def setup_method(self): + clear_db_runs() + + def teardown_method(self): + clear_db_runs() + + def test_dag_run_clear(self, client, session, dag_maker): + dag_id = "test_trigger_dag_run_clear" + run_id = "test_run_id" + + with dag_maker(dag_id=dag_id, session=session, serialized=True): + EmptyOperator(task_id="test_task") + + dag_maker.create_dagrun(run_id=run_id, state=DagRunState.SUCCESS) + + session.commit() + + response = client.post(f"/execution/dag-runs/{dag_id}/{run_id}/clear") + + assert response.status_code == 204 + + session.expire_all() + dag_run = session.query(DagRun).filter(DagRun.run_id == run_id).one() + assert dag_run.state == DagRunState.QUEUED + + def test_dag_run_import_error(self, client, session, dag_maker): + """Test that a DAG with import errors cannot be cleared.""" + + dag_id = "test_trigger_dag_run_import_error" + run_id = "test_run_id" + + with dag_maker(dag_id=dag_id, session=session, serialized=True): + EmptyOperator(task_id="test_task") + + session.query(DagModel).filter(DagModel.dag_id == dag_id).update({"has_import_errors": True}) + + session.commit() + + response = client.post(f"/execution/dag-runs/{dag_id}/{run_id}/clear") + + assert response.status_code == 400 + assert response.json() == { + "detail": { + "message": "DAG with dag_id: 'test_trigger_dag_run_import_error' has import errors and cannot be triggered", + "reason": "import_errors", + } + } + + def test_dag_run_not_found(self, client): + """Test that a DAG that does not exist cannot be cleared.""" + dag_id = "dag_not_found" + run_id = "test_run_id" + + response = client.post(f"/execution/dag-runs/{dag_id}/{run_id}/clear") + + assert response.status_code == 404 + + +class TestDagRunState: + def setup_method(self): + clear_db_runs() + + def teardown_method(self): + clear_db_runs() + + def test_get_state(self, client, session, dag_maker): + dag_id = "test_get_state" + run_id = "test_run_id" + + with dag_maker(dag_id=dag_id, session=session, serialized=True): + EmptyOperator(task_id="test_task") + + dag_maker.create_dagrun(run_id=run_id, state=DagRunState.SUCCESS) + + session.commit() + + response = client.get(f"/execution/dag-runs/{dag_id}/{run_id}/state") + + assert response.status_code == 200 + assert response.json() == {"state": "success"} + + def test_dag_run_not_found(self, client): + dag_id = "dag_not_found" + run_id = "test_run_id" + + response = client.post(f"/execution/dag-runs/{dag_id}/{run_id}/clear") + + assert response.status_code == 404