This is an automated email from the ASF dual-hosted git repository.
dabla 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 54171ff7cc8 Fix MSGraphSensor timeout being reset on every re-defer
(#75774)
54171ff7cc8 is described below
commit 54171ff7cc8e532f4488cce74aacbc8d3cd28027
Author: Mathieu Monet <[email protected]>
AuthorDate: Wed Oct 28 17:53:51 2026 +0100
Fix MSGraphSensor timeout being reset on every re-defer (#75774)
* Fix MSGraphSensor timeout being reset on every re-defer
A defer timeout becomes an absolute deadline of now + timeout, and
MSGraphSensor re-deferred with the full sensor timeout between polls
(while the request defer carried none), so the sensor never timed out.
Carry the first execute's start time in next_kwargs, give each request
defer the remaining time, and raise AirflowSensorTimeout after a
not-done poll once the timeout has elapsed, so soft_fail and never_fail
apply as for core sensors.
Stop passing the sensor timeout to MSGraphTrigger, where it was used as
the HTTP request timeout.
Co-Authored-By: Claude Opus 5.5 <[email protected]>
* Keep a deadline on the MSGraphSensor sleep defer
Give the TimeDeltaTrigger sleep a defer timeout of the remaining time
plus retry_delay, so a lost sleep trigger still ends in a sensor
timeout. Drop the unreachable timeout=None branches and use started_at
consistently.
Co-Authored-By: Claude Opus 5.5 <[email protected]>
---------
Co-authored-by: Claude Opus 5.5 <[email protected]>
---
.../providers/microsoft/azure/sensors/msgraph.py | 44 ++++--
.../unit/microsoft/azure/sensors/test_msgraph.py | 161 +++++++++++++++++++--
2 files changed, 181 insertions(+), 24 deletions(-)
diff --git
a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/sensors/msgraph.py
b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/sensors/msgraph.py
index 103fb1d1b45..beada87b7a1 100644
---
a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/sensors/msgraph.py
+++
b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/sensors/msgraph.py
@@ -21,7 +21,12 @@ from collections.abc import Callable, Sequence
from datetime import timedelta
from typing import TYPE_CHECKING, Any
-from airflow.providers.common.compat.sdk import AirflowException,
BaseSensorOperator
+from airflow.providers.common.compat.sdk import (
+ AirflowException,
+ AirflowSensorTimeout,
+ BaseSensorOperator,
+ timezone,
+)
from airflow.providers.common.compat.standard.triggers import TimeDeltaTrigger
from airflow.providers.microsoft.azure.hooks.msgraph import
KiotaRequestAdapterHook
from airflow.providers.microsoft.azure.operators.msgraph import
execute_callable
@@ -108,6 +113,14 @@ class MSGraphSensor(BaseSensorOperator):
self.serializer = serializer()
def execute(self, context: Context):
+ self._defer_request(timezone.utcnow().timestamp(), self.timeout)
+
+ def retry_execute(self, context: Context, sensor_started_at: float | None
= None, **kwargs) -> Any:
+ started_at = self._started_at(sensor_started_at)
+ remaining = max(self._remaining_timeout(started_at),
self.retry_delay.total_seconds())
+ self._defer_request(started_at, remaining)
+
+ def _defer_request(self, started_at: float, remaining: float) -> None:
self.defer(
trigger=MSGraphTrigger(
url=self.url,
@@ -119,26 +132,28 @@ class MSGraphSensor(BaseSensorOperator):
headers=self.headers,
data=self.data,
conn_id=self.conn_id,
- timeout=self.timeout,
proxies=self.proxies,
scopes=self.scopes,
api_version=self.api_version,
serializer=type(self.serializer),
),
method_name=self.execute_complete.__name__,
+ kwargs={"sensor_started_at": started_at},
+ timeout=timedelta(seconds=remaining),
)
- def retry_execute(
- self,
- context: Context,
- **kwargs,
- ) -> Any:
- self.execute(context=context)
+ @staticmethod
+ def _started_at(sensor_started_at: float | None) -> float:
+ return sensor_started_at if sensor_started_at is not None else
timezone.utcnow().timestamp()
+
+ def _remaining_timeout(self, started_at: float) -> float:
+ return self.timeout - max(timezone.utcnow().timestamp() - started_at,
0)
def execute_complete(
self,
context: Context,
event: dict[Any, Any] | None = None,
+ sensor_started_at: float | None = None,
) -> Any:
"""
Execute callback when MSGraphSensor finishes execution.
@@ -183,11 +198,18 @@ class MSGraphSensor(BaseSensorOperator):
return result
- # Re-defer with timeout so Airflow enforces the sensor timeout
natively
+ started_at = self._started_at(sensor_started_at)
+ remaining = self._remaining_timeout(started_at)
+ if remaining <= 0:
+ raise AirflowSensorTimeout(
+ f"Sensor has timed out; run duration of {self.timeout
- remaining} seconds "
+ f"exceeds the specified timeout of {self.timeout}."
+ )
self.defer(
- trigger=TimeDeltaTrigger(self.retry_delay),
+ trigger=TimeDeltaTrigger(min(self.retry_delay,
timedelta(seconds=remaining))),
method_name=self.retry_execute.__name__,
- timeout=timedelta(seconds=self.timeout) if self.timeout is
not None else None,
+ kwargs={"sensor_started_at": started_at},
+ timeout=timedelta(seconds=remaining) + self.retry_delay,
)
return None
diff --git
a/providers/microsoft/azure/tests/unit/microsoft/azure/sensors/test_msgraph.py
b/providers/microsoft/azure/tests/unit/microsoft/azure/sensors/test_msgraph.py
index 66f439ba722..2dd951c1e99 100644
---
a/providers/microsoft/azure/tests/unit/microsoft/azure/sensors/test_msgraph.py
+++
b/providers/microsoft/azure/tests/unit/microsoft/azure/sensors/test_msgraph.py
@@ -17,20 +17,48 @@
from __future__ import annotations
import json
-from datetime import datetime, timedelta
+from datetime import UTC, datetime, timedelta
from os.path import dirname
from unittest.mock import patch
import pytest
+import time_machine
from airflow.exceptions import AirflowProviderDeprecationWarning
+from airflow.providers.common.compat.sdk import AirflowSensorTimeout,
AirflowSkipException, TaskDeferred
+from airflow.providers.common.compat.standard.triggers import TimeDeltaTrigger
from airflow.providers.microsoft.azure.sensors.msgraph import MSGraphSensor
+from airflow.providers.microsoft.azure.triggers.msgraph import MSGraphTrigger
from airflow.triggers.base import TriggerEvent
from tests_common.test_utils.file_loading import load_json_from_resources
from tests_common.test_utils.operators.run_deferrable import execute_operator
+from tests_common.test_utils.version_compat import AIRFLOW_V_3_1_PLUS
from unit.microsoft.azure.test_utils import mock_json_response,
patch_hook_and_request_adapter
+T0 = datetime(2026, 1, 1, tzinfo=UTC)
+NOT_DONE = {"status": "success", "type": "builtins.dict", "response":
json.dumps({"status": "Running"})}
+DONE = {
+ "status": "success",
+ "type": "builtins.dict",
+ "response": json.dumps({"status": "Succeeded", "id": "x"}),
+}
+
+
+def _sensor(**kwargs) -> MSGraphSensor:
+ return MSGraphSensor(
+ task_id="check",
+ conn_id="powerbi",
+ url="myorg/admin/workspaces/scanStatus/{scanId}",
+ timeout=100,
+ retry_delay=30,
+ **kwargs,
+ )
+
+
+def _at(seconds: float):
+ return time_machine.travel(T0 + timedelta(seconds=seconds), tick=False)
+
class TestMSGraphSensor:
def test_execute_with_result_processor_with_old_signature(self):
@@ -143,17 +171,124 @@ class TestMSGraphSensor:
for template_field in MSGraphSensor.template_fields:
getattr(sensor, template_field)
- def test_execute_complete_passes_timeout_to_defer(self):
- sensor = MSGraphSensor(
- task_id="check_timeout",
- conn_id="powerbi",
- url="myorg/admin/workspaces/scanStatus/{scanId}",
- timeout=10,
- )
+ def test_execute_defers_with_timeout_and_started_at(self):
+ sensor = _sensor()
+
+ with _at(0), pytest.raises(TaskDeferred) as exc:
+ sensor.execute(context={})
+
+ assert isinstance(exc.value.trigger, MSGraphTrigger)
+ assert exc.value.trigger.timeout is None
+ assert exc.value.method_name == "execute_complete"
+ assert exc.value.timeout == timedelta(seconds=100)
+ assert exc.value.kwargs == {"sensor_started_at": T0.timestamp()}
+
+ def test_execute_complete_redefers_with_capped_sleep_and_started_at(self):
+ sensor = _sensor()
+
+ with _at(10), pytest.raises(TaskDeferred) as exc:
+ sensor.execute_complete(context={}, event=NOT_DONE,
sensor_started_at=T0.timestamp())
+
+ assert isinstance(exc.value.trigger, TimeDeltaTrigger)
+ assert exc.value.trigger.moment == T0 + timedelta(seconds=40)
+ assert exc.value.method_name == "retry_execute"
+ assert exc.value.timeout == timedelta(seconds=120)
+ assert exc.value.kwargs == {"sensor_started_at": T0.timestamp()}
+
+ @pytest.mark.parametrize("elapsed", [10, 40, 70])
+ def test_retry_execute_deadline_does_not_slide(self, elapsed):
+ sensor = _sensor()
+
+ with _at(elapsed), pytest.raises(TaskDeferred) as exc:
+ sensor.retry_execute(context={}, event=T0,
sensor_started_at=T0.timestamp())
+
+ assert isinstance(exc.value.trigger, MSGraphTrigger)
+ assert exc.value.method_name == "execute_complete"
+ assert exc.value.timeout == timedelta(seconds=100 - elapsed)
+ assert exc.value.kwargs == {"sensor_started_at": T0.timestamp()}
+
+ def test_sleep_capped_to_remaining(self):
+ sensor = _sensor()
+
+ with _at(90), pytest.raises(TaskDeferred) as exc:
+ sensor.execute_complete(context={}, event=NOT_DONE,
sensor_started_at=T0.timestamp())
- with patch.object(sensor, "defer") as mock_defer:
- sensor.execute_complete(
- context={}, event={"status": "success", "response":
json.dumps({"status": "running"})}
+ assert exc.value.trigger.moment == T0 + timedelta(seconds=100)
+ assert exc.value.timeout == timedelta(seconds=40)
+
+ @pytest.mark.parametrize("elapsed", [100, 150])
+ def test_past_deadline_raises_sensor_timeout(self, elapsed):
+ sensor = _sensor()
+
+ with (
+ _at(elapsed),
+ patch.object(sensor, "defer") as mock_defer,
+ pytest.raises(AirflowSensorTimeout, match="Sensor has timed out"),
+ ):
+ sensor.execute_complete(context={}, event=NOT_DONE,
sensor_started_at=T0.timestamp())
+
+ mock_defer.assert_not_called()
+
+ @pytest.mark.parametrize("elapsed", [80, 99.7, 100, 150])
+ def test_retry_execute_near_or_past_deadline_polls_once_more(self,
elapsed):
+ sensor = _sensor()
+
+ with _at(elapsed), pytest.raises(TaskDeferred) as exc:
+ sensor.retry_execute(context={}, event=T0,
sensor_started_at=T0.timestamp())
+
+ assert isinstance(exc.value.trigger, MSGraphTrigger)
+ assert exc.value.method_name == "execute_complete"
+ assert exc.value.timeout == timedelta(seconds=30)
+ assert exc.value.kwargs == {"sensor_started_at": T0.timestamp()}
+
+ def test_done_after_deadline_returns_result(self):
+ sensor = _sensor(result_processor=lambda result, **context:
result["id"])
+
+ with _at(150):
+ assert sensor.execute_complete(context={}, event=DONE,
sensor_started_at=T0.timestamp()) == "x"
+
+ @pytest.mark.parametrize(
+ ("flags", "expected"),
+ [
+ pytest.param({}, AirflowSensorTimeout, id="no_flag"),
+ pytest.param({"soft_fail": True}, AirflowSkipException,
id="soft_fail"),
+ pytest.param(
+ {"never_fail": True},
+ AirflowSkipException,
+ id="never_fail",
+ marks=pytest.mark.skipif(not AIRFLOW_V_3_1_PLUS,
reason="never_fail on resume is 3.1+"),
+ ),
+ ],
+ )
+ def test_timeout_via_resume_execution(self, flags, expected):
+ sensor = _sensor(**flags)
+
+ with _at(150), pytest.raises(expected):
+ sensor.resume_execution(
+ "execute_complete", {"event": NOT_DONE, "sensor_started_at":
T0.timestamp()}, {}
)
- mock_defer.assert_called_once()
- assert mock_defer.call_args.kwargs["timeout"] ==
timedelta(seconds=10)
+
+ def test_missing_started_at_falls_back_to_now(self):
+ sensor = _sensor()
+
+ with _at(500), pytest.raises(TaskDeferred) as exc:
+ sensor.retry_execute(context={}, event=T0)
+
+ assert exc.value.timeout == timedelta(seconds=100)
+ assert exc.value.kwargs == {"sensor_started_at": (T0 +
timedelta(seconds=500)).timestamp()}
+
+ with _at(500), pytest.raises(TaskDeferred) as exc:
+ sensor.execute_complete(context={}, event=NOT_DONE)
+
+ assert exc.value.trigger.moment == T0 + timedelta(seconds=530)
+ assert exc.value.kwargs == {"sensor_started_at": (T0 +
timedelta(seconds=500)).timestamp()}
+
+ def test_future_started_at_clamped(self):
+ sensor = _sensor()
+
+ with _at(0), pytest.raises(TaskDeferred) as exc:
+ sensor.retry_execute(
+ context={}, event=T0, sensor_started_at=(T0 +
timedelta(seconds=10)).timestamp()
+ )
+
+ assert exc.value.timeout == timedelta(seconds=100)
