Support timezone-aware `time`s for `TimeSensor` and `TimeSensorAsync`.

[dwreeves]

Description

Right now, if you try doing something like the below, the timezone is completely ignored:

from pendulum import Time
from airflow.sensors.time_sensor import TimeSensorAsync
import pytz

wait_for_start_of_business = TimeSensorAsync(
    task_id="wait_for_business_hours",
    target_time=Time(9, 0, tzinfo=pytz.timezone("America/New_York")),
)

The timezone is completely ignored.

It would be nice if the timezone were not ignored in this case. I think the change that would need to occur in the code would be to replace this:

aware_time = timezone.coerce_datetime(
    datetime.datetime.combine(datetime.datetime.today(), self.target_time, self.dag.timezone)
)

with this:

tz = self.target_time.tzinfo or self.dag.timezone
aware_time = timezone.coerce_datetime(
    datetime.datetime.combine(datetime.datetime.now(tz=tz), self.target_time, tz)
)

This is backwards compatible so long as:

• The time does not have a tzinfo or alternatively the tzinfo is equal to self.dag.timezone
• The self.dag.timezone agrees with the system time. This is a weird one, as I would argue the previous behavior was a minor bug: e.g. datetime.datetime.today() running on Jan 2nd at 03:00:00+0600 (i.e. 6 hours ahead of UTC) will output Jan 2nd even though it is Jan 1st in UTC, so if the system is +0600 and the DAG is +0000, the sensor will potentially not fire (for example, target_time=time(20, 0) would not fire on Jan 1st, 20:00:00 UTC in this setup).

(While we are at it, TimeSensor is untyped, so some type annotations could be added during this PR.)

Use case/motivation

I have a Slack notification system that needs to wait until 9am EST before it's allowed to fire off notifications. What constitutes 9am EST in UTC time (which is our Airflow server's timezone) differs depending on daylight savings.

Related issues

No response

Are you willing to submit a PR?

• Yes I am willing to submit a PR!

Code of Conduct

• I agree to follow this project's Code of Conduct

[Taragolis]

The timezone is completely ignored.

Heh, that is what I exactly (or close) asked into the #contributors channel in Slack today: https://apache-airflow.slack.com/archives/C06K9Q5G2UA/p1710946977026229

I have a look on DTZ rules (ruff, flake8-datetimez) and somethings tells me that enable some of the rules or subsequent of it might be useful.
Especially with .today() which might return inconsistent result in case if server setup different timezone rather than defined into the https://airflow.apache.org/docs/apache-airflow/stable/configurations-ref.html#default-timezone
WDYT?

[dwreeves]

@Taragolis Good timing! The stars have aligned.

I've been thinking about all the time-related sensors for like, the past week, and decided to post this and #38336 back to back to make sure my thoughts are somewhere.

This time, I was looking at the source code, and I've been thinking there is a lot more than can be done. Not only just in terms of expanding the API, but simply and much more importantly, just "formalizing" the API that already exists. Across all the time-related control flow sensors, there are things like: lack of type annotations, inconsistent behavior with render_template_as_native_obj=True, and this issue with timezone awareness, which I'm surprised has never been caught before (perhaps you're surprised too!).

There are potentially more minor issues in these control-flow operators and sensors. This is a very old part of Airflow, and nobody wants to touch it...

---

For sake of completeness:

I will say, the "business case" I stated above is actually simplified. What I ended up doing was implementing a custom operator called TimeBetweenSensor, and it looks like this:

from pendulum import Time
from dags.utils.control_flow import TimeBetweenSensorAsync
import pytz

wait_for_start_of_business = TimeBetweenSensorAsync(
    task_id="wait_for_business_hours",
    start_time=Time(9, 0, tzinfo=pytz.timezone("America/New_York")),
    end_time=Time(17, 0, tzinfo=pytz.timezone("America/New_York")),
)

And so to make this work, you actually first need to check whether you've exceeded the end_time, if true, add a day to the start time, if false, then just use the start time.)

[eladkal]

@dwreeves if you have a suggestion how to address then simply open a PR it will be easier to review what you propose

[dwreeves]

Yeah sorry, I've just had my plate full for a while, and my local fork of the Airflow repo started encountering issues that I was not able to fully diagnose within a reasonable amount of time. Unsure when I'll get to this.

[dwreeves]

First of all, thank you @SML0127 for addressing this!! And thank you Wei Lee for reviewing!

@Lee-W, I don't think this is completely closed. The PR addresses half of the issue, which is that it changes the default behavior to look at the DAG's timezone. That's a great change! But it still ignores the timezone info of the target_time if it exists, and it doesn't seem like that should be the case.

The example of "only send during business hours" is an example of a case where this matters: if I'm in New York, for example, that's either UTC-4 or UTC-5 depending on time of year, and my Airflow instance is set to UTC as the default timezone for everything.

I can also open a new issue if you'd like, or simply just open a PR addressing this without creating a new issue. This really isn't that hard to resolve, I was just very slammed with other obligations when I wrote up the issue to get to it.

Thanks again to both of you!