diff --git a/dataretrieval/transport/pagination.py b/dataretrieval/transport/pagination.py index 6199138cd..523463c62 100644 --- a/dataretrieval/transport/pagination.py +++ b/dataretrieval/transport/pagination.py @@ -1,4 +1,8 @@ -"""Callback-driven cursor pagination independent of any service protocol.""" +"""Callback-driven cursor pagination independent of any service protocol. + +:func:`paginate` is the page walk itself; :func:`run_paginated` composes +it with the shared executor, so an adapter supplies only its strategies. +""" from __future__ import annotations @@ -17,8 +21,18 @@ _safe_elapsed, ) from dataretrieval.exceptions import DataRetrievalError, RateLimited + +# One-way: ``fanout`` does not import this module, so this edge cannot cycle. +from dataretrieval.transport.fanout import ( + _CONCURRENCY_DEFAULT, + FanOut, + _Finalize, + _passthrough_result, + active_client, +) from dataretrieval.transport.http import open_async_client from dataretrieval.transport.liveness import note_progress +from dataretrieval.transport.retry import RetryPolicy logger = logging.getLogger(__name__) _Cursor = TypeVar("_Cursor") @@ -129,3 +143,53 @@ def report_page(page: httpx.Response, frame: pd.DataFrame) -> None: if row_cap is not None: result = result.head(row_cap) return result, final_response + + +def run_paginated( + requests: list[httpx.Request], + *, + parse_response: Callable[[httpx.Response], tuple[pd.DataFrame, Any]], + follow_up: Callable[[Any, httpx.AsyncClient], Awaitable[httpx.Response]], + raise_for_status: Callable[[httpx.Response], None], + service: str, + finalize: _Finalize = _passthrough_result, + client: httpx.AsyncClient | None = None, + client_options: dict[str, Any] | None = None, + default_concurrent: int = _CONCURRENCY_DEFAULT, + canonical_url: str | None = None, +) -> tuple[pd.DataFrame, Any]: + """Drive one full page walk per request through the shared executor. + + The adapter supplies its strategies (``parse_response``, ``follow_up``, + ``raise_for_status``, and optionally ``finalize``); this driver owns the + composition three adapters used to copy -- each request paginated on the + client the executor publishes unless ``client`` is injected, the retry + policy, bounded concurrency, and the canonical URL the aggregate reports + (the first request's, unless overridden). + + Raw transport errors need no mapping in the strategies: the executor + retries them and normalizes a deterministic one into the typed + :class:`~dataretrieval.exceptions.NetworkError`. + """ + + async def fetch(request: httpx.Request) -> tuple[pd.DataFrame, httpx.Response]: + return await paginate( + request, + parse_response=parse_response, + follow_up=follow_up, + client=client if client is not None else active_client(), + raise_for_status=raise_for_status, + ) + + if canonical_url is None and requests: + canonical_url = str(requests[0].url) + return FanOut( + requests, + fetch, + RetryPolicy.from_env(), + finalize=finalize, + client_options=client_options, + default_concurrent=default_concurrent, + canonical_url=canonical_url, + service=service, + ).resume() diff --git a/dataretrieval/waterdata/ratings.py b/dataretrieval/waterdata/ratings.py index 5aa979b95..e00c949d1 100644 --- a/dataretrieval/waterdata/ratings.py +++ b/dataretrieval/waterdata/ratings.py @@ -28,7 +28,7 @@ default_headers as _default_headers, ) from dataretrieval.transport.links import resolve_next_url -from dataretrieval.transport.pagination import paginate +from dataretrieval.transport.pagination import run_paginated from dataretrieval.transport.retry import RetryPolicy from dataretrieval.waterdata.endpoints import STAC_URL @@ -244,11 +244,8 @@ def _search( STAC ``next`` link is followed until exhausted so a result set larger than one page isn't silently truncated. - The page walk is :func:`dataretrieval.transport.pagination.paginate` with - STAC strategies, driven as a one-item - :class:`~dataretrieval.transport.fanout.FanOut` -- the same executor and - semantics (retry, stall budget, progress line, resumable interruption) as - every other getter. Pages carry features rather than rows, so each page + The page walk is :func:`~dataretrieval.transport.pagination.run_paginated` + with STAC strategies. Pages carry features rather than rows, so each page frame wraps the raw feature dicts in a single ``feature`` column. """ query_params: dict[str, Any] = {"limit": min(limit, 10000)} @@ -284,24 +281,14 @@ def parse_response(resp: httpx.Response) -> tuple[pd.DataFrame, str | None]: async def follow_up(cursor: str, sess: httpx.AsyncClient) -> httpx.Response: return await sess.get(cursor, headers=_default_headers(cursor)) - async def fetch(request: httpx.Request) -> tuple[pd.DataFrame, httpx.Response]: - return await paginate( - request, - parse_response=parse_response, - follow_up=follow_up, - # Borrow the executor's shared client for every page. - client=active_client(), - raise_for_status=_raise_for_non_200, - ) - - df, _ = FanOut( + df, _ = run_paginated( [req], - fetch, - RetryPolicy.from_env(), + parse_response=parse_response, + follow_up=follow_up, + raise_for_status=_raise_for_non_200, client_options={"verify": ssl_check}, - canonical_url=str(req.url), service="ratings", - ).resume() + ) # Every page frame is built with a ``feature`` column, and the combine # helpers preserve it, so the empty case needs no special branch. return list(df["feature"]) diff --git a/dataretrieval/waterdata/stats.py b/dataretrieval/waterdata/stats.py index 88ff751e0..9239f324f 100644 --- a/dataretrieval/waterdata/stats.py +++ b/dataretrieval/waterdata/stats.py @@ -26,10 +26,8 @@ _empty_feature_frame, _geo_feature_frame, ) -from dataretrieval.transport.fanout import FanOut, active_client from dataretrieval.transport.http import default_headers -from dataretrieval.transport.pagination import paginate -from dataretrieval.transport.retry import RetryPolicy +from dataretrieval.transport.pagination import run_paginated from dataretrieval.waterdata.endpoints import STATISTICS_API_URL __all__ = ["get_data"] @@ -273,27 +271,14 @@ async def follow_up(cursor: str, sess: httpx.AsyncClient) -> httpx.Response: method, url=url, params={**args, "next_token": cursor}, headers=headers ) - async def _fetch(request: httpx.Request) -> tuple[pd.DataFrame, httpx.Response]: - return await paginate( - request, - parse_response=parse_response, - follow_up=follow_up, - # Borrow the executor's shared client unless a caller injected one. - client=client if client is not None else active_client(), - raise_for_status=_raise_for_non_200, - ) - - # A one-item fan-out. Statistics has no chunkable axis, so the plan is a - # single request -- but running it through the same executor as every other - # getter is what gives it retry, the resumable interruption taxonomy, and - # the progress line, instead of a private sync bridge that had none of them. - df, response = FanOut( + df, response = run_paginated( [req], - _fetch, - RetryPolicy.from_env(), - canonical_url=str(req.url), + parse_response=parse_response, + follow_up=follow_up, + raise_for_status=_raise_for_non_200, + client=client, service=service, - ).resume() + ) if expand_percentiles: df = _expand_percentiles(df) diff --git a/dataretrieval/wateruse.py b/dataretrieval/wateruse.py index 7fe69e591..39fc23f01 100644 --- a/dataretrieval/wateruse.py +++ b/dataretrieval/wateruse.py @@ -52,11 +52,9 @@ from dataretrieval._response_metadata import BaseMetadata from dataretrieval.codes.states import to_state from dataretrieval.exceptions import DataRetrievalError -from dataretrieval.transport.fanout import FanOut, active_client -from dataretrieval.transport.http import default_headers, network_error +from dataretrieval.transport.http import default_headers from dataretrieval.transport.links import resolve_next_url -from dataretrieval.transport.pagination import paginate -from dataretrieval.transport.retry import RetryPolicy +from dataretrieval.transport.pagination import run_paginated __all__ = [ "get_wateruse", @@ -341,18 +339,12 @@ def _fan_out( ) -> tuple[pd.DataFrame, BaseMetadata]: """Fetch every request (each paginated) over the shared fan-out executor. - Each request is paginated by :func:`dataretrieval.transport.pagination.paginate` - with NWDC strategies: parse a CSV page and read its ``Link`` header cursor - (``parse``), follow that cursor (``follow``), and raise the typed error - carrying the NWDC ``detail`` (``raise_for_status``). + This function is only the NWDC-specific half: parse a CSV page and read + its ``Link`` header cursor, follow that cursor, raise the typed error + carrying the NWDC ``detail``, and shape the result. + :func:`~dataretrieval.transport.pagination.run_paginated` owns the rest. - Everything else -- bounded concurrency, per-attempt retry, failure - precedence, progress, and resumable interruption -- belongs to - :class:`~dataretrieval.transport.fanout.FanOut`, which Water Data and NGWMN - drive too. This function is now only the NWDC-specific half: what a - chunk is, and how to read one. - - The plan is the request list itself. ``FanOut`` asks a plan only to be + The plan is the request list itself. The executor asks a plan only to be sized and iterable, and the NWDC accepts one ``location=`` per request, so the caller's locations arrive already separate -- there is nothing to divide and so nothing for a plan class to hold. @@ -371,43 +363,21 @@ async def follow(cursor: str, sess: httpx.AsyncClient) -> httpx.Response: def raise_for_status(response: httpx.Response) -> None: _raise_for_status(response, detail_from=_nwdc_error_detail) - async def fetch(request: httpx.Request) -> tuple[pd.DataFrame, httpx.Response]: - """One location's full page walk, over the executor's shared client. - - ``active_client()`` is the client :meth:`FanOut._run` published for this - run; borrowing it keeps every location's pages on one connection pool - instead of opening a client per location. - """ - try: - return await paginate( - request, - parse_response=parse, - follow_up=follow, - client=active_client(), - raise_for_status=raise_for_status, - ) - except httpx.TransportError as exc: - raise network_error(request.url, exc) from exc - def finalize( frame: pd.DataFrame, response: httpx.Response ) -> tuple[pd.DataFrame, BaseMetadata]: return frame, BaseMetadata(response) - return FanOut( + return run_paginated( requests, - fetch, - RetryPolicy.from_env(), + parse_response=parse, + follow_up=follow, + raise_for_status=raise_for_status, finalize=finalize, client_options={"verify": ssl_check}, default_concurrent=DEFAULT_CONCURRENT_REQUESTS, - # No single URL expresses "all of these locations" -- the service - # has no such request -- so the aggregate reports the first, - # matching what an un-fanned single-location call would show. - canonical_url=str(requests[0].url) if requests else None, - # Labels the progress line the executor opens for this drive. service="wateruse", - ).resume() + ) def _read_csv_page(response: httpx.Response) -> pd.DataFrame: diff --git a/tests/wateruse_test.py b/tests/wateruse_test.py index 5e88f0060..9e003e65b 100644 --- a/tests/wateruse_test.py +++ b/tests/wateruse_test.py @@ -762,7 +762,9 @@ async def fail(*_args, **_kwargs): failure.__context__ = resolution raise failure - monkeypatch.setattr(wateruse, "paginate", fail) + from dataretrieval.transport import pagination + + monkeypatch.setattr(pagination, "paginate", fail) with pytest.raises(dataretrieval.NetworkError) as excinfo: get_wateruse(model="wu-public-supply-wd", state="RI")