Add new datetime branch operator #11964

tomasfarias · 2020-10-30T08:57:40Z

closes: #11929

This PR includes a new datetime branching operator: the current date and time, as given by datetime.datetime.now is compared against target datetime attributes, like year or hour, to decide which task id branch to take.

No tests were written yet as this is intended as a POC to allow us to review the implementation first.

^ Add meaningful description above

Read the Pull Request Guidelines for more information.
In case of fundamental code change, Airflow Improvement Proposal (AIP) is needed.
In case of a new dependency, check compliance with the ASF 3rd Party License Policy.
In case of backwards incompatible changes please leave a note in UPDATING.md.

boring-cyborg · 2020-10-30T08:57:42Z

Congratulations on your first Pull Request and welcome to the Apache Airflow community! If you have any issues or are unsure about any anything please check our Contribution Guide (https://github.com/apache/airflow/blob/master/CONTRIBUTING.rst)
Here are some useful points:

Pay attention to the quality of your code (flake8, pylint and type annotations). Our pre-commits will help you with that.
In case of a new feature add useful documentation (in docstrings or in docs/ directory). Adding a new operator? Check this short guide Consider adding an example DAG that shows how users should use it.
Consider using Breeze environment for testing locally, it’s a heavy docker but it ships with a working Airflow and a lot of integrations.
Be patient and persistent. It might take some time to get a review or get the final approval from Committers.
Please follow ASF Code of Conduct for all communication including (but not limited to) comments on Pull Requests, Mailing list and Slack.
Be sure to read the Airflow Coding style.
Apache Airflow is a community-driven project and together we are making it better 🚀.
In case of doubts contact the developers at:
Mailing List: dev@airflow.apache.org
Slack: https://s.apache.org/airflow-slack

github-actions · 2020-10-30T11:06:44Z

The Workflow run is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks,^Build docs$,^Spell check docs$,^Backport packages$,^Provider packages,^Checks: Helm tests$,^Test OpenAPI*.

mik-laj · 2020-10-30T11:39:26Z

DateTimeSensor has a much simpler interface. Why do we need so many parameters here?

github-actions · 2020-10-30T12:26:27Z

The Workflow run is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks,^Build docs$,^Spell check docs$,^Backport packages$,^Provider packages,^Checks: Helm tests$,^Test OpenAPI*.

eladkal · 2020-10-30T13:59:02Z

I don't think you can actually compare exact specific datetime as you can't know when the task will be executed. For example branching based on specific minute will probably wont work as expected.

I think the concept should be more to get acceptable range : lower & upper. If current time is between the range then continue true branch otherwise continue false branch.
WDYT?

kaxil · 2020-11-04T00:10:39Z

Can you please rebased your PR on latest Master since we have applied Black and PyUpgrade on Master.

It will help if your squash your commits into single commit first so that there are less conflicts.

github-actions · 2020-11-04T01:17:47Z

The Workflow run is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks,^Build docs$,^Spell check docs$,^Backport packages$,^Provider packages,^Checks: Helm tests$,^Test OpenAPI*.

github-actions · 2020-11-04T02:19:11Z

The Workflow run is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks,^Build docs$,^Spell check docs$,^Backport packages$,^Provider packages,^Checks: Helm tests$,^Test OpenAPI*.

tomasfarias · 2020-11-04T03:11:18Z

@eladkal I agree, any unit smaller than hour is probably not something you'd be able to match reliably anyways 😓. I like your idea, and I've also allowed for either the upper bound or the lower bound to be set as None, but not both, to allow for unilateral comparisons, for example: if target_upper is None and the current datetime falls above target_lower then return true branch.

This also reduces the number of arguments, which was getting unnecessarily out of control as @mik-laj suggested.

Branch has been rebased as @kaxil asked.

I'm was having a bit of trouble getting the docs build to pass: now docs build passes locally, so I'm expecting it to pass with the latest commit.

Thanks all for reviews and input 💪

github-actions · 2020-11-04T03:48:13Z

The Workflow run is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks,^Build docs$,^Spell check docs$,^Backport packages$,^Provider packages,^Checks: Helm tests$,^Test OpenAPI*.

eladkal · 2020-11-10T13:35:25Z

airflow/operators/datetime_branch_operator.py

+        follow_task_ids_if_false: Union[str, Iterable[str]],
+        target_lower: Optional[datetime.datetime],
+        target_upper: Optional[datetime.datetime],
+        timezone: Optional[datetime.timezone] = None,


I think this param is redundant.
We can respect the default timezone set in airflow.cfg and allow overwrite if the DAG is timezone aware.
Similar to how TimeSensor was modified #9882 #9699

Thanks for pointing it out 👍 Latest commit should replicate the work done for TimeSensor.

can you add unit tests ?

Unit tests added 🙌

airflow/operators/datetime_branch_operator.py

RosterIn · 2020-11-16T15:11:25Z

airflow/operators/datetime_branch_operator.py

+        `datetime.datetime.now()` falls below target_lower or above `target_upper`.
+    :type follow_task_ids_if_false: str or list[str]
+    :param target_lower: target lower bound.
+    :type target_lower: Optional[datetime.datetime]


why this optional?

Both target_* arguments are optional to support unilateral comparisons, i.e. comparing the current time to a lower target but without any upper bound and vice-versa. We do need at least one of the two to be defined otherwise there is no target to compare the current date to, which is why that is checked below in __init__.

RosterIn · 2020-11-16T15:11:33Z

airflow/operators/datetime_branch_operator.py

+    :param target_lower: target lower bound.
+    :type target_lower: Optional[datetime.datetime]
+    :param target_upper: target upper bound.
+    :type target_upper: Optional[datetime.datetime]


same question here

Answered above ☝️

ashb · 2020-11-16T15:53:04Z

tests/operators/test_datetime_branch_operator.py

+        branch_op = DateTimeBranchOperator(
+            task_id='datetime_branch',
+            follow_task_ids_if_true='branch_1',
+            follow_task_ids_if_false='branch_2',
+            target_upper=datetime.datetime(2020, 7, 7, 11, 0, 0),
+            target_lower=None,
+            dag=self.dag,
+        )
+
+        self.branch_1.set_upstream(branch_op)
+        self.branch_2.set_upstream(branch_op)
+        self.dag.clear()
+
+        dr = self.dag.create_dagrun(
+            run_id='manual__', start_date=DEFAULT_DATE, execution_date=DEFAULT_DATE, state=State.RUNNING
+        )


This appears in every test with almost indentical content -- could you create a helper function to reduce this boilerplate please.

Good suggestion! Moved boilerplate code over to setUp

ashb · 2020-11-16T15:54:10Z

tests/operators/test_datetime_branch_operator.py

+    @mock.patch('airflow.operators.datetime_branch_operator.timezone.utcnow')
+    def test_datetime_branch_operator_lower_comparison_outside_range(self, mock_timezone):


Rather than mocking the datetime like this, please use freezegun

Thanks for the suggestion! I've used freezegun instead.

tomasfarias · 2020-11-20T22:15:21Z

Not sure why MySQL build is failing, all tests appear to be passing according to logs.

eladkal · 2020-11-21T20:24:46Z

airflow/operators/datetime_branch_operator.py

@@ -0,0 +1,77 @@
+# Licensed to the Apache Software Foundation (ASF) under one


the suffix of _operator in the file name should be removed (to comply with AIP 21)
see #11178

eladkal · 2020-11-26T15:11:14Z

tests/operators/test_datetime_branch_operator.py

@@ -0,0 +1,195 @@
+#


please also remove the _operator suffix from the test file

eladkal · 2020-11-26T15:15:30Z

airflow/operators/datetime_branch.py

+        target_lower: Optional[datetime.datetime],
+        target_upper: Optional[datetime.datetime],


I'm wondring about the use cases of datetime in general.
I would assume that it's more likely that user would like to branch based on times rather than datetime.
Do you think we will need a different TimeBranchOperator or the two use cases can be combined?

Branching of times should be possible by setting the target values to something like:

target_lower=datetime.datetime.combine( datetime.datetime.now(), datetime.time(hour=target_hour, minute=target_minute, second=target_second) )

But this is not as straightforward as passing a single datetime.time.

I think adding a TimeBranchOperator may be a bit redundant, instead, we could support targets of type Optional[Union[datetime.datetime, datetime.time]], and when encountering a datetime.time we combine it with the current date and proceed as it is right now:

def choose_branch(self, context: Dict) -> Union[str, Iterable[str]]: now = timezone.make_naive(timezone.utcnow(), self.dag.timezone) target_upper = datetime.datetime.combine(now, self.target_upper) if isinstance(self.target_upper, datetime.time) else self.target_upper

What do you think?

yeah I think the Optional[Union[datetime.datetime, datetime.time]] approach is better as it simplify the usage for the users but maybe worth asking for more opinions on this one.

note that there are edge cases for example lower is 23:00 and upper is 01:00 (the next day) so the duration of success is total of 2 hours yet it spread over two dates.

Also it will be good to add documentation about this operator. you can use previous PRs to see examples #11472

Figured out it may be easier to hear other opinions with a concrete example, so I implemented the idea mostly as described. The edge case you described is handled by pushing the smaller time to the next day. Should the same thing be done when both times are equal? Or should an exception be raised instead?

Gave it a shot at adding documentation, do let me know if it needs to be extended or rewritten.

github-actions · 2020-12-06T01:16:58Z

The Workflow run is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks,^Build docs$,^Spell check docs$,^Backport packages$,^Provider packages,^Checks: Helm tests$,^Test OpenAPI*.

eladkal · 2020-12-15T15:24:39Z

airflow/example_dags/example_datetime_branch_operator.py

+Example DAG demonstrating the usage of BranchPythonOperator with depends_on_past=True, where tasks may be run
+or skipped on alternating runs.


This is a description of another operator

Sorry about that, copy and paste blunder. Thanks for the review!

github-actions · 2020-12-20T22:29:55Z

The Workflow run is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks,^Build docs$,^Spell check docs$,^Backport packages$,^Provider packages,^Checks: Helm tests$,^Test OpenAPI*.

eladkal

LGTM
you will need to rebase to fix the tests/docs due to #13201

kaxil · 2021-02-23T11:29:51Z

airflow/operators/datetime_branch.py

+    True branch will be returned when `datetime.datetime.now()` falls below
+    `target_upper` and above `target_lower`.


What's the behaviour when backfilling or catching-up ?

How about adding something like:

airflow/airflow/operators/weekday.py

Lines 86 to 89 in e363ac1

if self.use_task_execution_day:

now = context["execution_date"]

else:

now = timezone.make_naive(timezone.utcnow(), self.dag.timezone)

cc @eladkal

Good call! Behavior when backfilling would be to still use system date, which is probably not what a user would expect. I added a use_task_execution_date argument very similar to the one you linked, with some minor adaptation since we are talking about dates instead of days of the week.

airflow/operators/datetime_branch.py

kaxil · 2021-02-23T11:32:14Z

airflow/operators/datetime_branch.py

+        self.follow_task_ids_if_false = follow_task_ids_if_false
+
+    def choose_branch(self, context: Dict) -> Union[str, Iterable[str]]:
+        now = timezone.make_naive(timezone.utcnow(), self.dag.timezone)


Same comment as https://github.com/apache/airflow/pull/11964/files#r580961797

Implemented your suggestion (see the other comment).

eladkal · 2021-03-09T20:55:15Z

@tomasfarias can you please address issues/questions raised by kaxil?

Co-authored-by: Kaxil Naik <kaxilnaik@gmail.com>

tomasfarias · 2021-03-10T19:50:04Z

@eladkal Been a bit busy with work the last few weeks, but all comments should be addressed! Thanks for pinging me 👍

github-actions · 2021-03-10T22:43:06Z

The PR most likely needs to run full matrix of tests because it modifies parts of the core of Airflow. However, committers might decide to merge it quickly and take the risk. If they don't merge it quickly - please rebase it to the latest master at your convenience, or amend the last commit of the PR, and push it with --force-with-lease.

boring-cyborg · 2021-03-10T22:44:12Z

Awesome work, congrats on your first merged pull request!

eladkal · 2021-03-11T11:38:16Z

while commenting on #11931 i notice that we might have two issues here:

The operator name DateTimeBranchOperator isn't consistent with other branch operators: BranchPythonOperator, BranchSQLOperator, BranchDayOfWeekOperator for consistnacy the operator should be renamed to BranchDateTimeOperator
Should datetime_branch.py be datetime.py?

Since the operator isn't released yet we can handle this easily without deprecation notice.
@kaxil WDYT?

I can handle this if @tomasfarias has no time.

tomasfarias · 2021-03-11T11:44:00Z

@eladkal Things are much quieter in my schedule now so I can push any changes to names if they are deemed appropriate.

I personally agree 100% with 1. as it plays nice with any IDE autocomplete where you can just get all the Branch* operator suggestions typing only that.

For 2. though I see we have both branch_operator.py as well as sql_branch_operator.py, so maybe we should settle for datetime_branch_operator.py? I'm fine either way on this one though.

eladkal · 2021-03-11T11:50:50Z

For 2. though I see we have both branch_operator.py as well as sql_branch_operator.py, so maybe we should settle for datetime_branch_operator.py? I'm fine either way on this one though.

They are deprecated :)
branch_operator.py -> branch.py
sql_branch_operator.py > sql.py

tomasfarias · 2021-03-11T12:13:53Z

Now that you mention that, we should also change the deprecated imports, e.g.:

from airflow.operators.branch_operator import BranchBaseOperator

To:

from airflow.operators.branch import BaseBranchOperator

I've pushed the changes to my fork, see this commit: tomasfarias@5fc2862. I can create a PR as soon as the changes are confirmed as needed 👍

Great feedback!

kaxil · 2021-03-11T12:16:33Z

Good point @eladkal -- agree to your suggestions.

And yes since they are not released we can rename them without deprecations

tomasfarias · 2021-03-11T12:54:59Z

I opened a PR with the changes suggested, linking here for reference: #14720

tomasfarias force-pushed the branch-datetime-operator branch from 277a146 to 144e6b8 Compare October 30, 2020 09:01

mik-laj requested a review from kaxil October 30, 2020 11:39

tomasfarias force-pushed the branch-datetime-operator branch from 08bbfaa to 911735f Compare November 4, 2020 00:58

tomasfarias force-pushed the branch-datetime-operator branch from 911735f to 044cf1f Compare November 4, 2020 01:56

tomasfarias force-pushed the branch-datetime-operator branch from 044cf1f to 14793c2 Compare November 4, 2020 03:11

tomasfarias marked this pull request as ready for review November 9, 2020 17:18

eladkal reviewed Nov 10, 2020

View reviewed changes

tomasfarias force-pushed the branch-datetime-operator branch 2 times, most recently from 0fe054c to f6ff87e Compare November 11, 2020 19:27

RosterIn reviewed Nov 16, 2020

View reviewed changes

ashb reviewed Nov 16, 2020

View reviewed changes

eladkal requested changes Nov 21, 2020

View reviewed changes

tomasfarias force-pushed the branch-datetime-operator branch from e184167 to 78ec08c Compare November 23, 2020 23:33

eladkal reviewed Nov 26, 2020

View reviewed changes

eladkal reviewed Dec 15, 2020

View reviewed changes

eladkal approved these changes Dec 21, 2020

View reviewed changes

kaxil reviewed Feb 23, 2021

View reviewed changes

airflow/operators/datetime_branch.py Outdated Show resolved Hide resolved

kaxil reviewed Feb 23, 2021

View reviewed changes

airflow/operators/datetime_branch.py Outdated Show resolved Hide resolved

kaxil reviewed Feb 23, 2021

View reviewed changes

airflow/operators/datetime_branch.py Outdated Show resolved Hide resolved

kaxil reviewed Feb 23, 2021

View reviewed changes

Add use_task_execution_date argument to DateTimeBranchOperator

817955b

tomasfarias requested review from mik-laj and vikramkoka as code owners March 10, 2021 19:46

Use double backticks in rst docs

8253b4a

Co-authored-by: Kaxil Naik <kaxilnaik@gmail.com>

tomasfarias force-pushed the branch-datetime-operator branch from 74b4d9f to 8253b4a Compare March 10, 2021 19:48

kaxil approved these changes Mar 10, 2021

View reviewed changes

github-actions bot added the full tests needed We need to run full set of tests for this PR to merge label Mar 10, 2021

kaxil merged commit 1e37a11 into apache:master Mar 10, 2021

kaxil added this to the Airflow 2.1 milestone Mar 10, 2021

tomasfarias mentioned this pull request Mar 11, 2021

Update names and paths to be consistent with other Branch operators #14720

Merged

		@mock.patch('airflow.operators.datetime_branch_operator.timezone.utcnow')
		def test_datetime_branch_operator_lower_comparison_outside_range(self, mock_timezone):

		@@ -0,0 +1,77 @@
		# Licensed to the Apache Software Foundation (ASF) under one

		target_lower: Optional[datetime.datetime],
		target_upper: Optional[datetime.datetime],

		Example DAG demonstrating the usage of BranchPythonOperator with depends_on_past=True, where tasks may be run
		or skipped on alternating runs.

		True branch will be returned when `datetime.datetime.now()` falls below
		`target_upper` and above `target_lower`.

	if self.use_task_execution_day:
	now = context["execution_date"]
	else:
	now = timezone.make_naive(timezone.utcnow(), self.dag.timezone)

Add new datetime branch operator #11964

Add new datetime branch operator #11964

Conversation

tomasfarias commented Oct 30, 2020

boring-cyborg bot commented Oct 30, 2020

github-actions bot commented Oct 30, 2020

mik-laj commented Oct 30, 2020

github-actions bot commented Oct 30, 2020

eladkal commented Oct 30, 2020

kaxil commented Nov 4, 2020

github-actions bot commented Nov 4, 2020

github-actions bot commented Nov 4, 2020

tomasfarias commented Nov 4, 2020

github-actions bot commented Nov 4, 2020

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

tomasfarias commented Nov 20, 2020

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

tomasfarias Nov 26, 2020 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

github-actions bot commented Dec 6, 2020

Choose a reason for hiding this comment

Choose a reason for hiding this comment

github-actions bot commented Dec 20, 2020

eladkal left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

tomasfarias Mar 10, 2021 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

eladkal commented Mar 9, 2021

tomasfarias commented Mar 10, 2021

github-actions bot commented Mar 10, 2021

boring-cyborg bot commented Mar 10, 2021

eladkal commented Mar 11, 2021

tomasfarias commented Mar 11, 2021

eladkal commented Mar 11, 2021

tomasfarias commented Mar 11, 2021

kaxil commented Mar 11, 2021 • edited Loading

tomasfarias commented Mar 11, 2021

tomasfarias Nov 26, 2020 •

edited

Loading

tomasfarias Mar 10, 2021 •

edited

Loading

kaxil commented Mar 11, 2021 •

edited

Loading