dabla commented on code in PR #75649:
URL: https://github.com/apache/airflow/pull/75649#discussion_r4411827574
##########
providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/msgraph.py:
##########
@@ -676,6 +676,59 @@ async def assert_allowed_host(self, url: str | None) ->
None:
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 it is
+ fetched. ``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
+
+ yield response
Review Comment:
Good catch, and it was worse than a docstring matter: with the page yielded
first, a consumer that changes it (say `page.pop("value")` after writing the
items of a counted `$top`/`@odata.count` query) made `default_pagination` see
an empty page and end the collection after the first request, with no error.
Rather than asking callers not to touch the page,
9e6b1d8b19f699937af738cd437f377361576095 makes the generator not depend on it:
the next page is resolved and host-checked **before** the page is yielded, so a
yielded page is the caller's and is never read again. `paginated_run` keeps
handing `responses` with the current page included by collecting the page in
its wrapper, since the pagination function now runs before the consumer sees
it. One visible consequence: a cross-host link now raises before the page that
carried it is yielded, instead of after it.
Tests added: a consumer that empties every page still gets all four items
from three requests; a custom `pagination_function` sees the arguments it
expects; and one that declares `responses` is rejected by `iter_paginated_run`
rather than fed an empty list. Docstring updated as well.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
