You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by GitBox <gi...@apache.org> on 2020/10/30 08:57:52 UTC
[GitHub] [airflow] tomasfarias opened a new pull request #11964: Add new datetime branch operator
tomasfarias opened a new pull request #11964:
URL: https://github.com/apache/airflow/pull/11964
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](https://github.com/apache/airflow/blob/master/CONTRIBUTING.rst#pull-request-guidelines)** for more information.
In case of fundamental code change, Airflow Improvement Proposal ([AIP](https://cwiki.apache.org/confluence/display/AIRFLOW/Airflow+Improvements+Proposals)) is needed.
In case of a new dependency, check compliance with the [ASF 3rd Party License Policy](https://www.apache.org/legal/resolved.html#category-x).
In case of backwards incompatible changes please leave a note in [UPDATING.md](https://github.com/apache/airflow/blob/master/UPDATING.md).
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] boring-cyborg[bot] commented on pull request #11964: Add new datetime branch operator
Posted by GitBox <gi...@apache.org>.
boring-cyborg[bot] commented on pull request #11964:
URL: https://github.com/apache/airflow/pull/11964#issuecomment-719429860
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]( https://github.com/apache/airflow/blob/master/STATIC_CODE_CHECKS.rst#prerequisites-for-pre-commit-hooks) 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](https://github.com/apache/airflow/blob/master/docs/howto/custom-operator.rst) Consider adding an example DAG that shows how users should use it.
- Consider using [Breeze environment](https://github.com/apache/airflow/blob/master/BREEZE.rst) 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](https://www.apache.org/foundation/policies/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]( https://github.com/apache/airflow/blob/master/CONTRIBUTING.rst#coding-style-and-best-practices).
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
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] tomasfarias commented on a change in pull request #11964: Add new datetime branch operator
Posted by GitBox <gi...@apache.org>.
tomasfarias commented on a change in pull request #11964:
URL: https://github.com/apache/airflow/pull/11964#discussion_r521589482
##########
File path: airflow/operators/datetime_branch_operator.py
##########
@@ -0,0 +1,80 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import datetime
+from typing import Dict, Iterable, Optional, Union
+
+from airflow.exceptions import AirflowException
+from airflow.operators.branch_operator import BaseBranchOperator
+from airflow.utils.decorators import apply_defaults
+
+
+class DateTimeBranchOperator(BaseBranchOperator):
+ """
+ Branches into one of two lists of tasks depending on the current datetime.
+
+ True branch will be returned when `datetime.datetime.now()` falls below
+ `target_upper` and above `target_lower`.
+
+ :param follow_task_ids_if_true: task id or task ids to follow if
+ `datetime.datetime.now()` falls above target_lower and below `target_upper`.
+ :type follow_task_ids_if_true: str or list[str]
+ :param follow_task_ids_if_false: task id or task ids to follow if
+ `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]
+ :param target_upper: target upper bound.
+ :type target_upper: Optional[datetime.datetime]
+ :param timezone: timezone to override local time.
+ :type timezone: Optional[datetime.timezone]
+ """
+
+ @apply_defaults
+ def __init__(
+ self,
+ *,
+ follow_task_ids_if_true: Union[str, Iterable[str]],
+ follow_task_ids_if_false: Union[str, Iterable[str]],
+ target_lower: Optional[datetime.datetime],
+ target_upper: Optional[datetime.datetime],
+ timezone: Optional[datetime.timezone] = None,
Review comment:
Unit tests added π
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] boring-cyborg[bot] commented on pull request #11964: Add new datetime branch operator
Posted by GitBox <gi...@apache.org>.
boring-cyborg[bot] commented on pull request #11964:
URL: https://github.com/apache/airflow/pull/11964#issuecomment-796255556
Awesome work, congrats on your first merged pull request!
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] kaxil edited a comment on pull request #11964: Add new datetime branch operator
Posted by GitBox <gi...@apache.org>.
kaxil edited a comment on pull request #11964:
URL: https://github.com/apache/airflow/pull/11964#issuecomment-796694915
Good point @eladkal -- agree to your suggestions.
And yes since they are not released we can rename them without deprecations
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] tomasfarias commented on a change in pull request #11964: Add new datetime branch operator
Posted by GitBox <gi...@apache.org>.
tomasfarias commented on a change in pull request #11964:
URL: https://github.com/apache/airflow/pull/11964#discussion_r591822333
##########
File path: airflow/operators/datetime_branch.py
##########
@@ -0,0 +1,101 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import datetime
+from typing import Dict, Iterable, Union
+
+from airflow.exceptions import AirflowException
+from airflow.operators.branch_operator import BaseBranchOperator
+from airflow.utils import timezone
+from airflow.utils.decorators import apply_defaults
+
+
+class DateTimeBranchOperator(BaseBranchOperator):
+ """
+ Branches into one of two lists of tasks depending on the current datetime.
+
+ True branch will be returned when `datetime.datetime.now()` falls below
+ `target_upper` and above `target_lower`.
Review comment:
Good call! Behaviour 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.
##########
File path: airflow/operators/datetime_branch.py
##########
@@ -0,0 +1,101 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import datetime
+from typing import Dict, Iterable, Union
+
+from airflow.exceptions import AirflowException
+from airflow.operators.branch_operator import BaseBranchOperator
+from airflow.utils import timezone
+from airflow.utils.decorators import apply_defaults
+
+
+class DateTimeBranchOperator(BaseBranchOperator):
+ """
+ Branches into one of two lists of tasks depending on the current datetime.
+
+ True branch will be returned when `datetime.datetime.now()` falls below
+ `target_upper` and above `target_lower`.
Review comment:
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.
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] ashb commented on a change in pull request #11964: Add new datetime branch operator
Posted by GitBox <gi...@apache.org>.
ashb commented on a change in pull request #11964:
URL: https://github.com/apache/airflow/pull/11964#discussion_r524370257
##########
File path: airflow/operators/datetime_branch_operator.py
##########
@@ -0,0 +1,77 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import datetime
+from typing import Dict, Iterable, Optional, Union
+
+from airflow.exceptions import AirflowException
+from airflow.operators.branch_operator import BaseBranchOperator
+from airflow.utils import timezone
+from airflow.utils.decorators import apply_defaults
+
+
+class DateTimeBranchOperator(BaseBranchOperator):
+ """
+ Branches into one of two lists of tasks depending on the current datetime.
+
+ True branch will be returned when `datetime.datetime.now()` falls below
+ `target_upper` and above `target_lower`.
+
+ :param follow_task_ids_if_true: task id or task ids to follow if
+ `datetime.datetime.now()` falls above target_lower and below `target_upper`.
+ :type follow_task_ids_if_true: str or list[str]
+ :param follow_task_ids_if_false: task id or task ids to follow if
+ `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]
+ :param target_upper: target upper bound.
+ :type target_upper: Optional[datetime.datetime]
+ """
+
+ @apply_defaults
+ def __init__(
+ self,
+ *,
+ follow_task_ids_if_true: Union[str, Iterable[str]],
+ follow_task_ids_if_false: Union[str, Iterable[str]],
+ target_lower: Optional[datetime.datetime],
+ target_upper: Optional[datetime.datetime],
+ **kwargs,
+ ) -> None:
+ super().__init__(**kwargs)
+ if target_lower is None and target_upper is None:
+ raise AirflowException(
+ "Both target_upper and target_lower are None. At least one "
+ "must be defined to be compared to the current datetime"
+ )
+
+ self.target_lower = target_lower
+ self.target_upper = target_upper
+ self.follow_task_ids_if_true = follow_task_ids_if_true
+ 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)
Review comment:
Probably not, no. Doing this based on "wall clock" dates is right.
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] eladkal commented on pull request #11964: Add new datetime branch operator
Posted by GitBox <gi...@apache.org>.
eladkal commented on pull request #11964:
URL: https://github.com/apache/airflow/pull/11964#issuecomment-796680972
> 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](https://github.com/apache/airflow/blob/master/airflow/operators/branch_operator.py) -> [branch.py](https://github.com/apache/airflow/blob/master/airflow/operators/branch.py)
[sql_branch_operator.py](https://github.com/apache/airflow/blob/master/airflow/operators/sql_branch_operator.py) > [sql.py](https://github.com/apache/airflow/blob/master/airflow/operators/sql.py)
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] kaxil commented on a change in pull request #11964: Add new datetime branch operator
Posted by GitBox <gi...@apache.org>.
kaxil commented on a change in pull request #11964:
URL: https://github.com/apache/airflow/pull/11964#discussion_r580960901
##########
File path: airflow/operators/datetime_branch.py
##########
@@ -0,0 +1,101 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import datetime
+from typing import Dict, Iterable, Union
+
+from airflow.exceptions import AirflowException
+from airflow.operators.branch_operator import BaseBranchOperator
+from airflow.utils import timezone
+from airflow.utils.decorators import apply_defaults
+
+
+class DateTimeBranchOperator(BaseBranchOperator):
+ """
+ Branches into one of two lists of tasks depending on the current datetime.
+
+ True branch will be returned when `datetime.datetime.now()` falls below
+ `target_upper` and above `target_lower`.
Review comment:
What's the behaviour when backfilling or catching-up ?
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] github-actions[bot] commented on pull request #11964: Add new datetime branch operator
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #11964:
URL: https://github.com/apache/airflow/pull/11964#issuecomment-748679932
[The Workflow run](https://github.com/apache/airflow/actions/runs/434545527) 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*.
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] tomasfarias commented on a change in pull request #11964: Add new datetime branch operator
Posted by GitBox <gi...@apache.org>.
tomasfarias commented on a change in pull request #11964:
URL: https://github.com/apache/airflow/pull/11964#discussion_r527962973
##########
File path: airflow/operators/datetime_branch_operator.py
##########
@@ -0,0 +1,77 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import datetime
+from typing import Dict, Iterable, Optional, Union
+
+from airflow.exceptions import AirflowException
+from airflow.operators.branch_operator import BaseBranchOperator
+from airflow.utils import timezone
+from airflow.utils.decorators import apply_defaults
+
+
+class DateTimeBranchOperator(BaseBranchOperator):
+ """
+ Branches into one of two lists of tasks depending on the current datetime.
+
+ True branch will be returned when `datetime.datetime.now()` falls below
+ `target_upper` and above `target_lower`.
+
+ :param follow_task_ids_if_true: task id or task ids to follow if
+ `datetime.datetime.now()` falls above target_lower and below `target_upper`.
+ :type follow_task_ids_if_true: str or list[str]
+ :param follow_task_ids_if_false: task id or task ids to follow if
+ `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]
Review comment:
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__`.
##########
File path: airflow/operators/datetime_branch_operator.py
##########
@@ -0,0 +1,77 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import datetime
+from typing import Dict, Iterable, Optional, Union
+
+from airflow.exceptions import AirflowException
+from airflow.operators.branch_operator import BaseBranchOperator
+from airflow.utils import timezone
+from airflow.utils.decorators import apply_defaults
+
+
+class DateTimeBranchOperator(BaseBranchOperator):
+ """
+ Branches into one of two lists of tasks depending on the current datetime.
+
+ True branch will be returned when `datetime.datetime.now()` falls below
+ `target_upper` and above `target_lower`.
+
+ :param follow_task_ids_if_true: task id or task ids to follow if
+ `datetime.datetime.now()` falls above target_lower and below `target_upper`.
+ :type follow_task_ids_if_true: str or list[str]
+ :param follow_task_ids_if_false: task id or task ids to follow if
+ `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]
+ :param target_upper: target upper bound.
+ :type target_upper: Optional[datetime.datetime]
Review comment:
Answered above βοΈ
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] tomasfarias commented on a change in pull request #11964: Add new datetime branch operator
Posted by GitBox <gi...@apache.org>.
tomasfarias commented on a change in pull request #11964:
URL: https://github.com/apache/airflow/pull/11964#discussion_r527963197
##########
File path: tests/operators/test_datetime_branch_operator.py
##########
@@ -0,0 +1,281 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import datetime
+import unittest
+from unittest import mock
+
+from airflow.exceptions import AirflowException
+from airflow.models import DAG, DagRun, TaskInstance as TI
+from airflow.operators.datetime_branch_operator import DateTimeBranchOperator
+from airflow.operators.dummy_operator import DummyOperator
+from airflow.utils import timezone
+from airflow.utils.session import create_session
+from airflow.utils.state import State
+
+DEFAULT_DATE = timezone.datetime(2016, 1, 1)
+INTERVAL = datetime.timedelta(hours=12)
+
+
+class TestDateTimeBranchOperator(unittest.TestCase):
+ @classmethod
+ def setUpClass(cls):
+ super().setUpClass()
+
+ with create_session() as session:
+ session.query(DagRun).delete()
+ session.query(TI).delete()
+
+ def setUp(self):
+ self.dag = DAG(
+ 'datetime_branch_operator_test',
+ default_args={'owner': 'airflow', 'start_date': DEFAULT_DATE},
+ schedule_interval=INTERVAL,
+ )
+
+ self.branch_1 = DummyOperator(task_id='branch_1', dag=self.dag)
+ self.branch_2 = DummyOperator(task_id='branch_2', dag=self.dag)
+
+ def tearDown(self):
+ super().tearDown()
+
+ with create_session() as session:
+ session.query(DagRun).delete()
+ session.query(TI).delete()
+
+ def test_no_target_time(self):
+ """Check if DateTimeBranchOperator raises exception on missing target"""
+ with self.assertRaises(AirflowException):
+ DateTimeBranchOperator(
+ task_id='datetime_branch',
+ follow_task_ids_if_true='branch_1',
+ follow_task_ids_if_false='branch_2',
+ target_upper=None,
+ target_lower=None,
+ dag=self.dag,
+ )
+
+ @mock.patch('airflow.operators.datetime_branch_operator.timezone.utcnow')
+ def test_datetime_branch_operator_falls_within_range(self, mock_timezone):
+ """Check DateTimeBranchOperator branch operation"""
+ 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=datetime.datetime(2020, 7, 7, 10, 0, 0),
+ 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
+ )
+
+ mock_timezone.return_value = datetime.datetime(2020, 7, 7, 10, 54, 5, tzinfo=datetime.timezone.utc)
+ branch_op.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE)
+
+ tis = dr.get_task_instances()
+ for ti in tis:
+ if ti.task_id == 'datetime_branch':
+ self.assertEqual(ti.state, State.SUCCESS)
+ elif ti.task_id == 'branch_1':
+ self.assertEqual(ti.state, State.NONE)
+ elif ti.task_id == 'branch_2':
+ self.assertEqual(ti.state, State.SKIPPED)
+ else:
+ raise ValueError(f'Invalid task id {ti.task_id} found!')
+
+ @mock.patch('airflow.operators.datetime_branch_operator.timezone.utcnow')
+ def test_datetime_branch_operator_falls_outside_range(self, mock_timezone):
+ """Check DateTimeBranchOperator branch operation"""
+ 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=datetime.datetime(2020, 7, 7, 10, 0, 0),
+ 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
+ )
+
+ dates = [
+ datetime.datetime(2020, 7, 7, 12, 0, 0, tzinfo=datetime.timezone.utc),
+ datetime.datetime(2020, 6, 7, tzinfo=datetime.timezone.utc),
+ ]
+
+ for date in dates:
+ mock_timezone.return_value = date
+ branch_op.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE)
+
+ tis = dr.get_task_instances()
+ for ti in tis:
+ if ti.task_id == 'datetime_branch':
+ self.assertEqual(ti.state, State.SUCCESS)
+ elif ti.task_id == 'branch_1':
+ self.assertEqual(ti.state, State.SKIPPED)
+ elif ti.task_id == 'branch_2':
+ self.assertEqual(ti.state, State.NONE)
+ else:
+ raise ValueError(f'Invalid task id {ti.task_id} found!')
+
+ @mock.patch('airflow.operators.datetime_branch_operator.timezone.utcnow')
+ def test_datetime_branch_operator_upper_comparison_within_range(self, mock_timezone):
+ """Check DateTimeBranchOperator branch operation"""
+ 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
+ )
+
+ mock_timezone.return_value = datetime.datetime(2020, 7, 7, 10, 54, 5, tzinfo=datetime.timezone.utc)
+ branch_op.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE)
+
+ tis = dr.get_task_instances()
+ for ti in tis:
+ if ti.task_id == 'datetime_branch':
+ self.assertEqual(ti.state, State.SUCCESS)
+ elif ti.task_id == 'branch_1':
+ self.assertEqual(ti.state, State.NONE)
+ elif ti.task_id == 'branch_2':
+ self.assertEqual(ti.state, State.SKIPPED)
+ else:
+ raise ValueError(f'Invalid task id {ti.task_id} found!')
+
+ @mock.patch('airflow.operators.datetime_branch_operator.timezone.utcnow')
+ def test_datetime_branch_operator_lower_comparison_within_range(self, mock_timezone):
+ """Check DateTimeBranchOperator branch operation"""
+ branch_op = DateTimeBranchOperator(
+ task_id='datetime_branch',
+ follow_task_ids_if_true='branch_1',
+ follow_task_ids_if_false='branch_2',
+ target_upper=None,
+ target_lower=datetime.datetime(2020, 7, 7, 10, 0, 0),
+ 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
+ )
+
+ mock_timezone.return_value = datetime.datetime(2020, 7, 7, 10, 54, 5, tzinfo=datetime.timezone.utc)
+ branch_op.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE)
+
+ tis = dr.get_task_instances()
+ for ti in tis:
+ if ti.task_id == 'datetime_branch':
+ self.assertEqual(ti.state, State.SUCCESS)
+ elif ti.task_id == 'branch_1':
+ self.assertEqual(ti.state, State.NONE)
+ elif ti.task_id == 'branch_2':
+ self.assertEqual(ti.state, State.SKIPPED)
+ else:
+ raise ValueError(f'Invalid task id {ti.task_id} found!')
+
+ @mock.patch('airflow.operators.datetime_branch_operator.timezone.utcnow')
+ def test_datetime_branch_operator_upper_comparison_outside_range(self, mock_timezone):
+ """Check DateTimeBranchOperator branch operation"""
+ 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
+ )
Review comment:
Good suggestion! Moved boilerplate code over to `setUp`
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] tomasfarias commented on a change in pull request #11964: Add new datetime branch operator
Posted by GitBox <gi...@apache.org>.
tomasfarias commented on a change in pull request #11964:
URL: https://github.com/apache/airflow/pull/11964#discussion_r520992673
##########
File path: airflow/operators/datetime_branch_operator.py
##########
@@ -0,0 +1,80 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import datetime
+from typing import Dict, Iterable, Optional, Union
+
+from airflow.exceptions import AirflowException
+from airflow.operators.branch_operator import BaseBranchOperator
+from airflow.utils.decorators import apply_defaults
+
+
+class DateTimeBranchOperator(BaseBranchOperator):
+ """
+ Branches into one of two lists of tasks depending on the current datetime.
+
+ True branch will be returned when `datetime.datetime.now()` falls below
+ `target_upper` and above `target_lower`.
+
+ :param follow_task_ids_if_true: task id or task ids to follow if
+ `datetime.datetime.now()` falls above target_lower and below `target_upper`.
+ :type follow_task_ids_if_true: str or list[str]
+ :param follow_task_ids_if_false: task id or task ids to follow if
+ `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]
+ :param target_upper: target upper bound.
+ :type target_upper: Optional[datetime.datetime]
+ :param timezone: timezone to override local time.
+ :type timezone: Optional[datetime.timezone]
+ """
+
+ @apply_defaults
+ def __init__(
+ self,
+ *,
+ follow_task_ids_if_true: Union[str, Iterable[str]],
+ follow_task_ids_if_false: Union[str, Iterable[str]],
+ target_lower: Optional[datetime.datetime],
+ target_upper: Optional[datetime.datetime],
+ timezone: Optional[datetime.timezone] = None,
Review comment:
Thanks for pointing it out π Latest commit should replicate the work done for `TimeSensor`.
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] tomasfarias commented on pull request #11964: Add new datetime branch operator
Posted by GitBox <gi...@apache.org>.
tomasfarias commented on pull request #11964:
URL: https://github.com/apache/airflow/pull/11964#issuecomment-795988509
@eladkal Been a bit busy with work the last few weeks, but all comments should be addressed! Thanks for pinging me π
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] eladkal commented on a change in pull request #11964: Add new datetime branch operator
Posted by GitBox <gi...@apache.org>.
eladkal commented on a change in pull request #11964:
URL: https://github.com/apache/airflow/pull/11964#discussion_r528238842
##########
File path: airflow/operators/datetime_branch_operator.py
##########
@@ -0,0 +1,77 @@
+# Licensed to the Apache Software Foundation (ASF) under one
Review comment:
the suffix of `_operator` in the file name should be removed (to comply with AIP 21)
see https://github.com/apache/airflow/issues/11178
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] kaxil commented on a change in pull request #11964: Add new datetime branch operator
Posted by GitBox <gi...@apache.org>.
kaxil commented on a change in pull request #11964:
URL: https://github.com/apache/airflow/pull/11964#discussion_r580961797
##########
File path: airflow/operators/datetime_branch.py
##########
@@ -0,0 +1,101 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import datetime
+from typing import Dict, Iterable, Union
+
+from airflow.exceptions import AirflowException
+from airflow.operators.branch_operator import BaseBranchOperator
+from airflow.utils import timezone
+from airflow.utils.decorators import apply_defaults
+
+
+class DateTimeBranchOperator(BaseBranchOperator):
+ """
+ Branches into one of two lists of tasks depending on the current datetime.
+
+ True branch will be returned when `datetime.datetime.now()` falls below
+ `target_upper` and above `target_lower`.
Review comment:
How about adding something like:
https://github.com/apache/airflow/blob/e363ac1f33cd82188fb0923791f5cb747629d501/airflow/operators/weekday.py#L86-L89
cc @eladkal
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] tomasfarias commented on a change in pull request #11964: Add new datetime branch operator
Posted by GitBox <gi...@apache.org>.
tomasfarias commented on a change in pull request #11964:
URL: https://github.com/apache/airflow/pull/11964#discussion_r537015330
##########
File path: airflow/operators/datetime_branch.py
##########
@@ -0,0 +1,77 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import datetime
+from typing import Dict, Iterable, Optional, Union
+
+from airflow.exceptions import AirflowException
+from airflow.operators.branch_operator import BaseBranchOperator
+from airflow.utils import timezone
+from airflow.utils.decorators import apply_defaults
+
+
+class DateTimeBranchOperator(BaseBranchOperator):
+ """
+ Branches into one of two lists of tasks depending on the current datetime.
+
+ True branch will be returned when `datetime.datetime.now()` falls below
+ `target_upper` and above `target_lower`.
+
+ :param follow_task_ids_if_true: task id or task ids to follow if
+ `datetime.datetime.now()` falls above target_lower and below `target_upper`.
+ :type follow_task_ids_if_true: str or list[str]
+ :param follow_task_ids_if_false: task id or task ids to follow if
+ `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]
+ :param target_upper: target upper bound.
+ :type target_upper: Optional[datetime.datetime]
+ """
+
+ @apply_defaults
+ def __init__(
+ self,
+ *,
+ follow_task_ids_if_true: Union[str, Iterable[str]],
+ follow_task_ids_if_false: Union[str, Iterable[str]],
+ target_lower: Optional[datetime.datetime],
+ target_upper: Optional[datetime.datetime],
Review comment:
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.
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] github-actions[bot] commented on pull request #11964: Add new datetime branch operator
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #11964:
URL: https://github.com/apache/airflow/pull/11964#issuecomment-721458558
[The Workflow run](https://github.com/apache/airflow/actions/runs/344666665) 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*.
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] github-actions[bot] commented on pull request #11964: Add new datetime branch operator
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #11964:
URL: https://github.com/apache/airflow/pull/11964#issuecomment-721496197
[The Workflow run](https://github.com/apache/airflow/actions/runs/344817513) 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*.
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] tomasfarias commented on a change in pull request #11964: Add new datetime branch operator
Posted by GitBox <gi...@apache.org>.
tomasfarias commented on a change in pull request #11964:
URL: https://github.com/apache/airflow/pull/11964#discussion_r531244619
##########
File path: airflow/operators/datetime_branch.py
##########
@@ -0,0 +1,77 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import datetime
+from typing import Dict, Iterable, Optional, Union
+
+from airflow.exceptions import AirflowException
+from airflow.operators.branch_operator import BaseBranchOperator
+from airflow.utils import timezone
+from airflow.utils.decorators import apply_defaults
+
+
+class DateTimeBranchOperator(BaseBranchOperator):
+ """
+ Branches into one of two lists of tasks depending on the current datetime.
+
+ True branch will be returned when `datetime.datetime.now()` falls below
+ `target_upper` and above `target_lower`.
+
+ :param follow_task_ids_if_true: task id or task ids to follow if
+ `datetime.datetime.now()` falls above target_lower and below `target_upper`.
+ :type follow_task_ids_if_true: str or list[str]
+ :param follow_task_ids_if_false: task id or task ids to follow if
+ `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]
+ :param target_upper: target upper bound.
+ :type target_upper: Optional[datetime.datetime]
+ """
+
+ @apply_defaults
+ def __init__(
+ self,
+ *,
+ follow_task_ids_if_true: Union[str, Iterable[str]],
+ follow_task_ids_if_false: Union[str, Iterable[str]],
+ target_lower: Optional[datetime.datetime],
+ target_upper: Optional[datetime.datetime],
Review comment:
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)
)
```
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?
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] eladkal commented on pull request #11964: Add new datetime branch operator
Posted by GitBox <gi...@apache.org>.
eladkal commented on pull request #11964:
URL: https://github.com/apache/airflow/pull/11964#issuecomment-796674191
while commenting on https://github.com/apache/airflow/issues/11931 i notice that we might have two issues here:
1. The operator name `DateTimeBranchOperator` isn't consistent with other branch operators: `BranchPythonOperator`, `BranchSQLOperator`, `BranchDayOfWeekOperator` for consistnacy the operator should be renamed to `BranchDateTimeOperator`
2. 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.
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] tomasfarias commented on pull request #11964: Add new datetime branch operator
Posted by GitBox <gi...@apache.org>.
tomasfarias commented on pull request #11964:
URL: https://github.com/apache/airflow/pull/11964#issuecomment-796677397
@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.
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] kaxil commented on pull request #11964: Add new datetime branch operator
Posted by GitBox <gi...@apache.org>.
kaxil commented on pull request #11964:
URL: https://github.com/apache/airflow/pull/11964#issuecomment-796694915
Good point @eladkal -- agree to your suggestions
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] eladkal commented on a change in pull request #11964: Add new datetime branch operator
Posted by GitBox <gi...@apache.org>.
eladkal commented on a change in pull request #11964:
URL: https://github.com/apache/airflow/pull/11964#discussion_r543441471
##########
File path: airflow/example_dags/example_datetime_branch_operator.py
##########
@@ -0,0 +1,83 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""
+Example DAG demonstrating the usage of BranchPythonOperator with depends_on_past=True, where tasks may be run
+or skipped on alternating runs.
Review comment:
This is a description of another operator
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] kaxil merged pull request #11964: Add new datetime branch operator
Posted by GitBox <gi...@apache.org>.
kaxil merged pull request #11964:
URL: https://github.com/apache/airflow/pull/11964
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] ashb commented on a change in pull request #11964: Add new datetime branch operator
Posted by GitBox <gi...@apache.org>.
ashb commented on a change in pull request #11964:
URL: https://github.com/apache/airflow/pull/11964#discussion_r524372464
##########
File path: tests/operators/test_datetime_branch_operator.py
##########
@@ -0,0 +1,281 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import datetime
+import unittest
+from unittest import mock
+
+from airflow.exceptions import AirflowException
+from airflow.models import DAG, DagRun, TaskInstance as TI
+from airflow.operators.datetime_branch_operator import DateTimeBranchOperator
+from airflow.operators.dummy_operator import DummyOperator
+from airflow.utils import timezone
+from airflow.utils.session import create_session
+from airflow.utils.state import State
+
+DEFAULT_DATE = timezone.datetime(2016, 1, 1)
+INTERVAL = datetime.timedelta(hours=12)
+
+
+class TestDateTimeBranchOperator(unittest.TestCase):
+ @classmethod
+ def setUpClass(cls):
+ super().setUpClass()
+
+ with create_session() as session:
+ session.query(DagRun).delete()
+ session.query(TI).delete()
+
+ def setUp(self):
+ self.dag = DAG(
+ 'datetime_branch_operator_test',
+ default_args={'owner': 'airflow', 'start_date': DEFAULT_DATE},
+ schedule_interval=INTERVAL,
+ )
+
+ self.branch_1 = DummyOperator(task_id='branch_1', dag=self.dag)
+ self.branch_2 = DummyOperator(task_id='branch_2', dag=self.dag)
+
+ def tearDown(self):
+ super().tearDown()
+
+ with create_session() as session:
+ session.query(DagRun).delete()
+ session.query(TI).delete()
+
+ def test_no_target_time(self):
+ """Check if DateTimeBranchOperator raises exception on missing target"""
+ with self.assertRaises(AirflowException):
+ DateTimeBranchOperator(
+ task_id='datetime_branch',
+ follow_task_ids_if_true='branch_1',
+ follow_task_ids_if_false='branch_2',
+ target_upper=None,
+ target_lower=None,
+ dag=self.dag,
+ )
+
+ @mock.patch('airflow.operators.datetime_branch_operator.timezone.utcnow')
+ def test_datetime_branch_operator_falls_within_range(self, mock_timezone):
+ """Check DateTimeBranchOperator branch operation"""
+ 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=datetime.datetime(2020, 7, 7, 10, 0, 0),
+ 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
+ )
+
+ mock_timezone.return_value = datetime.datetime(2020, 7, 7, 10, 54, 5, tzinfo=datetime.timezone.utc)
+ branch_op.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE)
+
+ tis = dr.get_task_instances()
+ for ti in tis:
+ if ti.task_id == 'datetime_branch':
+ self.assertEqual(ti.state, State.SUCCESS)
+ elif ti.task_id == 'branch_1':
+ self.assertEqual(ti.state, State.NONE)
+ elif ti.task_id == 'branch_2':
+ self.assertEqual(ti.state, State.SKIPPED)
+ else:
+ raise ValueError(f'Invalid task id {ti.task_id} found!')
+
+ @mock.patch('airflow.operators.datetime_branch_operator.timezone.utcnow')
+ def test_datetime_branch_operator_falls_outside_range(self, mock_timezone):
+ """Check DateTimeBranchOperator branch operation"""
+ 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=datetime.datetime(2020, 7, 7, 10, 0, 0),
+ 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
+ )
+
+ dates = [
+ datetime.datetime(2020, 7, 7, 12, 0, 0, tzinfo=datetime.timezone.utc),
+ datetime.datetime(2020, 6, 7, tzinfo=datetime.timezone.utc),
+ ]
+
+ for date in dates:
+ mock_timezone.return_value = date
+ branch_op.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE)
+
+ tis = dr.get_task_instances()
+ for ti in tis:
+ if ti.task_id == 'datetime_branch':
+ self.assertEqual(ti.state, State.SUCCESS)
+ elif ti.task_id == 'branch_1':
+ self.assertEqual(ti.state, State.SKIPPED)
+ elif ti.task_id == 'branch_2':
+ self.assertEqual(ti.state, State.NONE)
+ else:
+ raise ValueError(f'Invalid task id {ti.task_id} found!')
+
+ @mock.patch('airflow.operators.datetime_branch_operator.timezone.utcnow')
+ def test_datetime_branch_operator_upper_comparison_within_range(self, mock_timezone):
+ """Check DateTimeBranchOperator branch operation"""
+ 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
+ )
+
+ mock_timezone.return_value = datetime.datetime(2020, 7, 7, 10, 54, 5, tzinfo=datetime.timezone.utc)
+ branch_op.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE)
+
+ tis = dr.get_task_instances()
+ for ti in tis:
+ if ti.task_id == 'datetime_branch':
+ self.assertEqual(ti.state, State.SUCCESS)
+ elif ti.task_id == 'branch_1':
+ self.assertEqual(ti.state, State.NONE)
+ elif ti.task_id == 'branch_2':
+ self.assertEqual(ti.state, State.SKIPPED)
+ else:
+ raise ValueError(f'Invalid task id {ti.task_id} found!')
+
+ @mock.patch('airflow.operators.datetime_branch_operator.timezone.utcnow')
+ def test_datetime_branch_operator_lower_comparison_within_range(self, mock_timezone):
+ """Check DateTimeBranchOperator branch operation"""
+ branch_op = DateTimeBranchOperator(
+ task_id='datetime_branch',
+ follow_task_ids_if_true='branch_1',
+ follow_task_ids_if_false='branch_2',
+ target_upper=None,
+ target_lower=datetime.datetime(2020, 7, 7, 10, 0, 0),
+ 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
+ )
+
+ mock_timezone.return_value = datetime.datetime(2020, 7, 7, 10, 54, 5, tzinfo=datetime.timezone.utc)
+ branch_op.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE)
+
+ tis = dr.get_task_instances()
+ for ti in tis:
+ if ti.task_id == 'datetime_branch':
+ self.assertEqual(ti.state, State.SUCCESS)
+ elif ti.task_id == 'branch_1':
+ self.assertEqual(ti.state, State.NONE)
+ elif ti.task_id == 'branch_2':
+ self.assertEqual(ti.state, State.SKIPPED)
+ else:
+ raise ValueError(f'Invalid task id {ti.task_id} found!')
+
+ @mock.patch('airflow.operators.datetime_branch_operator.timezone.utcnow')
+ def test_datetime_branch_operator_upper_comparison_outside_range(self, mock_timezone):
+ """Check DateTimeBranchOperator branch operation"""
+ 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
+ )
Review comment:
This appears in every test with almost indentical content -- could you create a helper function to reduce this boilerplate please.
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] tomasfarias commented on pull request #11964: Add new datetime branch operator
Posted by GitBox <gi...@apache.org>.
tomasfarias commented on pull request #11964:
URL: https://github.com/apache/airflow/pull/11964#issuecomment-796693453
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: https://github.com/tomasfarias/airflow/commit/5fc2862b5cddb1c427fc4727b09ce1490eb91c6a. I can create a PR as soon as the changes are confirmed as needed π
Great feedback!
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] kaxil commented on a change in pull request #11964: Add new datetime branch operator
Posted by GitBox <gi...@apache.org>.
kaxil commented on a change in pull request #11964:
URL: https://github.com/apache/airflow/pull/11964#discussion_r580961942
##########
File path: airflow/operators/datetime_branch.py
##########
@@ -0,0 +1,101 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import datetime
+from typing import Dict, Iterable, Union
+
+from airflow.exceptions import AirflowException
+from airflow.operators.branch_operator import BaseBranchOperator
+from airflow.utils import timezone
+from airflow.utils.decorators import apply_defaults
+
+
+class DateTimeBranchOperator(BaseBranchOperator):
+ """
+ Branches into one of two lists of tasks depending on the current datetime.
+
+ True branch will be returned when `datetime.datetime.now()` falls below
+ `target_upper` and above `target_lower`.
Review comment:
```suggestion
``target_upper`` and above ``target_lower``.
```
since it is rst
##########
File path: airflow/operators/datetime_branch.py
##########
@@ -0,0 +1,101 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import datetime
+from typing import Dict, Iterable, Union
+
+from airflow.exceptions import AirflowException
+from airflow.operators.branch_operator import BaseBranchOperator
+from airflow.utils import timezone
+from airflow.utils.decorators import apply_defaults
+
+
+class DateTimeBranchOperator(BaseBranchOperator):
+ """
+ Branches into one of two lists of tasks depending on the current datetime.
+
+ True branch will be returned when `datetime.datetime.now()` falls below
+ `target_upper` and above `target_lower`.
+
+ :param follow_task_ids_if_true: task id or task ids to follow if
+ `datetime.datetime.now()` falls above target_lower and below `target_upper`.
Review comment:
```suggestion
``datetime.datetime.now()`` falls above target_lower and below ``target_upper``.
```
##########
File path: airflow/operators/datetime_branch.py
##########
@@ -0,0 +1,101 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import datetime
+from typing import Dict, Iterable, Union
+
+from airflow.exceptions import AirflowException
+from airflow.operators.branch_operator import BaseBranchOperator
+from airflow.utils import timezone
+from airflow.utils.decorators import apply_defaults
+
+
+class DateTimeBranchOperator(BaseBranchOperator):
+ """
+ Branches into one of two lists of tasks depending on the current datetime.
+
+ True branch will be returned when `datetime.datetime.now()` falls below
+ `target_upper` and above `target_lower`.
+
+ :param follow_task_ids_if_true: task id or task ids to follow if
+ `datetime.datetime.now()` falls above target_lower and below `target_upper`.
+ :type follow_task_ids_if_true: str or list[str]
+ :param follow_task_ids_if_false: task id or task ids to follow if
+ `datetime.datetime.now()` falls below target_lower or above `target_upper`.
Review comment:
```suggestion
``datetime.datetime.now()`` falls below target_lower or above ``target_upper``.
```
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] tomasfarias commented on a change in pull request #11964: Add new datetime branch operator
Posted by GitBox <gi...@apache.org>.
tomasfarias commented on a change in pull request #11964:
URL: https://github.com/apache/airflow/pull/11964#discussion_r528210968
##########
File path: tests/operators/test_datetime_branch_operator.py
##########
@@ -0,0 +1,281 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import datetime
+import unittest
+from unittest import mock
+
+from airflow.exceptions import AirflowException
+from airflow.models import DAG, DagRun, TaskInstance as TI
+from airflow.operators.datetime_branch_operator import DateTimeBranchOperator
+from airflow.operators.dummy_operator import DummyOperator
+from airflow.utils import timezone
+from airflow.utils.session import create_session
+from airflow.utils.state import State
+
+DEFAULT_DATE = timezone.datetime(2016, 1, 1)
+INTERVAL = datetime.timedelta(hours=12)
+
+
+class TestDateTimeBranchOperator(unittest.TestCase):
+ @classmethod
+ def setUpClass(cls):
+ super().setUpClass()
+
+ with create_session() as session:
+ session.query(DagRun).delete()
+ session.query(TI).delete()
+
+ def setUp(self):
+ self.dag = DAG(
+ 'datetime_branch_operator_test',
+ default_args={'owner': 'airflow', 'start_date': DEFAULT_DATE},
+ schedule_interval=INTERVAL,
+ )
+
+ self.branch_1 = DummyOperator(task_id='branch_1', dag=self.dag)
+ self.branch_2 = DummyOperator(task_id='branch_2', dag=self.dag)
+
+ def tearDown(self):
+ super().tearDown()
+
+ with create_session() as session:
+ session.query(DagRun).delete()
+ session.query(TI).delete()
+
+ def test_no_target_time(self):
+ """Check if DateTimeBranchOperator raises exception on missing target"""
+ with self.assertRaises(AirflowException):
+ DateTimeBranchOperator(
+ task_id='datetime_branch',
+ follow_task_ids_if_true='branch_1',
+ follow_task_ids_if_false='branch_2',
+ target_upper=None,
+ target_lower=None,
+ dag=self.dag,
+ )
+
+ @mock.patch('airflow.operators.datetime_branch_operator.timezone.utcnow')
+ def test_datetime_branch_operator_falls_within_range(self, mock_timezone):
+ """Check DateTimeBranchOperator branch operation"""
+ 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=datetime.datetime(2020, 7, 7, 10, 0, 0),
+ 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
+ )
+
+ mock_timezone.return_value = datetime.datetime(2020, 7, 7, 10, 54, 5, tzinfo=datetime.timezone.utc)
+ branch_op.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE)
+
+ tis = dr.get_task_instances()
+ for ti in tis:
+ if ti.task_id == 'datetime_branch':
+ self.assertEqual(ti.state, State.SUCCESS)
+ elif ti.task_id == 'branch_1':
+ self.assertEqual(ti.state, State.NONE)
+ elif ti.task_id == 'branch_2':
+ self.assertEqual(ti.state, State.SKIPPED)
+ else:
+ raise ValueError(f'Invalid task id {ti.task_id} found!')
+
+ @mock.patch('airflow.operators.datetime_branch_operator.timezone.utcnow')
+ def test_datetime_branch_operator_falls_outside_range(self, mock_timezone):
+ """Check DateTimeBranchOperator branch operation"""
+ 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=datetime.datetime(2020, 7, 7, 10, 0, 0),
+ 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
+ )
+
+ dates = [
+ datetime.datetime(2020, 7, 7, 12, 0, 0, tzinfo=datetime.timezone.utc),
+ datetime.datetime(2020, 6, 7, tzinfo=datetime.timezone.utc),
+ ]
+
+ for date in dates:
+ mock_timezone.return_value = date
+ branch_op.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE)
+
+ tis = dr.get_task_instances()
+ for ti in tis:
+ if ti.task_id == 'datetime_branch':
+ self.assertEqual(ti.state, State.SUCCESS)
+ elif ti.task_id == 'branch_1':
+ self.assertEqual(ti.state, State.SKIPPED)
+ elif ti.task_id == 'branch_2':
+ self.assertEqual(ti.state, State.NONE)
+ else:
+ raise ValueError(f'Invalid task id {ti.task_id} found!')
+
+ @mock.patch('airflow.operators.datetime_branch_operator.timezone.utcnow')
+ def test_datetime_branch_operator_upper_comparison_within_range(self, mock_timezone):
+ """Check DateTimeBranchOperator branch operation"""
+ 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
+ )
+
+ mock_timezone.return_value = datetime.datetime(2020, 7, 7, 10, 54, 5, tzinfo=datetime.timezone.utc)
+ branch_op.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE)
+
+ tis = dr.get_task_instances()
+ for ti in tis:
+ if ti.task_id == 'datetime_branch':
+ self.assertEqual(ti.state, State.SUCCESS)
+ elif ti.task_id == 'branch_1':
+ self.assertEqual(ti.state, State.NONE)
+ elif ti.task_id == 'branch_2':
+ self.assertEqual(ti.state, State.SKIPPED)
+ else:
+ raise ValueError(f'Invalid task id {ti.task_id} found!')
+
+ @mock.patch('airflow.operators.datetime_branch_operator.timezone.utcnow')
+ def test_datetime_branch_operator_lower_comparison_within_range(self, mock_timezone):
+ """Check DateTimeBranchOperator branch operation"""
+ branch_op = DateTimeBranchOperator(
+ task_id='datetime_branch',
+ follow_task_ids_if_true='branch_1',
+ follow_task_ids_if_false='branch_2',
+ target_upper=None,
+ target_lower=datetime.datetime(2020, 7, 7, 10, 0, 0),
+ 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
+ )
+
+ mock_timezone.return_value = datetime.datetime(2020, 7, 7, 10, 54, 5, tzinfo=datetime.timezone.utc)
+ branch_op.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE)
+
+ tis = dr.get_task_instances()
+ for ti in tis:
+ if ti.task_id == 'datetime_branch':
+ self.assertEqual(ti.state, State.SUCCESS)
+ elif ti.task_id == 'branch_1':
+ self.assertEqual(ti.state, State.NONE)
+ elif ti.task_id == 'branch_2':
+ self.assertEqual(ti.state, State.SKIPPED)
+ else:
+ raise ValueError(f'Invalid task id {ti.task_id} found!')
+
+ @mock.patch('airflow.operators.datetime_branch_operator.timezone.utcnow')
+ def test_datetime_branch_operator_upper_comparison_outside_range(self, mock_timezone):
+ """Check DateTimeBranchOperator branch operation"""
+ 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
+ )
+
+ mock_timezone.return_value = datetime.datetime(2020, 7, 7, 12, 0, 0, tzinfo=datetime.timezone.utc)
+ branch_op.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE)
+
+ tis = dr.get_task_instances()
+ for ti in tis:
+ if ti.task_id == 'datetime_branch':
+ self.assertEqual(ti.state, State.SUCCESS)
+ elif ti.task_id == 'branch_1':
+ self.assertEqual(ti.state, State.SKIPPED)
+ elif ti.task_id == 'branch_2':
+ self.assertEqual(ti.state, State.NONE)
+ else:
+ raise ValueError(f'Invalid task id {ti.task_id} found!')
+
+ @mock.patch('airflow.operators.datetime_branch_operator.timezone.utcnow')
+ def test_datetime_branch_operator_lower_comparison_outside_range(self, mock_timezone):
Review comment:
Thanks for the suggestion! I've used `freezegun` instead.
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] tomasfarias commented on a change in pull request #11964: Add new datetime branch operator
Posted by GitBox <gi...@apache.org>.
tomasfarias commented on a change in pull request #11964:
URL: https://github.com/apache/airflow/pull/11964#discussion_r531244619
##########
File path: airflow/operators/datetime_branch.py
##########
@@ -0,0 +1,77 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import datetime
+from typing import Dict, Iterable, Optional, Union
+
+from airflow.exceptions import AirflowException
+from airflow.operators.branch_operator import BaseBranchOperator
+from airflow.utils import timezone
+from airflow.utils.decorators import apply_defaults
+
+
+class DateTimeBranchOperator(BaseBranchOperator):
+ """
+ Branches into one of two lists of tasks depending on the current datetime.
+
+ True branch will be returned when `datetime.datetime.now()` falls below
+ `target_upper` and above `target_lower`.
+
+ :param follow_task_ids_if_true: task id or task ids to follow if
+ `datetime.datetime.now()` falls above target_lower and below `target_upper`.
+ :type follow_task_ids_if_true: str or list[str]
+ :param follow_task_ids_if_false: task id or task ids to follow if
+ `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]
+ :param target_upper: target upper bound.
+ :type target_upper: Optional[datetime.datetime]
+ """
+
+ @apply_defaults
+ def __init__(
+ self,
+ *,
+ follow_task_ids_if_true: Union[str, Iterable[str]],
+ follow_task_ids_if_false: Union[str, Iterable[str]],
+ target_lower: Optional[datetime.datetime],
+ target_upper: Optional[datetime.datetime],
Review comment:
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?
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] eladkal commented on pull request #11964: Add new datetime branch operator
Posted by GitBox <gi...@apache.org>.
eladkal commented on pull request #11964:
URL: https://github.com/apache/airflow/pull/11964#issuecomment-794442086
@tomasfarias can you please address issues/questions raised by kaxil?
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] tomasfarias commented on pull request #11964: Add new datetime branch operator
Posted by GitBox <gi...@apache.org>.
tomasfarias commented on pull request #11964:
URL: https://github.com/apache/airflow/pull/11964#issuecomment-796715445
I opened a PR with the changes suggested, linking here for reference: https://github.com/apache/airflow/pull/14720
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] github-actions[bot] commented on pull request #11964: Add new datetime branch operator
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #11964:
URL: https://github.com/apache/airflow/pull/11964#issuecomment-721474221
[The Workflow run](https://github.com/apache/airflow/actions/runs/344729593) 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*.
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] kaxil commented on a change in pull request #11964: Add new datetime branch operator
Posted by GitBox <gi...@apache.org>.
kaxil commented on a change in pull request #11964:
URL: https://github.com/apache/airflow/pull/11964#discussion_r580962260
##########
File path: airflow/operators/datetime_branch.py
##########
@@ -0,0 +1,101 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import datetime
+from typing import Dict, Iterable, Union
+
+from airflow.exceptions import AirflowException
+from airflow.operators.branch_operator import BaseBranchOperator
+from airflow.utils import timezone
+from airflow.utils.decorators import apply_defaults
+
+
+class DateTimeBranchOperator(BaseBranchOperator):
+ """
+ Branches into one of two lists of tasks depending on the current datetime.
+
+ True branch will be returned when `datetime.datetime.now()` falls below
+ `target_upper` and above `target_lower`.
+
+ :param follow_task_ids_if_true: task id or task ids to follow if
+ `datetime.datetime.now()` falls above target_lower and below `target_upper`.
+ :type follow_task_ids_if_true: str or list[str]
+ :param follow_task_ids_if_false: task id or task ids to follow if
+ `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]
+ :param target_upper: target upper bound.
+ :type target_upper: Optional[datetime.datetime]
+ """
+
+ @apply_defaults
+ def __init__(
+ self,
+ *,
+ follow_task_ids_if_true: Union[str, Iterable[str]],
+ follow_task_ids_if_false: Union[str, Iterable[str]],
+ target_lower: Union[datetime.datetime, datetime.time, None],
+ target_upper: Union[datetime.datetime, datetime.time, None],
+ **kwargs,
+ ) -> None:
+ super().__init__(**kwargs)
+ if target_lower is None and target_upper is None:
+ raise AirflowException(
+ "Both target_upper and target_lower are None. At least one "
+ "must be defined to be compared to the current datetime"
+ )
+
+ self.target_lower = target_lower
+ self.target_upper = target_upper
+ self.follow_task_ids_if_true = follow_task_ids_if_true
+ 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)
Review comment:
Same comment as https://github.com/apache/airflow/pull/11964/files#r580961797
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] github-actions[bot] commented on pull request #11964: Add new datetime branch operator
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #11964:
URL: https://github.com/apache/airflow/pull/11964#issuecomment-796253918
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.
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] tomasfarias commented on a change in pull request #11964: Add new datetime branch operator
Posted by GitBox <gi...@apache.org>.
tomasfarias commented on a change in pull request #11964:
URL: https://github.com/apache/airflow/pull/11964#discussion_r591823237
##########
File path: airflow/operators/datetime_branch.py
##########
@@ -0,0 +1,101 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import datetime
+from typing import Dict, Iterable, Union
+
+from airflow.exceptions import AirflowException
+from airflow.operators.branch_operator import BaseBranchOperator
+from airflow.utils import timezone
+from airflow.utils.decorators import apply_defaults
+
+
+class DateTimeBranchOperator(BaseBranchOperator):
+ """
+ Branches into one of two lists of tasks depending on the current datetime.
+
+ True branch will be returned when `datetime.datetime.now()` falls below
+ `target_upper` and above `target_lower`.
+
+ :param follow_task_ids_if_true: task id or task ids to follow if
+ `datetime.datetime.now()` falls above target_lower and below `target_upper`.
+ :type follow_task_ids_if_true: str or list[str]
+ :param follow_task_ids_if_false: task id or task ids to follow if
+ `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]
+ :param target_upper: target upper bound.
+ :type target_upper: Optional[datetime.datetime]
+ """
+
+ @apply_defaults
+ def __init__(
+ self,
+ *,
+ follow_task_ids_if_true: Union[str, Iterable[str]],
+ follow_task_ids_if_false: Union[str, Iterable[str]],
+ target_lower: Union[datetime.datetime, datetime.time, None],
+ target_upper: Union[datetime.datetime, datetime.time, None],
+ **kwargs,
+ ) -> None:
+ super().__init__(**kwargs)
+ if target_lower is None and target_upper is None:
+ raise AirflowException(
+ "Both target_upper and target_lower are None. At least one "
+ "must be defined to be compared to the current datetime"
+ )
+
+ self.target_lower = target_lower
+ self.target_upper = target_upper
+ self.follow_task_ids_if_true = follow_task_ids_if_true
+ 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)
Review comment:
Implemented your suggestion (see the other comment).
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] tomasfarias commented on pull request #11964: Add new datetime branch operator
Posted by GitBox <gi...@apache.org>.
tomasfarias commented on pull request #11964:
URL: https://github.com/apache/airflow/pull/11964#issuecomment-721487330
@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 πͺ
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] tomasfarias commented on a change in pull request #11964: Add new datetime branch operator
Posted by GitBox <gi...@apache.org>.
tomasfarias commented on a change in pull request #11964:
URL: https://github.com/apache/airflow/pull/11964#discussion_r546430798
##########
File path: airflow/example_dags/example_datetime_branch_operator.py
##########
@@ -0,0 +1,83 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""
+Example DAG demonstrating the usage of BranchPythonOperator with depends_on_past=True, where tasks may be run
+or skipped on alternating runs.
Review comment:
Sorry about that, copy and paste blunder. Thanks for the review!
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] tomasfarias commented on pull request #11964: Add new datetime branch operator
Posted by GitBox <gi...@apache.org>.
tomasfarias commented on pull request #11964:
URL: https://github.com/apache/airflow/pull/11964#issuecomment-731434999
Not sure why MySQL build is failing, all tests appear to be passing according to logs.
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] ashb commented on a change in pull request #11964: Add new datetime branch operator
Posted by GitBox <gi...@apache.org>.
ashb commented on a change in pull request #11964:
URL: https://github.com/apache/airflow/pull/11964#discussion_r524373311
##########
File path: tests/operators/test_datetime_branch_operator.py
##########
@@ -0,0 +1,281 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import datetime
+import unittest
+from unittest import mock
+
+from airflow.exceptions import AirflowException
+from airflow.models import DAG, DagRun, TaskInstance as TI
+from airflow.operators.datetime_branch_operator import DateTimeBranchOperator
+from airflow.operators.dummy_operator import DummyOperator
+from airflow.utils import timezone
+from airflow.utils.session import create_session
+from airflow.utils.state import State
+
+DEFAULT_DATE = timezone.datetime(2016, 1, 1)
+INTERVAL = datetime.timedelta(hours=12)
+
+
+class TestDateTimeBranchOperator(unittest.TestCase):
+ @classmethod
+ def setUpClass(cls):
+ super().setUpClass()
+
+ with create_session() as session:
+ session.query(DagRun).delete()
+ session.query(TI).delete()
+
+ def setUp(self):
+ self.dag = DAG(
+ 'datetime_branch_operator_test',
+ default_args={'owner': 'airflow', 'start_date': DEFAULT_DATE},
+ schedule_interval=INTERVAL,
+ )
+
+ self.branch_1 = DummyOperator(task_id='branch_1', dag=self.dag)
+ self.branch_2 = DummyOperator(task_id='branch_2', dag=self.dag)
+
+ def tearDown(self):
+ super().tearDown()
+
+ with create_session() as session:
+ session.query(DagRun).delete()
+ session.query(TI).delete()
+
+ def test_no_target_time(self):
+ """Check if DateTimeBranchOperator raises exception on missing target"""
+ with self.assertRaises(AirflowException):
+ DateTimeBranchOperator(
+ task_id='datetime_branch',
+ follow_task_ids_if_true='branch_1',
+ follow_task_ids_if_false='branch_2',
+ target_upper=None,
+ target_lower=None,
+ dag=self.dag,
+ )
+
+ @mock.patch('airflow.operators.datetime_branch_operator.timezone.utcnow')
+ def test_datetime_branch_operator_falls_within_range(self, mock_timezone):
+ """Check DateTimeBranchOperator branch operation"""
+ 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=datetime.datetime(2020, 7, 7, 10, 0, 0),
+ 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
+ )
+
+ mock_timezone.return_value = datetime.datetime(2020, 7, 7, 10, 54, 5, tzinfo=datetime.timezone.utc)
+ branch_op.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE)
+
+ tis = dr.get_task_instances()
+ for ti in tis:
+ if ti.task_id == 'datetime_branch':
+ self.assertEqual(ti.state, State.SUCCESS)
+ elif ti.task_id == 'branch_1':
+ self.assertEqual(ti.state, State.NONE)
+ elif ti.task_id == 'branch_2':
+ self.assertEqual(ti.state, State.SKIPPED)
+ else:
+ raise ValueError(f'Invalid task id {ti.task_id} found!')
+
+ @mock.patch('airflow.operators.datetime_branch_operator.timezone.utcnow')
+ def test_datetime_branch_operator_falls_outside_range(self, mock_timezone):
+ """Check DateTimeBranchOperator branch operation"""
+ 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=datetime.datetime(2020, 7, 7, 10, 0, 0),
+ 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
+ )
+
+ dates = [
+ datetime.datetime(2020, 7, 7, 12, 0, 0, tzinfo=datetime.timezone.utc),
+ datetime.datetime(2020, 6, 7, tzinfo=datetime.timezone.utc),
+ ]
+
+ for date in dates:
+ mock_timezone.return_value = date
+ branch_op.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE)
+
+ tis = dr.get_task_instances()
+ for ti in tis:
+ if ti.task_id == 'datetime_branch':
+ self.assertEqual(ti.state, State.SUCCESS)
+ elif ti.task_id == 'branch_1':
+ self.assertEqual(ti.state, State.SKIPPED)
+ elif ti.task_id == 'branch_2':
+ self.assertEqual(ti.state, State.NONE)
+ else:
+ raise ValueError(f'Invalid task id {ti.task_id} found!')
+
+ @mock.patch('airflow.operators.datetime_branch_operator.timezone.utcnow')
+ def test_datetime_branch_operator_upper_comparison_within_range(self, mock_timezone):
+ """Check DateTimeBranchOperator branch operation"""
+ 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
+ )
+
+ mock_timezone.return_value = datetime.datetime(2020, 7, 7, 10, 54, 5, tzinfo=datetime.timezone.utc)
+ branch_op.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE)
+
+ tis = dr.get_task_instances()
+ for ti in tis:
+ if ti.task_id == 'datetime_branch':
+ self.assertEqual(ti.state, State.SUCCESS)
+ elif ti.task_id == 'branch_1':
+ self.assertEqual(ti.state, State.NONE)
+ elif ti.task_id == 'branch_2':
+ self.assertEqual(ti.state, State.SKIPPED)
+ else:
+ raise ValueError(f'Invalid task id {ti.task_id} found!')
+
+ @mock.patch('airflow.operators.datetime_branch_operator.timezone.utcnow')
+ def test_datetime_branch_operator_lower_comparison_within_range(self, mock_timezone):
+ """Check DateTimeBranchOperator branch operation"""
+ branch_op = DateTimeBranchOperator(
+ task_id='datetime_branch',
+ follow_task_ids_if_true='branch_1',
+ follow_task_ids_if_false='branch_2',
+ target_upper=None,
+ target_lower=datetime.datetime(2020, 7, 7, 10, 0, 0),
+ 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
+ )
+
+ mock_timezone.return_value = datetime.datetime(2020, 7, 7, 10, 54, 5, tzinfo=datetime.timezone.utc)
+ branch_op.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE)
+
+ tis = dr.get_task_instances()
+ for ti in tis:
+ if ti.task_id == 'datetime_branch':
+ self.assertEqual(ti.state, State.SUCCESS)
+ elif ti.task_id == 'branch_1':
+ self.assertEqual(ti.state, State.NONE)
+ elif ti.task_id == 'branch_2':
+ self.assertEqual(ti.state, State.SKIPPED)
+ else:
+ raise ValueError(f'Invalid task id {ti.task_id} found!')
+
+ @mock.patch('airflow.operators.datetime_branch_operator.timezone.utcnow')
+ def test_datetime_branch_operator_upper_comparison_outside_range(self, mock_timezone):
+ """Check DateTimeBranchOperator branch operation"""
+ 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
+ )
+
+ mock_timezone.return_value = datetime.datetime(2020, 7, 7, 12, 0, 0, tzinfo=datetime.timezone.utc)
+ branch_op.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE)
+
+ tis = dr.get_task_instances()
+ for ti in tis:
+ if ti.task_id == 'datetime_branch':
+ self.assertEqual(ti.state, State.SUCCESS)
+ elif ti.task_id == 'branch_1':
+ self.assertEqual(ti.state, State.SKIPPED)
+ elif ti.task_id == 'branch_2':
+ self.assertEqual(ti.state, State.NONE)
+ else:
+ raise ValueError(f'Invalid task id {ti.task_id} found!')
+
+ @mock.patch('airflow.operators.datetime_branch_operator.timezone.utcnow')
+ def test_datetime_branch_operator_lower_comparison_outside_range(self, mock_timezone):
Review comment:
Rather than mocking the datetime like this, please use `freezegun`
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] kaxil commented on pull request #11964: Add new datetime branch operator
Posted by GitBox <gi...@apache.org>.
kaxil commented on pull request #11964:
URL: https://github.com/apache/airflow/pull/11964#issuecomment-721440649
Can you please rebased your PR on latest Master since we have applied [Black](https://github.com/apache/airflow/commit/4e8f9cc8d02b29c325b8a5a76b4837671bdf5f68) and [PyUpgrade](https://github.com/apache/airflow/commit/8c42cf1b00c90f0d7f11b8a3a455381de8e003c5) on Master.
It will help if your squash your commits into single commit first so that there are less conflicts.
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] eladkal commented on a change in pull request #11964: Add new datetime branch operator
Posted by GitBox <gi...@apache.org>.
eladkal commented on a change in pull request #11964:
URL: https://github.com/apache/airflow/pull/11964#discussion_r532194095
##########
File path: airflow/operators/datetime_branch.py
##########
@@ -0,0 +1,77 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import datetime
+from typing import Dict, Iterable, Optional, Union
+
+from airflow.exceptions import AirflowException
+from airflow.operators.branch_operator import BaseBranchOperator
+from airflow.utils import timezone
+from airflow.utils.decorators import apply_defaults
+
+
+class DateTimeBranchOperator(BaseBranchOperator):
+ """
+ Branches into one of two lists of tasks depending on the current datetime.
+
+ True branch will be returned when `datetime.datetime.now()` falls below
+ `target_upper` and above `target_lower`.
+
+ :param follow_task_ids_if_true: task id or task ids to follow if
+ `datetime.datetime.now()` falls above target_lower and below `target_upper`.
+ :type follow_task_ids_if_true: str or list[str]
+ :param follow_task_ids_if_false: task id or task ids to follow if
+ `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]
+ :param target_upper: target upper bound.
+ :type target_upper: Optional[datetime.datetime]
+ """
+
+ @apply_defaults
+ def __init__(
+ self,
+ *,
+ follow_task_ids_if_true: Union[str, Iterable[str]],
+ follow_task_ids_if_false: Union[str, Iterable[str]],
+ target_lower: Optional[datetime.datetime],
+ target_upper: Optional[datetime.datetime],
Review comment:
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](https://github.com/apache/airflow/tree/master/docs/howto/operator) about this operator. you can use previous PRs to see examples https://github.com/apache/airflow/pull/11472
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] RosterIn commented on a change in pull request #11964: Add new datetime branch operator
Posted by GitBox <gi...@apache.org>.
RosterIn commented on a change in pull request #11964:
URL: https://github.com/apache/airflow/pull/11964#discussion_r524339719
##########
File path: airflow/operators/datetime_branch_operator.py
##########
@@ -0,0 +1,77 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import datetime
+from typing import Dict, Iterable, Optional, Union
+
+from airflow.exceptions import AirflowException
+from airflow.operators.branch_operator import BaseBranchOperator
+from airflow.utils import timezone
+from airflow.utils.decorators import apply_defaults
+
+
+class DateTimeBranchOperator(BaseBranchOperator):
+ """
+ Branches into one of two lists of tasks depending on the current datetime.
+
+ True branch will be returned when `datetime.datetime.now()` falls below
+ `target_upper` and above `target_lower`.
+
+ :param follow_task_ids_if_true: task id or task ids to follow if
+ `datetime.datetime.now()` falls above target_lower and below `target_upper`.
+ :type follow_task_ids_if_true: str or list[str]
+ :param follow_task_ids_if_false: task id or task ids to follow if
+ `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]
Review comment:
why this optional?
##########
File path: airflow/operators/datetime_branch_operator.py
##########
@@ -0,0 +1,77 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import datetime
+from typing import Dict, Iterable, Optional, Union
+
+from airflow.exceptions import AirflowException
+from airflow.operators.branch_operator import BaseBranchOperator
+from airflow.utils import timezone
+from airflow.utils.decorators import apply_defaults
+
+
+class DateTimeBranchOperator(BaseBranchOperator):
+ """
+ Branches into one of two lists of tasks depending on the current datetime.
+
+ True branch will be returned when `datetime.datetime.now()` falls below
+ `target_upper` and above `target_lower`.
+
+ :param follow_task_ids_if_true: task id or task ids to follow if
+ `datetime.datetime.now()` falls above target_lower and below `target_upper`.
+ :type follow_task_ids_if_true: str or list[str]
+ :param follow_task_ids_if_false: task id or task ids to follow if
+ `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]
+ :param target_upper: target upper bound.
+ :type target_upper: Optional[datetime.datetime]
+ """
+
+ @apply_defaults
+ def __init__(
+ self,
+ *,
+ follow_task_ids_if_true: Union[str, Iterable[str]],
+ follow_task_ids_if_false: Union[str, Iterable[str]],
+ target_lower: Optional[datetime.datetime],
+ target_upper: Optional[datetime.datetime],
+ **kwargs,
+ ) -> None:
+ super().__init__(**kwargs)
+ if target_lower is None and target_upper is None:
+ raise AirflowException(
+ "Both target_upper and target_lower are None. At least one "
+ "must be defined to be compared to the current datetime"
+ )
+
+ self.target_lower = target_lower
+ self.target_upper = target_upper
+ self.follow_task_ids_if_true = follow_task_ids_if_true
+ 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)
Review comment:
is https://github.com/apache/airflow/pull/12353 effect this as well?
##########
File path: airflow/operators/datetime_branch_operator.py
##########
@@ -0,0 +1,77 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import datetime
+from typing import Dict, Iterable, Optional, Union
+
+from airflow.exceptions import AirflowException
+from airflow.operators.branch_operator import BaseBranchOperator
+from airflow.utils import timezone
+from airflow.utils.decorators import apply_defaults
+
+
+class DateTimeBranchOperator(BaseBranchOperator):
+ """
+ Branches into one of two lists of tasks depending on the current datetime.
+
+ True branch will be returned when `datetime.datetime.now()` falls below
+ `target_upper` and above `target_lower`.
+
+ :param follow_task_ids_if_true: task id or task ids to follow if
+ `datetime.datetime.now()` falls above target_lower and below `target_upper`.
+ :type follow_task_ids_if_true: str or list[str]
+ :param follow_task_ids_if_false: task id or task ids to follow if
+ `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]
+ :param target_upper: target upper bound.
+ :type target_upper: Optional[datetime.datetime]
Review comment:
same question here
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] eladkal commented on a change in pull request #11964: Add new datetime branch operator
Posted by GitBox <gi...@apache.org>.
eladkal commented on a change in pull request #11964:
URL: https://github.com/apache/airflow/pull/11964#discussion_r520565235
##########
File path: airflow/operators/datetime_branch_operator.py
##########
@@ -0,0 +1,80 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import datetime
+from typing import Dict, Iterable, Optional, Union
+
+from airflow.exceptions import AirflowException
+from airflow.operators.branch_operator import BaseBranchOperator
+from airflow.utils.decorators import apply_defaults
+
+
+class DateTimeBranchOperator(BaseBranchOperator):
+ """
+ Branches into one of two lists of tasks depending on the current datetime.
+
+ True branch will be returned when `datetime.datetime.now()` falls below
+ `target_upper` and above `target_lower`.
+
+ :param follow_task_ids_if_true: task id or task ids to follow if
+ `datetime.datetime.now()` falls above target_lower and below `target_upper`.
+ :type follow_task_ids_if_true: str or list[str]
+ :param follow_task_ids_if_false: task id or task ids to follow if
+ `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]
+ :param target_upper: target upper bound.
+ :type target_upper: Optional[datetime.datetime]
+ :param timezone: timezone to override local time.
+ :type timezone: Optional[datetime.timezone]
+ """
+
+ @apply_defaults
+ def __init__(
+ self,
+ *,
+ follow_task_ids_if_true: Union[str, Iterable[str]],
+ follow_task_ids_if_false: Union[str, Iterable[str]],
+ target_lower: Optional[datetime.datetime],
+ target_upper: Optional[datetime.datetime],
+ timezone: Optional[datetime.timezone] = None,
Review comment:
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 https://github.com/apache/airflow/pull/9882 https://github.com/apache/airflow/pull/9699
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] eladkal commented on a change in pull request #11964: Add new datetime branch operator
Posted by GitBox <gi...@apache.org>.
eladkal commented on a change in pull request #11964:
URL: https://github.com/apache/airflow/pull/11964#discussion_r531091730
##########
File path: tests/operators/test_datetime_branch_operator.py
##########
@@ -0,0 +1,195 @@
+#
Review comment:
please also remove the` _operator` suffix from the test file
##########
File path: airflow/operators/datetime_branch.py
##########
@@ -0,0 +1,77 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import datetime
+from typing import Dict, Iterable, Optional, Union
+
+from airflow.exceptions import AirflowException
+from airflow.operators.branch_operator import BaseBranchOperator
+from airflow.utils import timezone
+from airflow.utils.decorators import apply_defaults
+
+
+class DateTimeBranchOperator(BaseBranchOperator):
+ """
+ Branches into one of two lists of tasks depending on the current datetime.
+
+ True branch will be returned when `datetime.datetime.now()` falls below
+ `target_upper` and above `target_lower`.
+
+ :param follow_task_ids_if_true: task id or task ids to follow if
+ `datetime.datetime.now()` falls above target_lower and below `target_upper`.
+ :type follow_task_ids_if_true: str or list[str]
+ :param follow_task_ids_if_false: task id or task ids to follow if
+ `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]
+ :param target_upper: target upper bound.
+ :type target_upper: Optional[datetime.datetime]
+ """
+
+ @apply_defaults
+ def __init__(
+ self,
+ *,
+ follow_task_ids_if_true: Union[str, Iterable[str]],
+ follow_task_ids_if_false: Union[str, Iterable[str]],
+ target_lower: Optional[datetime.datetime],
+ target_upper: Optional[datetime.datetime],
Review comment:
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?
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] eladkal commented on a change in pull request #11964: Add new datetime branch operator
Posted by GitBox <gi...@apache.org>.
eladkal commented on a change in pull request #11964:
URL: https://github.com/apache/airflow/pull/11964#discussion_r521369948
##########
File path: airflow/operators/datetime_branch_operator.py
##########
@@ -0,0 +1,80 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import datetime
+from typing import Dict, Iterable, Optional, Union
+
+from airflow.exceptions import AirflowException
+from airflow.operators.branch_operator import BaseBranchOperator
+from airflow.utils.decorators import apply_defaults
+
+
+class DateTimeBranchOperator(BaseBranchOperator):
+ """
+ Branches into one of two lists of tasks depending on the current datetime.
+
+ True branch will be returned when `datetime.datetime.now()` falls below
+ `target_upper` and above `target_lower`.
+
+ :param follow_task_ids_if_true: task id or task ids to follow if
+ `datetime.datetime.now()` falls above target_lower and below `target_upper`.
+ :type follow_task_ids_if_true: str or list[str]
+ :param follow_task_ids_if_false: task id or task ids to follow if
+ `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]
+ :param target_upper: target upper bound.
+ :type target_upper: Optional[datetime.datetime]
+ :param timezone: timezone to override local time.
+ :type timezone: Optional[datetime.timezone]
+ """
+
+ @apply_defaults
+ def __init__(
+ self,
+ *,
+ follow_task_ids_if_true: Union[str, Iterable[str]],
+ follow_task_ids_if_false: Union[str, Iterable[str]],
+ target_lower: Optional[datetime.datetime],
+ target_upper: Optional[datetime.datetime],
+ timezone: Optional[datetime.timezone] = None,
Review comment:
can you add unit tests ?
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] tomasfarias commented on pull request #11964: Add new datetime branch operator
Posted by GitBox <gi...@apache.org>.
tomasfarias commented on pull request #11964:
URL: https://github.com/apache/airflow/pull/11964#issuecomment-755018681
@eladkal Thanks for the review, I've rebased the branch as instructed. I apologize for the delay, I took a couple of weeks of vacation for the holidays, but I'm now back in full force.
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] github-actions[bot] commented on pull request #11964: Add new datetime branch operator
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #11964:
URL: https://github.com/apache/airflow/pull/11964#issuecomment-739439758
[The Workflow run](https://github.com/apache/airflow/actions/runs/403294483) 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*.
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org