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 77cf38b1bc4 Keep MSGraph pagination offset across paginated pages
(#71986)
77cf38b1bc4 is described below
commit 77cf38b1bc47376d588cac8aaf765f2d93f5494a
Author: PoAn Yang <[email protected]>
AuthorDate: Wed Sep 9 00:53:35 2026 +0900
Keep MSGraph pagination offset across paginated pages (#71986)
Signed-off-by: PoAn Yang <[email protected]>
---
.../providers/microsoft/azure/hooks/msgraph.py | 11 ++++--
.../providers/microsoft/azure/operators/msgraph.py | 10 +++++-
.../unit/microsoft/azure/hooks/test_msgraph.py | 38 +++++++++++++++++++++
.../unit/microsoft/azure/operators/test_msgraph.py | 39 ++++++++++++++++++++--
.../unit/microsoft/azure/resources/messages.json | 2 +-
.../microsoft/azure/resources/second_messages.json | 1 +
.../{next_messages.json => third_messages.json} | 2 +-
7 files changed, 95 insertions(+), 8 deletions(-)
diff --git
a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/msgraph.py
b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/msgraph.py
index b81f1aacba5..0a9cca43d0d 100644
---
a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/msgraph.py
+++
b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/msgraph.py
@@ -576,6 +576,13 @@ class KiotaRequestAdapterHook(BaseHook):
query_parameters: dict[str, Any] | None = None,
responses: Callable[[], list[dict[str, Any]] | None] = lambda: [],
) -> tuple[Any, dict[str, Any] | None]:
+ """
+ Resolve the url and query parameters of the page following
``response``.
+
+ The ``$skip`` offset is derived from ``query_parameters`` rather than
from ``responses``:
+ callers accumulate whatever their own callbacks produced, so the
entries are not guaranteed
+ to be the raw pages this offset would have to be counted from.
+ """
if isinstance(response, dict):
odata_count = response.get("@odata.count")
if odata_count and query_parameters:
@@ -583,9 +590,7 @@ class KiotaRequestAdapterHook(BaseHook):
if top and odata_count:
if len(response.get("value", [])) == top:
- results = responses()
- skip = sum([len(result["value"]) for result in
results]) + top if results else top # type: ignore
- query_parameters["$skip"] = skip
+ query_parameters["$skip"] =
(query_parameters.get("$skip") or 0) + top
return url, query_parameters
return response.get("@odata.nextLink"), query_parameters
return None, query_parameters
diff --git
a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/operators/msgraph.py
b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/operators/msgraph.py
index 049c7004477..e60dd36bdc4 100644
---
a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/operators/msgraph.py
+++
b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/operators/msgraph.py
@@ -187,14 +187,22 @@ class MSGraphAsyncOperator(BaseOperator):
self,
context: Context,
event: dict[Any, Any] | None = None,
+ query_parameters: dict[str, Any] | None = None,
) -> Any:
"""
Execute callback when MSGraphTrigger finishes execution.
This method gets executed automatically when MSGraphTrigger completes
its execution.
+
+ :param query_parameters: The query parameters the completed page was
requested with, passed
+ back by :meth:`trigger_next_link`. The operator is rebuilt from
the serialized Dag on
+ every page, so without them ``self.query_parameters`` would still
describe the first page.
"""
self.log.debug("context: %s", context)
+ if query_parameters is not None:
+ self.query_parameters = query_parameters
+
if event:
self.log.debug("%s completed with %s: %s", self.task_id,
event.get("status"), event)
@@ -315,7 +323,6 @@ class MSGraphAsyncOperator(BaseOperator):
response=response,
url=operator.url,
query_parameters=operator.query_parameters,
- responses=lambda: operator.pull_xcom(context),
)
def trigger_next_link(self, response, method_name: str, context: Context)
-> None:
@@ -355,4 +362,5 @@ class MSGraphAsyncOperator(BaseOperator):
pagination_link=True,
),
method_name=method_name,
+ kwargs={"query_parameters": query_parameters},
)
diff --git
a/providers/microsoft/azure/tests/unit/microsoft/azure/hooks/test_msgraph.py
b/providers/microsoft/azure/tests/unit/microsoft/azure/hooks/test_msgraph.py
index 0a8dead772b..c2fe69ff9f3 100644
--- a/providers/microsoft/azure/tests/unit/microsoft/azure/hooks/test_msgraph.py
+++ b/providers/microsoft/azure/tests/unit/microsoft/azure/hooks/test_msgraph.py
@@ -442,6 +442,44 @@ class TestKiotaRequestAdapterHook:
assert isinstance(actual, list)
assert actual == [users, next_users]
+ @pytest.mark.asyncio
+ @pytest.mark.parametrize(
+ ("query_parameters", "expected_skips"),
+ [
+ pytest.param(
+ {"$top": 12, "$count": True},
+ ["", "&%24skip=12", "&%24skip=24"],
+ id="from_the_start",
+ ),
+ pytest.param(
+ {"$top": 12, "$count": True, "$skip": 100},
+ ["&%24skip=100", "&%24skip=112", "&%24skip=124"],
+ id="from_a_user_supplied_offset",
+ ),
+ ],
+ )
+ async def test_paginated_run_advances_the_skip_offset_by_a_single_page(
+ self, query_parameters, expected_skips
+ ):
+ messages = load_json_from_resources(dirname(__file__), "..",
"resources", "messages.json")
+ second_messages = load_json_from_resources(
+ dirname(__file__), "..", "resources", "second_messages.json"
+ )
+ third_messages = load_json_from_resources(dirname(__file__), "..",
"resources", "third_messages.json")
+ response = mock_json_response(200, messages, second_messages,
third_messages)
+
+ with patch_hook_and_request_adapter(response) as mocks:
+ mock_get_http_response = mocks[-1]
+ hook = KiotaRequestAdapterHook(conn_id="msgraph_api")
+
+ # paginated_run mutates the query parameters it is given, so hand
it a copy rather than
+ # the dict pytest built once at collection time.
+ await hook.paginated_run(url="users/messages",
query_parameters=dict(query_parameters))
+
+ urls = [call.args[0].url for call in
mock_get_http_response.call_args_list]
+
+ assert urls == [f"users/messages?%24top=12&%24count=true{skip}" for
skip in expected_skips]
+
@pytest.mark.asyncio
async def test_paginated_run_refuses_cross_host_next_link(self):
first_page = {
diff --git
a/providers/microsoft/azure/tests/unit/microsoft/azure/operators/test_msgraph.py
b/providers/microsoft/azure/tests/unit/microsoft/azure/operators/test_msgraph.py
index cdb1715b7e1..e2d14a67ff9 100644
---
a/providers/microsoft/azure/tests/unit/microsoft/azure/operators/test_msgraph.py
+++
b/providers/microsoft/azure/tests/unit/microsoft/azure/operators/test_msgraph.py
@@ -356,10 +356,44 @@ class TestMSGraphAsyncOperator:
assert trigger.url ==
"users/{user_id}/mailFolders/{mailFolder_id}/messages"
assert trigger.path_parameters == path_parameters
+ def
test_execute_complete_advances_the_skip_offset_it_was_resumed_with(self):
+ # The operator is rebuilt from the serialized Dag on every page, so
its own query parameters
+ # describe the first request; only the resume kwargs know where the
completed page came from.
+ operator = MSGraphAsyncOperator(
+ task_id="messages",
+ conn_id="msgraph_api",
+ url="users/messages",
+ query_parameters={"$top": 12, "$count": True},
+ )
+ context = mock_context(task=operator)
+ second_messages = load_json_from_resources(
+ dirname(__file__), "..", "resources", "second_messages.json"
+ )
+ event = {"status": "success", "type": "builtins.dict", "response":
json.dumps(second_messages)}
+
+ with mock.patch.object(operator, "defer") as mock_defer:
+ operator.execute_complete(
+ context=context,
+ event=event,
+ query_parameters={"$top": 12, "$count": True, "$skip": 12},
+ )
+
+ assert mock_defer.call_args.kwargs["trigger"].query_parameters == {
+ "$top": 12,
+ "$count": True,
+ "$skip": 24,
+ }
+ assert mock_defer.call_args.kwargs["kwargs"] == {
+ "query_parameters": {"$top": 12, "$count": True, "$skip": 24}
+ }
+
def test_skip_pagination_expands_the_url_template_on_every_page(self):
messages = load_json_from_resources(dirname(__file__), "..",
"resources", "messages.json")
- next_messages = load_json_from_resources(dirname(__file__), "..",
"resources", "next_messages.json")
- response = mock_json_response(200, messages, next_messages)
+ second_messages = load_json_from_resources(
+ dirname(__file__), "..", "resources", "second_messages.json"
+ )
+ third_messages = load_json_from_resources(dirname(__file__), "..",
"resources", "third_messages.json")
+ response = mock_json_response(200, messages, second_messages,
third_messages)
with patch_hook_and_request_adapter(response) as (*_,
mock_get_http_response):
operator = MSGraphAsyncOperator(
@@ -378,6 +412,7 @@ class TestMSGraphAsyncOperator:
assert urls == [
"users/48d31887-5fad-4d73-a9f5-3c356e68a038/mailFolders/inbox/messages?%24top=12&%24count=true",
"users/48d31887-5fad-4d73-a9f5-3c356e68a038/mailFolders/inbox/messages?%24top=12&%24count=true&%24skip=12",
+
"users/48d31887-5fad-4d73-a9f5-3c356e68a038/mailFolders/inbox/messages?%24top=12&%24count=true&%24skip=24",
]
def test_pagination_issues_every_page_with_the_configured_request(self):
diff --git
a/providers/microsoft/azure/tests/unit/microsoft/azure/resources/messages.json
b/providers/microsoft/azure/tests/unit/microsoft/azure/resources/messages.json
index e8994c021c9..c45e0457d99 100644
---
a/providers/microsoft/azure/tests/unit/microsoft/azure/resources/messages.json
+++
b/providers/microsoft/azure/tests/unit/microsoft/azure/resources/messages.json
@@ -1 +1 @@
-{"@odata.context":
"https://graph.microsoft.com/v1.0/$metadata#users('48d31887-5fad-4d73-a9f5-3c356e68a038')/mailFolders('inbox')/messages",
"@odata.count": 18, "@odata.nextLink":
"https://graph.microsoft.com/v1.0/users('48d31887-5fad-4d73-a9f5-3c356e68a038')/mailFolders('inbox')/messages?%24top=12&%24count=true&%24skip=12",
"value": [{"@odata.etag": "W/\"CQAAABYAAADHcgC8Hl9tRZ/hc1wEUs1TAAAwR4Hg\"",
"id": "AAMkAGUAAAwTW09AAA=", "subject": "Weekly status report",
"receivedDateTime": "2026 [...]
+{"@odata.context":
"https://graph.microsoft.com/v1.0/$metadata#users('48d31887-5fad-4d73-a9f5-3c356e68a038')/mailFolders('inbox')/messages",
"@odata.count": 30, "@odata.nextLink":
"https://graph.microsoft.com/v1.0/users('48d31887-5fad-4d73-a9f5-3c356e68a038')/mailFolders('inbox')/messages?%24top=12&%24count=true&%24skip=12",
"value": [{"@odata.etag": "W/\"CQAAABYAAADHcgC8Hl9tRZ/hc1wEUs1TAAAwR4Hg\"",
"id": "AAMkAGUAAAwTW09AAA=", "subject": "Weekly status report",
"receivedDateTime": "2026 [...]
diff --git
a/providers/microsoft/azure/tests/unit/microsoft/azure/resources/second_messages.json
b/providers/microsoft/azure/tests/unit/microsoft/azure/resources/second_messages.json
new file mode 100644
index 00000000000..48c4f7f6721
--- /dev/null
+++
b/providers/microsoft/azure/tests/unit/microsoft/azure/resources/second_messages.json
@@ -0,0 +1 @@
+{"@odata.context":
"https://graph.microsoft.com/v1.0/$metadata#users('48d31887-5fad-4d73-a9f5-3c356e68a038')/mailFolders('inbox')/messages",
"@odata.count": 30, "@odata.nextLink":
"https://graph.microsoft.com/v1.0/users('48d31887-5fad-4d73-a9f5-3c356e68a038')/mailFolders('inbox')/messages?%24top=12&%24count=true&%24skip=24",
"value": [{"@odata.etag": "W/\"CQAAABYAAADHcgC8Hl9tRZ/hc1wEUs1TAAAwTA01\"",
"id": "AAMkAGUAAAwUW09AAA=", "subject": "Re: Reminder: timesheet due",
"receivedDateTime" [...]
diff --git
a/providers/microsoft/azure/tests/unit/microsoft/azure/resources/next_messages.json
b/providers/microsoft/azure/tests/unit/microsoft/azure/resources/third_messages.json
similarity index 95%
rename from
providers/microsoft/azure/tests/unit/microsoft/azure/resources/next_messages.json
rename to
providers/microsoft/azure/tests/unit/microsoft/azure/resources/third_messages.json
index 81470210d56..e2f8bb14649 100644
---
a/providers/microsoft/azure/tests/unit/microsoft/azure/resources/next_messages.json
+++
b/providers/microsoft/azure/tests/unit/microsoft/azure/resources/third_messages.json
@@ -1 +1 @@
-{"@odata.context":
"https://graph.microsoft.com/v1.0/$metadata#users('48d31887-5fad-4d73-a9f5-3c356e68a038')/mailFolders('inbox')/messages",
"@odata.count": 18, "value": [{"@odata.etag":
"W/\"CQAAABYAAADHcgC8Hl9tRZ/hc1wEUs1TAAAwSp4B\"", "id": "AAMkAGUAAAwTXCMAAA=",
"subject": "Re: Deployment window moved", "receivedDateTime":
"2026-08-17T09:18:47Z", "isRead": true}, {"@odata.etag":
"W/\"CQAAABYAAADHcgC8Hl9tRZ/hc1wEUs1TAAAwStQm\"", "id": "AAMkAGUAAAwTXDNAAA=",
"subject": "Offsite logistic [...]
+{"@odata.context":
"https://graph.microsoft.com/v1.0/$metadata#users('48d31887-5fad-4d73-a9f5-3c356e68a038')/mailFolders('inbox')/messages",
"@odata.count": 30, "value": [{"@odata.etag":
"W/\"CQAAABYAAADHcgC8Hl9tRZ/hc1wEUs1TAAAwSp4B\"", "id": "AAMkAGUAAAwTXCMAAA=",
"subject": "Re: Deployment window moved", "receivedDateTime":
"2026-08-17T09:18:47Z", "isRead": true}, {"@odata.etag":
"W/\"CQAAABYAAADHcgC8Hl9tRZ/hc1wEUs1TAAAwStQm\"", "id": "AAMkAGUAAAwTXDNAAA=",
"subject": "Offsite logistic [...]
