diff --git a/providers/standard/src/airflow/providers/standard/operators/datetime.py b/providers/standard/src/airflow/providers/standard/operators/datetime.py index 589e106fe31ba..26f90efe14008 100644 --- a/providers/standard/src/airflow/providers/standard/operators/datetime.py +++ b/providers/standard/src/airflow/providers/standard/operators/datetime.py @@ -71,10 +71,16 @@ def __init__( def choose_branch(self, context: Context) -> str | Iterable[str]: if self.use_task_logical_date: - now = context.get("logical_date") - if not now: - dag_run = context.get("dag_run") - now = dag_run.run_after # type: ignore[union-attr, assignment] + logical_date = context.get("logical_date") + dag_run = context.get("dag_run") + + if not (logical_date or (dag_run and dag_run.run_after)): + raise ValueError( + "Either `logical_date` or `run_after` should be provided in the task context when " + "`use_task_logical_date` is True" + ) + + now = logical_date or dag_run.run_after # type: ignore[union-attr] else: now = timezone.coerce_datetime(timezone.utcnow()) if TYPE_CHECKING: @@ -111,4 +117,4 @@ def target_times_as_dates( if upper < lower: upper += datetime.timedelta(days=1) - return lower, upper + return lower, upper \ No newline at end of file diff --git a/providers/standard/src/airflow/providers/standard/operators/weekday.py b/providers/standard/src/airflow/providers/standard/operators/weekday.py index e663bc7e4f0f1..0012755b3f917 100644 --- a/providers/standard/src/airflow/providers/standard/operators/weekday.py +++ b/providers/standard/src/airflow/providers/standard/operators/weekday.py @@ -111,13 +111,19 @@ def __init__( def choose_branch(self, context: Context) -> str | Iterable[str]: if self.use_task_logical_date: - now = context.get("logical_date") - if not now: - dag_run = context.get("dag_run") - now = dag_run.run_after # type: ignore[union-attr, assignment] + logical_date = context.get("logical_date") + dag_run = context.get("dag_run") + + if not (logical_date or (dag_run and dag_run.run_after)): + raise ValueError( + "Either `logical_date` or `run_after` should be provided in the task context when " + "`use_task_logical_date` is True" + ) + + now = logical_date or dag_run.run_after # type: ignore[union-attr] else: now = timezone.make_naive(timezone.utcnow(), self.dag.timezone) if now.isoweekday() in self._week_day_num: # type: ignore[union-attr] return self.follow_task_ids_if_true - return self.follow_task_ids_if_false + return self.follow_task_ids_if_false \ No newline at end of file diff --git a/providers/standard/tests/unit/standard/operators/test_datetime.py b/providers/standard/tests/unit/standard/operators/test_datetime.py index c9e38e6c9602f..78cca83f27f1b 100644 --- a/providers/standard/tests/unit/standard/operators/test_datetime.py +++ b/providers/standard/tests/unit/standard/operators/test_datetime.py @@ -344,3 +344,32 @@ def test_choose_branch_should_use_run_after_when_logical_date_none(self, dag_mak **{"run_after": timezone.datetime(2020, 8, 7)}, ) assert branch_op.choose_branch(context={"dag_run": dr}) == "branch_1" + + @pytest.mark.skipif(not AIRFLOW_V_3_0_PLUS, reason="Skip on Airflow < 3.0") + @time_machine.travel("2020-12-01 09:00:00") + def test_choose_branch_should_raise_when_neither_logical_date_nor_run_after(self, dag_maker): + with dag_maker( + "branch_datetime_operator_raises_without_date", + default_args={"owner": "airflow", "start_date": DEFAULT_DATE}, + schedule=INTERVAL, + serialized=True, + ): + branch_1 = EmptyOperator(task_id="branch_1") + branch_2 = EmptyOperator(task_id="branch_2") + + branch_op = BranchDateTimeOperator( + task_id="datetime_branch", + follow_task_ids_if_true="branch_1", + follow_task_ids_if_false="branch_2", + target_upper=datetime.datetime(2020, 9, 7, 11, 0, 0), + target_lower=datetime.datetime(2020, 6, 7, 10, 0, 0), + use_task_logical_date=True, + ) + branch_1.set_upstream(branch_op) + branch_2.set_upstream(branch_op) + + with pytest.raises( + ValueError, + match="Either `logical_date` or `run_after` should be provided in the task context", + ): + branch_op.choose_branch(context={}) diff --git a/providers/standard/tests/unit/standard/operators/test_weekday.py b/providers/standard/tests/unit/standard/operators/test_weekday.py index 957c94e7295ef..636b57e713030 100644 --- a/providers/standard/tests/unit/standard/operators/test_weekday.py +++ b/providers/standard/tests/unit/standard/operators/test_weekday.py @@ -206,6 +206,30 @@ def test_choose_branch_should_use_run_after_when_logical_date_none(self, dag_mak ) assert branch_op.choose_branch(context={"dag_run": dr}) == "branch_1" + @pytest.mark.skipif(not AIRFLOW_V_3_0_PLUS, reason="Skip on Airflow < 3.0") + @time_machine.travel("2021-01-25") # Monday + def test_choose_branch_should_raise_when_neither_logical_date_nor_run_after(self, dag_maker): + with dag_maker( + "branch_day_of_week_operator_test", start_date=DEFAULT_DATE, schedule=INTERVAL, serialized=True + ): + branch_op = BranchDayOfWeekOperator( + task_id="make_choice", + follow_task_ids_if_true="branch_1", + follow_task_ids_if_false="branch_2", + week_day="Wednesday", + use_task_logical_date=True, + ) + branch_1 = EmptyOperator(task_id="branch_1") + branch_2 = EmptyOperator(task_id="branch_2") + branch_1.set_upstream(branch_op) + branch_2.set_upstream(branch_op) + + with pytest.raises( + ValueError, + match="Either `logical_date` or `run_after` should be provided in the task context", + ): + branch_op.choose_branch(context={}) + @time_machine.travel("2021-01-25") # Monday def test_branch_follow_false(self, dag_maker): """Checks if BranchDayOfWeekOperator follow false branch"""