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 e99f9af45a8 Add iter_paginated_run to KiotaRequestAdapterHook (#75649)
e99f9af45a8 is described below
commit e99f9af45a8d3286e29418d337c3b598ae286b53
Author: David Blain <[email protected]>
AuthorDate: Thu Nov 12 17:09:05 2026 +0100
Add iter_paginated_run to KiotaRequestAdapterHook (#75649)
* Add iter_paginated_run to KiotaRequestAdapterHook
paginated_run collects every page of a Microsoft Graph request in one
list before it returns, so a task that extracts a large collection (all
users of a tenant, their group memberships) holds the whole dataset in
memory even when it only writes each page away.
iter_paginated_run yields the pages one at a time, with the same
pagination function and the same allowed-host check on every link, and
paginated_run now collects what it yields, so there is a single
pagination loop. The pagination function of paginated_run keeps
receiving the pages collected so far; the generator has nothing to hand
it, since it keeps no page it already yielded.
We have run this as a monkey patch in production Dags that stream
Entra ID objects to S3 page by page.
* Resolve the next page before iter_paginated_run yields the current one
The generator yielded a page and read it again afterwards to resolve the
next request, so a consumer that changed the page (popping ``value`` after
writing its items, for instance) made ``default_pagination`` end a counted
query after the first request, silently.
The next page is now resolved and host-checked before the page is yielded,
so a yielded page belongs to the consumer and is never read again.
``paginated_run`` collects the page inside its pagination wrapper, which
keeps ``responses()`` including the page being paginated. A cross-host link
therefore raises before the page that carried it is yielded.
Tests: a consumer that empties every page still receives every item, a
custom ``pagination_function`` receives the expected arguments, and one
that declares ``responses`` is rejected by ``iter_paginated_run``.
---
.../providers/microsoft/azure/hooks/msgraph.py | 120 +++++++++++------
.../unit/microsoft/azure/hooks/test_msgraph.py | 148 +++++++++++++++++++++
2 files changed, 229 insertions(+), 39 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 adbf499e305..b0e62b85d45 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
@@ -33,7 +33,7 @@ import re
import warnings
from ast import literal_eval
from base64 import b64encode
-from collections.abc import Callable
+from collections.abc import AsyncIterator, Callable
from contextlib import suppress
from email.utils import parseaddr
from http import HTTPStatus
@@ -676,6 +676,60 @@ class KiotaRequestAdapterHook(BaseHook):
f"Microsoft Graph endpoints {sorted(allowed_netlocs)}."
)
+ async def iter_paginated_run(
+ self,
+ url: str = "",
+ response_type: str | None = None,
+ path_parameters: dict[str, Any] | None = None,
+ method: str = "GET",
+ query_parameters: dict[str, Any] | None = None,
+ headers: dict[str, str] | None = None,
+ data: dict[str, Any] | str | BytesIO | None = None,
+ pagination_function: PaginationCallable | None = None,
+ ) -> AsyncIterator[Any]:
+ """
+ Yield the pages of a paginated request one at a time.
+
+ The caller handles each page before the next one is requested, so a
collection of any size passes
+ through in page-sized steps instead of the single list
``paginated_run`` builds. The next page is
+ resolved by ``pagination_function`` (``default_pagination`` by
default: ``@odata.nextLink``, or
+ ``$skip`` on counted queries) and its host is checked with
``assert_allowed_host`` before the page
+ is yielded, so a yielded page is the caller's to change: it is not
read again. ``pagination_function``
+ receives the same arguments as in ``paginated_run`` except
``responses``: nothing keeps the pages
+ already yielded.
+ """
+ if pagination_function is None:
+ pagination_function = self.default_pagination
+
+ while url:
+ response = await self.run(
+ url=url,
+ response_type=response_type,
+ path_parameters=path_parameters,
+ method=method,
+ query_parameters=query_parameters,
+ headers=headers,
+ data=data,
+ )
+
+ if not response:
+ return
+
+ url, query_parameters = execute_callable(
+ pagination_function,
+ response=response,
+ url=url,
+ response_type=response_type,
+ path_parameters=path_parameters,
+ method=method,
+ query_parameters=query_parameters,
+ headers=headers,
+ data=data,
+ )
+ await self.assert_allowed_host(url)
+
+ yield response
+
async def paginated_run(
self,
url: str = "",
@@ -686,49 +740,37 @@ class KiotaRequestAdapterHook(BaseHook):
headers: dict[str, str] | None = None,
data: dict[str, Any] | str | BytesIO | None = None,
pagination_function: PaginationCallable | None = None,
- ):
+ ) -> list[Any]:
+ """
+ Collect the pages of a paginated request in a list.
+
+ The pages come from ``iter_paginated_run``; use that directly when the
collection should be handled
+ page by page. ``pagination_function`` additionally receives
``responses``, a callable returning the
+ pages collected so far, the one being paginated included.
+ """
if pagination_function is None:
pagination_function = self.default_pagination
- responses: list[dict] = []
+ responses: list[Any] = []
- async def run(
- url: str = "",
- query_parameters: dict[str, Any] | None = None,
- ):
- while url:
- response = await self.run(
- url=url,
- response_type=response_type,
- path_parameters=path_parameters,
- method=method,
- query_parameters=query_parameters,
- headers=headers,
- data=data,
- )
+ def paginate(response: Any, **kwargs: Any) -> tuple[str, dict[str,
Any] | None]:
+ # Called before the page is yielded, so the page is collected here
for ``responses`` to include it.
+ responses.append(response)
+ return execute_callable(
+ pagination_function, response=response, responses=lambda:
responses, **kwargs
+ )
- if response:
- responses.append(response)
-
- if pagination_function:
- next_url, query_parameters = execute_callable(
- pagination_function,
- response=response,
- url=url,
- response_type=response_type,
- path_parameters=path_parameters,
- method=method,
- query_parameters=query_parameters,
- headers=headers,
- data=data,
- responses=lambda: responses,
- )
- await self.assert_allowed_host(next_url)
- url = next_url
- else:
- break
-
- await run(url=url, query_parameters=query_parameters)
+ async for _ in self.iter_paginated_run(
+ url=url,
+ response_type=response_type,
+ path_parameters=path_parameters,
+ method=method,
+ query_parameters=query_parameters,
+ headers=headers,
+ data=data,
+ pagination_function=paginate,
+ ):
+ pass
return responses
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 27e4e0ca06e..bc46d6e2495 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
@@ -504,6 +504,154 @@ class TestKiotaRequestAdapterHook:
# token is never sent to the attacker host.
assert mock_get_http_response.call_count == 1
+ @pytest.mark.asyncio
+ async def
test_paginated_run_hands_the_pages_collected_so_far_to_the_pagination_function(self):
+ users = load_json_from_resources(dirname(__file__), "..", "resources",
"users.json")
+ next_users = load_json_from_resources(dirname(__file__), "..",
"resources", "next_users.json")
+ response = mock_json_response(200, users, next_users)
+ collected: list[list] = []
+
+ def pagination_function(response, responses):
+ collected.append(list(responses()))
+ return response.get("@odata.nextLink"), None
+
+ with patch_hook_and_request_adapter(response):
+ hook = KiotaRequestAdapterHook(conn_id="msgraph_api")
+
+ actual = await hook.paginated_run(url="users",
pagination_function=pagination_function)
+
+ assert actual == [users, next_users]
+ assert collected == [[users], [users, next_users]]
+
+ @pytest.mark.asyncio
+ async def test_iter_paginated_run(self):
+ users = load_json_from_resources(dirname(__file__), "..", "resources",
"users.json")
+ next_users = load_json_from_resources(dirname(__file__), "..",
"resources", "next_users.json")
+ response = mock_json_response(200, users, next_users)
+
+ with patch_hook_and_request_adapter(response):
+ hook = KiotaRequestAdapterHook(conn_id="msgraph_api")
+
+ actual = [page async for page in
hook.iter_paginated_run(url="users")]
+
+ assert actual == [users, next_users]
+
+ @pytest.mark.asyncio
+ async def
test_iter_paginated_run_requests_the_next_page_only_when_asked(self):
+ users = load_json_from_resources(dirname(__file__), "..", "resources",
"users.json")
+ next_users = load_json_from_resources(dirname(__file__), "..",
"resources", "next_users.json")
+ response = mock_json_response(200, users, next_users)
+
+ with patch_hook_and_request_adapter(response) as mocks:
+ mock_get_http_response = mocks[-1]
+ hook = KiotaRequestAdapterHook(conn_id="msgraph_api")
+ pages = hook.iter_paginated_run(url="users")
+
+ assert await anext(pages) == users
+ assert mock_get_http_response.call_count == 1
+
+ assert await anext(pages) == next_users
+ assert mock_get_http_response.call_count == 2
+
+ await pages.aclose()
+
+ @pytest.mark.asyncio
+ async def test_iter_paginated_run_stops_on_an_empty_response(self):
+ users = load_json_from_resources(dirname(__file__), "..", "resources",
"users.json")
+ response = mock_json_response(200, users, None)
+
+ with patch_hook_and_request_adapter(response) as mocks:
+ mock_get_http_response = mocks[-1]
+ hook = KiotaRequestAdapterHook(conn_id="msgraph_api")
+
+ actual = [page async for page in
hook.iter_paginated_run(url="users")]
+
+ assert actual == [users]
+ assert mock_get_http_response.call_count == 2
+
+ @pytest.mark.asyncio
+ async def test_iter_paginated_run_refuses_cross_host_next_link(self):
+ first_page = {
+ "@odata.nextLink":
"https://attacker.example/v1.0/users?$skiptoken=steal";,
+ "value": [{"id": "1"}],
+ }
+ response = mock_json_response(200, first_page)
+
+ with patch_hook_and_request_adapter(response) as mocks:
+ mock_get_http_response = mocks[-1]
+ hook = KiotaRequestAdapterHook(conn_id="msgraph_api")
+
+ with pytest.raises(ValueError, match="attacker.example"):
+ [page async for page in hook.iter_paginated_run(url="users")]
+
+ # The off-host pagination link is refused before it is fetched, so
the bearer
+ # token is never sent to the attacker host.
+ assert mock_get_http_response.call_count == 1
+
+ @pytest.mark.asyncio
+ async def test_iter_paginated_run_does_not_read_a_yielded_page_again(self):
+ first_page = {"@odata.count": 4, "value": [{"id": "1"}, {"id": "2"}]}
+ second_page = {"@odata.count": 4, "value": [{"id": "3"}, {"id": "4"}]}
+ response = mock_json_response(200, first_page, second_page, None)
+
+ with patch_hook_and_request_adapter(response) as mocks:
+ mock_get_http_response = mocks[-1]
+ hook = KiotaRequestAdapterHook(conn_id="msgraph_api")
+ ids = []
+
+ async for page in hook.iter_paginated_run(url="users",
query_parameters={"$top": 2}):
+ # The consumer owns the page once it is yielded: emptying it
must not end the pagination,
+ # which ``default_pagination`` decides from ``value`` and
``@odata.count``.
+ ids.extend(item["id"] for item in page.pop("value"))
+
+ assert ids == ["1", "2", "3", "4"]
+ assert mock_get_http_response.call_count == 3
+
+ @pytest.mark.asyncio
+ async def test_iter_paginated_run_with_a_custom_pagination_function(self):
+ users = load_json_from_resources(dirname(__file__), "..", "resources",
"users.json")
+ response = mock_json_response(200, users, users)
+ seen: list[tuple[str, dict | None]] = []
+
+ def skip_once(response, url, query_parameters):
+ seen.append((url, query_parameters))
+ if query_parameters and query_parameters.get("$skip"):
+ return None, query_parameters
+ return url, {"$top": 1, "$skip": 1}
+
+ with patch_hook_and_request_adapter(response) as mocks:
+ mock_get_http_response = mocks[-1]
+ hook = KiotaRequestAdapterHook(conn_id="msgraph_api")
+
+ actual = [
+ page
+ async for page in hook.iter_paginated_run(
+ url="users", query_parameters={"$top": 1},
pagination_function=skip_once
+ )
+ ]
+
+ assert actual == [users, users]
+ assert seen == [("users", {"$top": 1}), ("users", {"$top": 1,
"$skip": 1})]
+ assert mock_get_http_response.call_count == 2
+
+ @pytest.mark.asyncio
+ async def
test_iter_paginated_run_does_not_pass_responses_to_the_pagination_function(self):
+ users = load_json_from_resources(dirname(__file__), "..", "resources",
"users.json")
+
+ def needs_responses(response, responses):
+ return None, None
+
+ with patch_hook_and_request_adapter(mock_json_response(200, users)):
+ hook = KiotaRequestAdapterHook(conn_id="msgraph_api")
+
+ with pytest.raises(TypeError, match="responses"):
+ [
+ page
+ async for page in hook.iter_paginated_run(
+ url="users", pagination_function=needs_responses
+ )
+ ]
+
@pytest.mark.asyncio
async def test_assert_allowed_host_refuses_another_host(self):
with patch_hook_and_request_adapter(mock_json_response(200, {})):
