dabla commented on code in PR #73928:
URL: https://github.com/apache/airflow/pull/73928#discussion_r4179003733
##########
providers/standard/src/airflow/providers/standard/sensors/date_time.py:
##########
@@ -115,39 +117,53 @@ 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(
- method_name="execute_complete",
- trigger=DateTimeTrigger(
+
+ if AIRFLOW_V_3_3_PLUS:
+ trigger = DateTimeTrigger(
+ target_time=self.target_time,
Review Comment:
**[blocker]** `execute()` branches on the wrong condition: it should check
`self.start_from_trigger`, not the Airflow version, and keep building the
trigger from `moment`.
There are two separate paths:
- **Triggerer path** (`start_from_trigger=True`): the scheduler defers the
task directly from `start_trigger_args`, and the triggerer renders the raw
`target_time` from `trigger_kwargs`. `execute()` is not involved.
- **Worker path** (`start_from_trigger=False`): the worker has already
rendered `self.target_time` when `execute()` runs, so the trigger can be built
from `moment=self._moment` exactly as on `main`.
Selecting `target_time=` in `execute()` on 3.3+ applies the triggerer-path
mechanism to the worker path, which has two costs:
- It ships the rendered value with whatever type it has. A `datetime`
(passed directly, or rendered natively) reaches the trigger as `target_time`
and fails there, see my comment on `DateTimeTrigger.moment`.
- Plain deferral now depends on the triggerer running this provider version.
An older `DateTimeTrigger` rejects the unknown `target_time` kwarg, while
`moment` works with every version.
Suggestion:
```python
def execute(self, context: Context) -> None:
if not self.start_from_trigger:
if AIRFLOW_V_3_0_PLUS:
trigger = DateTimeTrigger(moment=self._moment,
end_from_trigger=self.end_from_trigger)
else:
trigger = DateTimeTrigger(moment=self._moment)
self.defer(method_name="execute_complete", trigger=trigger)
```
With that, the changes to `test_supervisor.py` and `test_task_runner.py` are
no longer needed.
##########
providers/standard/src/airflow/providers/standard/triggers/temporal.py:
##########
@@ -31,31 +31,55 @@ 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). ``target_time`` is listed in ``template_fields`` so that, with
``start_from_trigger``,
+ the triggerer renders it in place before ``run()`` and it is 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:
+ template_fields = ("target_time",)
+
+ 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.
+ if not isinstance(self.target_time, str) or not self.target_time:
Review Comment:
**[blocker]** A `datetime` `target_time` is accepted by `__init__` but makes
the trigger fail at run time.
The signature allows `target_time: datetime.datetime | str | None`, but
`moment` only resolves strings. Reproduced against this branch:
```python
trigger = DateTimeTrigger(target_time=datetime.datetime(2020, 1, 1,
tzinfo=datetime.timezone.utc))
await trigger.run().__anext__()
# TypeError: DateTimeTrigger has neither a 'moment' nor a usable
'target_time'
```
On Airflow 3.3+ the sensor now sends a `datetime` this way in cases that
work on `main`:
- `DateTimeSensorAsync(target_time=<datetime>, start_from_trigger=True)`:
`__init__` puts the raw value in `trigger_kwargs`. The updated
`test_async_start_from_trigger_localizes_naive_datetime` asserts exactly that
(a naive `datetime` under `target_time`).
- `DateTimeSensorAsync(target_time=<datetime>)` through `execute()`: the
updated expectation in `test_supervisor.py` shows a `datetime.datetime` being
sent as `target_time`.
- A template that renders to a native `datetime`
(`render_template_as_native_obj=True`) would end up in the same place.
`DateTimeSensor._moment` already handles both types, so the trigger can
mirror it:
```python
@property
def moment(self) -> pendulum.DateTime:
if self._moment is None:
target_time = 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(f"Expected str or datetime.datetime type for
target_time. Got {type(target_time)}")
self._moment = timezone.convert_to_utc(timezone.parse(target_time))
return self._moment
```
`timezone.parse` also localizes naive values, which restores the behaviour
the test name promises. Please add trigger tests in `triggers/test_temporal.py`
for `target_time` as `str`, aware `datetime` and naive `datetime`, for the
"exactly one of" validation, and for the `serialize()` round trip.
##########
providers/standard/src/airflow/providers/standard/triggers/temporal.py:
##########
@@ -31,31 +31,55 @@ 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). ``target_time`` is listed in ``template_fields`` so that, with
``start_from_trigger``,
+ the triggerer renders it in place before ``run()`` and it is 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:
+ template_fields = ("target_time",)
Review Comment:
**[nit]** This class attribute has no effect, and the docstring above
describes a mechanism that is not the one in use.
`BaseTrigger.__init__` sets `self.template_fields = ()` on the instance,
which shadows the class attribute:
```python
>>> DateTimeTrigger(target_time="{{ ds }}").template_fields
()
```
The fields that get rendered are set by the `task_instance` setter in
`BaseTrigger`: the operator's `template_fields` that are also keys of
`start_trigger_args.trigger_kwargs` and attributes of the trigger. So rendering
works here because `DateTimeSensor.template_fields` contains `target_time` and
the sensor puts `target_time` in `trigger_kwargs`. I would drop this line and
reword the docstring accordingly.
##########
providers/standard/tests/unit/standard/sensors/test_date_time.py:
##########
@@ -134,19 +135,60 @@ 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.2")
+ async def test_full_run_worker_path_templated_past_target_time(self):
Review Comment:
**[warning]** This test exercises a flow that does not occur in production,
so it does not cover the new feature.
It calls `op.execute(ctx)` without rendering the operator first (the worker
always renders before `execute`), then renders the template on the trigger
through a `SimpleNamespace` task instance. With `start_from_trigger=True` the
scheduler defers the task directly and `execute()` is never called; without it,
the triggerer does not render anything.
A test closer to the real path would build the trigger the way the triggerer
does:
```python
trigger = DateTimeTrigger(**op.start_trigger_args.trigger_kwargs)
trigger.task_instance = ... # object exposing task_id and task=op
trigger.render_template_fields(ctx)
event = await trigger.run().__anext__()
```
Smaller points in this file:
- The `skipif` reasons do not match the conditions: `not AIRFLOW_V_3_3_PLUS`
with "Test only for AF < 3.2" skips on < 3.3, and the `AIRFLOW_V_3_3_PLUS` one
below should read "< 3.3".
- `TaskDeferred` and `SimpleNamespace` should be imported at the top of the
module.
- `test_async_start_from_trigger_molocalment` looks like an accidental
rename.
- The new `UserWarning` on Airflow < 3.3 has no test (`pytest.warns`).
##########
providers/standard/tests/unit/standard/sensors/test_date_time.py:
##########
@@ -156,8 +198,13 @@ 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)
Review Comment:
**[warning]** This assertion pins the bug described on
`DateTimeTrigger.moment`.
The test is named `..._localizes_naive_datetime`, but on 3.3+ it now asserts
that the naive `datetime` is passed through unchanged. The trigger then rejects
it at run time. Once the trigger handles `datetime`, I would extend this test
to check the resolved value, for example
`DateTimeTrigger(**op.start_trigger_args.trigger_kwargs).moment ==
pendulum.datetime(2020, 1, 1, tz="UTC")`.
##########
task-sdk/tests/task_sdk/execution_time/test_supervisor.py:
##########
@@ -810,15 +811,22 @@ def mock_monotonic():
# Validate calls to the client
mock_client.task_instances.start.assert_called_once_with(ti.id,
mocker.ANY, mocker.ANY)
mock_client.task_instances.heartbeat.assert_called_once_with(ti.id,
pid=mocker.ANY)
+
+ trigger_kwargs_key_name = "moment"
+ trigger_kwargs_class_name = "pendulum.datetime.DateTime"
+ if AIRFLOW_V_3_3_PLUS:
Review Comment:
**[warning]** A change in the standard provider should not need edits in the
task-sdk tests.
These two tests only use `DateTimeSensorAsync` as a sample deferrable task
and assert the trigger kwargs it sends. Their expectations had to change
because `execute()` now sends `target_time` instead of `moment`. Once
`execute()` builds the trigger from `moment` again (see my comment there), both
`test_supervisor.py` and `test_task_runner.py` can be reverted.
The `AIRFLOW_V_3_3_PLUS` branches added here are also dead code: the
task-sdk tests only run against the current `main` (3.4.0), never against older
Airflow versions.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
