From cb5a7b776212f3097e9c95a5dba41c051dd9de56 Mon Sep 17 00:00:00 2001 From: Beata Kossakowska Date: Tue, 31 Jan 2023 07:53:18 +0000 Subject: [PATCH 01/12] Add deferrable mode to ExternalTaskSensor --- .../example_external_task_marker_dag.py | 1 - airflow/sensors/external_task.py | 88 +++++++++++++++++++ .../example_external_task_child_deferrable.py | 41 +++++++++ ...example_external_task_parent_deferrable.py | 66 ++++++++++++++ 4 files changed, 195 insertions(+), 1 deletion(-) create mode 100644 tests/system/providers/core/example_external_task_child_deferrable.py create mode 100644 tests/system/providers/core/example_external_task_parent_deferrable.py diff --git a/airflow/example_dags/example_external_task_marker_dag.py b/airflow/example_dags/example_external_task_marker_dag.py index 1b4f5d3ffd8f2..7158a5401f845 100644 --- a/airflow/example_dags/example_external_task_marker_dag.py +++ b/airflow/example_dags/example_external_task_marker_dag.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 diff --git a/airflow/sensors/external_task.py b/airflow/sensors/external_task.py index b5c1c1321f7ed..585fcef879066 100644 --- a/airflow/sensors/external_task.py +++ b/airflow/sensors/external_task.py @@ -33,6 +33,7 @@ from airflow.models.taskinstance import TaskInstance from airflow.operators.empty import EmptyOperator from airflow.sensors.base import BaseSensorOperator +from airflow.triggers.external_task import ExternalTaskTrigger from airflow.utils.file import correct_maybe_zipped from airflow.utils.helpers import build_airflow_url_with_query from airflow.utils.session import NEW_SESSION, provide_session @@ -428,6 +429,93 @@ def _handle_execution_date_fn(self, context) -> Any: return kwargs_callable(logical_date, **kwargs) +class ExternalTaskAsyncSensor(ExternalTaskSensor): + """ + Waits for a different DAG, task group, or task to complete for a specific logical date. + + If both `external_task_group_id` and `external_task_id` are ``None`` (default), the sensor + waits for the DAG. + Values for `external_task_group_id` and `external_task_id` can't be set at the same time. + + By default, the ExternalTaskAsyncSensor will wait for the external task to + succeed, at which point it will also succeed. However, by default it will + *not* fail if the external task fails, but will continue to check the status + until the sensor times out (thus giving you time to retry the external task + without also having to clear the sensor). + + It is possible to alter the default behavior by setting states which + cause the sensor to fail, e.g. by setting ``allowed_states=[State.FAILED]`` + and ``failed_states=[State.SUCCESS]`` you will flip the behaviour to get a + sensor which goes green when the external task *fails* and immediately goes + red if the external task *succeeds*! + + Note that ``soft_fail`` is respected when examining the failed_states. Thus + if the external task enters a failed state and ``soft_fail == True`` the + sensor will _skip_ rather than fail. As a result, setting ``soft_fail=True`` + and ``failed_states=[State.SKIPPED]`` will result in the sensor skipping if + the external task skips. + + :param external_dag_id: The dag_id that contains the task you want to + wait for + :param external_task_id: The task_id that contains the task you want to + wait for. + :param external_task_ids: The list of task_ids that you want to wait for. + If ``None`` (default value) the sensor waits for the DAG. Either + external_task_id or external_task_ids can be passed to + ExternalTaskSensor, but not both. + :param allowed_states: Iterable of allowed states, default is ``['success']`` + :param failed_states: Iterable of failed or dis-allowed states, default is ``None`` + :param execution_delta: time difference with the previous execution to + look at, the default is the same logical date as the current task or DAG. + For yesterday, use [positive!] datetime.timedelta(days=1). Either + execution_delta or execution_date_fn can be passed to + ExternalTaskSensor, but not both. + :param execution_date_fn: function that receives the current execution's logical date as the first + positional argument and optionally any number of keyword arguments available in the + context dictionary, and returns the desired logical dates to query. + Either execution_delta or execution_date_fn can be passed to ExternalTaskSensor, + but not both. + :param check_existence: Set to `True` to check if the external task exists (when + external_task_id is not None) or check if the DAG to wait for exists (when + external_task_id is None), and immediately cease waiting if the external task + or DAG does not exist (default value: False). + :param polling_interval_seconds: Time (seconds) to wait between calls to check the run status. + """ + + def __init__( + self, + external_task_id: str | None = None, + external_task_ids: Collection[str] | None = None, + polling_interval_seconds: int = 10, + **kwargs, + ) -> None: + super().__init__(**kwargs) + self.polling_interval_seconds = polling_interval_seconds + if external_task_id is not None: + external_task_ids = [external_task_id] + self.external_task_id = external_task_id + self.external_task_ids = external_task_ids + + def execute(self, context: Context) -> None: + """Airflow runs this method on the worker and defers using the trigger.""" + self.defer( + timeout=datetime.timedelta(seconds=self.timeout), + trigger=ExternalTaskTrigger( + external_dag_id=self.external_dag_id, + dttm=self._get_dttm_filter(context), + external_task_ids=self.external_task_ids, + allowed_states=self.allowed_states, + failed_states=self.failed_states, + polling_interval_seconds=self.polling_interval_seconds, + ), + method_name="execute_complete", + ) + + def execute_complete(self, context, event=None): + """Callback for when the trigger fires - returns immediately.""" + return None + + class ExternalTaskMarker(EmptyOperator): """ Use this operator to indicate that a task on a different DAG depends on this task. diff --git a/tests/system/providers/core/example_external_task_child_deferrable.py b/tests/system/providers/core/example_external_task_child_deferrable.py new file mode 100644 index 0000000000000..7318e91ce093e --- /dev/null +++ b/tests/system/providers/core/example_external_task_child_deferrable.py @@ -0,0 +1,41 @@ +# 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 + +from datetime import datetime + +from airflow import DAG +from airflow.sensors.external_task import ExternalTaskMarker + +with DAG( + dag_id="wait_for_dag_child", + start_date=datetime(2022, 1, 1), + schedule="@once", + catchup=False, + tags=["example", "async", "core"], +) as dag: + wait_for_task = ExternalTaskMarker( + task_id="wait_for_task", + external_dag_id="example_external_task", + external_task_id="wait_for_task", + ) + + +from tests.system.utils import get_test_run + +# Needed to run the example DAG with pytest (see: tests/system/README.md#run_via_pytest) +test_run = get_test_run(dag) diff --git a/tests/system/providers/core/example_external_task_parent_deferrable.py b/tests/system/providers/core/example_external_task_parent_deferrable.py new file mode 100644 index 0000000000000..b65d9b29a9a1a --- /dev/null +++ b/tests/system/providers/core/example_external_task_parent_deferrable.py @@ -0,0 +1,66 @@ +# 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 + +from airflow import DAG +from airflow.operators.dummy import DummyOperator +from airflow.operators.trigger_dagrun import TriggerDagRunOperator +from airflow.sensors.external_task import ExternalTaskAsyncSensor +from airflow.utils.timezone import datetime + +with DAG( + dag_id="example_external_task", + start_date=datetime(2022, 1, 1), + schedule="@once", + catchup=False, + tags=["example", "async", "core"], +) as dag: + start = DummyOperator(task_id="start") + + # [START howto_external_task_async_sensor] + wait_for_task = ExternalTaskAsyncSensor( + task_id="wait_for_task", + external_task_id="wait_for_task_child", + external_dag_id="wait_for_dag_child", + ) + # [END howto_external_task_async_sensor] + + external_task = TriggerDagRunOperator( + task_id="external_task", + trigger_dag_id="wait_for_dag_child", + allowed_states=["success", "failed", "skipped"], + execution_date="{{execution_date}}", + poke_interval=5, + reset_dag_run=True, + wait_for_completion=True, + ) + + end = DummyOperator(task_id="end") + + start >> [external_task, wait_for_task] >> end + + from tests.system.utils.watcher import watcher + + # This test needs watcher in order to properly mark success/failure + # when "teardown" task with trigger rule is part of the DAG + list(dag.tasks) >> watcher() + + +from tests.system.utils import get_test_run + +# Needed to run the example DAG with pytest (see: tests/system/README.md#run_via_pytest) +test_run = get_test_run(dag) From f88e363211f404e18168165cfdd033ad221df456 Mon Sep 17 00:00:00 2001 From: VladaZakharova Date: Fri, 3 Feb 2023 15:47:00 +0100 Subject: [PATCH 02/12] Rebase main and change ExternalTaskAsyncSensor implementation --- airflow/sensors/external_task.py | 15 +++- airflow/triggers/external_task.py | 68 ++++++++++++++++--- tests/system/providers/core/__init__.py | 0 .../example_external_task_child_deferrable.py | 11 ++- ...example_external_task_parent_deferrable.py | 16 ++--- tests/triggers/test_external_task.py | 6 ++ 6 files changed, 89 insertions(+), 27 deletions(-) create mode 100644 tests/system/providers/core/__init__.py diff --git a/airflow/sensors/external_task.py b/airflow/sensors/external_task.py index 585fcef879066..d55f4a77ba661 100644 --- a/airflow/sensors/external_task.py +++ b/airflow/sensors/external_task.py @@ -33,11 +33,12 @@ from airflow.models.taskinstance import TaskInstance from airflow.operators.empty import EmptyOperator from airflow.sensors.base import BaseSensorOperator -from airflow.triggers.external_task import ExternalTaskTrigger +from airflow.triggers.external_task import TaskStateTrigger from airflow.utils.file import correct_maybe_zipped from airflow.utils.helpers import build_airflow_url_with_query from airflow.utils.session import NEW_SESSION, provide_session from airflow.utils.sqlalchemy import tuple_in_condition +from airflow.utils.timezone import utcnow from airflow.utils.state import State, TaskInstanceState if TYPE_CHECKING: @@ -484,6 +485,7 @@ class ExternalTaskAsyncSensor(ExternalTaskSensor): def __init__( self, + *, external_task_id: str | None = None, external_task_ids: Collection[str] | None = None, polling_interval_seconds: int = 10, @@ -513,7 +515,16 @@ def execute(self, context: Context) -> None: def execute_complete(self, context, event=None): """Callback for when the trigger fires - returns immediately.""" - return None + if event["status"] == "success": + self.log.info("External task %s has executed successfully", self.external_task_id) + return None + elif event["status"] == "timeout": + raise AirflowException("took longer then expected") + else: + raise AirflowException( + "Error occurred while trying to retrieve task status. Please, check the " + "name of executed task and Dag" + ) class ExternalTaskMarker(EmptyOperator): diff --git a/airflow/triggers/external_task.py b/airflow/triggers/external_task.py index e739c7a7cb4e8..461dd5feca264 100644 --- a/airflow/triggers/external_task.py +++ b/airflow/triggers/external_task.py @@ -19,6 +19,7 @@ import asyncio import datetime import typing +from datetime import datetime from asgiref.sync import sync_to_async from sqlalchemy import func @@ -28,6 +29,7 @@ from airflow.triggers.base import BaseTrigger, TriggerEvent from airflow.utils.session import NEW_SESSION, provide_session from airflow.utils.state import DagRunState +from airflow.utils.timezone import utcnow class TaskStateTrigger(BaseTrigger): @@ -38,18 +40,24 @@ class TaskStateTrigger(BaseTrigger): :param task_id: The task_id that contains the task you want to wait for. If ``None`` (default value) the sensor waits for the DAG :param states: allowed states, default is ``['success']`` - :param execution_dates: + :param execution_dates: task execution time interval :param poll_interval: The time interval in seconds to check the state. The default value is 5 sec. + :param trigger_start_time: time in Datetime format when the trigger was started. Is used + to control the execution of trigger to prevent infinite loop in case if specified name + of the dag does not exist in database. It will wait period of time equals _timeout_sec parameter + from the time, when the trigger was started and if the execution lasts more time than expected, + the trigger will terminate with 'timeout' status. """ def __init__( self, dag_id: str, - task_id: str, states: list[str], - execution_dates: list[datetime.datetime], - poll_interval: float = 5.0, + execution_dates: list[datetime], + trigger_start_time: datetime, + task_id: str | None = None, + poll_interval: float = 2.0, ): super().__init__() self.dag_id = dag_id @@ -57,6 +65,8 @@ def __init__( self.states = states self.execution_dates = execution_dates self.poll_interval = poll_interval + self.trigger_start_time = trigger_start_time + self._timeout_sec = 60 def serialize(self) -> tuple[str, dict[str, typing.Any]]: """Serializes TaskStateTrigger arguments and classpath.""" @@ -68,17 +78,50 @@ def serialize(self) -> tuple[str, dict[str, typing.Any]]: "states": self.states, "execution_dates": self.execution_dates, "poll_interval": self.poll_interval, + "trigger_start_time": self.trigger_start_time, }, ) async def run(self) -> typing.AsyncIterator[TriggerEvent]: - """Checks periodically in the database to see if the task exists and has hit one of the states.""" + """ + Checks periodically in the database to see if the dag exists and is in the running state. If found, + wait until the task specified will reach one of the expected states. If dag with specified name was + not in the running state after _timeout_sec seconds after starting execution process of the trigger, + terminate with status 'timeout'. + """ while True: - # mypy confuses typing here - num_tasks = await self.count_tasks() # type: ignore[call-arg] - if num_tasks == len(self.execution_dates): - yield TriggerEvent(True) - await asyncio.sleep(self.poll_interval) + try: + delta = utcnow() - self.trigger_start_time + if delta.total_seconds() < self._timeout_sec: + if await self.count_running_dags() == 0: + self.log.info("Waiting for DAG to start execution...") + await asyncio.sleep(self.poll_interval) + else: + yield TriggerEvent({"status": "timeout"}) + return + if await self.count_tasks() == len(self.execution_dates): + yield TriggerEvent({"status": "success"}) + return + self.log.info("Task is still running, sleeping for %s seconds...", self.poll_interval) + await asyncio.sleep(self.poll_interval) + except Exception: + yield TriggerEvent({"status": "failed"}) + return + + @sync_to_async + @provide_session + def count_running_dags(self, session: Session): + """Count how many dag instances in running state in the database.""" + dags = ( + session.query(func.count("*")) + .filter( + TaskInstance.dag_id == self.dag_id, + TaskInstance.execution_date.in_(self.execution_dates), + TaskInstance.state.in_(["running", "success"]), + ) + .scalar() + ) + return dags @sync_to_async @provide_session @@ -134,7 +177,10 @@ def serialize(self) -> tuple[str, dict[str, typing.Any]]: ) async def run(self) -> typing.AsyncIterator[TriggerEvent]: - """Checks periodically in the database to see if the dag run exists and has hit one of the states.""" + """ + Checks periodically in the database to see if the dag run exists, and has + hit one of the states yet, or not. + """ while True: # mypy confuses typing here num_dags = await self.count_dags() # type: ignore[call-arg] diff --git a/tests/system/providers/core/__init__.py b/tests/system/providers/core/__init__.py new file mode 100644 index 0000000000000..e69de29bb2d1d diff --git a/tests/system/providers/core/example_external_task_child_deferrable.py b/tests/system/providers/core/example_external_task_child_deferrable.py index 7318e91ce093e..f75eb4f23479f 100644 --- a/tests/system/providers/core/example_external_task_child_deferrable.py +++ b/tests/system/providers/core/example_external_task_child_deferrable.py @@ -19,19 +19,18 @@ from datetime import datetime from airflow import DAG -from airflow.sensors.external_task import ExternalTaskMarker +from airflow.operators.bash import BashOperator with DAG( - dag_id="wait_for_dag_child", + dag_id="child_dag", start_date=datetime(2022, 1, 1), schedule="@once", catchup=False, tags=["example", "async", "core"], ) as dag: - wait_for_task = ExternalTaskMarker( - task_id="wait_for_task", - external_dag_id="example_external_task", - external_task_id="wait_for_task", + dummy_task = BashOperator( + task_id="child_task", + bash_command="echo 1; sleep 1; echo 2; sleep 2; echo 3; sleep 3", ) diff --git a/tests/system/providers/core/example_external_task_parent_deferrable.py b/tests/system/providers/core/example_external_task_parent_deferrable.py index b65d9b29a9a1a..af4f43edaa6d1 100644 --- a/tests/system/providers/core/example_external_task_parent_deferrable.py +++ b/tests/system/providers/core/example_external_task_parent_deferrable.py @@ -32,16 +32,16 @@ start = DummyOperator(task_id="start") # [START howto_external_task_async_sensor] - wait_for_task = ExternalTaskAsyncSensor( - task_id="wait_for_task", - external_task_id="wait_for_task_child", - external_dag_id="wait_for_dag_child", + external_task_sensor = ExternalTaskAsyncSensor( + task_id="parent_task_sensor", + external_task_id="child_task", + external_dag_id="child_dag", ) # [END howto_external_task_async_sensor] - external_task = TriggerDagRunOperator( - task_id="external_task", - trigger_dag_id="wait_for_dag_child", + trigger_child_task = TriggerDagRunOperator( + task_id="trigger_child_task", + trigger_dag_id="child_dag", allowed_states=["success", "failed", "skipped"], execution_date="{{execution_date}}", poke_interval=5, @@ -51,7 +51,7 @@ end = DummyOperator(task_id="end") - start >> [external_task, wait_for_task] >> end + start >> [trigger_child_task, external_task_sensor] >> end from tests.system.utils.watcher import watcher diff --git a/tests/triggers/test_external_task.py b/tests/triggers/test_external_task.py index e8a4a67ba5d57..a8569a9c558d1 100644 --- a/tests/triggers/test_external_task.py +++ b/tests/triggers/test_external_task.py @@ -26,6 +26,7 @@ from airflow.triggers.external_task import DagStateTrigger, TaskStateTrigger from airflow.utils import timezone from airflow.utils.state import DagRunState, TaskInstanceState +from airflow.utils.timezone import utcnow class TestTaskStateTrigger: @@ -40,6 +41,7 @@ async def test_task_state_trigger(self, session): Asserts that the TaskStateTrigger only goes off on or after a TaskInstance reaches an allowed state (i.e. SUCCESS). """ + trigger_start_time = utcnow() dag = DAG(self.DAG_ID, start_date=timezone.datetime(2022, 1, 1)) dag_run = DagRun( dag_id=dag.dag_id, @@ -61,6 +63,7 @@ async def test_task_state_trigger(self, session): states=self.STATES, execution_dates=[timezone.datetime(2022, 1, 1)], poll_interval=0.2, + trigger_start_time=trigger_start_time, ) task = asyncio.create_task(trigger.run().__anext__()) @@ -83,12 +86,14 @@ def test_serialization(self): Asserts that the TaskStateTrigger correctly serializes its arguments and classpath. """ + trigger_start_time = utcnow() trigger = TaskStateTrigger( dag_id=self.DAG_ID, task_id=self.TASK_ID, states=self.STATES, execution_dates=[timezone.datetime(2022, 1, 1)], poll_interval=5, + trigger_start_time=trigger_start_time, ) classpath, kwargs = trigger.serialize() assert classpath == "airflow.triggers.external_task.TaskStateTrigger" @@ -98,6 +103,7 @@ def test_serialization(self): "states": self.STATES, "execution_dates": [timezone.datetime(2022, 1, 1)], "poll_interval": 5, + "trigger_start_time": trigger_start_time, } From 4f6040f66e29ff37a9af57c57a90518409e1e353 Mon Sep 17 00:00:00 2001 From: VladaZakharova Date: Fri, 3 Feb 2023 16:16:10 +0100 Subject: [PATCH 03/12] Update __init__.py --- tests/system/providers/core/__init__.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/tests/system/providers/core/__init__.py b/tests/system/providers/core/__init__.py index e69de29bb2d1d..13a83393a9124 100644 --- a/tests/system/providers/core/__init__.py +++ b/tests/system/providers/core/__init__.py @@ -0,0 +1,16 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. From bee2d6c6a11c25a2da3cdca1ab04243aa2e4b418 Mon Sep 17 00:00:00 2001 From: Ulada Zakharava Date: Wed, 8 Feb 2023 15:51:36 +0000 Subject: [PATCH 04/12] Add unit tests for the AsyncSensor --- .../example_external_task_marker_dag.py | 1 + airflow/sensors/external_task.py | 2 +- tests/sensors/test_external_task_sensor.py | 75 ++++++++++++++++++- 3 files changed, 75 insertions(+), 3 deletions(-) diff --git a/airflow/example_dags/example_external_task_marker_dag.py b/airflow/example_dags/example_external_task_marker_dag.py index 7158a5401f845..1b4f5d3ffd8f2 100644 --- a/airflow/example_dags/example_external_task_marker_dag.py +++ b/airflow/example_dags/example_external_task_marker_dag.py @@ -1,3 +1,4 @@ +# # 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 diff --git a/airflow/sensors/external_task.py b/airflow/sensors/external_task.py index d55f4a77ba661..6839747d296fc 100644 --- a/airflow/sensors/external_task.py +++ b/airflow/sensors/external_task.py @@ -519,7 +519,7 @@ def execute_complete(self, context, event=None): self.log.info("External task %s has executed successfully", self.external_task_id) return None elif event["status"] == "timeout": - raise AirflowException("took longer then expected") + raise AirflowException("Dag was not started within 1 minute, assuming fail") else: raise AirflowException( "Error occurred while trying to retrieve task status. Please, check the " diff --git a/tests/sensors/test_external_task_sensor.py b/tests/sensors/test_external_task_sensor.py index a5259084b13f6..def19f9ad1278 100644 --- a/tests/sensors/test_external_task_sensor.py +++ b/tests/sensors/test_external_task_sensor.py @@ -22,12 +22,13 @@ import tempfile import zipfile from datetime import time, timedelta +from unittest import mock import pytest from airflow import exceptions, settings from airflow.decorators import task as task_deco -from airflow.exceptions import AirflowException, AirflowSensorTimeout +from airflow.exceptions import AirflowException, AirflowSensorTimeout, TaskDeferred from airflow.models import DagBag, DagRun, TaskInstance from airflow.models.dag import DAG from airflow.models.serialized_dag import SerializedDagModel @@ -35,9 +36,15 @@ from airflow.operators.bash import BashOperator from airflow.operators.empty import EmptyOperator from airflow.operators.python import PythonOperator -from airflow.sensors.external_task import ExternalTaskMarker, ExternalTaskSensor, ExternalTaskSensorLink +from airflow.sensors.external_task import ( + ExternalTaskAsyncSensor, + ExternalTaskMarker, + ExternalTaskSensor, + ExternalTaskSensorLink, +) from airflow.sensors.time_sensor import TimeSensor from airflow.serialization.serialized_objects import SerializedBaseOperator +from airflow.triggers.external_task import TaskStateTrigger from airflow.utils.hashlib_wrapper import md5 from airflow.utils.session import create_session, provide_session from airflow.utils.state import DagRunState, State, TaskInstanceState @@ -829,6 +836,70 @@ def test_external_task_group_when_there_is_no_TIs(self): ) +class TestExternalTaskAsyncSensor: + TASK_ID = "external_task_sensor_check" + EXTERNAL_DAG_ID = "child_dag" # DAG the external task sensor is waiting on + EXTERNAL_TASK_ID = "child_task" # Task the external task sensor is waiting on + + def test_defer_and_fire_task_state_trigger(self): + """ + Asserts that a task is deferred and TaskStateTrigger will be fired + when the ExternalTaskAsyncSensor is provided with all required arguments + (i.e. including the external_task_id). + """ + sensor = ExternalTaskAsyncSensor( + task_id=self.TASK_ID, + external_task_id=self.EXTERNAL_TASK_ID, + external_dag_id=self.EXTERNAL_DAG_ID, + ) + + with pytest.raises(TaskDeferred) as exc: + sensor.execute(context=mock.MagicMock()) + + assert isinstance(exc.value.trigger, TaskStateTrigger), "Trigger is not a TaskStateTrigger" + + def test_defer_and_fire_failed_state_trigger(self): + """Tests that an AirflowException is raised in case of error event""" + sensor = ExternalTaskAsyncSensor( + task_id=self.TASK_ID, + external_task_id=self.EXTERNAL_TASK_ID, + external_dag_id=self.EXTERNAL_DAG_ID, + ) + + with pytest.raises(AirflowException): + sensor.execute_complete( + context=mock.MagicMock(), event={"status": "error", "message": "test failure message"} + ) + + def test_defer_and_fire_timeout_state_trigger(self): + """Tests that an AirflowException is raised in case of timeout event""" + sensor = ExternalTaskAsyncSensor( + task_id=self.TASK_ID, + external_task_id=self.EXTERNAL_TASK_ID, + external_dag_id=self.EXTERNAL_DAG_ID, + ) + + with pytest.raises(AirflowException): + sensor.execute_complete( + context=mock.MagicMock(), event={"status": "timeout", "message": "Dag was not started"} + ) + + def test_defer_execute_check_correct_logging(self): + """Asserts that logging occurs as expected""" + sensor = ExternalTaskAsyncSensor( + task_id=self.TASK_ID, + external_task_id=self.EXTERNAL_TASK_ID, + external_dag_id=self.EXTERNAL_DAG_ID, + ) + + with mock.patch.object(sensor.log, "info") as mock_log_info: + sensor.execute_complete( + context=mock.MagicMock(), + event={"status": "success"}, + ) + mock_log_info.assert_called_with("External task %s has executed successfully", self.EXTERNAL_TASK_ID) + + def test_external_task_sensor_check_zipped_dag_existence(dag_zip_maker): with dag_zip_maker("test_external_task_sensor_check_existense.py") as dagbag: with create_session() as session: From 2d32bc6a683346cb934be3109003c880b5e925a1 Mon Sep 17 00:00:00 2001 From: Ulada Zakharava Date: Thu, 9 Feb 2023 07:08:14 +0000 Subject: [PATCH 05/12] Update message --- airflow/sensors/external_task.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/airflow/sensors/external_task.py b/airflow/sensors/external_task.py index 6839747d296fc..02d00ccf2e005 100644 --- a/airflow/sensors/external_task.py +++ b/airflow/sensors/external_task.py @@ -519,7 +519,7 @@ def execute_complete(self, context, event=None): self.log.info("External task %s has executed successfully", self.external_task_id) return None elif event["status"] == "timeout": - raise AirflowException("Dag was not started within 1 minute, assuming fail") + raise AirflowException("Dag was not started within 1 minute, assuming fail.") else: raise AirflowException( "Error occurred while trying to retrieve task status. Please, check the " From 99bf4600bcaf077fa4c16289279be236112e0926 Mon Sep 17 00:00:00 2001 From: Ulada Zakharava Date: Thu, 9 Feb 2023 10:07:49 +0000 Subject: [PATCH 06/12] Update message --- airflow/sensors/external_task.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/airflow/sensors/external_task.py b/airflow/sensors/external_task.py index 02d00ccf2e005..14784933386b3 100644 --- a/airflow/sensors/external_task.py +++ b/airflow/sensors/external_task.py @@ -523,7 +523,7 @@ def execute_complete(self, context, event=None): else: raise AirflowException( "Error occurred while trying to retrieve task status. Please, check the " - "name of executed task and Dag" + "name of executed task and Dag." ) From 9bc8d5ceba4fd96bd2713b687043406705a5f1ad Mon Sep 17 00:00:00 2001 From: Ulada Zakharava Date: Mon, 13 Feb 2023 07:35:12 +0000 Subject: [PATCH 07/12] Update message --- airflow/sensors/external_task.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/airflow/sensors/external_task.py b/airflow/sensors/external_task.py index 14784933386b3..d72cf05ab62a5 100644 --- a/airflow/sensors/external_task.py +++ b/airflow/sensors/external_task.py @@ -516,7 +516,7 @@ def execute(self, context: Context) -> None: def execute_complete(self, context, event=None): """Callback for when the trigger fires - returns immediately.""" if event["status"] == "success": - self.log.info("External task %s has executed successfully", self.external_task_id) + self.log.info("External task %s has executed successfully.", self.external_task_id) return None elif event["status"] == "timeout": raise AirflowException("Dag was not started within 1 minute, assuming fail.") From 226ff697a1089fde7c4676fd6e3aa03e52126528 Mon Sep 17 00:00:00 2001 From: Ulada Zakharava Date: Mon, 13 Feb 2023 10:03:36 +0000 Subject: [PATCH 08/12] Update test message --- tests/sensors/test_external_task_sensor.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/tests/sensors/test_external_task_sensor.py b/tests/sensors/test_external_task_sensor.py index def19f9ad1278..30542bdeff8af 100644 --- a/tests/sensors/test_external_task_sensor.py +++ b/tests/sensors/test_external_task_sensor.py @@ -881,7 +881,10 @@ def test_defer_and_fire_timeout_state_trigger(self): with pytest.raises(AirflowException): sensor.execute_complete( - context=mock.MagicMock(), event={"status": "timeout", "message": "Dag was not started"} + context=mock.MagicMock(), event={ + "status": "timeout", + "message": "Dag was not started within 1 minute, assuming fail." + } ) def test_defer_execute_check_correct_logging(self): @@ -897,7 +900,7 @@ def test_defer_execute_check_correct_logging(self): context=mock.MagicMock(), event={"status": "success"}, ) - mock_log_info.assert_called_with("External task %s has executed successfully", self.EXTERNAL_TASK_ID) + mock_log_info.assert_called_with("External task %s has executed successfully.", self.EXTERNAL_TASK_ID) def test_external_task_sensor_check_zipped_dag_existence(dag_zip_maker): From d2987bd3276de8772bd7d8cf4a4da357e25ef279 Mon Sep 17 00:00:00 2001 From: Ulada Zakharava Date: Mon, 13 Feb 2023 11:29:58 +0000 Subject: [PATCH 09/12] update test --- tests/sensors/test_external_task_sensor.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/tests/sensors/test_external_task_sensor.py b/tests/sensors/test_external_task_sensor.py index 30542bdeff8af..3346241b878f9 100644 --- a/tests/sensors/test_external_task_sensor.py +++ b/tests/sensors/test_external_task_sensor.py @@ -881,10 +881,8 @@ def test_defer_and_fire_timeout_state_trigger(self): with pytest.raises(AirflowException): sensor.execute_complete( - context=mock.MagicMock(), event={ - "status": "timeout", - "message": "Dag was not started within 1 minute, assuming fail." - } + context=mock.MagicMock(), + event={"status": "timeout", "message": "Dag was not started within 1 minute, assuming fail."}, ) def test_defer_execute_check_correct_logging(self): From 994aa43e3da443340854c1af0581583acc313cf8 Mon Sep 17 00:00:00 2001 From: Beata Kossakowska Date: Wed, 17 May 2023 08:16:32 +0000 Subject: [PATCH 10/12] Merge ExternalTaskAsyncSensor into ExternalTaskSensor --- airflow/sensors/external_task.py | 137 +++++------------- .../howto/operator/external_task_sensor.rst | 9 ++ tests/sensors/test_external_task_sensor.py | 42 +++--- ...example_external_task_parent_deferrable.py | 5 +- 4 files changed, 76 insertions(+), 117 deletions(-) diff --git a/airflow/sensors/external_task.py b/airflow/sensors/external_task.py index d72cf05ab62a5..3e282a0b3a37c 100644 --- a/airflow/sensors/external_task.py +++ b/airflow/sensors/external_task.py @@ -38,6 +38,7 @@ from airflow.utils.helpers import build_airflow_url_with_query from airflow.utils.session import NEW_SESSION, provide_session from airflow.utils.sqlalchemy import tuple_in_condition +from airflow.utils.state import State from airflow.utils.timezone import utcnow from airflow.utils.state import State, TaskInstanceState @@ -128,6 +129,8 @@ class ExternalTaskSensor(BaseSensorOperator): external_task_id is not None) or check if the DAG to wait for exists (when external_task_id is None), and immediately cease waiting if the external task or DAG does not exist (default value: False). + :param poll_interval: polling period in seconds to check for the status + :param deferrable: Run sensor in deferrable mode """ template_fields = ["external_dag_id", "external_task_id", "external_task_ids", "external_task_group_id"] @@ -147,6 +150,8 @@ def __init__( execution_delta: datetime.timedelta | None = None, execution_date_fn: Callable | None = None, check_existence: bool = False, + poll_interval: float = 2.0, + deferrable: bool = False, **kwargs, ): super().__init__(**kwargs) @@ -213,6 +218,8 @@ def __init__( self.external_task_group_id = external_task_group_id self.check_existence = check_existence self._has_checked_existence = False + self.deferrable = deferrable + self.poll_interval = poll_interval def _get_dttm_filter(self, context): if self.execution_delta: @@ -320,6 +327,39 @@ def poke(self, context: Context, session: Session = NEW_SESSION) -> bool: count_allowed = self.get_count(dttm_filter, session, self.allowed_states) return count_allowed == len(dttm_filter) + def execute(self, context: Context) -> None: + """ + Airflow runs this method on the worker and defers using the triggers + if deferrable is set to True. + """ + if not self.deferrable: + super().execute(context) + else: + self.defer( + trigger=TaskStateTrigger( + dag_id=self.external_dag_id, + task_id=self.external_task_id, + execution_dates=self._get_dttm_filter(context), + states=self.allowed_states, + trigger_start_time=utcnow(), + poll_interval=self.poll_interval, + ), + method_name="execute_complete", + ) + + def execute_complete(self, context, event=None): + """Callback for when the trigger fires - returns immediately.""" + if event["status"] == "success": + self.log.info("External task %s has executed successfully.", self.external_task_id) + return None + elif event["status"] == "timeout": + raise AirflowException("Dag was not started within 1 minute, assuming fail.") + else: + raise AirflowException( + "Error occurred while trying to retrieve task status. Please, check the " + "name of executed task and Dag." + ) + def _check_for_existence(self, session) -> None: dag_to_wait = DagModel.get_current(self.external_dag_id, session) @@ -430,103 +470,6 @@ def _handle_execution_date_fn(self, context) -> Any: return kwargs_callable(logical_date, **kwargs) -class ExternalTaskAsyncSensor(ExternalTaskSensor): - """ - Waits for a different DAG, task group, or task to complete for a specific logical date. - - If both `external_task_group_id` and `external_task_id` are ``None`` (default), the sensor - waits for the DAG. - Values for `external_task_group_id` and `external_task_id` can't be set at the same time. - - By default, the ExternalTaskAsyncSensor will wait for the external task to - succeed, at which point it will also succeed. However, by default it will - *not* fail if the external task fails, but will continue to check the status - until the sensor times out (thus giving you time to retry the external task - without also having to clear the sensor). - - It is possible to alter the default behavior by setting states which - cause the sensor to fail, e.g. by setting ``allowed_states=[State.FAILED]`` - and ``failed_states=[State.SUCCESS]`` you will flip the behaviour to get a - sensor which goes green when the external task *fails* and immediately goes - red if the external task *succeeds*! - - Note that ``soft_fail`` is respected when examining the failed_states. Thus - if the external task enters a failed state and ``soft_fail == True`` the - sensor will _skip_ rather than fail. As a result, setting ``soft_fail=True`` - and ``failed_states=[State.SKIPPED]`` will result in the sensor skipping if - the external task skips. - - :param external_dag_id: The dag_id that contains the task you want to - wait for - :param external_task_id: The task_id that contains the task you want to - wait for. - :param external_task_ids: The list of task_ids that you want to wait for. - If ``None`` (default value) the sensor waits for the DAG. Either - external_task_id or external_task_ids can be passed to - ExternalTaskSensor, but not both. - :param allowed_states: Iterable of allowed states, default is ``['success']`` - :param failed_states: Iterable of failed or dis-allowed states, default is ``None`` - :param execution_delta: time difference with the previous execution to - look at, the default is the same logical date as the current task or DAG. - For yesterday, use [positive!] datetime.timedelta(days=1). Either - execution_delta or execution_date_fn can be passed to - ExternalTaskSensor, but not both. - :param execution_date_fn: function that receives the current execution's logical date as the first - positional argument and optionally any number of keyword arguments available in the - context dictionary, and returns the desired logical dates to query. - Either execution_delta or execution_date_fn can be passed to ExternalTaskSensor, - but not both. - :param check_existence: Set to `True` to check if the external task exists (when - external_task_id is not None) or check if the DAG to wait for exists (when - external_task_id is None), and immediately cease waiting if the external task - or DAG does not exist (default value: False). - :param polling_interval_seconds: Time (seconds) to wait between calls to check the run status. - """ - - def __init__( - self, - *, - external_task_id: str | None = None, - external_task_ids: Collection[str] | None = None, - polling_interval_seconds: int = 10, - **kwargs, - ) -> None: - super().__init__(**kwargs) - self.polling_interval_seconds = polling_interval_seconds - if external_task_id is not None: - external_task_ids = [external_task_id] - self.external_task_id = external_task_id - self.external_task_ids = external_task_ids - - def execute(self, context: Context) -> None: - """Airflow runs this method on the worker and defers using the trigger.""" - self.defer( - timeout=datetime.timedelta(seconds=self.timeout), - trigger=ExternalTaskTrigger( - external_dag_id=self.external_dag_id, - dttm=self._get_dttm_filter(context), - external_task_ids=self.external_task_ids, - allowed_states=self.allowed_states, - failed_states=self.failed_states, - polling_interval_seconds=self.polling_interval_seconds, - ), - method_name="execute_complete", - ) - - def execute_complete(self, context, event=None): - """Callback for when the trigger fires - returns immediately.""" - if event["status"] == "success": - self.log.info("External task %s has executed successfully.", self.external_task_id) - return None - elif event["status"] == "timeout": - raise AirflowException("Dag was not started within 1 minute, assuming fail.") - else: - raise AirflowException( - "Error occurred while trying to retrieve task status. Please, check the " - "name of executed task and Dag." - ) - - class ExternalTaskMarker(EmptyOperator): """ Use this operator to indicate that a task on a different DAG depends on this task. diff --git a/docs/apache-airflow/howto/operator/external_task_sensor.rst b/docs/apache-airflow/howto/operator/external_task_sensor.rst index 923f8ec3d1161..f6f53f87e55fa 100644 --- a/docs/apache-airflow/howto/operator/external_task_sensor.rst +++ b/docs/apache-airflow/howto/operator/external_task_sensor.rst @@ -53,6 +53,15 @@ via ``allowed_states`` and ``failed_states`` parameters. :start-after: [START howto_operator_external_task_sensor] :end-before: [END howto_operator_external_task_sensor] +Also for this action you can use sensor in the deferrable mode: + +.. exampleinclude:: /../../tests/system/providers/core/example_external_task_parent_deferrable.py + :language: python + :dedent: 4 + :start-after: [START howto_external_task_async_sensor] + :end-before: [END howto_external_task_async_sensor] + + ExternalTaskSensor with task_group dependency --------------------------------------------- In Addition, we can also use the :class:`~airflow.sensors.external_task.ExternalTaskSensor` to make tasks on a DAG diff --git a/tests/sensors/test_external_task_sensor.py b/tests/sensors/test_external_task_sensor.py index 3346241b878f9..e84b3f69f48e0 100644 --- a/tests/sensors/test_external_task_sensor.py +++ b/tests/sensors/test_external_task_sensor.py @@ -37,7 +37,6 @@ from airflow.operators.empty import EmptyOperator from airflow.operators.python import PythonOperator from airflow.sensors.external_task import ( - ExternalTaskAsyncSensor, ExternalTaskMarker, ExternalTaskSensor, ExternalTaskSensorLink, @@ -61,6 +60,9 @@ TEST_TASK_ID_ALTERNATE = "time_sensor_check_alternate" TEST_TASK_GROUP_ID = "time_sensor_group_id" DEV_NULL = "/dev/null" +TASK_ID = "external_task_sensor_check" +EXTERNAL_DAG_ID = "child_dag" # DAG the external task sensor is waiting on +EXTERNAL_TASK_ID = "child_task" # Task the external task sensor is waiting on @pytest.fixture(autouse=True) @@ -847,10 +849,11 @@ def test_defer_and_fire_task_state_trigger(self): when the ExternalTaskAsyncSensor is provided with all required arguments (i.e. including the external_task_id). """ - sensor = ExternalTaskAsyncSensor( - task_id=self.TASK_ID, - external_task_id=self.EXTERNAL_TASK_ID, - external_dag_id=self.EXTERNAL_DAG_ID, + sensor = ExternalTaskSensor( + task_id=TASK_ID, + external_task_id=EXTERNAL_TASK_ID, + external_dag_id=EXTERNAL_DAG_ID, + deferrable=True, ) with pytest.raises(TaskDeferred) as exc: @@ -860,10 +863,11 @@ def test_defer_and_fire_task_state_trigger(self): def test_defer_and_fire_failed_state_trigger(self): """Tests that an AirflowException is raised in case of error event""" - sensor = ExternalTaskAsyncSensor( - task_id=self.TASK_ID, - external_task_id=self.EXTERNAL_TASK_ID, - external_dag_id=self.EXTERNAL_DAG_ID, + sensor = ExternalTaskSensor( + task_id=TASK_ID, + external_task_id=EXTERNAL_TASK_ID, + external_dag_id=EXTERNAL_DAG_ID, + deferrable=True, ) with pytest.raises(AirflowException): @@ -873,10 +877,11 @@ def test_defer_and_fire_failed_state_trigger(self): def test_defer_and_fire_timeout_state_trigger(self): """Tests that an AirflowException is raised in case of timeout event""" - sensor = ExternalTaskAsyncSensor( - task_id=self.TASK_ID, - external_task_id=self.EXTERNAL_TASK_ID, - external_dag_id=self.EXTERNAL_DAG_ID, + sensor = ExternalTaskSensor( + task_id=TASK_ID, + external_task_id=EXTERNAL_TASK_ID, + external_dag_id=EXTERNAL_DAG_ID, + deferrable=True, ) with pytest.raises(AirflowException): @@ -887,10 +892,11 @@ def test_defer_and_fire_timeout_state_trigger(self): def test_defer_execute_check_correct_logging(self): """Asserts that logging occurs as expected""" - sensor = ExternalTaskAsyncSensor( - task_id=self.TASK_ID, - external_task_id=self.EXTERNAL_TASK_ID, - external_dag_id=self.EXTERNAL_DAG_ID, + sensor = ExternalTaskSensor( + task_id=TASK_ID, + external_task_id=EXTERNAL_TASK_ID, + external_dag_id=EXTERNAL_DAG_ID, + deferrable=True, ) with mock.patch.object(sensor.log, "info") as mock_log_info: @@ -898,7 +904,7 @@ def test_defer_execute_check_correct_logging(self): context=mock.MagicMock(), event={"status": "success"}, ) - mock_log_info.assert_called_with("External task %s has executed successfully.", self.EXTERNAL_TASK_ID) + mock_log_info.assert_called_with("External task %s has executed successfully.", EXTERNAL_TASK_ID) def test_external_task_sensor_check_zipped_dag_existence(dag_zip_maker): diff --git a/tests/system/providers/core/example_external_task_parent_deferrable.py b/tests/system/providers/core/example_external_task_parent_deferrable.py index af4f43edaa6d1..7ae405674828e 100644 --- a/tests/system/providers/core/example_external_task_parent_deferrable.py +++ b/tests/system/providers/core/example_external_task_parent_deferrable.py @@ -19,7 +19,7 @@ from airflow import DAG from airflow.operators.dummy import DummyOperator from airflow.operators.trigger_dagrun import TriggerDagRunOperator -from airflow.sensors.external_task import ExternalTaskAsyncSensor +from airflow.sensors.external_task import ExternalTaskSensor from airflow.utils.timezone import datetime with DAG( @@ -32,10 +32,11 @@ start = DummyOperator(task_id="start") # [START howto_external_task_async_sensor] - external_task_sensor = ExternalTaskAsyncSensor( + external_task_sensor = ExternalTaskSensor( task_id="parent_task_sensor", external_task_id="child_task", external_dag_id="child_dag", + deferrable=True, ) # [END howto_external_task_async_sensor] From e801b7db4b6cc99f6924adc4563dfafb75fc6bd8 Mon Sep 17 00:00:00 2001 From: Beata Kossakowska Date: Thu, 1 Jun 2023 11:28:40 +0000 Subject: [PATCH 11/12] Fix static checks --- airflow/triggers/external_task.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/airflow/triggers/external_task.py b/airflow/triggers/external_task.py index 461dd5feca264..602799d9b79d7 100644 --- a/airflow/triggers/external_task.py +++ b/airflow/triggers/external_task.py @@ -93,13 +93,15 @@ async def run(self) -> typing.AsyncIterator[TriggerEvent]: try: delta = utcnow() - self.trigger_start_time if delta.total_seconds() < self._timeout_sec: - if await self.count_running_dags() == 0: + # mypy confuses typing here + if await self.count_running_dags() == 0: # type: ignore[call-arg] self.log.info("Waiting for DAG to start execution...") await asyncio.sleep(self.poll_interval) else: yield TriggerEvent({"status": "timeout"}) return - if await self.count_tasks() == len(self.execution_dates): + # mypy confuses typing here + if await self.count_tasks() == len(self.execution_dates): # type: ignore[call-arg] yield TriggerEvent({"status": "success"}) return self.log.info("Task is still running, sleeping for %s seconds...", self.poll_interval) From 9c314d1d7aa25a4548153c3f2393c99d98ad86e0 Mon Sep 17 00:00:00 2001 From: Beata Kossakowska Date: Wed, 12 Jul 2023 10:58:07 +0000 Subject: [PATCH 12/12] Rebase main and fix. --- airflow/sensors/external_task.py | 7 ++++--- airflow/triggers/external_task.py | 10 +++++----- .../core/example_external_task_parent_deferrable.py | 5 ++++- 3 files changed, 13 insertions(+), 9 deletions(-) diff --git a/airflow/sensors/external_task.py b/airflow/sensors/external_task.py index 3e282a0b3a37c..5e42820ffe997 100644 --- a/airflow/sensors/external_task.py +++ b/airflow/sensors/external_task.py @@ -25,6 +25,7 @@ import attr from sqlalchemy import func +from airflow.configuration import conf from airflow.exceptions import AirflowException, AirflowSkipException, RemovedInAirflow3Warning from airflow.models.baseoperator import BaseOperatorLink from airflow.models.dag import DagModel @@ -38,9 +39,8 @@ from airflow.utils.helpers import build_airflow_url_with_query from airflow.utils.session import NEW_SESSION, provide_session from airflow.utils.sqlalchemy import tuple_in_condition -from airflow.utils.state import State -from airflow.utils.timezone import utcnow from airflow.utils.state import State, TaskInstanceState +from airflow.utils.timezone import utcnow if TYPE_CHECKING: from sqlalchemy.orm import Query, Session @@ -151,10 +151,11 @@ def __init__( execution_date_fn: Callable | None = None, check_existence: bool = False, poll_interval: float = 2.0, - deferrable: bool = False, + deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False), **kwargs, ): super().__init__(**kwargs) + self.allowed_states = list(allowed_states) if allowed_states else [TaskInstanceState.SUCCESS.value] self.skipped_states = list(skipped_states) if skipped_states else [] self.failed_states = list(failed_states) if failed_states else [] diff --git a/airflow/triggers/external_task.py b/airflow/triggers/external_task.py index 602799d9b79d7..f179cba259ca4 100644 --- a/airflow/triggers/external_task.py +++ b/airflow/triggers/external_task.py @@ -17,7 +17,6 @@ from __future__ import annotations import asyncio -import datetime import typing from datetime import datetime @@ -28,7 +27,7 @@ from airflow.models import DagRun, TaskInstance from airflow.triggers.base import BaseTrigger, TriggerEvent from airflow.utils.session import NEW_SESSION, provide_session -from airflow.utils.state import DagRunState +from airflow.utils.state import DagRunState, TaskInstanceState from airflow.utils.timezone import utcnow @@ -38,7 +37,7 @@ class TaskStateTrigger(BaseTrigger): :param dag_id: The dag_id that contains the task you want to wait for :param task_id: The task_id that contains the task you want to - wait for. If ``None`` (default value) the sensor waits for the DAG + wait for. :param states: allowed states, default is ``['success']`` :param execution_dates: task execution time interval :param poll_interval: The time interval in seconds to check the state. @@ -53,9 +52,9 @@ class TaskStateTrigger(BaseTrigger): def __init__( self, dag_id: str, - states: list[str], execution_dates: list[datetime], trigger_start_time: datetime, + states: list[str] | None = None, task_id: str | None = None, poll_interval: float = 2.0, ): @@ -66,6 +65,7 @@ def __init__( self.execution_dates = execution_dates self.poll_interval = poll_interval self.trigger_start_time = trigger_start_time + self.states = states if states else [TaskInstanceState.SUCCESS.value] self._timeout_sec = 60 def serialize(self) -> tuple[str, dict[str, typing.Any]]: @@ -157,7 +157,7 @@ def __init__( self, dag_id: str, states: list[DagRunState], - execution_dates: list[datetime.datetime], + execution_dates: list[datetime], poll_interval: float = 5.0, ): super().__init__() diff --git a/tests/system/providers/core/example_external_task_parent_deferrable.py b/tests/system/providers/core/example_external_task_parent_deferrable.py index 7ae405674828e..7cec2ce13815a 100644 --- a/tests/system/providers/core/example_external_task_parent_deferrable.py +++ b/tests/system/providers/core/example_external_task_parent_deferrable.py @@ -43,7 +43,10 @@ trigger_child_task = TriggerDagRunOperator( task_id="trigger_child_task", trigger_dag_id="child_dag", - allowed_states=["success", "failed", "skipped"], + allowed_states=[ + "success", + "failed", + ], execution_date="{{execution_date}}", poke_interval=5, reset_dag_run=True,