DateTimeSensorAsync support start_from_trigger and templated target_time

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(

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]:
         """

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,19 +193,37 @@ 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 }}",
             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
+        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": "",