diff --git a/providers/standard/src/airflow/providers/standard/sensors/date_time.py b/providers/standard/src/airflow/providers/standard/sensors/date_time.py index cc3485a906d7d..4455c49db9335 100644 --- a/providers/standard/src/airflow/providers/standard/sensors/date_time.py +++ b/providers/standard/src/airflow/providers/standard/sensors/date_time.py @@ -19,12 +19,13 @@ import dataclasses import datetime +import warnings from collections.abc import Sequence from typing import TYPE_CHECKING, Any, NoReturn from airflow.providers.common.compat.sdk import BaseSensorOperator, timezone from airflow.providers.standard.triggers.temporal import DateTimeTrigger -from airflow.providers.standard.version_compat import AIRFLOW_V_3_0_PLUS +from airflow.providers.standard.version_compat import AIRFLOW_V_3_0_PLUS, AIRFLOW_V_3_3_PLUS from airflow.triggers.base import StartTriggerArgs if TYPE_CHECKING: @@ -64,6 +65,7 @@ class DateTimeSensor(BaseSensorOperator): def __init__(self, *, target_time: str | datetime.datetime, **kwargs) -> None: super().__init__(**kwargs) + self.target_time = target_time def poke(self, context: Context) -> bool: @@ -89,8 +91,8 @@ class DateTimeSensorAsync(DateTimeSensor): :param target_time: datetime after which the job succeeds. (templated) :param start_from_trigger: Start the task directly from the triggerer without going into the worker. - Ignored when ``target_time`` is a Jinja template: templates are only rendered on the worker, so the - task defers from the worker instead. + It's ignored on Airflow < 3.3 when ``target_time`` is a Jinja template + :param trigger_kwargs: The keyword arguments passed to the trigger when start_from_trigger is set to True during dynamic task mapping. This argument is not used in standard usage. :param end_from_trigger: End the task directly from the triggerer without going into the worker. @@ -115,28 +117,35 @@ def __init__( ) -> None: super().__init__(**kwargs) self.end_from_trigger = end_from_trigger - - # A templated target is rendered after Dag parsing, so it cannot be used to - # construct the trigger arguments at task initialization time. - if ( - start_from_trigger - and isinstance(self.target_time, str) - and any(delimiter in self.target_time for delimiter in ("{{", "{%", "{#")) - ): - start_from_trigger = False - self.start_from_trigger = start_from_trigger + if self.start_from_trigger: - # Replaced rather than mutated: ``start_trigger_args`` is a class attribute, so - # assigning through it would overwrite the arguments of every other task built - # from this operator. - self.start_trigger_args = dataclasses.replace( - self.start_trigger_args, - trigger_kwargs=dict( - moment=self._moment, - end_from_trigger=self.end_from_trigger, - ), - ) + if AIRFLOW_V_3_3_PLUS: + self.start_trigger_args = dataclasses.replace( + self.start_trigger_args, + trigger_kwargs={ + "target_time": self.target_time, + "end_from_trigger": self.end_from_trigger, + }, + ) + elif isinstance(self.target_time, str) and any( + delimiter in self.target_time for delimiter in ("{{", "{%", "{#") + ): + warnings.warn( + f"start_from_trigger=True with a templated target_time ({self.target_time!r}) requires " + "Airflow >= 3.3. Disabling start_from_trigger, the task will defer from the worker instead.", + UserWarning, + stacklevel=2, + ) + self.start_from_trigger = False + else: + self.start_trigger_args = dataclasses.replace( + self.start_trigger_args, + trigger_kwargs={ + "moment": self._moment, + "end_from_trigger": self.end_from_trigger, + }, + ) def execute(self, context: Context) -> NoReturn: self.defer( diff --git a/providers/standard/src/airflow/providers/standard/triggers/temporal.py b/providers/standard/src/airflow/providers/standard/triggers/temporal.py index cfc206ff3131d..f1d01b5401f79 100644 --- a/providers/standard/src/airflow/providers/standard/triggers/temporal.py +++ b/providers/standard/src/airflow/providers/standard/triggers/temporal.py @@ -31,31 +31,57 @@ class DateTimeTrigger(BaseTrigger): """ Trigger based on a datetime. - A trigger that fires exactly once, at the given datetime, give or take - a few seconds. - - The provided datetime MUST be in UTC. + Pass either ``moment`` (a tz-aware datetime) or ``target_time`` (a string, possibly a Jinja + template, or a datetime). With ``start_from_trigger``, the operator lists ``target_time`` in its + own ``template_fields`` and puts it in ``start_trigger_args.trigger_kwargs``, so the triggerer + renders it in place before ``run()``; it is then parsed into ``moment`` on first use. :param moment: when to yield event + :param target_time: raw (possibly templated) datetime string, an alternative to ``moment`` :param end_from_trigger: whether the trigger should mark the task successful after time condition reached or resume the task after time condition reached. """ - def __init__(self, moment: datetime.datetime, *, end_from_trigger: bool = False) -> None: + def __init__( + self, + moment: datetime.datetime | None = None, + *, + target_time: datetime.datetime | str | None = None, + end_from_trigger: bool = False, + ) -> None: super().__init__() - if not isinstance(moment, datetime.datetime): - raise TypeError(f"Expected datetime.datetime type for moment. Got {type(moment)}") - # Make sure it's in UTC - if moment.tzinfo is None: - raise ValueError("You cannot pass naive datetimes") - self.moment: pendulum.DateTime = timezone.convert_to_utc(moment) + if (moment is None) == (target_time is None): + raise TypeError("DateTimeTrigger requires exactly one of 'moment' or 'target_time'") + self.target_time = target_time + self._moment: pendulum.DateTime | None = None + if moment is not None: + if not isinstance(moment, datetime.datetime): + raise TypeError(f"Expected datetime.datetime type for moment. Got {type(moment)}") + # Make sure it's in UTC + if moment.tzinfo is None: + raise ValueError("You cannot pass naive datetimes") + self._moment = timezone.convert_to_utc(moment) self.end_from_trigger = end_from_trigger + @property + def moment(self) -> pendulum.DateTime: + if self._moment is None: + # Resolved lazily: by now the triggerer has rendered target_time in place. + target_time: Any = self.target_time + if isinstance(target_time, datetime.datetime): + target_time = target_time.isoformat() + if not isinstance(target_time, str) or not target_time: + raise TypeError("DateTimeTrigger has neither a 'moment' nor a usable 'target_time'") + self._moment = timezone.convert_to_utc(timezone.parse(target_time)) + return self._moment + def serialize(self) -> tuple[str, dict[str, Any]]: - return ( - "airflow.providers.standard.triggers.temporal.DateTimeTrigger", - {"moment": self.moment, "end_from_trigger": self.end_from_trigger}, - ) + if self._moment is None: + kwargs: dict[str, Any] = {"target_time": self.target_time} + else: + kwargs = {"moment": self._moment} + kwargs["end_from_trigger"] = self.end_from_trigger + return ("airflow.providers.standard.triggers.temporal.DateTimeTrigger", kwargs) async def run(self) -> AsyncIterator[TriggerEvent]: """ diff --git a/providers/standard/tests/unit/standard/sensors/test_date_time.py b/providers/standard/tests/unit/standard/sensors/test_date_time.py index 12818e27ea7c4..c860b524e184a 100644 --- a/providers/standard/tests/unit/standard/sensors/test_date_time.py +++ b/providers/standard/tests/unit/standard/sensors/test_date_time.py @@ -17,7 +17,9 @@ # under the License. from __future__ import annotations +import asyncio import datetime +from types import SimpleNamespace from unittest.mock import patch import pendulum @@ -26,8 +28,9 @@ from airflow import macros from airflow.models.dag import DAG from airflow.providers.standard.sensors.date_time import DateTimeSensor, DateTimeSensorAsync +from airflow.providers.standard.triggers.temporal import DateTimeTrigger -from tests_common.test_utils.version_compat import timezone +from tests_common.test_utils.version_compat import AIRFLOW_V_3_3_PLUS, timezone DEFAULT_DATE = timezone.datetime(2015, 1, 1) @@ -134,19 +137,53 @@ def test_poke_with_natively_rendered_datetime(self, mock_utcnow): assert isinstance(op.target_time, datetime.datetime) assert op.poke(None) is True + @pytest.mark.asyncio + @pytest.mark.skipif(not AIRFLOW_V_3_3_PLUS, reason="Test only for AF >= 3.3") + async def test_full_run_triggerer_path_templated_past_target_time(self): + """Simulates the triggerer path: start_trigger_args carries the raw target_time, which + the triggerer renders in place (never going through the worker's execute()).""" + dag = DAG( + dag_id="full_run_triggerer_path_dag", + start_date=pendulum.datetime(2020, 1, 1, tz="UTC"), + schedule=None, + ) + op = DateTimeSensorAsync( + task_id="full_run_triggerer_path", + target_time="{{ data_interval_end }}", + start_from_trigger=True, + dag=dag, + ) + ctx = {"data_interval_end": pendulum.datetime(2020, 1, 1, tz="UTC")} + + trigger = DateTimeTrigger(**op.start_trigger_args.trigger_kwargs) + trigger.task_instance = SimpleNamespace(task_id=op.task_id, task=op) + trigger.render_template_fields(ctx) + + assert trigger.target_time == "2020-01-01 00:00:00+00:00" + + # Triggerer: run the trigger the task deferred to. The past moment fires immediately. + event = await asyncio.wait_for(trigger.run().__anext__(), timeout=5) + + assert op.execute_complete(context=ctx, event=event.payload) is None + def test_moment_localizes_naive_datetime(self): """A naive datetime target_time is localized to UTC via _moment (mirrors old isoformat()).""" op = DateTimeSensor(task_id="naive", target_time=datetime.datetime(2020, 1, 1), dag=self.dag) assert op._moment == pendulum.datetime(2020, 1, 1, tz="UTC") - def test_async_start_from_trigger_moment(self): + def test_async_start_from_trigger_sets_trigger_kwargs(self): op = DateTimeSensorAsync( task_id="async", target_time="2020-01-01T00:00:00+00:00", start_from_trigger=True, dag=self.dag, ) - assert op.start_trigger_args.trigger_kwargs["moment"] == pendulum.parse("2020-01-01T00:00:00+00:00") + if AIRFLOW_V_3_3_PLUS: + assert op.start_trigger_args.trigger_kwargs["target_time"] == "2020-01-01T00:00:00+00:00" + else: + assert op.start_trigger_args.trigger_kwargs["moment"] == pendulum.parse( + "2020-01-01T00:00:00+00:00" + ) def test_async_start_from_trigger_localizes_naive_datetime(self): """DateTimeSensorAsync never pokes, so _moment must still localize a naive datetime.""" @@ -156,9 +193,28 @@ def test_async_start_from_trigger_localizes_naive_datetime(self): start_from_trigger=True, dag=self.dag, ) - assert op.start_trigger_args.trigger_kwargs["moment"] == pendulum.datetime(2020, 1, 1, tz="UTC") + if AIRFLOW_V_3_3_PLUS: + assert op.start_trigger_args.trigger_kwargs["target_time"] == datetime.datetime(2020, 1, 1, 0, 0) + else: + assert op.start_trigger_args.trigger_kwargs["moment"] == pendulum.datetime(2020, 1, 1, tz="UTC") + + @pytest.mark.skipif(AIRFLOW_V_3_3_PLUS, reason="Test only for AF < 3.3") def test_async_start_from_trigger_skips_templated_target_time(self): + with pytest.warns(UserWarning, match="requires Airflow >= 3.3"): + op = DateTimeSensorAsync( + task_id="async_templated", + target_time="{{ data_interval_end }}", + start_from_trigger=True, + dag=self.dag, + ) + + assert op.start_from_trigger is False + # The class attribute must be left alone, not replaced with a copy built from the raw template. + assert op.start_trigger_args is DateTimeSensorAsync.start_trigger_args + + @pytest.mark.skipif(not AIRFLOW_V_3_3_PLUS, reason="Test only for AF >= 3.3") + def test_async_start_from_trigger_templated_target_time(self): op = DateTimeSensorAsync( task_id="async_templated", target_time="{{ data_interval_end }}", @@ -166,9 +222,8 @@ def test_async_start_from_trigger_skips_templated_target_time(self): dag=self.dag, ) - assert op.start_from_trigger is False - # The class attribute must be left alone, not replaced with a copy built from the raw template. - assert op.start_trigger_args is DateTimeSensorAsync.start_trigger_args + assert op.start_from_trigger is True + assert op.start_trigger_args.trigger_kwargs["target_time"] == "{{ data_interval_end }}" def test_start_trigger_args_are_not_shared_between_tasks(self): """Each task must carry its own trigger arguments. @@ -190,12 +245,17 @@ def test_start_trigger_args_are_not_shared_between_tasks(self): ) assert first.start_trigger_args is not second.start_trigger_args - assert first.start_trigger_args.trigger_kwargs["moment"] == pendulum.parse( - "2030-01-01T00:00:00+00:00" - ) - assert second.start_trigger_args.trigger_kwargs["moment"] == pendulum.parse( - "2040-06-06T00:00:00+00:00" - ) + if AIRFLOW_V_3_3_PLUS: + assert first.start_trigger_args.trigger_kwargs["target_time"] == "2030-01-01T00:00:00+00:00" + assert second.start_trigger_args.trigger_kwargs["target_time"] == "2040-06-06T00:00:00+00:00" + else: + assert first.start_trigger_args.trigger_kwargs["moment"] == pendulum.parse( + "2030-01-01T00:00:00+00:00" + ) + assert second.start_trigger_args.trigger_kwargs["moment"] == pendulum.parse( + "2040-06-06T00:00:00+00:00" + ) + # the class level template must survive untouched for the next task built from it assert DateTimeSensorAsync.start_trigger_args.trigger_kwargs == { "moment": "",