diff --git a/airflow-core/src/airflow/jobs/triggerer_job_runner.py b/airflow-core/src/airflow/jobs/triggerer_job_runner.py index 520699e9e5093..be7fe23004b1c 100644 --- a/airflow-core/src/airflow/jobs/triggerer_job_runner.py +++ b/airflow-core/src/airflow/jobs/triggerer_job_runner.py @@ -50,8 +50,10 @@ GetConnection, GetDagRunState, GetDRCount, + GetTICount, GetVariable, GetXCom, + TICount, VariableResult, XComResult, ) @@ -222,6 +224,7 @@ class TriggerStateSync(BaseModel): XComResult, DagRunStateResult, DRCount, + TICount, ErrorResponse, ], Field(discriminator="type"), @@ -233,7 +236,15 @@ class TriggerStateSync(BaseModel): ToTriggerSupervisor = Annotated[ - Union[messages.TriggerStateChanges, GetConnection, GetVariable, GetXCom, GetDagRunState, GetDRCount], + Union[ + messages.TriggerStateChanges, + GetConnection, + GetVariable, + GetXCom, + GetTICount, + GetDagRunState, + GetDRCount, + ], Field(discriminator="type"), ] """ @@ -411,6 +422,16 @@ def _handle_request(self, msg: ToTriggerSupervisor, log: FilteringBoundLogger) - 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() + elif isinstance(msg, GetTICount): + ti_count = self.client.task_instances.get_count( + dag_id=msg.dag_id, + task_ids=msg.task_ids, + task_group_id=msg.task_group_id, + logical_dates=msg.logical_dates, + run_ids=msg.run_ids, + states=msg.states, + ) + resp = ti_count.model_dump_json().encode() else: raise ValueError(f"Unknown message type {type(msg)}") diff --git a/airflow-core/tests/unit/jobs/test_triggerer_job.py b/airflow-core/tests/unit/jobs/test_triggerer_job.py index af197d6f35bfa..50320681df1de 100644 --- a/airflow-core/tests/unit/jobs/test_triggerer_job.py +++ b/airflow-core/tests/unit/jobs/test_triggerer_job.py @@ -767,3 +767,103 @@ async def test_trigger_can_fetch_trigger_dag_run_count_in_deferrable(session, da assert task_instance.state == TaskInstanceState.SCHEDULED assert task_instance.next_method != "__fail__" assert task_instance.next_kwargs == {"event": {"count": 1}} + + +class CustomTriggerWorkflowStateTrigger(BaseTrigger): + """Custom Trigger to check the triggerer can access the get_ti_count and get_dr_count.""" + + def __init__(self, external_dag_id, execution_dates, external_task_ids, allowed_states, run_ids): + self.external_dag_id = external_dag_id + self.execution_dates = execution_dates + self.external_task_ids = external_task_ids + self.allowed_states = allowed_states + self.run_ids = run_ids + + def serialize(self) -> tuple[str, dict[str, Any]]: + return ( + f"{type(self).__module__}.{type(self).__qualname__}", + { + "external_dag_id": self.external_dag_id, + "execution_dates": self.execution_dates, + "external_task_ids": self.external_task_ids, + "allowed_states": self.allowed_states, + "run_ids": self.run_ids, + }, + ) + + async def run(self, **args) -> AsyncIterator[TriggerEvent]: + from airflow.sdk.execution_time.task_runner import RuntimeTaskInstance + + ti_count = await sync_to_async(RuntimeTaskInstance.get_ti_count)( + dag_id=self.external_dag_id, + task_ids=self.external_task_ids, + task_group_id=None, + run_ids=self.run_ids, + logical_dates=self.execution_dates, + states=self.allowed_states, + ) + dr_count = await sync_to_async(RuntimeTaskInstance.get_dr_count)( + dag_id=self.external_dag_id, + run_ids=self.run_ids, + logical_dates=self.execution_dates, + states=["running"], + ) + yield TriggerEvent({"ti_count": ti_count, "dr_count": dr_count}) + + +@pytest.mark.xfail( + reason="We know that test is flaky and have no time to fix it before 3.0. " + "We should fix it later. TODO: AIP-72" +) +@pytest.mark.asyncio +@pytest.mark.flaky(reruns=2, reruns_delay=10) +@pytest.mark.execution_timeout(30) +async def test_trigger_can_fetch_dag_run_count_ti_count_in_deferrable(session, dag_maker): + """Checks that the trigger will successfully fetch the count of trigger DAG runs.""" + # Create the test DAG and task + with dag_maker(dag_id="parent_dag", session=session): + EmptyOperator(task_id="parent_task") + parent_dag_run = dag_maker.create_dagrun() + parent_task = parent_dag_run.task_instances[0] + parent_task.state = TaskInstanceState.SUCCESS + + with dag_maker(dag_id="trigger_can_fetch_dag_run_count_ti_count_in_deferrable", session=session): + EmptyOperator(task_id="dummy1") + dr = dag_maker.create_dagrun() + task_instance = dr.task_instances[0] + task_instance.state = TaskInstanceState.DEFERRED + + # Use the same dag run with states deferred to fetch the count + trigger = CustomTriggerWorkflowStateTrigger( + external_dag_id=parent_task.dag_id, + execution_dates=[parent_task.logical_date], + external_task_ids=[parent_task.task_id], + allowed_states=[State.SUCCESS], + run_ids=[parent_task.run_id], + ) + trigger_orm = Trigger( + classpath=trigger.serialize()[0], + kwargs={ + "external_dag_id": parent_dag_run.dag_id, + "execution_dates": [parent_dag_run.logical_date], + "external_task_ids": [parent_task.task_id], + "allowed_states": [State.SUCCESS], + "run_ids": [parent_dag_run.run_id], + }, + ) + session.add(trigger_orm) + session.commit() + task_instance.trigger_id = trigger_orm.id + + job = Job() + session.add(job) + session.commit() + + supervisor = DummyTriggerRunnerSupervisor.start(job=job, capacity=1, logger=None) + supervisor.run() + + parent_task.refresh_from_db() + task_instance.refresh_from_db() + assert task_instance.state == TaskInstanceState.SCHEDULED + assert task_instance.next_method != "__fail__" + assert task_instance.next_kwargs == {"event": {"ti_count": 1, "dr_count": 1}} diff --git a/providers/standard/src/airflow/providers/standard/sensors/external_task.py b/providers/standard/src/airflow/providers/standard/sensors/external_task.py index dd5ad6c4e7a3a..e64eb0d6763b0 100644 --- a/providers/standard/src/airflow/providers/standard/sensors/external_task.py +++ b/providers/standard/src/airflow/providers/standard/sensors/external_task.py @@ -346,8 +346,6 @@ def _calculate_count(self, count: int, dttm_filter: list[datetime.datetime]) -> """Calculate the normalized count based on the type of check.""" if self.external_task_ids: return count / len(self.external_task_ids) - elif self.external_task_group_id: - return count / len(dttm_filter) else: return count @@ -421,16 +419,22 @@ def execute(self, context: Context) -> None: if not self.deferrable: super().execute(context) else: + dttm_filter = self._get_dttm_filter(context) + logical_or_execution_dates = ( + {"logical_dates": dttm_filter} if AIRFLOW_V_3_0_PLUS else {"execution_date": dttm_filter} + ) self.defer( timeout=self.execution_timeout, trigger=WorkflowTrigger( external_dag_id=self.external_dag_id, external_task_group_id=self.external_task_group_id, external_task_ids=self.external_task_ids, - logical_dates=self._get_dttm_filter(context), allowed_states=self.allowed_states, + failed_states=self.failed_states, + skipped_states=self.skipped_states, poke_interval=self.poll_interval, soft_fail=self.soft_fail, + **logical_or_execution_dates, ), method_name="execute_complete", ) diff --git a/providers/standard/src/airflow/providers/standard/triggers/external_task.py b/providers/standard/src/airflow/providers/standard/triggers/external_task.py index adf3a535e3ef2..a5f8b67f54972 100644 --- a/providers/standard/src/airflow/providers/standard/triggers/external_task.py +++ b/providers/standard/src/airflow/providers/standard/triggers/external_task.py @@ -50,6 +50,7 @@ class WorkflowTrigger(BaseTrigger): :param allowed_states: States considered as successful for external tasks. :param poke_interval: The interval (in seconds) for poking the external tasks. :param soft_fail: If True, the trigger will not fail the entire dag on external task failure. + :param logical_dates: A list of logical dates for the external dag. """ def __init__( @@ -57,6 +58,7 @@ def __init__( external_dag_id: str, run_ids: list[str] | None = None, execution_dates: list[datetime] | None = None, + logical_dates: list[datetime] | None = None, external_task_ids: typing.Collection[str] | None = None, external_task_group_id: str | None = None, failed_states: typing.Iterable[str] | None = None, @@ -76,6 +78,7 @@ def __init__( self.poke_interval = poke_interval self.soft_fail = soft_fail self.execution_dates = execution_dates + self.logical_dates = logical_dates super().__init__(**kwargs) def serialize(self) -> tuple[str, dict[str, Any]]: @@ -92,6 +95,7 @@ def serialize(self) -> tuple[str, dict[str, Any]]: } if AIRFLOW_V_3_0_PLUS: data["run_ids"] = self.run_ids + data["logical_dates"] = self.logical_dates else: data["execution_dates"] = self.execution_dates @@ -99,9 +103,16 @@ def serialize(self) -> tuple[str, dict[str, Any]]: async def run(self) -> typing.AsyncIterator[TriggerEvent]: """Check periodically tasks, task group or dag status.""" + if AIRFLOW_V_3_0_PLUS: + get_count_func = self._get_count_af_3 + run_id_or_dates = (self.run_ids or self.logical_dates) or [] + else: + get_count_func = self._get_count + run_id_or_dates = self.execution_dates or [] + while True: if self.failed_states: - failed_count = await self._get_count(self.failed_states) + failed_count = await get_count_func(self.failed_states) if failed_count > 0: yield TriggerEvent({"status": "failed"}) return @@ -109,18 +120,43 @@ async def run(self) -> typing.AsyncIterator[TriggerEvent]: yield TriggerEvent({"status": "success"}) return if self.skipped_states: - skipped_count = await self._get_count(self.skipped_states) + skipped_count = await get_count_func(self.skipped_states) if skipped_count > 0: yield TriggerEvent({"status": "skipped"}) return - allowed_count = await self._get_count(self.allowed_states) - _dates = self.run_ids if AIRFLOW_V_3_0_PLUS else self.execution_dates - if allowed_count == len(_dates): # type: ignore[arg-type] + allowed_count = await get_count_func(self.allowed_states) + + if allowed_count == len(run_id_or_dates): # type: ignore[arg-type] yield TriggerEvent({"status": "success"}) return self.log.info("Sleeping for %s seconds", self.poke_interval) await asyncio.sleep(self.poke_interval) + async def _get_count_af_3(self, states): + from airflow.sdk.execution_time.task_runner import RuntimeTaskInstance + + if self.external_task_ids or self.external_task_group_id: + count = await sync_to_async(RuntimeTaskInstance.get_ti_count)( + dag_id=self.external_dag_id, + task_ids=self.external_task_ids, + task_group_id=self.external_task_group_id, + logical_dates=self.logical_dates, + run_ids=self.run_ids, + states=states, + ) + else: + count = await sync_to_async(RuntimeTaskInstance.get_dr_count)( + dag_id=self.external_dag_id, + logical_dates=self.logical_dates, + run_ids=self.run_ids, + states=states, + ) + + if self.external_task_ids: + return count / len(self.external_task_ids) + else: + return count + @sync_to_async def _get_count(self, states: typing.Iterable[str] | None) -> int: """ diff --git a/providers/standard/tests/unit/standard/triggers/test_external_task.py b/providers/standard/tests/unit/standard/triggers/test_external_task.py index 247085a407ff4..3aa581257d79e 100644 --- a/providers/standard/tests/unit/standard/triggers/test_external_task.py +++ b/providers/standard/tests/unit/standard/triggers/test_external_task.py @@ -39,11 +39,315 @@ key, value = next(iter(_DATES.items())) +@pytest.mark.skipif(not AIRFLOW_V_3_0_PLUS, reason="Test only for Airflow 3") class TestWorkflowTrigger: DAG_ID = "external_task" TASK_ID = "external_task_op" RUN_ID = "external_task_run_id" STATES = ["success", "fail"] + LOGICAL_DATE = timezone.datetime(2022, 1, 1) + + @pytest.mark.flaky(reruns=5) + @mock.patch("airflow.sdk.execution_time.task_runner.RuntimeTaskInstance.get_ti_count") + @pytest.mark.asyncio + async def test_task_workflow_trigger_success(self, mock_get_count): + """check the db count get called correctly.""" + mock_get_count.side_effect = mocked_get_count + + trigger = WorkflowTrigger( + external_dag_id=self.DAG_ID, + logical_dates=[self.LOGICAL_DATE], + external_task_ids=[self.TASK_ID], + allowed_states=self.STATES, + poke_interval=0.2, + ) + + gen = trigger.run() + trigger_task = asyncio.create_task(gen.__anext__()) + fake_task = asyncio.create_task(fake_async_fun()) + await trigger_task + await fake_task + assert fake_task.done() # confirm that get_count is done in an async fashion + assert trigger_task.done() + result = trigger_task.result() + assert result.payload == {"status": "success"} + + mock_get_count.assert_called_once_with( + dag_id="external_task", + task_ids=["external_task_op"], + task_group_id=None, + logical_dates=[self.LOGICAL_DATE], + run_ids=None, + states=["success", "fail"], + ) + # test that it returns after yielding + with pytest.raises(StopAsyncIteration): + await gen.__anext__() + + @pytest.mark.flaky(reruns=5) + @mock.patch("airflow.sdk.execution_time.task_runner.RuntimeTaskInstance.get_ti_count") + @pytest.mark.asyncio + async def test_task_workflow_trigger_failed(self, mock_get_count): + mock_get_count.side_effect = mocked_get_count + + trigger = WorkflowTrigger( + external_dag_id=self.DAG_ID, + logical_dates=[self.LOGICAL_DATE], + run_ids=[self.RUN_ID], + external_task_ids=[self.TASK_ID], + failed_states=self.STATES, + poke_interval=0.2, + ) + + gen = trigger.run() + trigger_task = asyncio.create_task(gen.__anext__()) + fake_task = asyncio.create_task(fake_async_fun()) + await trigger_task + await fake_task + assert fake_task.done() # confirm that get_count is done in an async fashion + assert trigger_task.done() + result = trigger_task.result() + assert isinstance(result, TriggerEvent) + assert result.payload == {"status": "failed"} + mock_get_count.assert_called_once_with( + dag_id="external_task", + task_ids=["external_task_op"], + task_group_id=None, + logical_dates=[self.LOGICAL_DATE], + run_ids=[self.RUN_ID], + states=["success", "fail"], + ) + # test that it returns after yielding + with pytest.raises(StopAsyncIteration): + await gen.__anext__() + + @pytest.mark.asyncio + @mock.patch("airflow.sdk.execution_time.task_runner.RuntimeTaskInstance.get_ti_count") + async def test_task_workflow_trigger_fail_count_eq_0(self, mock_get_count): + mock_get_count.return_value = 0 + + trigger = WorkflowTrigger( + external_dag_id=self.DAG_ID, + logical_dates=[self.LOGICAL_DATE], + run_ids=[self.RUN_ID], + external_task_ids=[self.TASK_ID], + failed_states=self.STATES, + poke_interval=0.2, + ) + + gen = trigger.run() + trigger_task = asyncio.create_task(gen.__anext__()) + await trigger_task + assert trigger_task.done() + result = trigger_task.result() + assert isinstance(result, TriggerEvent) + assert result.payload == {"status": "success"} + mock_get_count.assert_called_once_with( + dag_id="external_task", + task_ids=["external_task_op"], + task_group_id=None, + logical_dates=[self.LOGICAL_DATE], + run_ids=[self.RUN_ID], + states=["success", "fail"], + ) + # test that it returns after yielding + with pytest.raises(StopAsyncIteration): + await gen.__anext__() + + @pytest.mark.flaky(reruns=5) + @mock.patch("airflow.sdk.execution_time.task_runner.RuntimeTaskInstance.get_ti_count") + @pytest.mark.asyncio + async def test_task_workflow_trigger_skipped(self, mock_get_count): + mock_get_count.side_effect = mocked_get_count + + trigger = WorkflowTrigger( + external_dag_id=self.DAG_ID, + logical_dates=[self.LOGICAL_DATE], + external_task_ids=[self.TASK_ID], + skipped_states=self.STATES, + poke_interval=0.2, + ) + + gen = trigger.run() + trigger_task = asyncio.create_task(gen.__anext__()) + fake_task = asyncio.create_task(fake_async_fun()) + await trigger_task + await fake_task + assert fake_task.done() # confirm that get_count is done in an async fashion + assert trigger_task.done() + result = trigger_task.result() + assert isinstance(result, TriggerEvent) + assert result.payload == {"status": "skipped"} + mock_get_count.assert_called_once_with( + dag_id="external_task", + task_ids=["external_task_op"], + task_group_id=None, + logical_dates=[self.LOGICAL_DATE], + run_ids=None, + states=["success", "fail"], + ) + + @mock.patch("airflow.sdk.execution_time.task_runner.RuntimeTaskInstance.get_ti_count") + @mock.patch("asyncio.sleep") + @pytest.mark.asyncio + async def test_task_workflow_trigger_sleep_success(self, mock_sleep, mock_get_count): + mock_get_count.side_effect = [0, 1] + + trigger = WorkflowTrigger( + external_dag_id=self.DAG_ID, + logical_dates=[self.LOGICAL_DATE], + external_task_ids=[self.TASK_ID], + poke_interval=0.2, + ) + + gen = trigger.run() + trigger_task = asyncio.create_task(gen.__anext__()) + await trigger_task + assert trigger_task.done() + result = trigger_task.result() + assert isinstance(result, TriggerEvent) + assert result.payload == {"status": "success"} + mock_get_count.assert_called() + assert mock_get_count.call_count == 2 + + # test that it returns after yielding + with pytest.raises(StopAsyncIteration): + await gen.__anext__() + + mock_sleep.assert_awaited() + assert mock_sleep.await_count == 1 + + @pytest.mark.parametrize( + "task_ids, task_group_id, states, logical_dates, mock_ti_count, mock_dag_count, expected", + [ + ( + ["task_id_one", "task_id_two"], + None, + ["success"], + [ + timezone.datetime(2020, 7, 6, 13, tzinfo=timezone.utc), + timezone.datetime(2020, 7, 6, 13, tzinfo=timezone.utc), + ], + 4, + 2, + 2, + ), + ( + [], + "task_group_id", + ["success"], + [ + timezone.datetime(2020, 7, 6, 13, tzinfo=timezone.utc), + timezone.datetime(2020, 7, 6, 13, tzinfo=timezone.utc), + ], + 2, + 2, + 2, + ), + ( + [], + None, + ["success"], + [ + timezone.datetime(2020, 7, 6, 13, tzinfo=timezone.utc), + timezone.datetime(2020, 7, 6, 13, tzinfo=timezone.utc), + ], + 2, + 2, + 2, + ), + ], + ids=[ + "with_task_ids", + "with task_group_id only", + "no task_ids or task_group_id", + ], + ) + @mock.patch("airflow.sdk.execution_time.task_runner.RuntimeTaskInstance.get_ti_count") + @mock.patch("airflow.sdk.execution_time.task_runner.RuntimeTaskInstance.get_dr_count") + @pytest.mark.asyncio + async def test_get_count_af_3( + self, + mock_get_dr_count, + mock_get_ti_count, + task_ids, + task_group_id, + states, + logical_dates, + mock_ti_count, + mock_dag_count, + expected, + ): + """ + case1: when provided two task_ids, and two dag runs, the get_ti_count should return 4(each dag run returns two tasks) + and normalized count becomes 2 + case2: when provided task_group_id, and two dag runs, the get_ti_count should return 2(each dag run returns 1 task group) + case3: when not provided any task_ids or task_group_id, the get_dr_count should return 2(total dag runs 2) + """ + + mock_get_ti_count.return_value = mock_ti_count + mock_get_dr_count.return_value = mock_dag_count + + trigger = WorkflowTrigger( + external_dag_id=self.DAG_ID, + logical_dates=logical_dates, + external_task_ids=task_ids, + external_task_group_id=task_group_id, + allowed_states=states, + poke_interval=0.2, + ) + + get_count_af_3 = await trigger._get_count_af_3(states) + assert get_count_af_3 == expected + + if task_ids or task_group_id: + mock_get_ti_count.assert_called_once() + assert mock_get_ti_count.call_count == 1 + mock_get_dr_count.assert_not_called() + assert mock_get_dr_count.call_count == 0 + + if not task_ids and not task_group_id: + mock_get_dr_count.assert_called_once() + assert mock_get_dr_count.call_count == 1 + mock_get_ti_count.assert_not_called() + assert mock_get_ti_count.call_count == 0 + + def test_serialization(self): + """ + Asserts that the WorkflowTrigger correctly serializes its arguments and classpath. + """ + trigger = WorkflowTrigger( + external_dag_id=self.DAG_ID, + logical_dates=[self.LOGICAL_DATE], + run_ids=[self.RUN_ID], + failed_states=["failed"], + skipped_states=["skipped"], + external_task_ids=[self.TASK_ID], + allowed_states=self.STATES, + poke_interval=5, + ) + classpath, kwargs = trigger.serialize() + assert classpath == "airflow.providers.standard.triggers.external_task.WorkflowTrigger" + assert kwargs == { + "external_dag_id": self.DAG_ID, + "logical_dates": [self.LOGICAL_DATE], + "run_ids": [self.RUN_ID], + "external_task_ids": [self.TASK_ID], + "external_task_group_id": None, + "failed_states": ["failed"], + "skipped_states": ["skipped"], + "allowed_states": self.STATES, + "poke_interval": 5, + "soft_fail": False, + } + + +@pytest.mark.skipif(AIRFLOW_V_3_0_PLUS, reason="Test only for Airflow 2") +class TestWorkflowTriggerAF2: + DAG_ID = "external_task" + TASK_ID = "external_task_op" + RUN_ID = "external_task_run_id" + STATES = ["success", "fail"] @pytest.mark.flaky(reruns=5) @mock.patch("airflow.providers.standard.triggers.external_task._get_count") 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 9cecdc5547d44..2d5a3add593bd 100644 --- a/task-sdk/src/airflow/sdk/execution_time/task_runner.py +++ b/task-sdk/src/airflow/sdk/execution_time/task_runner.py @@ -413,18 +413,19 @@ def get_ti_count( """Return the number of task instances matching the given criteria.""" log = structlog.get_logger(logger_name="task") - SUPERVISOR_COMMS.send_request( - log=log, - msg=GetTICount( - dag_id=dag_id, - task_ids=task_ids, - task_group_id=task_group_id, - logical_dates=logical_dates, - run_ids=run_ids, - states=states, - ), - ) - response = SUPERVISOR_COMMS.get_message() + with SUPERVISOR_COMMS.lock: + SUPERVISOR_COMMS.send_request( + log=log, + msg=GetTICount( + dag_id=dag_id, + task_ids=task_ids, + task_group_id=task_group_id, + logical_dates=logical_dates, + run_ids=run_ids, + states=states, + ), + ) + response = SUPERVISOR_COMMS.get_message() if TYPE_CHECKING: assert isinstance(response, TICount) @@ -441,16 +442,17 @@ def get_dr_count( """Return the number of DAG runs matching the given criteria.""" log = structlog.get_logger(logger_name="task") - SUPERVISOR_COMMS.send_request( - log=log, - msg=GetDRCount( - dag_id=dag_id, - logical_dates=logical_dates, - run_ids=run_ids, - states=states, - ), - ) - response = SUPERVISOR_COMMS.get_message() + with SUPERVISOR_COMMS.lock: + SUPERVISOR_COMMS.send_request( + log=log, + msg=GetDRCount( + dag_id=dag_id, + logical_dates=logical_dates, + run_ids=run_ids, + states=states, + ), + ) + response = SUPERVISOR_COMMS.get_message() if TYPE_CHECKING: assert isinstance(response, DRCount)