This is an automated email from the ASF dual-hosted git repository.
gopidesu pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 4761082189e Fix key error in _handle_execution_date_fn for
ExternalTaskSensor (#53728)
4761082189e is described below
commit 4761082189e4a70de474c634128cd5bfa98ad8ce
Author: GPK <[email protected]>
AuthorDate: Sat Jul 26 13:20:20 2025 +0100
Fix key error in _handle_execution_date_fn for ExternalTaskSensor (#53728)
* fix(external_task_sensor): fallback to execution_date key for
execution_date_fn context - Try context['logical_date'], then
context['execution_date'], then dag_run.run_after - Prevents KeyError when
users run in environments without 'logical_date' Closes #52236
fix(external_task_sensor): fallback to execution_date key for
execution_date_fn context
- Try context['logical_date'], then context['execution_date'], then
dag_run.run_after
- Prevents KeyError when users run in environments without 'logical_date'
Closes #52236
* test: alias stdlib datetime to dt to avoid shadowing timezone datetime
* chore: apply blank-line fix to docstring and add logical_date tests
Signed-off-by: yeswanth-bf <[email protected]>
* Update external_task.py
* Fix key error in _handle_execution_date_fn for ExternalTaskSensor
* Fix tests
---------
Signed-off-by: yeswanth-bf <[email protected]>
Co-authored-by: yeswanth-bf <[email protected]>
Co-authored-by: Yeswanth Surampudi
<[email protected]>
---
.../providers/standard/sensors/external_task.py | 33 +++++++++---
.../standard/sensors/test_external_task_sensor.py | 63 +++++++++++++++++++++-
2 files changed, 87 insertions(+), 9 deletions(-)
diff --git
a/providers/standard/src/airflow/providers/standard/sensors/external_task.py
b/providers/standard/src/airflow/providers/standard/sensors/external_task.py
index 7506a88ec0c..ffff3f85517 100644
--- a/providers/standard/src/airflow/providers/standard/sensors/external_task.py
+++ b/providers/standard/src/airflow/providers/standard/sensors/external_task.py
@@ -252,20 +252,15 @@ class ExternalTaskSensor(BaseSensorOperator):
self.poll_interval = poll_interval
def _get_dttm_filter(self, context):
- logical_date = context.get("logical_date")
- if AIRFLOW_V_3_0_PLUS:
- if logical_date is None:
- dag_run = context.get("dag_run")
- if TYPE_CHECKING:
- assert dag_run
+ logical_date = self._get_logical_date(context)
- logical_date = dag_run.run_after
if self.execution_delta:
dttm = logical_date - self.execution_delta
elif self.execution_date_fn:
dttm = self._handle_execution_date_fn(context=context)
else:
dttm = logical_date
+
return dttm if isinstance(dttm, list) else [dttm]
def poke(self, context: Context) -> bool:
@@ -522,6 +517,28 @@ class ExternalTaskSensor(BaseSensorOperator):
dttm_filter, self.external_task_group_id, self.external_dag_id,
session
)
+ def _get_logical_date(self, context) -> datetime.datetime:
+ """
+ Handle backwards- and forwards-compatible retrieval of the date.
+
+ to pass as the positional argument to execution_date_fn.
+ """
+ # Airflow 3.x: contexts define "logical_date" (or fall back to
dag_run.run_after).
+ if AIRFLOW_V_3_0_PLUS:
+ 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 `dag_run.run_after` must be
provided in the context"
+ )
+ return logical_date or dag_run.run_after
+
+ # Airflow 2.x and earlier: contexts used "execution_date"
+ execution_date = context.get("execution_date")
+ if not execution_date:
+ raise ValueError("Either `execution_date` must be provided in the
context`")
+ return execution_date
+
def _handle_execution_date_fn(self, context) -> Any:
"""
Handle backward compatibility.
@@ -534,7 +551,7 @@ class ExternalTaskSensor(BaseSensorOperator):
from airflow.utils.operator_helpers import make_kwargs_callable
# Remove "logical_date" because it is already a mandatory positional
argument
- logical_date = context["logical_date"]
+ logical_date = self._get_logical_date(context)
kwargs = {k: v for k, v in context.items() if k not in
{"execution_date", "logical_date"}}
# Add "context" in the kwargs for backward compatibility (because
context used to be
# an acceptable argument of execution_date_fn)
diff --git
a/providers/standard/tests/unit/standard/sensors/test_external_task_sensor.py
b/providers/standard/tests/unit/standard/sensors/test_external_task_sensor.py
index 978c617b4de..8c2435c16c4 100644
---
a/providers/standard/tests/unit/standard/sensors/test_external_task_sensor.py
+++
b/providers/standard/tests/unit/standard/sensors/test_external_task_sensor.py
@@ -1024,7 +1024,7 @@ exit 0
@pytest.mark.execution_timeout(10)
def test_external_task_sensor_deferrable(self, dag_maker):
- context = {}
+ context = {"execution_date": DEFAULT_DATE}
with dag_maker() as dag:
op = ExternalTaskSensor(
task_id="test_external_task_sensor_check",
@@ -1047,6 +1047,30 @@ exit 0
assert exc.value.trigger.external_task_ids == ["test_task"]
assert exc.value.trigger.execution_dates == [DEFAULT_DATE]
+ def test_get_logical_date(self):
+ """For AF 2, we check for execution_date in context."""
+ context = {"execution_date": DEFAULT_DATE}
+ op = ExternalTaskSensor(
+ task_id="test_external_task_sensor_check",
+ external_dag_id="test_dag_parent",
+ external_task_id="test_task",
+ )
+ assert op._get_logical_date(context) == DEFAULT_DATE
+
+ def test_handle_execution_date_fn(self):
+ def func(dt, context):
+ assert context["execution_date"] == dt
+ return dt + timedelta(0)
+
+ op = ExternalTaskSensor(
+ task_id="test_external_task_sensor_check",
+ external_dag_id="test_dag_parent",
+ external_task_id="test_task",
+ execution_date_fn=func,
+ )
+ context = {"execution_date": DEFAULT_DATE}
+ assert op._handle_execution_date_fn(context) == DEFAULT_DATE
+
@pytest.mark.skipif(not AIRFLOW_V_3_0_PLUS, reason="Different test for AF 2")
@pytest.mark.usefixtures("testing_dag_bundle")
@@ -1335,6 +1359,43 @@ class TestExternalTaskSensorV3:
task_group_id="test_group",
)
+ def test_get_logical_date(self):
+ """For AF 3, we check for logical date or dag_run.run_after in
context."""
+
+ context = {"logical_date": DEFAULT_DATE}
+ op = ExternalTaskSensor(
+ task_id="test_external_task_sensor_check",
+ external_dag_id="test_dag_parent",
+ external_task_id="test_task",
+ )
+ assert op._get_logical_date(context) == DEFAULT_DATE
+
+ def test_get_logical_date_with_dag_run_after(self):
+ """For AF 3, we check for logical date or dag_run.run_after in
context."""
+ op = ExternalTaskSensor(
+ task_id="test_external_task_sensor_check",
+ external_dag_id="test_dag_parent",
+ external_task_id="test_task",
+ )
+ mock_dag_run = mock.MagicMock()
+ mock_dag_run.run_after = DEFAULT_DATE
+ context = {"dag_run": mock_dag_run}
+ assert op._get_logical_date(context) == DEFAULT_DATE
+
+ def test_handle_execution_date_fn(self):
+ def func(dt, context):
+ assert context["logical_date"] == dt
+ return dt + timedelta(0)
+
+ op = ExternalTaskSensor(
+ task_id="test_external_task_sensor_check",
+ external_dag_id="test_dag_parent",
+ external_task_id="test_task",
+ execution_date_fn=func,
+ )
+ context = {"logical_date": DEFAULT_DATE}
+ assert op._handle_execution_date_fn(context) == DEFAULT_DATE
+
class TestExternalTaskAsyncSensor:
TASK_ID = "external_task_sensor_check"
