From cb8d6a9bf22c9640d2560223f74490932d424cdf Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tom=C3=A1s=20Far=C3=ADas=20Santana?= Date: Fri, 30 Oct 2020 09:33:40 +0100 Subject: [PATCH 01/15] Add new datetime branch operator --- airflow/operators/datetime_branch_operator.py | 81 +++++++++++++++++++ 1 file changed, 81 insertions(+) create mode 100644 airflow/operators/datetime_branch_operator.py diff --git a/airflow/operators/datetime_branch_operator.py b/airflow/operators/datetime_branch_operator.py new file mode 100644 index 0000000000000..1906b5a31cd24 --- /dev/null +++ b/airflow/operators/datetime_branch_operator.py @@ -0,0 +1,81 @@ +# 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 operator +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, + **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.timezone = timezone + 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 = datetime.datetime.now(self.timezone) + + if self.target_upper is not None and self.target_upper < now: + return self.follow_task_ids_if_false + + if self.target_lower is not None and self.target_lower > now: + return self.follow_task_ids_if_false + + return self.follow_task_ids_if_true From d0175b6252d3e9081d40d55b284cad50addef487 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tom=C3=A1s=20Far=C3=ADas=20Santana?= Date: Wed, 4 Nov 2020 09:52:09 +0100 Subject: [PATCH 02/15] Remove unused import --- airflow/operators/datetime_branch_operator.py | 1 - 1 file changed, 1 deletion(-) diff --git a/airflow/operators/datetime_branch_operator.py b/airflow/operators/datetime_branch_operator.py index 1906b5a31cd24..65db98d3c26ef 100644 --- a/airflow/operators/datetime_branch_operator.py +++ b/airflow/operators/datetime_branch_operator.py @@ -16,7 +16,6 @@ # under the License. import datetime -import operator from typing import Dict, Iterable, Optional, Union from airflow.exceptions import AirflowException From afab35d4ed993e9773338e7b4e99e60e4698e0b3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tom=C3=A1s=20Far=C3=ADas=20Santana?= Date: Wed, 11 Nov 2020 01:43:00 +0100 Subject: [PATCH 03/15] Consider timezone-aware DAGs and default_timezone config --- airflow/operators/datetime_branch_operator.py | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/airflow/operators/datetime_branch_operator.py b/airflow/operators/datetime_branch_operator.py index 65db98d3c26ef..31f92a1b287d1 100644 --- a/airflow/operators/datetime_branch_operator.py +++ b/airflow/operators/datetime_branch_operator.py @@ -20,6 +20,7 @@ from airflow.exceptions import AirflowException from airflow.operators.branch_operator import BaseBranchOperator +from airflow.utils import timezone from airflow.utils.decorators import apply_defaults @@ -40,8 +41,6 @@ class DateTimeBranchOperator(BaseBranchOperator): :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 @@ -52,7 +51,6 @@ def __init__( follow_task_ids_if_false: Union[str, Iterable[str]], target_lower: Optional[datetime.datetime], target_upper: Optional[datetime.datetime], - timezone: Optional[datetime.timezone] = None, **kwargs, ) -> None: super().__init__(**kwargs) @@ -64,12 +62,11 @@ def __init__( self.target_lower = target_lower self.target_upper = target_upper - self.timezone = timezone 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 = datetime.datetime.now(self.timezone) + now = timezone.make_naive(timezone.utcnow(), self.dag.timezone) if self.target_upper is not None and self.target_upper < now: return self.follow_task_ids_if_false From 7b2ef5314c6574eb16cd7969b13009cae75dd36e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tom=C3=A1s=20Far=C3=ADas=20Santana?= Date: Wed, 11 Nov 2020 18:13:07 +0100 Subject: [PATCH 04/15] Add unit tests for DateTimeBranchOperator --- .../test_datetime_branch_operator.py | 281 ++++++++++++++++++ 1 file changed, 281 insertions(+) create mode 100644 tests/operators/test_datetime_branch_operator.py diff --git a/tests/operators/test_datetime_branch_operator.py b/tests/operators/test_datetime_branch_operator.py new file mode 100644 index 0000000000000..ac674c37cafbe --- /dev/null +++ b/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): + """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, 9, 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!') From 144ce663eaeb442b24fcf358d6467878ac9dd0f6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tom=C3=A1s=20Far=C3=ADas=20Santana?= Date: Fri, 20 Nov 2020 21:28:49 +0100 Subject: [PATCH 05/15] Include test operator and test dagrun in setUp function --- .../test_datetime_branch_operator.py | 143 ++++-------------- 1 file changed, 29 insertions(+), 114 deletions(-) diff --git a/tests/operators/test_datetime_branch_operator.py b/tests/operators/test_datetime_branch_operator.py index ac674c37cafbe..2db505892f8bb 100644 --- a/tests/operators/test_datetime_branch_operator.py +++ b/tests/operators/test_datetime_branch_operator.py @@ -51,6 +51,23 @@ def setUp(self): self.branch_1 = DummyOperator(task_id='branch_1', dag=self.dag) self.branch_2 = DummyOperator(task_id='branch_2', dag=self.dag) + self.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(self.branch_op) + self.branch_2.set_upstream(self.branch_op) + self.dag.clear() + + self.dr = self.dag.create_dagrun( + run_id='manual__', start_date=DEFAULT_DATE, execution_date=DEFAULT_DATE, state=State.RUNNING + ) + def tearDown(self): super().tearDown() @@ -73,27 +90,10 @@ def test_no_target_time(self): @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) + self.branch_op.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE) - tis = dr.get_task_instances() + tis = self.dr.get_task_instances() for ti in tis: if ti.task_id == 'datetime_branch': self.assertEqual(ti.state, State.SUCCESS) @@ -107,23 +107,6 @@ def test_datetime_branch_operator_falls_within_range(self, mock_timezone): @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), @@ -131,9 +114,9 @@ def test_datetime_branch_operator_falls_outside_range(self, mock_timezone): for date in dates: mock_timezone.return_value = date - branch_op.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE) + self.branch_op.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE) - tis = dr.get_task_instances() + tis = self.dr.get_task_instances() for ti in tis: if ti.task_id == 'datetime_branch': self.assertEqual(ti.state, State.SUCCESS) @@ -147,27 +130,10 @@ def test_datetime_branch_operator_falls_outside_range(self, mock_timezone): @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) + self.branch_op.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE) - tis = dr.get_task_instances() + tis = self.dr.get_task_instances() for ti in tis: if ti.task_id == 'datetime_branch': self.assertEqual(ti.state, State.SUCCESS) @@ -181,27 +147,10 @@ def test_datetime_branch_operator_upper_comparison_within_range(self, mock_timez @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) + self.branch_op.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE) - tis = dr.get_task_instances() + tis = self.dr.get_task_instances() for ti in tis: if ti.task_id == 'datetime_branch': self.assertEqual(ti.state, State.SUCCESS) @@ -215,27 +164,10 @@ def test_datetime_branch_operator_lower_comparison_within_range(self, mock_timez @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) + self.branch_op.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE) - tis = dr.get_task_instances() + tis = self.dr.get_task_instances() for ti in tis: if ti.task_id == 'datetime_branch': self.assertEqual(ti.state, State.SUCCESS) @@ -249,27 +181,10 @@ def test_datetime_branch_operator_upper_comparison_outside_range(self, mock_time @mock.patch('airflow.operators.datetime_branch_operator.timezone.utcnow') def test_datetime_branch_operator_lower_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=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, 9, 0, 0, tzinfo=datetime.timezone.utc) - branch_op.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE) + self.branch_op.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE) - tis = dr.get_task_instances() + tis = self.dr.get_task_instances() for ti in tis: if ti.task_id == 'datetime_branch': self.assertEqual(ti.state, State.SUCCESS) From ce2ed0378d0aacdb741d76aaf7597240f5347808 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tom=C3=A1s=20Far=C3=ADas=20Santana?= Date: Fri, 20 Nov 2020 21:43:03 +0100 Subject: [PATCH 06/15] Use freezegun instead of mocking dates --- .../test_datetime_branch_operator.py | 39 +++++++++---------- 1 file changed, 19 insertions(+), 20 deletions(-) diff --git a/tests/operators/test_datetime_branch_operator.py b/tests/operators/test_datetime_branch_operator.py index 2db505892f8bb..5e20bc638dcb3 100644 --- a/tests/operators/test_datetime_branch_operator.py +++ b/tests/operators/test_datetime_branch_operator.py @@ -18,7 +18,8 @@ import datetime import unittest -from unittest import mock + +import freezegun from airflow.exceptions import AirflowException from airflow.models import DAG, DagRun, TaskInstance as TI @@ -87,10 +88,9 @@ def test_no_target_time(self): dag=self.dag, ) - @mock.patch('airflow.operators.datetime_branch_operator.timezone.utcnow') - def test_datetime_branch_operator_falls_within_range(self, mock_timezone): + @freezegun.freeze_time("2020-07-07 10:54:05") + def test_datetime_branch_operator_falls_within_range(self): """Check DateTimeBranchOperator branch operation""" - mock_timezone.return_value = datetime.datetime(2020, 7, 7, 10, 54, 5, tzinfo=datetime.timezone.utc) self.branch_op.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE) tis = self.dr.get_task_instances() @@ -104,8 +104,7 @@ def test_datetime_branch_operator_falls_within_range(self, mock_timezone): 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): + def test_datetime_branch_operator_falls_outside_range(self): """Check DateTimeBranchOperator branch operation""" dates = [ datetime.datetime(2020, 7, 7, 12, 0, 0, tzinfo=datetime.timezone.utc), @@ -113,8 +112,8 @@ def test_datetime_branch_operator_falls_outside_range(self, mock_timezone): ] for date in dates: - mock_timezone.return_value = date - self.branch_op.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE) + with freezegun.freeze_time(date): + self.branch_op.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE) tis = self.dr.get_task_instances() for ti in tis: @@ -127,10 +126,10 @@ def test_datetime_branch_operator_falls_outside_range(self, mock_timezone): 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): + @freezegun.freeze_time("2020-07-07 10:54:05") + def test_datetime_branch_operator_upper_comparison_within_range(self): """Check DateTimeBranchOperator branch operation""" - mock_timezone.return_value = datetime.datetime(2020, 7, 7, 10, 54, 5, tzinfo=datetime.timezone.utc) + self.branch_op.target_lower = None self.branch_op.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE) tis = self.dr.get_task_instances() @@ -144,10 +143,10 @@ def test_datetime_branch_operator_upper_comparison_within_range(self, mock_timez 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): + @freezegun.freeze_time("2020-07-07 10:54:05") + def test_datetime_branch_operator_lower_comparison_within_range(self): """Check DateTimeBranchOperator branch operation""" - mock_timezone.return_value = datetime.datetime(2020, 7, 7, 10, 54, 5, tzinfo=datetime.timezone.utc) + self.branch_op.target_upper = None self.branch_op.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE) tis = self.dr.get_task_instances() @@ -161,10 +160,10 @@ def test_datetime_branch_operator_lower_comparison_within_range(self, mock_timez 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): + @freezegun.freeze_time("2020-07-07 12:00:00") + def test_datetime_branch_operator_upper_comparison_outside_range(self): """Check DateTimeBranchOperator branch operation""" - mock_timezone.return_value = datetime.datetime(2020, 7, 7, 12, 0, 0, tzinfo=datetime.timezone.utc) + self.branch_op.target_lower = None self.branch_op.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE) tis = self.dr.get_task_instances() @@ -178,10 +177,10 @@ def test_datetime_branch_operator_upper_comparison_outside_range(self, mock_time 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): + @freezegun.freeze_time("2020-07-07 09:00:00") + def test_datetime_branch_operator_lower_comparison_outside_range(self): """Check DateTimeBranchOperator branch operation""" - mock_timezone.return_value = datetime.datetime(2020, 7, 7, 9, 0, 0, tzinfo=datetime.timezone.utc) + self.branch_op.target_upper = None self.branch_op.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE) tis = self.dr.get_task_instances() From 60c1ad28f10c96a8c38fadb691237bc3e23aceec Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tom=C3=A1s=20Far=C3=ADas=20Santana?= Date: Mon, 23 Nov 2020 23:46:30 +0100 Subject: [PATCH 07/15] Remove _operator suffix from filename --- .../{datetime_branch_operator.py => datetime_branch.py} | 0 tests/operators/test_datetime_branch_operator.py | 2 +- 2 files changed, 1 insertion(+), 1 deletion(-) rename airflow/operators/{datetime_branch_operator.py => datetime_branch.py} (100%) diff --git a/airflow/operators/datetime_branch_operator.py b/airflow/operators/datetime_branch.py similarity index 100% rename from airflow/operators/datetime_branch_operator.py rename to airflow/operators/datetime_branch.py diff --git a/tests/operators/test_datetime_branch_operator.py b/tests/operators/test_datetime_branch_operator.py index 5e20bc638dcb3..65a3122a907d8 100644 --- a/tests/operators/test_datetime_branch_operator.py +++ b/tests/operators/test_datetime_branch_operator.py @@ -23,7 +23,7 @@ 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.datetime_branch import DateTimeBranchOperator from airflow.operators.dummy_operator import DummyOperator from airflow.utils import timezone from airflow.utils.session import create_session From fb03427d767471681390a43865f9fa73b2276be2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tom=C3=A1s=20Far=C3=ADas=20Santana?= Date: Thu, 26 Nov 2020 23:45:34 +0100 Subject: [PATCH 08/15] Rename datetime branch tests to remove _operator suffix --- .../{test_datetime_branch_operator.py => test_datetime_branch.py} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename tests/operators/{test_datetime_branch_operator.py => test_datetime_branch.py} (100%) diff --git a/tests/operators/test_datetime_branch_operator.py b/tests/operators/test_datetime_branch.py similarity index 100% rename from tests/operators/test_datetime_branch_operator.py rename to tests/operators/test_datetime_branch.py From bf4a4346e9de4f0a66db92e3a305aeca2fceb3a2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tom=C3=A1s=20Far=C3=ADas=20Santana?= Date: Sat, 5 Dec 2020 20:51:00 +0100 Subject: [PATCH 09/15] Support for datetime.time in DatetimeBranch operator --- airflow/operators/datetime_branch.py | 34 ++++- tests/operators/test_datetime_branch.py | 177 ++++++++++++++---------- 2 files changed, 133 insertions(+), 78 deletions(-) diff --git a/airflow/operators/datetime_branch.py b/airflow/operators/datetime_branch.py index 31f92a1b287d1..9a2fd962d0837 100644 --- a/airflow/operators/datetime_branch.py +++ b/airflow/operators/datetime_branch.py @@ -16,7 +16,7 @@ # under the License. import datetime -from typing import Dict, Iterable, Optional, Union +from typing import Dict, Iterable, Union from airflow.exceptions import AirflowException from airflow.operators.branch_operator import BaseBranchOperator @@ -49,8 +49,8 @@ def __init__( *, 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], + target_lower: Union[datetime.datetime, datetime.time, None], + target_upper: Union[datetime.datetime, datetime.time, None], **kwargs, ) -> None: super().__init__(**kwargs) @@ -68,10 +68,34 @@ def __init__( def choose_branch(self, context: Dict) -> Union[str, Iterable[str]]: now = timezone.make_naive(timezone.utcnow(), self.dag.timezone) - if self.target_upper is not None and self.target_upper < now: + lower, upper = target_times_as_dates(now, self.target_lower, self.target_upper) + + if upper is not None and upper < now: return self.follow_task_ids_if_false - if self.target_lower is not None and self.target_lower > now: + if lower is not None and lower > now: return self.follow_task_ids_if_false return self.follow_task_ids_if_true + + +def target_times_as_dates( + base_date: datetime.datetime, + lower: Union[datetime.datetime, datetime.time, None], + upper: Union[datetime.datetime, datetime.time, None], +): + """Ensures upper and lower time targets are datetimes by combining them with base_date""" + if isinstance(lower, datetime.datetime) and isinstance(upper, datetime.datetime): + return lower, upper + + if lower is not None and isinstance(lower, datetime.time): + lower = datetime.datetime.combine(base_date, lower) + if upper is not None and isinstance(upper, datetime.time): + upper = datetime.datetime.combine(base_date, upper) + + if any(date is None for date in (lower, upper)): + return lower, upper + + if upper < lower: + upper += datetime.timedelta(days=1) + return lower, upper diff --git a/tests/operators/test_datetime_branch.py b/tests/operators/test_datetime_branch.py index 65a3122a907d8..69e8081e7ae42 100644 --- a/tests/operators/test_datetime_branch.py +++ b/tests/operators/test_datetime_branch.py @@ -42,6 +42,14 @@ def setUpClass(cls): session.query(DagRun).delete() session.query(TI).delete() + cls.targets = [ + (datetime.datetime(2020, 7, 7, 10, 0, 0), datetime.datetime(2020, 7, 7, 11, 0, 0)), + (datetime.time(10, 0, 0), datetime.time(11, 0, 0)), + (datetime.datetime(2020, 7, 7, 10, 0, 0), datetime.time(11, 0, 0)), + (datetime.time(10, 0, 0), datetime.datetime(2020, 7, 7, 11, 0, 0)), + (datetime.time(11, 0, 0), datetime.time(10, 0, 0)), + ] + def setUp(self): self.dag = DAG( 'datetime_branch_operator_test', @@ -76,6 +84,21 @@ def tearDown(self): session.query(DagRun).delete() session.query(TI).delete() + def _assert_task_ids_match_states(self, task_ids_to_states): + """Helper that asserts task instances with a given id are in a given state""" + tis = self.dr.get_task_instances() + for ti in tis: + try: + expected_state = task_ids_to_states[ti.task_id] + except KeyError: + raise ValueError(f'Invalid task id {ti.task_id} found!') + else: + self.assertEqual( + ti.state, + expected_state, + f"Task {ti.task_id} has state {ti.state} instead of expected {expected_state}", + ) + def test_no_target_time(self): """Check if DateTimeBranchOperator raises exception on missing target""" with self.assertRaises(AirflowException): @@ -91,105 +114,113 @@ def test_no_target_time(self): @freezegun.freeze_time("2020-07-07 10:54:05") def test_datetime_branch_operator_falls_within_range(self): """Check DateTimeBranchOperator branch operation""" - self.branch_op.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE) + for target_lower, target_upper in self.targets: + with self.subTest(target_lower=target_lower, target_upper=target_upper): + print(f"{target_lower}, {target_upper}") + self.branch_op.target_lower = target_lower + self.branch_op.target_upper = target_upper + self.branch_op.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE) - tis = self.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!') + self._assert_task_ids_match_states( + { + 'datetime_branch': State.SUCCESS, + 'branch_1': State.NONE, + 'branch_2': State.SKIPPED, + } + ) def test_datetime_branch_operator_falls_outside_range(self): """Check DateTimeBranchOperator branch operation""" dates = [ datetime.datetime(2020, 7, 7, 12, 0, 0, tzinfo=datetime.timezone.utc), - datetime.datetime(2020, 6, 7, tzinfo=datetime.timezone.utc), + datetime.datetime(2020, 6, 7, 12, 0, 0, tzinfo=datetime.timezone.utc), ] - for date in dates: - with freezegun.freeze_time(date): - self.branch_op.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE) + for target_lower, target_upper in self.targets: + with self.subTest(target_lower=target_lower, target_upper=target_upper): + self.branch_op.target_lower = target_lower + self.branch_op.target_upper = target_upper - tis = self.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!') + for date in dates: + with freezegun.freeze_time(date): + self.branch_op.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE) + + self._assert_task_ids_match_states( + { + 'datetime_branch': State.SUCCESS, + 'branch_1': State.SKIPPED, + 'branch_2': State.NONE, + } + ) @freezegun.freeze_time("2020-07-07 10:54:05") def test_datetime_branch_operator_upper_comparison_within_range(self): """Check DateTimeBranchOperator branch operation""" - self.branch_op.target_lower = None - self.branch_op.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE) + for _, target_upper in self.targets: + with self.subTest(target_upper=target_upper): + self.branch_op.target_upper = target_upper + self.branch_op.target_lower = None - tis = self.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!') + self.branch_op.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE) + + self._assert_task_ids_match_states( + { + 'datetime_branch': State.SUCCESS, + 'branch_1': State.NONE, + 'branch_2': State.SKIPPED, + } + ) @freezegun.freeze_time("2020-07-07 10:54:05") def test_datetime_branch_operator_lower_comparison_within_range(self): """Check DateTimeBranchOperator branch operation""" - self.branch_op.target_upper = None - self.branch_op.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE) + for target_lower, _ in self.targets: + with self.subTest(target_lower=target_lower): + self.branch_op.target_lower = target_lower + self.branch_op.target_upper = None - tis = self.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!') + self.branch_op.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE) + + self._assert_task_ids_match_states( + { + 'datetime_branch': State.SUCCESS, + 'branch_1': State.NONE, + 'branch_2': State.SKIPPED, + } + ) @freezegun.freeze_time("2020-07-07 12:00:00") def test_datetime_branch_operator_upper_comparison_outside_range(self): """Check DateTimeBranchOperator branch operation""" - self.branch_op.target_lower = None - self.branch_op.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE) + for _, target_upper in self.targets: + with self.subTest(target_upper=target_upper): + self.branch_op.target_upper = target_upper + self.branch_op.target_lower = None - tis = self.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!') + self.branch_op.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE) + + self._assert_task_ids_match_states( + { + 'datetime_branch': State.SUCCESS, + 'branch_1': State.SKIPPED, + 'branch_2': State.NONE, + } + ) @freezegun.freeze_time("2020-07-07 09:00:00") def test_datetime_branch_operator_lower_comparison_outside_range(self): """Check DateTimeBranchOperator branch operation""" - self.branch_op.target_upper = None - self.branch_op.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE) + for target_lower, _ in self.targets: + with self.subTest(target_lower=target_lower): + self.branch_op.target_lower = target_lower + self.branch_op.target_upper = None - tis = self.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!') + self.branch_op.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE) + + self._assert_task_ids_match_states( + { + 'datetime_branch': State.SUCCESS, + 'branch_1': State.SKIPPED, + 'branch_2': State.NONE, + } + ) From 904d506bd37a23ec4c1e88c046a7e58e52879ab4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tom=C3=A1s=20Far=C3=ADas=20Santana?= Date: Sat, 5 Dec 2020 23:36:55 +0100 Subject: [PATCH 10/15] Add docs with example --- .../example_datetime_branch_operator.py | 72 +++++++++++++++++++ .../howto/operator/datetime_branch.rst | 39 ++++++++++ 2 files changed, 111 insertions(+) create mode 100644 airflow/example_dags/example_datetime_branch_operator.py create mode 100644 docs/apache-airflow/howto/operator/datetime_branch.rst diff --git a/airflow/example_dags/example_datetime_branch_operator.py b/airflow/example_dags/example_datetime_branch_operator.py new file mode 100644 index 0000000000000..0c1db7ba7e979 --- /dev/null +++ b/airflow/example_dags/example_datetime_branch_operator.py @@ -0,0 +1,72 @@ +# +# 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. +""" +import datetime + +from airflow import DAG +from airflow.operators.datetime_branch import DateTimeBranchOperator +from airflow.operators.dummy_operator import DummyOperator +from airflow.utils.dates import days_ago + +args = { + "owner": "airflow", +} + +dag = DAG( + dag_id="example_datetime_branch_operator", + start_date=days_ago(2), + default_args=args, + tags=["example"], + schedule_interval="@daily", +) + +# [START howto_operator_datetime_branch] +dummy_task_1 = DummyOperator(task_id='date_in_range', dag=dag) +dummy_task_2 = DummyOperator(task_id='date_outside_range', dag=dag) + +cond1 = DateTimeBranchOperator( + task_id='datetime_branch', + follow_task_ids_if_true=['date_in_range'], + follow_task_ids_if_false=['date_outside_range'], + target_upper=datetime.datetime(2020, 10, 10, 15, 0, 0), + target_lower=datetime.datetime(2020, 10, 10, 14, 0, 0), + dag=dag, +) + +# Run dummy_task_1 if cond1 executes between 2020-10-10 14:00:00 and 2020-10-10 15:00:00 +cond1 >> [dummy_task_1, dummy_task_2] +# [END howto_operator_datetime_branch] + +# [START howto_operator_datetime_branch_next_day] +cond2 = DateTimeBranchOperator( + task_id='datetime_branch', + follow_task_ids_if_true=['date_in_range'], + follow_task_ids_if_false=['date_outside_range'], + target_upper=datetime.time(0, 0, 0), + target_lower=datetime.time(15, 0, 0), + dag=dag, +) + +# Since target_lower happens after target_upper, target_upper will be moved to the following day +# Run dummy_task_1 if cond2 executes between 15:00:00 and 00:00:00 of the following day +cond2 >> [dummy_task_1, dummy_task_2] +# [END howto_operator_datetime_branch_next_day] diff --git a/docs/apache-airflow/howto/operator/datetime_branch.rst b/docs/apache-airflow/howto/operator/datetime_branch.rst new file mode 100644 index 0000000000000..b798468d23623 --- /dev/null +++ b/docs/apache-airflow/howto/operator/datetime_branch.rst @@ -0,0 +1,39 @@ + .. 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. + + + +.. _howto/operator:DatetimeBranch: + +DatetimeBranchOperator +====================== + +Use the :class:`~airflow.operators.datetime_branch.DatetimeBranchOperator` to branch into one of two execution paths depending on whether the date and/or time of execution falls into the range given by two target arguments. + +.. exampleinclude:: /../../airflow/example_dags/example_datetime_branch_operator.py + :language: python + :start-after: [START howto_operator_datetime_branch] + :end-before: [END howto_operator_datetime_branch] + +The target parameters, ``target_upper`` and ``target_lower``, can receive a ``datetime.datetime``, a ``datetime.time``, or ``None``. When a ``datetime.time`` object is used, it will be combined with the current date in order to allow comparisons with it. In the event that ``target_upper`` is set to a ``datetime.time`` that occurs before the given ``target_lower``, a day will be added to ``target_upper``. This is done to allow for time periods that span over two dates. + +.. exampleinclude:: /../../airflow/example_dags/example_datetime_branch_operator.py + :language: python + :start-after: [START howto_operator_datetime_branch_next_day] + :end-before: [END howto_operator_datetime_branch_next_day] + +If a target parameter is set to ``None``, the operator will perform a unilateral comparison using only the non-``None`` target. Setting both ``target_upper`` and ``target_lower`` to ``None`` will raise an exception. From 2b18f438b33e2a009ab5857c80d3010bc979149f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tom=C3=A1s=20Far=C3=ADas=20Santana?= Date: Sun, 6 Dec 2020 01:36:52 +0100 Subject: [PATCH 11/15] Fix example_dag tests --- .../example_dags/example_datetime_branch_operator.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/airflow/example_dags/example_datetime_branch_operator.py b/airflow/example_dags/example_datetime_branch_operator.py index 0c1db7ba7e979..01a0ee872e145 100644 --- a/airflow/example_dags/example_datetime_branch_operator.py +++ b/airflow/example_dags/example_datetime_branch_operator.py @@ -56,7 +56,18 @@ cond1 >> [dummy_task_1, dummy_task_2] # [END howto_operator_datetime_branch] + +dag = DAG( + dag_id="example_datetime_branch_operator_2", + start_date=days_ago(2), + default_args=args, + tags=["example"], + schedule_interval="@daily", +) # [START howto_operator_datetime_branch_next_day] +dummy_task_1 = DummyOperator(task_id='date_in_range', dag=dag) +dummy_task_2 = DummyOperator(task_id='date_outside_range', dag=dag) + cond2 = DateTimeBranchOperator( task_id='datetime_branch', follow_task_ids_if_true=['date_in_range'], From 437b6d1304d49dde81630bc2f35485c8e6ab4961 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tom=C3=A1s=20Far=C3=ADas=20Santana?= Date: Sun, 6 Dec 2020 04:00:25 +0100 Subject: [PATCH 12/15] Add operator to index in docs --- airflow/example_dags/example_datetime_branch_operator.py | 2 +- docs/apache-airflow/howto/operator/index.rst | 1 + 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/airflow/example_dags/example_datetime_branch_operator.py b/airflow/example_dags/example_datetime_branch_operator.py index 01a0ee872e145..0bf0910287d31 100644 --- a/airflow/example_dags/example_datetime_branch_operator.py +++ b/airflow/example_dags/example_datetime_branch_operator.py @@ -78,6 +78,6 @@ ) # Since target_lower happens after target_upper, target_upper will be moved to the following day -# Run dummy_task_1 if cond2 executes between 15:00:00 and 00:00:00 of the following day +# Run dummy_task_1 if cond2 executes between 15:00:00, and 00:00:00 of the following day cond2 >> [dummy_task_1, dummy_task_2] # [END howto_operator_datetime_branch_next_day] diff --git a/docs/apache-airflow/howto/operator/index.rst b/docs/apache-airflow/howto/operator/index.rst index 89ad5f3e5b884..4f615ccb96a21 100644 --- a/docs/apache-airflow/howto/operator/index.rst +++ b/docs/apache-airflow/howto/operator/index.rst @@ -32,5 +32,6 @@ determine what actually executes when your DAG runs. :maxdepth: 2 bash + datetime_branch python external_task_sensor From caa8c38a94c42b9ed6a2ea8f1f8c50a3baad61b2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tom=C3=A1s=20Far=C3=ADas=20Santana?= Date: Wed, 6 Jan 2021 01:50:42 +0100 Subject: [PATCH 13/15] Update example dag description --- airflow/example_dags/example_datetime_branch_operator.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/airflow/example_dags/example_datetime_branch_operator.py b/airflow/example_dags/example_datetime_branch_operator.py index 0bf0910287d31..58b1a0a5fc579 100644 --- a/airflow/example_dags/example_datetime_branch_operator.py +++ b/airflow/example_dags/example_datetime_branch_operator.py @@ -17,8 +17,8 @@ # 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. +Example DAG demonstrating the usage of DateTimeBranchOperator with datetime as well as time objects as +targets. """ import datetime From 817955ba2f4bcc672e3a2ea728d5d084863f0076 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tom=C3=A1s=20Far=C3=ADas=20Santana?= Date: Wed, 10 Mar 2021 20:44:51 +0100 Subject: [PATCH 14/15] Add use_task_execution_date argument to DateTimeBranchOperator --- airflow/operators/datetime_branch.py | 11 ++++++++-- tests/operators/test_datetime_branch.py | 27 ++++++++++++++++++++++++- 2 files changed, 35 insertions(+), 3 deletions(-) diff --git a/airflow/operators/datetime_branch.py b/airflow/operators/datetime_branch.py index 9a2fd962d0837..d5ce8c925e790 100644 --- a/airflow/operators/datetime_branch.py +++ b/airflow/operators/datetime_branch.py @@ -41,6 +41,9 @@ class DateTimeBranchOperator(BaseBranchOperator): :type target_lower: Optional[datetime.datetime] :param target_upper: target upper bound. :type target_upper: Optional[datetime.datetime] + :param use_task_execution_date: If ``True``, uses task's execution day to compare with targets. + Execution date is useful for backfilling. If ``False``, uses system's date. + :type use_task_execution_date: bool """ @apply_defaults @@ -51,6 +54,7 @@ def __init__( 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], + use_task_execution_date: bool = False, **kwargs, ) -> None: super().__init__(**kwargs) @@ -64,12 +68,15 @@ def __init__( 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 + self.use_task_execution_date = use_task_execution_date def choose_branch(self, context: Dict) -> Union[str, Iterable[str]]: - now = timezone.make_naive(timezone.utcnow(), self.dag.timezone) + if self.use_task_execution_date is True: + now = timezone.make_naive(context["execution_date"], self.dag.timezone) + else: + now = timezone.make_naive(timezone.utcnow(), self.dag.timezone) lower, upper = target_times_as_dates(now, self.target_lower, self.target_upper) - if upper is not None and upper < now: return self.follow_task_ids_if_false diff --git a/tests/operators/test_datetime_branch.py b/tests/operators/test_datetime_branch.py index 69e8081e7ae42..cfc2a86851002 100644 --- a/tests/operators/test_datetime_branch.py +++ b/tests/operators/test_datetime_branch.py @@ -116,7 +116,6 @@ def test_datetime_branch_operator_falls_within_range(self): """Check DateTimeBranchOperator branch operation""" for target_lower, target_upper in self.targets: with self.subTest(target_lower=target_lower, target_upper=target_upper): - print(f"{target_lower}, {target_upper}") self.branch_op.target_lower = target_lower self.branch_op.target_upper = target_upper self.branch_op.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE) @@ -224,3 +223,29 @@ def test_datetime_branch_operator_lower_comparison_outside_range(self): 'branch_2': State.NONE, } ) + + @freezegun.freeze_time("2020-12-01 09:00:00") + def test_datetime_branch_operator_use_task_execution_date(self): + """Check if DateTimeBranchOperator uses task execution date""" + in_between_date = timezone.datetime(2020, 7, 7, 10, 30, 0) + self.branch_op.use_task_execution_date = True + self.dr = self.dag.create_dagrun( + run_id='manual_exec_date__', + start_date=in_between_date, + execution_date=in_between_date, + state=State.RUNNING, + ) + + for target_lower, target_upper in self.targets: + with self.subTest(target_lower=target_lower, target_upper=target_upper): + self.branch_op.target_lower = target_lower + self.branch_op.target_upper = target_upper + self.branch_op.run(start_date=in_between_date, end_date=in_between_date) + + self._assert_task_ids_match_states( + { + 'datetime_branch': State.SUCCESS, + 'branch_1': State.NONE, + 'branch_2': State.SKIPPED, + } + ) From 8253b4a895b1334e5423e6551c89c52bacbc56f4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tom=C3=A1s=20Far=C3=ADas=20Santana?= Date: Wed, 10 Mar 2021 20:47:20 +0100 Subject: [PATCH 15/15] Use double backticks in rst docs Co-authored-by: Kaxil Naik --- airflow/operators/datetime_branch.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/airflow/operators/datetime_branch.py b/airflow/operators/datetime_branch.py index d5ce8c925e790..868bd23279de5 100644 --- a/airflow/operators/datetime_branch.py +++ b/airflow/operators/datetime_branch.py @@ -29,13 +29,13 @@ 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`. + ``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`. + ``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`. + ``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]