From d0e56c07ce568f66216e86e10b8c583242b7c646 Mon Sep 17 00:00:00 2001 From: thodson-usgs Date: Fri, 14 Aug 2026 14:21:36 -0500 Subject: [PATCH] feat(transport)!: page Water Data results in parallel Fetch large Water Data result sets in bounded offset waves while preserving cursor fallback, exact row limits, request state, and byte-driven chunking. BREAKING CHANGE: remove the parallel_chunks API and ChunkPlan.max_chunks. --- .importlinter | 7 +- NEWS.md | 2 + README.md | 89 ++- dataretrieval/__init__.py | 7 - dataretrieval/ogc/chunking.py | 161 +---- dataretrieval/ogc/engine.py | 106 +++- dataretrieval/ogc/planning.py | 120 +--- dataretrieval/ogc/policy.py | 8 + dataretrieval/ogc/requests.py | 24 + dataretrieval/transport/fanout.py | 25 +- dataretrieval/transport/offsets.py | 552 ++++++++++++++++++ dataretrieval/waterdata/__init__.py | 2 - dataretrieval/waterdata/utils.py | 6 + .../0006-service-neutral-transport.rst | 31 +- .../decisions/0008-fan-out-execution.rst | 5 +- docs/source/userguide/errors.rst | 69 ++- tests/architecture_test.py | 2 +- tests/contracts/public_api_test.py | 1 - tests/transport_test.py | 258 ++++++++ tests/utils_test.py | 18 +- tests/waterdata_chunking_test.py | 321 ++-------- tests/waterdata_offset_paging_test.py | 336 +++++++++++ 22 files changed, 1513 insertions(+), 637 deletions(-) create mode 100644 dataretrieval/transport/offsets.py create mode 100644 tests/waterdata_offset_paging_test.py diff --git a/.importlinter b/.importlinter index 0b993b9c..68967c40 100644 --- a/.importlinter +++ b/.importlinter @@ -54,17 +54,12 @@ type = protected ; The root ``dataretrieval`` package is deliberately NOT an allowed importer. ; ``allowed_importers`` is matched with ``as_packages``, so naming the root here ; would make every module in the distribution an allowed importer and the -; contract could never fail. Its two real imports are listed as explicit -; exceptions below instead -- narrow, visible, and they fail if they go stale. +; contract could never fail. protected_modules = dataretrieval.ogc allowed_importers = dataretrieval.ngwmn dataretrieval.waterdata -ignore_imports = -; The package __init__ re-exports the parallel-chunks context manager; it is -; part of the documented public surface, not a service reaching into OGC. - dataretrieval -> dataretrieval.ogc.chunking [importlinter:contract:ogc-facade] name = Facade-only OGC consumers, never its internals (ADR 0007) diff --git a/NEWS.md b/NEWS.md index b4abd5ec..4f6e2080 100644 --- a/NEWS.md +++ b/NEWS.md @@ -8,6 +8,8 @@ **08/06/2026:** Fan-out execution is now shared across services. Chunking is how a query is divided structurally (a Water Data/NGWMN URL over the byte limit); fan-out is how the pieces are distributed operationally. Only the first is protocol-specific, so the executor moved to `dataretrieval.transport.fanout` (`FanOut`, over a three-member `FanOutPlan` protocol) while chunk planning stays in `dataretrieval.ogc`. Water Use no longer re-implements the fan-out gather and inherits resume, progress reporting, and `API_USGS_CONCURRENT`: a multi-location pull interrupted by a rate limit now raises a resumable interruption whose `.call.resume()` re-issues only the locations that did not finish, instead of discarding every completed one. The interruption taxonomy moved to the `dataretrieval.interruptions` leaf and its base class is now `FanOutInterrupted`; **`ChunkInterrupted` is a permanent alias of the same class**, so `except ChunkInterrupted` keeps working. **Breaking change:** a Water Use fan-out interrupted by a 5xx, 429, or recoverable connection failure now raises `ServiceInterrupted`/`QuotaExhausted` rather than `ServiceUnavailable`/`RateLimited`/`NetworkError` — all remain `DataRetrievalError`, so broad handlers are unaffected, but narrow handlers around a Water Use call must widen. **Breaking change:** `wateruse.MAX_CONCURRENT_REQUESTS` is removed; set `API_USGS_CONCURRENT` (which now outranks any service default) or read `wateruse.DEFAULT_CONCURRENT_REQUESTS`. +**08/06/2026:** Large Water Data pulls are now paged in parallel automatically. A multi-page result previously cost one round trip per page, because a cursor page's URL only exists once the previous page has been parsed. Where the API honors `offset`, every page's URL is computable up front, so the pages are fetched concurrently — the request count is unchanged, only their timing, and since the USGS quota is volume-based the speedup costs no extra quota. Measured 2.1× on 8 sites × 2 years (16,000 rows) and 3.1–4.2× on a single site's full daily history (~19,000 rows). **Breaking change:** `parallel_chunks(n)` is removed (along with `ChunkPlan.max_chunks`), because it bought parallelism the other way — splitting a request that already fit into more sub-requests, which spent extra quota and did nothing at all for a single-site query, the case with no multi-value axis to split. Delete the `with parallel_chunks(...):` wrapper; the pages inside it are now overlapped without it. Byte-driven chunking is unchanged and still required for correctness (the ~8 KB URL limit is real). `API_USGS_CONCURRENT` now also caps the page-fetch wave width; set it to `1` to page strictly sequentially. Two fallbacks keep results correct rather than merely fast: a server that ignores `offset` (a non-standard extension, so ignoring it is conventional) is detected before any rows are returned and the query re-runs via standard `next`-link paging, and the API's hard `offset` ceiling of 40,000 hands the remainder off to the sequential walk, so an arbitrarily deep pull is still returned in full. + **08/03/2026:** Split the Water Data implementation into focused time-series, metadata, measurements, reference, samples, and CQL collection-family modules behind the unchanged `waterdata.api` facade. Active service modules now declare explicit exports; public Water Data imports, signatures, function identities, deprecations, and return contracts are protected by executable contract snapshots. OGC ambient context and schema/queryables execution are separated from request construction, adapter-to-adapter reach-through is prohibited by architecture tests, and service-specific output shapes are documented rather than forced into one model. **08/02/2026:** Added an internal API-neutral transport layer for guarded HTTP clients, host-scoped authentication, cursor pagination, bounded retry, response aggregation, progress, and sync-over-async dispatch. Water Use and the non-OGC Statistics API now consume transport directly instead of private OGC execution helpers; WQP, NLDI, and StreamStats opt into bounded transient retry while deprecated NWIS behavior remains unchanged. OGC retains CQL2, request construction, feature shaping, chunk planning, resumable calls, and interruption types, with compatibility imports at previous private paths. Failed pagination and fan-out still raise rather than returning partial data, and no public signatures or return shapes changed. diff --git a/README.md b/README.md index 41b3299c..edcab069 100644 --- a/README.md +++ b/README.md @@ -107,57 +107,52 @@ df, metadata = waterdata.get_continuous( print(f"Retrieved {len(df)} continuous gage height measurements") ``` -#### Speeding up large downloads with `parallel_chunks` - -By default the getters split a multi-value request only as far as the server's -~8 KB URL limit forces — the fewest sub-requests. For a **large, paginated** -pull, that default is needlessly conservative: every sub-request pages through -its own results, so dividing the query into more, smaller sub-requests lets -those pages be fetched **in parallel**. `parallel_chunks(n)` opts a single call -into that finer split, fanning it out into `n` sub-requests. The finer split -pays off only when the result is large enough to span many pages *and* the query -has a multi-value argument to divide, such as a list of monitoring locations. On -a small query — or one with nothing to split — it only adds requests, so -`parallel_chunks` is a deliberate, scoped `with` block, never the default. +#### Large Water Data downloads are paged in parallel automatically + +A large result arrives one page at a time. Standard OGC cursor pagination is +inherently sequential: page *N+1*'s URL is revealed only by page *N*. The Water +Data API also accepts `offset`, which makes page URLs computable up front +(`offset = i * limit`). `dataretrieval` therefore overlaps pages automatically; +there is no context manager to enable. ```python from dataretrieval import waterdata -# All stream gages in Ohio, then 20 years of their daily discharge — large -# enough to span many pages, so it profits from a finer split. -sites, _ = waterdata.get_monitoring_locations(state="Ohio", site_type_code="ST") - -with waterdata.parallel_chunks(32): # fan out into 32 sub-requests - df, md = waterdata.get_daily( - monitoring_location_id=sites["monitoring_location_id"], - parameter_code="00060", # discharge - time="2004-01-01/2023-12-31", - ) +# A deep daily-value history can span many pages even for one location. +df, md = waterdata.get_daily( + monitoring_location_id="USGS-01646500", + parameter_code="00060", # discharge + time="2004-01-01/2023-12-31", + limit=2000, +) ``` -`n` is the number of sub-requests to fan the call out into, capped by how many -values there are to split. Each sub-request costs a request against your hourly -[rate limit](https://api.waterdata.usgs.gov/signup/). How many run *at once* is -capped separately by `API_USGS_CONCURRENT` (default 32), so the useful range is -roughly `2` up to that value. - -Benchmark — a fixed 271-site subset of Ohio stream gages -(`get_daily`, `parameter_code="00060"`), with a small fixed page size -(`limit=250`) so every run fetches roughly the same number of pages (isolating -the effect of parallelism). Each `n` ran against its own cold 1-year time -window, so no run is served from the server's data-window cache: - -| `n` | parallelism | pages | wall-clock | speedup | -| ---- | ----------- | ----- | ----------------------- | ------- | -| off | 1 | ~30 | 9.5 s / 9.1 s (2 runs) | 1× | -| `8` | 8 | ~32 | 2.2 s / 1.9 s | ~4.5× | -| `32` | 32 | 54 | 1.2 s | ~8× | - -The gain comes from overlapping each sub-request's per-page latency and -server-side work. The exact multiplier therefore scales with how many pages the -pull spans: a larger pull (more pages) has more parallelism to exploit. The -extra sub-requests each cost quota, so reserve a large `n` for pulls you know -are large. +At a **fixed `limit`**, this changes timing rather than deliberately splitting +the query into extra requests. The unknown page count is discovered with +ramped speculative waves (1, 2, 4, ... requests), so the final wave can probe +past the end, but total requests remain below twice the pages needed. Reducing +`limit` creates more pages and therefore spends more hourly quota; the measured +speedups are not free if you shrink `limit` to obtain them. + +The default `limit` is 50,000 while Water Data accepts offsets only through +40,000. At that default there is no useful offset fan-out: the unbounded tail is +walked through standard cursors. Material speedups therefore require an +explicit paging-friendly `limit` below the offset ceiling, such as the `2000` +above. `API_USGS_CONCURRENT` (default 32) bounds the page-wave width; set it to +`1` to use cursor pagination strictly sequentially. + +Because `offset` is a Water Data extension rather than part of OGC API - +Features, the walk is defensive. It verifies that the server honors the +parameter before trusting rows and falls back to standard cursor pagination if +not. At the API's 40,000-row offset ceiling it rewinds one page and follows the +standard `next` links, keeping arbitrarily deep results complete without a gap +or duplicate seam. + +The removed `parallel_chunks(n)` API split a fitting multi-value query into more +chunks, spending extra quota and doing nothing for a single-location query. +Delete that wrapper when migrating. Byte-driven chunking remains automatic and +unchanged: queries above the service's ~8 KB request limit are still split for +correctness. Visit the [API Reference](https://doi-usgs.github.io/dataretrieval-python/reference/waterdata.html) @@ -176,8 +171,8 @@ logging.basicConfig(level=logging.DEBUG) ### National Ground-Water Monitoring Network (NGWMN) Access groundwater data aggregated from many state, federal, and local -agencies. NGWMN uses the same OGC engine as the Water Data API, -so chunking and pagination behave the same way: +agencies. NGWMN uses the same OGC engine and byte-driven chunker as Water Data. It +uses standard cursor pagination unless its API dialect declares offset support: ```python from dataretrieval import ngwmn diff --git a/dataretrieval/__init__.py b/dataretrieval/__init__.py index 2ed3d415..91c325bc 100644 --- a/dataretrieval/__init__.py +++ b/dataretrieval/__init__.py @@ -62,11 +62,6 @@ ServiceInterrupted, ) -# Parallel-chunks control (a context manager). Defined with the chunker in -# ``dataretrieval.ogc.chunking``; surfaced here for a stable public path -# ``from dataretrieval import parallel_chunks``. -from dataretrieval.ogc.chunking import parallel_chunks - from . import ( exceptions, ngwmn, @@ -109,7 +104,5 @@ "FanOutInterrupted", "QuotaExhausted", "ServiceInterrupted", - # parallel-chunks control (defined in ogc.chunking) - "parallel_chunks", "__version__", ] diff --git a/dataretrieval/ogc/chunking.py b/dataretrieval/ogc/chunking.py index c958b49d..81d8aa32 100644 --- a/dataretrieval/ogc/chunking.py +++ b/dataretrieval/ogc/chunking.py @@ -8,21 +8,16 @@ chunk URL under the budget. Requests that already fit get a trivial single-step plan — the executor has one code path either way. -This module owns the OGC-specific half: the byte budget, the -``parallel_chunks`` dial, and the ``multi_value_chunked`` decorator that -ties a plan to a fetcher. Driving the resulting chunks to -completion — bounded concurrency, retry, failure precedence, resume — is -API-neutral and belongs to -:class:`dataretrieval.transport.fanout.FanOut`, which this module hands -its plan to. :class:`~dataretrieval.ogc.planning.ChunkPlan` satisfies +This module owns the OGC-specific half: the byte budget and the +``multi_value_chunked`` decorator that ties a plan to a fetcher. The planner +splits only as far as the byte limit forces. Parallelism within each chunk is +the offset page walk's job, which overlaps pages without manufacturing extra +chunks. Driving chunks to completion — bounded concurrency, retry, failure +precedence, resume — is API-neutral and belongs to +:class:`dataretrieval.transport.fanout.FanOut`, which this module hands its plan +to. :class:`~dataretrieval.ogc.planning.ChunkPlan` satisfies :class:`~dataretrieval.transport.fanout.FanOutPlan` structurally. -Parallel chunks: the planner is conservative by default — it splits only as -far as the byte limit forces. A caller who knows their result is large can opt -into a finer split via the ``parallel_chunks(n)`` context manager, which fans -the query out into ``n`` parallel chunks. ``n`` drives -:meth:`ChunkPlan._refine`; see ``parallel_chunks`` for the why and the when. - Concurrency, retries, and interruption semantics are documented on :mod:`dataretrieval.transport.fanout`; ``API_USGS_CONCURRENT`` and ``API_USGS_RETRIES`` are read there. @@ -37,26 +32,25 @@ from __future__ import annotations import functools -from collections.abc import Callable, Iterator -from contextlib import contextmanager +from collections.abc import Callable from typing import Any import httpx import pandas as pd -from dataretrieval._ambient import Ambient from dataretrieval.transport.fanout import ( + _CONCURRENCY_DEFAULT, FanOut, _active_client, _Fetch, _Finalize, _passthrough_result, + _resolve_concurrency, active_client, ) from dataretrieval.transport.retry import RetryPolicy from .planning import ChunkPlan -from .policy import _require_positive_int # Compatibility aliases. ``ChunkedCall`` was this module's executor before it # moved down to transport as the API-neutral ``FanOut``; ``get_active_client`` @@ -78,116 +72,15 @@ _OGC_URL_BYTE_LIMIT = 8000 -# Parallel-chunks dial: opt-in to fan a query out *more finely* than the byte -# limit alone requires. Scoped to a ``with parallel_chunks(...):`` block (a -# ContextVar), deliberately NOT an env var (see :func:`parallel_chunks` for -# why). The ambient holds ``n`` — the requested cap on the plan's total -# chunk count; ``1`` (the default, outside any block) means "off — chunk -# only as much as the byte limit needs, no extra fan-out". -_parallel_chunks: Ambient[int] = Ambient("ogc_parallel_chunks", 1) - - -@contextmanager -def parallel_chunks(n: int) -> Iterator[None]: - """ - Fan the OGC getters' multi-value requests out into ``n`` parallel chunks. - - By default the Water Data / NGWMN getters chunk a request only as much as - the server's ~8 KB URL-byte limit forces — the fewest chunks that - fit. That is the safe default, but it can be *needlessly* conservative. - Because every chunk paginates, splitting a large result further costs - little or no extra quota *as long as each chunk still spans many - pages* — rows-per-chunk far exceeding the page size (ten states pulled as - one request page nearly as many times as ten per-state requests would). - When a split leaves each chunk only a page or two, its partial final - page is extra, so finer chunks do add some requests. This context manager - lets a caller who *knows* their pull is large ask for that finer split. The - trade is roughly the same pages for more, smaller chunks, which gives - smoother progress, more even concurrency, and a smaller unit of - retry/resume. - - This is a *deliberate* per-call knob rather than an automatic behavior or a - process-wide environment variable, because the library can't tell in - advance whether a query is large (ten states over a short window might fit - in a single page, where extra chunks would only burn quota). Scoping it to - a ``with`` block keeps an aggressive setting from leaking into unrelated - calls and accidentally spending quota. Outside any block the getters use - the conservative default. Only the OGC getters (Water Data, NGWMN) read - this; wrapping a legacy NWIS call in the block is a harmless no-op. +def page_concurrency() -> int: + """Return the bounded page-wave width for offset pagination. - Parameters - ---------- - n : int - The number of chunks to fan the whole call out into — a positive - integer such as ``2``, ``8``, or ``32``. It caps the plan's *total* - chunk count (the cartesian product across every multi-value - argument combined, not per argument), so several multi-value arguments - cannot multiply past it. The cap is a ceiling, never exceeded: the - actual count is bounded below by what the ~8 KB URL limit already - forces and above by ``n``. So an ``n`` larger than the input allows - simply yields one chunk per value, and with several multi-value - arguments the total may land somewhat below ``n`` because splits are - whole (the plan can't always divide evenly onto ``n``). ``n=1`` asks - for no extra fan-out. - - Each chunk fetches at least one page, so it costs at least one - request against your hourly rate limit — a larger ``n`` spends more - quota. How many chunks run *at once* is capped separately by - ``API_USGS_CONCURRENT`` (default 32), so an ``n`` beyond that adds - quota without adding parallelism; the useful range is roughly ``2`` - up to ``API_USGS_CONCURRENT``. - - Yields - ------ - None - - Raises - ------ - ValueError - If ``n`` is not a positive integer — raised on ``with`` entry, before - any request is issued, so a bad value fails loudly rather than silently - doing nothing. - - Notes - ----- - Fanning out carries the same consequences as the byte-limit chunking the - getters already do for oversized requests; opting in just brings them to a - request that would otherwise be a single call: - - - ``max_rows``: each chunk paginates up to ``max_rows`` rows - independently, then the combined result is sorted and truncated to - ``max_rows``. So a call with ``max_rows`` set returns a *different* - (though still valid and deterministically sorted) row set inside a - ``parallel_chunks`` block than without one. The cap is drawn from the - union of the chunks, not a single stream. Don't pair a tight - ``max_rows`` preview with ``parallel_chunks`` if you need exactly the - rows the un-fanned call would return. - - Resumability: a single request either fully succeeds or fully fails, - but a fanned-out call can fail partway (e.g. a mid-call rate-limit) and - raise a resumable :class:`~dataretrieval.ogc.interruptions.ChunkInterrupted` - (or ``QuotaExhausted``) carrying the completed chunks. Finish the - call with ``exc.call.resume()``. - - Cross-chunk de-duplication keys on the feature ``id``; features - with no ``id`` can't be deduped, so overlapping filter clauses split - across chunks may yield duplicate rows. - - Examples - -------- - >>> from dataretrieval import waterdata - >>> with waterdata.parallel_chunks(32): - ... df, md = waterdata.get_daily( - ... monitoring_location_id=many_sites, parameter_code="00060" - ... ) # doctest: +SKIP - - See Also - -------- - ChunkPlan._refine : the planning-side effect of ``n``. + ``API_USGS_CONCURRENT`` governs both chunk fan-out and page waves. The + ``unbounded`` chunk setting is clamped here because a speculative page wave + must remain finite. """ - # Fail loudly on a bad ``n`` at ``with`` entry, before any request. Shared - # rules with ``max_rows`` via the helper (accepts numpy ints, rejects bool). - _require_positive_int(n, "parallel_chunks(n)", examples="2, 8, 32") - with _parallel_chunks(n): - yield + resolved = _resolve_concurrency(_CONCURRENCY_DEFAULT) + return _CONCURRENCY_DEFAULT if resolved is None else resolved def multi_value_chunked( @@ -203,9 +96,9 @@ def multi_value_chunked( ``async def fetch(args) -> (df, response)``, and drives it to completion via :meth:`ChunkedCall.resume`. The plan splits multi-value list params and the cql-text filter so each chunk URL fits the - byte limit. An already-fitting request is a one-step plan, unless an - active :func:`parallel_chunks` block asks the plan to fan out more - finely. See the module docstring for the concurrency model. + byte limit. An already-fitting request is a one-step plan. Each chunk + then fetches its pages through the strategy selected by the OGC engine. + See the module docstring for the concurrency model. Parameters ---------- @@ -251,14 +144,7 @@ def wrapper( finalize: _Finalize = _passthrough_result, ) -> tuple[pd.DataFrame, Any]: limit = _OGC_URL_BYTE_LIMIT if url_limit is None else url_limit - # Read the parallel_chunks dial ``n`` from the ambient set by - # ``parallel_chunks`` (1 = off outside any such block; otherwise the - # requested total chunk cap). It only affects *planning*, done - # here up front, so a later resume — which re-issues the - # already-planned chunks — needs no snapshot. - plan = ChunkPlan( - args, build_request, limit, max_chunks=_parallel_chunks.get() - ) + plan = ChunkPlan(args, build_request, limit) retry_policy = RetryPolicy.from_env() # The concurrency cap is resolved inside ``resume()`` from # ``API_USGS_CONCURRENT``; ``1`` is a sequential gather, @@ -272,6 +158,9 @@ def wrapper( # The collection name, for the progress line the executor # opens. ``get_ogc_data`` puts it in ``args``. service=args.get("collection"), + # One chunk can fan out a page wave independently of the chunk + # semaphore, so the shared pool must cover both dimensions. + connection_multiplier=page_concurrency, ).resume() return wrapper diff --git a/dataretrieval/ogc/engine.py b/dataretrieval/ogc/engine.py index 67895f11..115cd1f7 100644 --- a/dataretrieval/ogc/engine.py +++ b/dataretrieval/ogc/engine.py @@ -49,10 +49,13 @@ _construct_cql_request, _switch_arg_id, _switch_properties_id, + page_limit, + with_offset, ) from dataretrieval.ogc.shaping import GEOPANDAS, _finalize_ogc, _get_resp_data from dataretrieval.transport.fanout import FanOut from dataretrieval.transport.links import resolve_next_url +from dataretrieval.transport.offsets import OffsetUnsupported, paginate_by_offset from dataretrieval.transport.pagination import paginate from dataretrieval.transport.retry import RetryPolicy @@ -116,6 +119,11 @@ def _next_req_url( return None +def _ogc_parse_page(resp: httpx.Response, *, geopd: bool) -> pd.DataFrame: + """Parse one OGC page for offset pagination, where no cursor is needed.""" + return _get_resp_data(resp, geopd=geopd) + + def _ogc_parse_response( resp: httpx.Response, *, geopd: bool ) -> tuple[pd.DataFrame, str | None]: @@ -198,6 +206,82 @@ async def follow_up(cursor: str, sess: httpx.AsyncClient) -> httpx.Response: ) +async def _walk_pages_by_offset( + geopd: bool, + req: httpx.Request, + client: httpx.AsyncClient | None = None, + *, + width: int, + max_offset: int | None, + row_cap: int | None = None, +) -> tuple[pd.DataFrame, httpx.Response]: + """Fetch computable offset pages concurrently, then cursor-walk the tail.""" + limit = page_limit(req) + if limit is None: + return await _walk_pages( + geopd=geopd, + req=req, + client=client, + row_cap=row_cap, + ) + + async def tail_walk( + resume_offset: int, rows_so_far: int, session: httpx.AsyncClient + ) -> tuple[pd.DataFrame, httpx.Response]: + remaining = None if row_cap is None else max(row_cap - rows_so_far, 0) + return await _walk_pages( + geopd=geopd, + req=with_offset(req, resume_offset), + client=session, + row_cap=remaining, + ) + + try: + return await paginate_by_offset( + build_page=functools.partial(with_offset, req), + parse_page=functools.partial(_ogc_parse_page, geopd=geopd), + raise_for_status=_raise_for_non_200, + client=client, + limit=limit, + width=width, + max_offset=max_offset, + row_cap=row_cap, + tail_walk=tail_walk, + ) + except OffsetUnsupported as exc: + logger.warning( + "Falling back to sequential pagination: %s This is slower but " + "uses only standard OGC API - Features paging.", + exc, + ) + return await _walk_pages( + geopd=geopd, + req=req, + client=client, + row_cap=row_cap, + ) + + +async def _walk_request( + geopd: bool, + req: httpx.Request, + *, + dialect: OgcDialect, + row_cap: int | None = None, +) -> tuple[pd.DataFrame, httpx.Response]: + """Select cursor or offset pagination once for one prepared request.""" + width = chunking.page_concurrency() + if dialect.max_offset is None or width <= 1: + return await _walk_pages(geopd=geopd, req=req, row_cap=row_cap) + return await _walk_pages_by_offset( + geopd, + req, + width=width, + max_offset=dialect.max_offset, + row_cap=row_cap, + ) + + def get_ogc_data( args: dict[str, Any], collection: str, @@ -271,7 +355,7 @@ def get_ogc_data( # Enforce a genuine positive integer up front: a float (even ``10.0``) or # ``bool`` would pass a bare ``< 1`` check and then crash deep in # ``pd.DataFrame.head`` with an opaque ``TypeError`` after HTTP I/O has - # already fired. Shared with ``parallel_chunks(n)`` via the helper. + # already fired. Validated by the shared OGC count helper. if max_rows is not None: _require_positive_int(max_rows, "max_rows") @@ -329,7 +413,12 @@ def get_ogc_data( # ``(df, BaseMetadata)`` shape rather than a raw response pair. return FanOut( [req], - functools.partial(_walk_pages, GEOPANDAS, row_cap=max_rows), + functools.partial( + _walk_request, + GEOPANDAS, + dialect=dialect, + row_cap=max_rows, + ), RetryPolicy.from_env(), finalize, canonical_url=str(req.url), @@ -346,7 +435,10 @@ def get_ogc_data( _construct_api_requests, base_url=base_url, dialect=dialect ) fetch = functools.partial( - _fetch_once, build_request=build_request, row_cap=max_rows + _fetch_once, + build_request=build_request, + dialect=dialect, + row_cap=max_rows, ) run = chunking.multi_value_chunked(build_request=build_request)(fetch) # No progress block here: the executor that emits the events owns the line @@ -358,6 +450,7 @@ async def _fetch_once( args: dict[str, Any], *, build_request: Callable[..., httpx.Request], + dialect: OgcDialect, row_cap: int | None = None, ) -> tuple[pd.DataFrame, httpx.Response]: """Send one prepared-args OGC request asynchronously; return (frame, response). @@ -375,4 +468,9 @@ async def _fetch_once( synchronously. The return shape is ``(frame, response)``. """ req = build_request(**args) - return await _walk_pages(geopd=GEOPANDAS, req=req, row_cap=row_cap) + return await _walk_request( + GEOPANDAS, + req, + dialect=dialect, + row_cap=row_cap, + ) diff --git a/dataretrieval/ogc/planning.py b/dataretrieval/ogc/planning.py index e5af5411..2f1bf082 100644 --- a/dataretrieval/ogc/planning.py +++ b/dataretrieval/ogc/planning.py @@ -243,12 +243,12 @@ def _extract_axes(args: dict[str, Any]) -> list[_Axis]: def _split_at(chunks: list[list[str]], idx: int) -> None: """Replace ``chunks[idx]`` in place with its two contiguous halves. - The single primitive both planning passes use to fan an axis out. It + The primitive the byte-driven planning pass uses to fan an axis out. It preserves the partition invariants every consumer relies on: *coverage* (each atom survives, exactly once) and *contiguous, deterministic order* (resume and :meth:`ChunkPlan.iter_chunk_args` depend on it). Kept in one - place so those invariants can't drift between :meth:`ChunkPlan._plan` - (byte-driven) and :meth:`ChunkPlan._refine` (fan-out-driven). + place so the partition invariants remain local to + :meth:`ChunkPlan._plan`. """ chunk = chunks[idx] mid = len(chunk) // 2 @@ -280,20 +280,6 @@ class ChunkPlan: url_limit : int Byte budget for the request (URL + body) — a hard ceiling every chunk must fit. - max_chunks : int, optional - Hard cap on the plan's total chunk count (default ``1`` = off). - ``1`` chunks only as much as ``url_limit`` requires — the most - conservative plan, fewest chunks — so a fitting request is a - passthrough. A cap of ``2`` or more fans the plan out to up to - ``max_chunks`` chunks overall (the cartesian product across axes, - never fewer than the byte budget already forces). The cap applies to - the plan as a whole, not per axis, so several multi-value axes can't - multiply past it. The plan never exceeds the cap and may land below it - when no whole split lands on it exactly. ``max_chunks`` is a - chunk count, so a value below ``1`` (``0`` or negative) is a - caller error and raises ``ValueError``. Set from the - :func:`~dataretrieval.ogc.chunking.parallel_chunks` ``n``; see - :meth:`_refine`. Attributes ---------- @@ -318,8 +304,6 @@ class ChunkPlan: Unchunkable If the request needs chunking but even the singleton plan doesn't fit ``url_limit``. - ValueError - If ``max_chunks`` is less than 1 (0 or negative). """ def __init__( @@ -327,19 +311,7 @@ def __init__( args: dict[str, Any], build_request: Callable[..., httpx.Request], url_limit: int, - max_chunks: int = 1, ) -> None: - if max_chunks < 1: - # ``max_chunks`` is a chunk *count*: the minimum is ``1`` - # (the ambient default outside any ``parallel_chunks`` block), - # which means "off — no extra fan-out". ``0`` or negative is a - # meaningless count and can only be a caller bug, so fail loudly - # rather than silently no-op. The public ``parallel_chunks(n)`` - # already rejects ``n < 1``; this guards direct construction. - raise ValueError( - f"max_chunks must be >= 1 (1 disables fan-out); got {max_chunks!r}." - ) - self.args = args self.axes: list[_Axis] = [] self.chunks: dict[str, list[list[str]]] = {} @@ -347,8 +319,8 @@ def __init__( axes = _extract_axes(args) if not axes: - # No chunkable axis: nothing to split, and ``parallel_chunks`` has - # nothing to act on either. If the single request fits, run it + # No chunkable axis: nothing to split. If the single request fits, + # run it # verbatim (the common passthrough). ``_safe_request_bytes`` treats # an un-constructable URL (httpx.InvalidURL, > 64 KB) as over budget. if _safe_request_bytes(build_request, args, url_limit) <= url_limit: @@ -388,25 +360,14 @@ def __init__( self.canonical_url = str(initial_request.url) fits = _request_bytes(initial_request) <= url_limit - # A request that already fits and hasn't opted into finer chunking is - # the common passthrough: leave ``axes``/``chunks`` empty so - # ``total == 1`` and ``iter_chunk_args`` yields the original args - # verbatim. ``max_chunks == 1`` (off / no extra fan-out) means - # "don't split", so it takes this path; only ``max_chunks >= 2`` asks - # for extra fan-out and sets the axes up to be refined below. - if fits and max_chunks <= 1: + # An already fitting request remains one chunk. Chunking exists to + # satisfy the server's byte ceiling, not as a parallelism dial. + if fits: return self.axes = axes self.chunks = {axis.arg_key: [list(axis.atoms)] for axis in axes} - if not fits: - # Hard pass: greedy-halve until every worst-case chunk fits - # the byte budget (may raise ``Unchunkable``). - self._plan(build_request, url_limit) - # Soft pass: optionally split further than the byte budget requires. - # Purely additive — never re-raises, and the byte budget stays - # satisfied; a no-op at ``max_chunks == 1``. - self._refine(max_chunks) + self._plan(build_request, url_limit) if self.canonical_url is None: # Original URL was un-constructable (httpx.InvalidURL); fall @@ -461,69 +422,6 @@ def _plan( ) _split_at(self.chunks[biggest_axis.arg_key], biggest_idx) - def _refine(self, max_chunks: int) -> None: - """ - Fan the plan out more finely than the byte budget alone requires. - - This is the ``parallel_chunks`` dial: see - :func:`~dataretrieval.ogc.chunking.parallel_chunks` for why a caller - would want this, and :class:`ChunkPlan`'s ``max_chunks`` parameter for - the cap's contract (total-not-per-axis, a hard ceiling that may land - below the cap). - - Implementation. Each split multiplies the plan by ``(k+1)/k`` for the - chosen axis (adding ``total // k`` chunks, not one), so a split - is taken only when it keeps :attr:`total` within the cap. When no - in-budget split remains, the plan stops *below* the cap rather than - overshooting (two even axes can reach 4 but not 5, so a cap of 5 yields - 4). Each split picks the single largest splittable chunk among the - in-budget axes (ties broken by axis-extraction order, then lowest - index), so growth is distributed round-robin rather than one axis - saturating before another is touched. Purely additive — only ever - *splits* existing chunks, so the byte pass's work and the ``url_limit`` - invariant are both preserved, and it never raises. A no-op at - ``max_chunks == 1``. - - Parameters - ---------- - max_chunks : int - The ``parallel_chunks(n)`` value; see :class:`ChunkPlan`'s - ``max_chunks`` parameter for the full contract. - """ - if max_chunks <= 1: - return - while True: - total = self.total - if total >= max_chunks: - return - # Largest splittable chunk among the axes whose split still fits the - # cap. Splitting any chunk of an axis with ``k`` chunks turns that - # ``k`` into ``k+1``, so it adds ``total // k`` chunks (the - # product of the other axes) regardless of which chunk. Hence the - # budget test is per axis, not per chunk. Skipping an over-budget - # axis makes ``max_chunks`` a true ceiling. The ranking key is atom - # count (``len``), not URL bytes like ``_plan`` — this pass balances - # work across chunks rather than fitting a byte budget. A - # chunk of size 1 can't be split further. Stable input order breaks - # ties by axis order, then lowest index within an axis. - candidate: tuple[_Axis, int] | None = None - candidate_size = -1 - for axis in self.axes: - axis_chunks = self.chunks[axis.arg_key] - if total + total // len(axis_chunks) > max_chunks: - continue # any split of this axis would overshoot the cap - for idx, chunk in enumerate(axis_chunks): - if len(chunk) <= 1: - continue - if len(chunk) > candidate_size: - candidate, candidate_size = (axis, idx), len(chunk) - if candidate is None: - # Every axis is saturated at one atom per chunk or would - # overshoot the cap; stop below it rather than exceed it. - return - axis, idx = candidate - _split_at(self.chunks[axis.arg_key], idx) - def _worst_case_args(self) -> dict[str, Any]: """ Args for the largest chunk the current partition will issue. diff --git a/dataretrieval/ogc/policy.py b/dataretrieval/ogc/policy.py index 0959e682..5fce9acf 100644 --- a/dataretrieval/ogc/policy.py +++ b/dataretrieval/ogc/policy.py @@ -60,6 +60,13 @@ class OgcDialect: Columns to sort the combined result by, in priority order. Sorting is applied only when the first (primary) column is present; any later columns also present are added as secondary keys. + max_offset : int or None + Largest ``offset`` the API accepts, or ``None`` when it rejects + ``offset`` entirely (the conservative default — offset-parallel page + fetching is only attempted when an API declares a ceiling here). + ``offset`` is *not* a standard OGC API - Features parameter: Part 1 + defines only ``limit`` and the ``next`` link relation, so support is a + per-server extension that has to be declared rather than assumed. """ cql2_services: frozenset[str] = field(default_factory=frozenset) @@ -67,6 +74,7 @@ class OgcDialect: time_cols: frozenset[str] = field(default_factory=frozenset) numerical_cols: frozenset[str] = field(default_factory=frozenset) sort_cols: tuple[str, ...] = field(default_factory=tuple) + max_offset: int | None = None # Default dialect: a plain OGC API with no CQL2-only collections and no diff --git a/dataretrieval/ogc/requests.py b/dataretrieval/ogc/requests.py index d3b76ae4..4fb73a60 100644 --- a/dataretrieval/ogc/requests.py +++ b/dataretrieval/ogc/requests.py @@ -95,6 +95,30 @@ def _ogc_query_params( return params +def with_offset(request: httpx.Request, offset: int) -> httpx.Request: + """Rebuild ``request`` with one absolute row ``offset`` applied.""" + url = request.url.copy_set_param("offset", str(offset)) + content = request.content if request.method == "POST" else None + return httpx.Request( + method=request.method, + url=url, + headers=request.headers, + content=content, + ) + + +def page_limit(request: httpx.Request) -> int | None: + """Return the positive page ``limit`` encoded on ``request``, if usable.""" + raw = request.url.params.get("limit") + if raw is None: + return None + try: + value = int(raw) + except (TypeError, ValueError): + return None + return value if value > 0 else None + + def _partition_request_params( params: dict[str, Any], *, use_cql2: bool ) -> tuple[dict[str, Any], dict[str, Any]]: diff --git a/dataretrieval/transport/fanout.py b/dataretrieval/transport/fanout.py index 53b46c8a..a4e03960 100644 --- a/dataretrieval/transport/fanout.py +++ b/dataretrieval/transport/fanout.py @@ -320,6 +320,7 @@ def __init__( finalize: _Finalize = _passthrough_result, client_options: dict[str, Any] | None = None, default_concurrent: int = _CONCURRENCY_DEFAULT, + connection_multiplier: Callable[[], int] | None = None, *, canonical_url: str | None = None, service: str | None = None, @@ -338,6 +339,11 @@ def __init__( # test's ``monkeypatch.setenv`` still applies. See # :func:`_resolve_concurrency` for why the env var outranks it. self.default_concurrent = default_concurrent + # Some adapters fan out work inside one chunk (OGC offset page waves). + # Let the adapter report that width without making generic transport + # depend on protocol code. It is resolved for each resume so an updated + # environment setting sizes both dimensions consistently. + self.connection_multiplier = connection_multiplier # Extra ``httpx.AsyncClient`` options merged into the shared client this # run opens (``verify`` for the Water Use ``ssl_check`` flag, say). The # executor owns client lifecycle, so an adapter with a per-call client @@ -570,9 +576,9 @@ async def _run(self, max_concurrent: int | None) -> tuple[pd.DataFrame, Any]: The gather dispatches *every* pending chunk at once, but an ``asyncio.Semaphore`` caps the number of concurrent fetches at ``N = max_concurrent`` — ``None`` lifts the cap, ``N=1`` runs them - one at a time. The connection pool is sized to the same ``N`` - (``httpx.Limits(max_connections=N, max_keepalive_connections=N)``) - so the in-flight fetches reuse keepalive connections. + one at a time. The connection pool is sized to ``N`` times the adapter-declared + per-chunk request multiplier, so nested page waves do not queue inside + httpx against the pool-acquire timeout. The semaphore, not the pool, is deliberately the throttle. If the pool throttled instead, the excess chunks would queue @@ -592,8 +598,9 @@ async def _run(self, max_concurrent: int | None) -> tuple[pd.DataFrame, Any]: Parameters ---------- max_concurrent : int or None - Maximum chunks in flight (the semaphore value, and the - connection-pool size). ``None`` lifts the cap entirely. + Maximum chunks in flight (the semaphore value). The connection pool + also covers any adapter-declared requests within one chunk. + ``None`` lifts the cap entirely. Returns ------- @@ -619,8 +626,14 @@ async def _run(self, max_concurrent: int | None) -> tuple[pd.DataFrame, Any]: # why the gate can't be the pool itself. ``unbounded`` # (``max_concurrent=None``) is a degenerate cap at the plan total — a # semaphore that can never block — so gated is the only code path. + multiplier = ( + self.connection_multiplier() + if self.connection_multiplier is not None + else 1 + ) + pool_size = None if max_concurrent is None else max_concurrent * multiplier limits = httpx.Limits( - max_connections=max_concurrent, max_keepalive_connections=max_concurrent + max_connections=pool_size, max_keepalive_connections=pool_size ) semaphore = asyncio.Semaphore( len(self.plan) if max_concurrent is None else max_concurrent diff --git a/dataretrieval/transport/offsets.py b/dataretrieval/transport/offsets.py new file mode 100644 index 00000000..a1771b7d --- /dev/null +++ b/dataretrieval/transport/offsets.py @@ -0,0 +1,552 @@ +"""Offset-parallel page fetching: overlap a page walk instead of serializing it. + +Cursor pagination is inherently sequential — page ``N+1``'s URL only exists once +page ``N`` has been parsed, so a 10-page result costs 10 round trips end to end. +When a service also honors ``offset``, every page's URL is computable up front +(``offset = i * limit``), so the same pages can be fetched concurrently. The +request *count* is unchanged; only their timing is. That distinction matters +because the USGS quota is volume-based (``x-ratelimit-limit``, default 1000/hr), +so overlapping pages costs no extra quota. + +This module owns the generic half of that strategy: given a page-request +builder and a page parser, drive a bounded, speculative, wave-by-wave fetch and +return the concatenated frames. It is service-neutral — no OGC or Water Data +knowledge — mirroring :mod:`dataretrieval.transport.pagination`, which owns the +sequential cursor walk this is an alternative to. + +Why waves, and why they ramp +---------------------------- +The size of the result is unknown before it is fetched. OGC API - Features +Part 1 makes ``numberMatched`` *optional* ("each page may include information +about the number of selected and returned features"), and the Water Data API +omits it — a page carries ``numberReturned`` but no total. So a client cannot +compute the page count in advance; it must probe. + +Probing is where a naive fan-out gets expensive. Issuing ``width`` requests +immediately means a *one-page* result costs ``width`` requests instead of one: +at ``width=32`` a small query would spend 32x the quota to discover it was +already done. Since the quota here is volume-based, that is a straight 32x tax +on exactly the queries that had nothing to gain from parallelism. + +So the wave width **ramps**: 1 request, then 2, then 4, doubling up to +``width``. The properties that buys: + +- A single-page result costs exactly **one** request — identical to the + sequential walk, so the common small query pays nothing for this feature. +- Total requests stay under **2x** the pages actually needed (doubling means + every prior wave summed is less than the current one), and approach ``width`` + overshoot only for results large enough to amortize it. +- Round trips are logarithmic in the page count rather than linear: a 10-page + result is 4 waves, not 10 round trips. + +That is the standard unbounded-search ramp, and it is the reason this walk can +claim to leave the request count essentially unchanged while still overlapping +pages. The ceiling clip in :func:`plan_offsets` is what bounds the final wave. + +Stop conditions +--------------- +A wave stops the walk when any of these holds — see :func:`_stop_index` for the +precedence, which is the single source of truth: + +1. **A short page.** A page with fewer than ``limit`` rows is the last page by + construction: the server had no more rows to give. This is the normal exit. +2. **An empty page.** Zero rows means the previous page ended exactly on a + ``limit`` boundary and this offset is past the end. +3. **The row cap.** ``row_cap`` (from ``max_rows``) is reached, so further + pages would be discarded anyway. +4. **The offset ceiling.** The service refuses offsets beyond ``max_offset`` + (Water Data: 40000). This is *not* an end-of-data signal, so it must not end + the walk: the caller supplies ``tail_walk``, a sequential continuation that + picks up where the offsets stop. Offsets have a ceiling; cursors don't, so + the hybrid is fast over the parallelizable prefix and complete over the rest. + + The seam needs care. The next offset the walk *would* need is by definition + past the ceiling, so it can't seed the continuation either — that request + would earn the same rejection. So the walk rewinds one page: it drops the + last page it fetched and re-seeds the cursor walk at the largest offset the + service still accepts. One page is re-fetched per deep query, in exchange + for a seam with neither a gap (missing rows) nor an overlap (duplicates). + +Ordering is preserved regardless of completion order: results are indexed by +wave position and concatenated in offset order, so the frame matches what a +sequential walk would have produced. +""" + +from __future__ import annotations + +import asyncio +import logging +from collections.abc import Awaitable, Callable +from dataclasses import dataclass, field +from datetime import timedelta + +import httpx +import pandas as pd + +from dataretrieval import progress as _progress +from dataretrieval.combining import ( + _QUOTA_HEADER, + _merge_response, + _safe_elapsed, +) +from dataretrieval.exceptions import DataRetrievalError +from dataretrieval.transport.liveness import note_progress +from dataretrieval.transport.pagination import _client_for, paginated_failure_message + +logger = logging.getLogger(__name__) + + +class OffsetUnsupported(Exception): + """The service does not honor ``offset``, so this strategy can't be used. + + Internal control-flow signal, not a user-facing error: the caller catches it + and re-runs the query through the sequential cursor walk, which needs no + non-standard parameters. Raised *before* any rows are returned, so a fallback + re-fetch can't produce a partial or double-counted result. + """ + + +# A page builder maps an absolute row offset to the request that fetches it. +PageRequest = Callable[[int], httpx.Request] + +# A page parser maps a response to its frame. Unlike the cursor walk's parser +# it returns no cursor — the offsets *are* the cursor, computed not discovered. +PageParser = Callable[[httpx.Response], pd.DataFrame] + +# The sequential continuation used past the offset ceiling: +# ``(resume_offset, rows_so_far, client) -> (frame, response)``. +TailWalk = Callable[ + [int, int, httpx.AsyncClient], "Awaitable[tuple[pd.DataFrame, httpx.Response]]" +] + + +def plan_offsets( + *, + limit: int, + width: int, + start: int, + max_offset: int | None, +) -> list[int]: + """Offsets for one wave, clipped to the service's offset ceiling. + + Returns up to ``width`` offsets spaced ``limit`` apart beginning at + ``start``, dropping any that would exceed ``max_offset``. An empty list + means the ceiling has been reached and the caller must stop (or fall back + to a cursor walk) rather than issue a request the service will reject. + + Parameters + ---------- + limit : int + Page size — the offset stride. + width : int + Maximum number of offsets to plan. + start : int + First offset in this wave. + max_offset : int or None + Largest offset the service accepts, or ``None`` for no ceiling. + + Returns + ------- + list of int + The planned offsets, ascending; possibly empty. + """ + offsets = [start + i * limit for i in range(width)] + if max_offset is not None: + offsets = [off for off in offsets if off <= max_offset] + return offsets + + +def _stop_index( + frames: list[pd.DataFrame], + *, + limit: int, + rows_before: int, + row_cap: int | None, +) -> int | None: + """Index of the page that ends the walk, or ``None`` to continue. + + Encodes the stop precedence documented in the module docstring. Pages are + inspected in offset order so the *earliest* terminal page wins: a short + page at index 2 ends the walk even if index 5 (fetched speculatively past + the end) also looks terminal. Returning the index — rather than a bool — + lets the caller discard the pages after it, which is what makes a + speculative overshoot harmless. + + Parameters + ---------- + frames : list of pandas.DataFrame + This wave's page frames, in offset order. + limit : int + The page size requested; a frame shorter than this is terminal. + rows_before : int + Rows already collected by earlier waves, for the ``row_cap`` test. + row_cap : int or None + Stop once this many rows are held, or ``None`` for uncapped. + + Returns + ------- + int or None + Index of the last page to keep, or ``None`` if the walk continues. + """ + running = rows_before + for i, frame in enumerate(frames): + n = len(frame) + running += n + # An empty page is past the end: keep everything before it. A short + # page is the genuine last page: keep it, including its rows. + if n == 0: + return i - 1 if i else -1 + if n < limit: + return i + if row_cap is not None and running >= row_cap: + return i + return None + + +def _offset_ignored(frames: list[pd.DataFrame]) -> bool: + """Whether the server appears to be ignoring ``offset``. + + ``offset`` is not a standard OGC API - Features parameter, and an + unrecognized query parameter is conventionally *ignored* rather than + rejected. A server that ignores it answers every offset with page 1, so the + walk would happily concatenate the same rows N times and report success — + silent duplication, the worst failure mode available to this design. + + The check: two full-length pages fetched at different offsets must not be + identical. Comparing the first two suffices — if the stride is being + honored at all, page 0 and page 1 hold different rows. This is a cheap + structural comparison on frames already in memory, run once on the first + wave, so it costs no extra request. + + False positives are possible in principle (two genuinely identical pages of + data), which is why the caller treats a positive as "fall back to the cursor + walk" rather than an error: the safe strategy always remains available. + """ + if len(frames) < 2: + return False + first, second = frames[0], frames[1] + if first.empty or len(first) != len(second): + return False + if list(first.columns) != list(second.columns): + return False + return bool(first.equals(second)) + + +async def _fetch_page( + build_page: PageRequest, + offset: int, + client: httpx.AsyncClient, + raise_for_status: Callable[[httpx.Response], None], + semaphore: asyncio.Semaphore | None, +) -> httpx.Response: + """Fetch one page at ``offset``, honoring the concurrency gate.""" + if semaphore is None: + response = await client.send(build_page(offset)) + else: + async with semaphore: + response = await client.send(build_page(offset)) + raise_for_status(response) + return response + + +async def _fetch_wave( + build_page: PageRequest, + offsets: list[int], + client: httpx.AsyncClient, + raise_for_status: Callable[[httpx.Response], None], +) -> list[httpx.Response]: + """Fetch one wave, cancelling and draining siblings on any failure.""" + tasks = [ + asyncio.create_task( + _fetch_page(build_page, offset, client, raise_for_status, None) + ) + for offset in offsets + ] + try: + return list(await asyncio.gather(*tasks)) + except BaseException: + # ``asyncio.gather`` propagates the first failure but deliberately + # leaves siblings running. A retry of the whole page walk would then + # overlap those abandoned requests, double-spend quota, and exceed the + # declared concurrency bound. Cancel and await them before control can + # return to retry policy; catching BaseException also cleans up when + # the caller cancels this walk. + for task in tasks: + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + raise + + +@dataclass +class _WalkState: + """Mutable state shared by the small steps of one offset page walk.""" + + frames: list[pd.DataFrame] = field(default_factory=list) + first_response: httpx.Response | None = None + last_response: httpx.Response | None = None + total_elapsed: timedelta = field(default_factory=timedelta) + offset: int = 0 + wave_width: int = 1 + offset_verified: bool = False + + +async def _fetch_and_parse_wave( + *, + build_page: PageRequest, + parse_page: PageParser, + raise_for_status: Callable[[httpx.Response], None], + session: httpx.AsyncClient, + offsets: list[int], + completed_pages: int, +) -> tuple[list[httpx.Response], list[pd.DataFrame]]: + """Fetch and parse one wave, adding standard pagination guidance.""" + try: + responses = await _fetch_wave( + build_page, + offsets, + session, + raise_for_status, + ) + except Exception as exc: # noqa: BLE001 + logger.warning("Offset-parallel page fetch failed at offsets %r.", offsets) + raise DataRetrievalError( + paginated_failure_message(completed_pages, exc) + ) from exc + + try: + return responses, [parse_page(response) for response in responses] + except Exception as exc: # noqa: BLE001 + logger.warning("Offset-parallel page parse failed.") + raise DataRetrievalError( + paginated_failure_message(completed_pages, exc) + ) from exc + + +def _verify_offset(wave: list[pd.DataFrame], prior: list[pd.DataFrame]) -> bool: + """Verify distinct offsets return distinct pages once two are available.""" + probe = wave if len(wave) >= 2 else [*prior[-1:], *wave] + if _offset_ignored(probe): + raise OffsetUnsupported( + "The service returned identical pages for different `offset` values, " + "so it appears to ignore `offset`." + ) + return len(probe) >= 2 + + +def _record_page( + response: httpx.Response, + frame: pd.DataFrame, + reporter: _progress.ProgressReporter | None, +) -> timedelta: + """Record progress for one kept page and return its elapsed duration.""" + note_progress() + if reporter is not None: + reporter.set_rate_remaining( + response.headers.get(_QUOTA_HEADER), + limit=response.headers.get("x-ratelimit-limit"), + ) + reporter.add_page(rows=len(frame)) + return _safe_elapsed(response) + + +def _accept_wave( + state: _WalkState, + *, + responses: list[httpx.Response], + wave: list[pd.DataFrame], + offsets: list[int], + limit: int, + width: int, + row_cap: int | None, + reporter: _progress.ProgressReporter | None, +) -> bool: + """Keep the useful prefix of one wave; return whether the walk is done.""" + if not state.offset_verified: + state.offset_verified = _verify_offset(wave, state.frames) + + if state.first_response is None: + state.first_response = responses[0] + if state.last_response is None: + # An empty first page keeps no frame but is still valid metadata. + state.last_response = responses[0] + + stop_at = _stop_index( + wave, + limit=limit, + rows_before=sum(len(frame) for frame in state.frames), + row_cap=row_cap, + ) + keep = wave if stop_at is None else wave[: stop_at + 1] + for response, frame in zip(responses[: len(keep)], keep, strict=False): + state.total_elapsed += _record_page(response, frame, reporter) + state.last_response = response + state.frames.extend(keep) + + if stop_at is not None: + return True + state.offset = offsets[-1] + limit + state.wave_width = min(state.wave_width * 2, width) + return False + + +async def _walk_offset_waves( + state: _WalkState, + *, + build_page: PageRequest, + parse_page: PageParser, + raise_for_status: Callable[[httpx.Response], None], + session: httpx.AsyncClient, + limit: int, + width: int, + max_offset: int | None, + row_cap: int | None, + reporter: _progress.ProgressReporter | None, +) -> bool: + """Drive ramped waves; return whether the offset ceiling was reached.""" + while True: + offsets = plan_offsets( + limit=limit, + width=state.wave_width, + start=state.offset, + max_offset=max_offset, + ) + if not offsets: + return True + responses, wave = await _fetch_and_parse_wave( + build_page=build_page, + parse_page=parse_page, + raise_for_status=raise_for_status, + session=session, + offsets=offsets, + completed_pages=len(state.frames), + ) + if _accept_wave( + state, + responses=responses, + wave=wave, + offsets=offsets, + limit=limit, + width=width, + row_cap=row_cap, + reporter=reporter, + ): + return False + + +async def _continue_after_ceiling( + state: _WalkState, + *, + session: httpx.AsyncClient, + tail_walk: TailWalk | None, + limit: int, + max_offset: int | None, +) -> None: + """Rewind one legal page and cursor-walk the unbounded tail.""" + if tail_walk is None: + logger.warning( + "Stopped at the service's offset ceiling (%s) with %d row(s) " + "collected; the result may be incomplete because no sequential " + "continuation was supplied.", + max_offset, + sum(len(frame) for frame in state.frames), + ) + return + + if state.frames: + state.offset -= limit + state.frames.pop() + rows_so_far = sum(len(frame) for frame in state.frames) + logger.debug( + "Offset ceiling (%s) reached after %d row(s); continuing sequentially " + "from offset %d.", + max_offset, + rows_so_far, + state.offset, + ) + tail_frame, tail_response = await tail_walk( + state.offset, + rows_so_far, + session, + ) + if len(tail_frame): + state.frames.append(tail_frame) + state.total_elapsed += _safe_elapsed(tail_response) + state.last_response = tail_response + + +def _finalize_walk( + state: _WalkState, + *, + max_offset: int | None, + row_cap: int | None, +) -> tuple[pd.DataFrame, httpx.Response]: + """Build the final frame and aggregate response from a completed walk.""" + if state.first_response is None or state.last_response is None: + raise DataRetrievalError( + "Offset-parallel pagination issued no requests; " + f"max_offset={max_offset!r} leaves no valid page offset." + ) + result = ( + pd.concat(state.frames, ignore_index=True) if state.frames else pd.DataFrame() + ) + if row_cap is not None: + result = result.head(row_cap) + return result, _merge_response( + state.first_response, + headers_from=state.last_response, + elapsed=state.total_elapsed, + ) + + +async def paginate_by_offset( + *, + build_page: PageRequest, + parse_page: PageParser, + raise_for_status: Callable[[httpx.Response], None], + client: httpx.AsyncClient | None = None, + limit: int, + width: int, + max_offset: int | None = None, + row_cap: int | None = None, + tail_walk: TailWalk | None = None, +) -> tuple[pd.DataFrame, httpx.Response]: + """Fetch pages concurrently in ramped waves until a stop condition. + + This is the offset-parallel counterpart to + :func:`dataretrieval.transport.pagination.paginate`. ``build_page`` maps an + absolute row offset to a request; ``parse_page`` maps a response to its + frame. ``limit`` is both page size and offset stride, while ``width`` caps + each speculative wave. ``row_cap`` stops and truncates early. + + ``max_offset`` is the largest offset the service accepts. Reaching it is + not end-of-data: when ``tail_walk`` is supplied, the last legal page is + rewound and the callback cursor-walks the unbounded remainder. Without a + callback, the partial result is returned with a warning. + + Any page failure raises :class:`DataRetrievalError` with the same recovery + guidance as cursor pagination. Sibling requests are cancelled and drained + before the error returns, so a retry cannot overlap abandoned work. + """ + async with _client_for(client) as session: + state = _WalkState() + ceiling_reached = await _walk_offset_waves( + state, + build_page=build_page, + parse_page=parse_page, + raise_for_status=raise_for_status, + session=session, + limit=limit, + width=width, + max_offset=max_offset, + row_cap=row_cap, + reporter=_progress.current(), + ) + if ceiling_reached: + await _continue_after_ceiling( + state, + session=session, + tail_walk=tail_walk, + limit=limit, + max_offset=max_offset, + ) + return _finalize_walk(state, max_offset=max_offset, row_cap=row_cap) diff --git a/dataretrieval/waterdata/__init__.py b/dataretrieval/waterdata/__init__.py index 48c4d9fb..34107ec8 100644 --- a/dataretrieval/waterdata/__init__.py +++ b/dataretrieval/waterdata/__init__.py @@ -8,7 +8,6 @@ from __future__ import annotations -from dataretrieval.ogc.chunking import parallel_chunks from dataretrieval.ogc.filters import FILTER_LANG # Public API exports @@ -50,7 +49,6 @@ "PROFILE_LOOKUP", "SERVICES", "WATERDATA_SERVICES", - "parallel_chunks", "get_channel", "get_codes", "get_combined_metadata", diff --git a/dataretrieval/waterdata/utils.py b/dataretrieval/waterdata/utils.py index e001b877..9d020454 100644 --- a/dataretrieval/waterdata/utils.py +++ b/dataretrieval/waterdata/utils.py @@ -94,6 +94,12 @@ } ), sort_cols=("time", "monitoring_location_id"), + # The API rejects ``offset > 40000`` with HTTP 400 ``InvalidQuery`` + # ("offset parameter must be less than or equal to 40000"), so + # offset-parallel page fetching can only cover a result's first 40k rows. + # Past that the walk hands off to the sequential cursor continuation, which + # has no ceiling, so a deep result is still returned in full. + max_offset=40_000, ) # The Water-Data-specific *extras* on top of the engine's own no-normalize set diff --git a/docs/source/architecture/decisions/0006-service-neutral-transport.rst b/docs/source/architecture/decisions/0006-service-neutral-transport.rst index 67b3c2c6..162a7f56 100644 --- a/docs/source/architecture/decisions/0006-service-neutral-transport.rst +++ b/docs/source/architecture/decisions/0006-service-neutral-transport.rst @@ -34,7 +34,8 @@ general. It owns: - synchronous and asynchronous HTTP client lifecycle and timeout defaults; - attaching the API key and stripping it at redirect time, over the predicate ``dataretrieval.credentials`` defines; -- callback-driven cursor pagination; +- callback-driven page walking in two strategies: sequential cursor + pagination and offset-parallel fetching for services that declare support; - bounded retry with exponential backoff, full jitter, capped ``Retry-After`` handling, and a no-progress budget bounding how long a call may receive nothing at all; and @@ -56,6 +57,15 @@ It must not import OGC modules or service adapters. Service adapters inject request construction, response parsing, cursor extraction, and API-specific error details. +Both page-walk strategies live in transport because both are HTTP execution +policy and neither is OGC-specific. Cursor pagination +(``transport.pagination``) discovers page *N+1* from page *N* and is the +standards-only fallback. Offset pagination (``transport.offsets``) computes +page requests in bounded ramped waves. Strategy selection remains adapter +policy: an ``OgcDialect`` declares the server's maximum accepted offset, or +``None`` to stay on cursors. At that ceiling, the offset walk invokes a cursor +continuation; the strategies compose rather than compete. + OGC retains its protocol concerns: dialects, CQL2, request construction, feature shaping, URL-byte chunk planning, resumable ``ChunkedCall`` state, and typed interruption handles. Thin imports at previous private OGC and utility paths @@ -86,6 +96,12 @@ Consequences - Water Use has no dependency on OGC implementation modules. - OGC and non-OGC adapters share authentication, timeout, retry, pagination, aggregation, progress, and sync-dispatch policy where their semantics match. +- A service can opt into faster page walks by declaring an offset ceiling; if it + stops honoring ``offset``, identical-page detection falls back to cursors + instead of returning duplicated rows. +- Parallelism belongs to the page walk, not the chunk planner. The planner + splits only to satisfy the URL byte ceiling, avoiding quota-positive splits + of requests that already fit. - Service-specific request and result contracts remain explicit instead of being forced into a universal adapter abstraction. - Retry can increase latency and quota consumption, so attempt counts, waits, @@ -95,10 +111,10 @@ Consequences service that cannot use an API key is not told to obtain one. - The transport package is internal infrastructure, not a new public API promise. -- Keeping presentation and frame assembly out means transport is roughly 570 - lines across five modules, each recognizably HTTP execution policy. Retry is - the one intricate module, and it is intricate because two independent bounds - are what make retry safe against a slow service. +- Keeping presentation and frame assembly out leaves each transport module + recognizably HTTP execution policy. Retry and the offset walk are intricate + for bounded reasons: retry enforces independent attempt/time bounds, while + the offset walk must discover an unknown page count without unbounded probes. Compliance ---------- @@ -110,4 +126,7 @@ transport, and that only ``dataretrieval.credentials`` names the API-key host. Component and adapter tests cover cursor termination, row caps, response aggregation, retry exhaustion, ``Retry-After`` limits, the no-progress budget, which failures are re-sent, cancellation, no-partial fan-out behavior, and -credential host scoping. +credential host scoping. ``tests/transport_test.py`` covers the offset walk's +stop conditions in isolation. ``tests/waterdata_offset_paging_test.py`` pins +query preservation, ignored-offset fallback, and a gap-free, duplicate-free +handoff at the offset ceiling through a real getter. diff --git a/docs/source/architecture/decisions/0008-fan-out-execution.rst b/docs/source/architecture/decisions/0008-fan-out-execution.rst index d5431f1e..01aaa936 100644 --- a/docs/source/architecture/decisions/0008-fan-out-execution.rst +++ b/docs/source/architecture/decisions/0008-fan-out-execution.rst @@ -78,8 +78,9 @@ planning and the OGC call site passes it through; Water Use passes its first location's URL, since the service has no request expressing "all of these". ``dataretrieval.ogc`` keeps chunk planning: the byte budget, the axis -partitioning, the CQL2 filter split, the ``parallel_chunks`` dial. Those are -division, and division is protocol-specific. +partitioning, and the CQL2 filter split. Those are division, and division is +protocol-specific. Page-level parallelism is transport execution policy and is +selected separately from chunk planning. The interruption taxonomy moves to ``dataretrieval.interruptions``, a top-level leaf, for the reason ADR 0006 gives for ``combining``, ``progress``, and diff --git a/docs/source/userguide/errors.rst b/docs/source/userguide/errors.rst index 046f7a5d..15c9fd2c 100644 --- a/docs/source/userguide/errors.rst +++ b/docs/source/userguide/errors.rst @@ -105,44 +105,49 @@ mid-stream, the work already completed is preserved: catch The same loop works for ``wateruse.get_wateruse`` with a list of states, counties, or HUCs. -Chunk a large request more finely -================================= - -By default the getters split an over-large request only as much as the -server's ~8 KB URL limit forces -- the fewest chunks. Because each -chunk paginates, splitting a large result further costs little or no -extra quota *as long as each chunk still spans many pages*. (Ten states -pulled as one request then page nearly as many times as ten per-state requests -would; a split that leaves each chunk only a page or two adds its partial -final page.) So if you *know* your pull is large, ask for a finer split with -``parallel_chunks(n)``: you trade roughly the same pages for more, smaller -chunks, which gives smoother progress, more even concurrency, and a -smaller unit of retry/resume. ``parallel_chunks`` is a scoped ``with`` block, so -an aggressive setting can't leak into unrelated calls and accidentally spend -quota: +Large Water Data pulls are paged in parallel +============================================== + +Water Data can compute page URLs with ``offset``, so pages of a large result are +fetched in ramped concurrent waves by default. There is no parallel-chunk knob +to enable: the old ``parallel_chunks(n)`` context manager was removed because +it split fitting queries into extra requests and could not help a single-site +query. + +At a fixed page ``limit`` the walk overlaps requests it was already going to +make, apart from bounded probes in the final speculative wave. Reducing +``limit`` creates more pages and consumes more quota. The default limit is +50,000, above Water Data's 40,000 offset ceiling, so material speedups require +an explicit smaller limit; for example: .. code-block:: python from dataretrieval import waterdata - with waterdata.parallel_chunks(32): - df, md = waterdata.get_daily( - monitoring_location_id=many_sites, parameter_code="00060" - ) + df, md = waterdata.get_daily( + monitoring_location_id="USGS-01646500", + parameter_code="00060", + limit=2000, + ) + +``API_USGS_CONCURRENT`` (default 32) bounds both chunk fan-out and the page-wave +width. Set it to ``1`` to use standard cursor pagination sequentially: + +.. code-block:: python + + import os + + os.environ["API_USGS_CONCURRENT"] = "1" + +``offset`` is a Water Data extension, not part of OGC API - Features. If a +server ignores it, the client detects identical pages before returning rows and +re-runs the query through standard ``next`` links. At Water Data's 40,000-row +offset ceiling, it rewinds one page and cursor-walks the tail so the result has +neither a gap nor duplicate rows. -``n`` is a positive integer (e.g. ``2``, ``8``, ``32``) -- the number of -chunks to fan the call out into; a non-integer or non-positive value -raises ``ValueError`` at the ``with``. ``n`` caps the *total* chunk count -across every multi-value argument combined (not per argument), bounded below by -what the byte limit already forces and above by how many values there are to -split. Several multi-value arguments therefore can't multiply past it, and -``n=1`` asks for no extra fan-out. Each chunk costs a request against your -hourly rate limit. How many run *at once* is capped separately by -``API_USGS_CONCURRENT`` (default 32), so an ``n`` beyond that adds quota without -adding parallelism -- the useful range is roughly ``2`` up to -``API_USGS_CONCURRENT``. There is no "off" level: don't enter the block -unless you already expect a large, multi-page result -- on a query that would -have fit in a single page, extra chunks only burn quota. +Byte-driven chunking is unchanged: a multi-value request above the service's +~8 KB request limit is still split for correctness. That division is separate +from page-level parallelism. The full taxonomy ================= diff --git a/tests/architecture_test.py b/tests/architecture_test.py index 88606b76..ecb9d100 100644 --- a/tests/architecture_test.py +++ b/tests/architecture_test.py @@ -38,7 +38,7 @@ #: How many request-building names ``ogc.engine`` currently needs. A ceiling #: rather than an exact list allows renames and deletions without weakening the #: rule that orchestration must not absorb request construction again. -_MAX_ENGINE_REQUEST_IMPORTS = 5 +_MAX_ENGINE_REQUEST_IMPORTS = 6 def _module_name(path: Path) -> str: diff --git a/tests/contracts/public_api_test.py b/tests/contracts/public_api_test.py index 53e28aeb..244f08df 100644 --- a/tests/contracts/public_api_test.py +++ b/tests/contracts/public_api_test.py @@ -24,7 +24,6 @@ "PROFILE_LOOKUP", "SERVICES", "WATERDATA_SERVICES", - "parallel_chunks", "get_channel", "get_codes", "get_combined_metadata", diff --git a/tests/transport_test.py b/tests/transport_test.py index dcd4b5c0..72463445 100644 --- a/tests/transport_test.py +++ b/tests/transport_test.py @@ -605,3 +605,261 @@ def test_exception_chain_walk_terminates_on_a_self_referencing_chain() -> None: ServiceInterrupted(completed_chunks=0, total_chunks=1, cause=first).status_code is None ) + + +# --------------------------------------------------------------------------- +# Offset-parallel pagination (``dataretrieval.transport.offsets``). Cursor +# paging is sequential by construction -- page N+1's URL is only revealed by +# page N -- so when a service also honors ``offset`` every page URL is +# computable up front and the pages can overlap. These tests pin the three +# things that makes correct: knowing when to stop, discarding a speculative +# overshoot, and refusing to trust a server that ignores ``offset``. +# --------------------------------------------------------------------------- + + +def _offset_server(total_rows: int, *, limit: int, ceiling: int | None = None): + """A fake paged service backed by ``total_rows`` sequential row ids. + + Returns ``(client, seen)`` where ``seen`` records the offsets requested, in + completion order, so a test can assert on the request *count* (the quota + cost) as well as the rows. + """ + seen: list[int] = [] + + async def send(request: httpx.Request) -> httpx.Response: + offset = int(request.url.params["offset"]) + seen.append(offset) + if ceiling is not None and offset > ceiling: + return httpx.Response(400, request=request) + rows = list(range(offset, min(offset + limit, total_rows))) + return httpx.Response( + 200, + json={"rows": rows}, + request=request, + headers={"x-ratelimit-limit": "1000"}, + ) + + client = mock.AsyncMock(spec=httpx.AsyncClient) + client.send.side_effect = send + return client, seen + + +def _parse_rows(response: httpx.Response) -> pd.DataFrame: + return pd.DataFrame({"row": response.json()["rows"]}) + + +def _walk(client, *, limit: int, width: int, **kwargs): + from dataretrieval.transport.offsets import paginate_by_offset + + return asyncio.run( + paginate_by_offset( + build_page=lambda offset: httpx.Request( + "GET", f"https://example.test/items?limit={limit}&offset={offset}" + ), + parse_page=_parse_rows, + raise_for_status=_raise_for_status, + client=client, + limit=limit, + width=width, + **kwargs, + ) + ) + + +def test_offset_walk_stops_on_a_short_page() -> None: + """The normal exit. 25 rows at ``limit=10`` is two full pages and a + 5-row third: the short page proves the server had no more rows, so the + walk ends there -- and its rows are kept, not discarded.""" + client, seen = _offset_server(25, limit=10) + frame, response = _walk(client, limit=10, width=4) + + assert frame["row"].tolist() == list(range(25)) + # The ramp costs nothing here: waves of 1 then 2 land exactly on the three + # pages that exist, so the short page ends the walk with no overshoot at all. + assert sorted(seen) == [0, 10, 20] + assert response.headers["x-ratelimit-limit"] == "1000" + + +def test_offset_walk_stops_on_an_empty_page_at_an_exact_boundary() -> None: + """20 rows at ``limit=10`` ends exactly on a page boundary, so no page is + short. The empty page at offset 20 is the only end-of-data signal available, + and it must not contribute a row.""" + client, _ = _offset_server(20, limit=10) + frame, _ = _walk(client, limit=10, width=3) + assert frame["row"].tolist() == list(range(20)) + + +def test_offset_walk_continues_across_waves() -> None: + """A result larger than one wave keeps going, and the next wave's offsets + continue where the last stopped -- no gap (missing rows) and no overlap + (duplicates), which a mis-computed stride would produce.""" + client, seen = _offset_server(95, limit=10, ceiling=None) + frame, _ = _walk(client, limit=10, width=4) + + assert frame["row"].tolist() == list(range(95)) + # Three waves: offsets 0-30, 40-70, 80-110. Every offset distinct. + assert len(seen) == len(set(seen)) + assert min(seen) == 0 + + +def test_offset_walk_honors_the_row_cap() -> None: + """``max_rows`` stops the walk once enough rows are held and truncates to + exactly the cap, so a preview doesn't page through a huge table.""" + client, _ = _offset_server(1000, limit=10, ceiling=None) + frame, _ = _walk(client, limit=10, width=4, row_cap=25) + assert frame["row"].tolist() == list(range(25)) + + +def test_offset_walk_hands_off_to_the_tail_walk_at_the_ceiling() -> None: + """The offset ceiling is NOT an end-of-data signal. Reaching it must hand + off to the sequential continuation -- otherwise a deep pull would silently + return a truncated result, the worst outcome available here. + + The seam is the subtle part. The next offset the walk *would* need (30 here) + is itself past the ceiling, so it can't seed the continuation either -- that + request would earn the same rejection. So the walk rewinds one page: it drops + the last page it fetched and re-seeds at 20, the largest offset the service + still accepts. One page is re-fetched; no row is missed or duplicated. + """ + client, seen = _offset_server(1000, limit=10, ceiling=None) + handoff: dict[str, object] = {} + + async def tail_walk(resume_offset, rows_so_far, session): + handoff["resume_offset"] = resume_offset + handoff["rows_so_far"] = rows_so_far + assert session is client + return ( + pd.DataFrame({"row": list(range(resume_offset, resume_offset + 15))}), + httpx.Response( + 200, request=httpx.Request("GET", "https://example.test/tail") + ), + ) + + frame, _ = _walk(client, limit=10, width=4, max_offset=25, tail_walk=tail_walk) + + # Offsets stop at the ceiling (0, 10, 20 -- 30 > 25 is never requested). + assert sorted(seen) == [0, 10, 20] + # Re-seeded at an offset the service accepts, and told how many rows are + # already held so it can rebase a remaining row cap onto the tail. + assert handoff == {"resume_offset": 20, "rows_so_far": 20} + # Seamless: rows 0-19 from offsets, 20-34 from the continuation. + assert frame["row"].tolist() == list(range(35)) + + +def test_offset_walk_warns_and_truncates_without_a_tail_walk(caplog) -> None: + """Without a continuation the ceiling result is knowingly partial, so it + must say so loudly rather than pass for a complete answer.""" + client, _ = _offset_server(1000, limit=10, ceiling=None) + with caplog.at_level("WARNING"): + frame, _ = _walk(client, limit=10, width=4, max_offset=25) + assert frame["row"].tolist() == list(range(30)) + assert "offset ceiling" in caplog.text.lower() + + +def test_offset_walk_refuses_a_server_that_ignores_offset() -> None: + """An unrecognized query parameter is conventionally *ignored*, not + rejected -- so a service that doesn't implement ``offset`` answers every + offset with page 1 and the walk would concatenate the same rows N times and + report success. That silent duplication is the design's worst failure mode, + so it is detected on the first wave and raises before any rows are + returned, leaving the caller free to re-run via cursors.""" + from dataretrieval.transport.offsets import OffsetUnsupported + + async def send(request: httpx.Request) -> httpx.Response: + # Same page regardless of the offset asked for. + return httpx.Response(200, json={"rows": list(range(10))}, request=request) + + client = mock.AsyncMock(spec=httpx.AsyncClient) + client.send.side_effect = send + + with pytest.raises(OffsetUnsupported): + _walk(client, limit=10, width=4) + + +def test_offset_walk_accepts_distinct_pages_of_equal_length() -> None: + """The ignore-detection compares page *contents*, not just their shape: two + full pages of the same length are the normal case and must not be mistaken + for a server echoing page 1.""" + client, _ = _offset_server(40, limit=10) + frame, _ = _walk(client, limit=10, width=4) + assert frame["row"].tolist() == list(range(40)) + + +def test_plan_offsets_clips_to_the_ceiling() -> None: + """The planner never proposes an offset the service would reject with a + 400: it clips to the ceiling, and an empty plan is the caller's signal to + hand off rather than to keep asking.""" + from dataretrieval.transport.offsets import plan_offsets + + assert plan_offsets(limit=10, width=4, start=0, max_offset=None) == [0, 10, 20, 30] + assert plan_offsets(limit=10, width=4, start=0, max_offset=25) == [0, 10, 20] + assert plan_offsets(limit=10, width=4, start=30, max_offset=25) == [] + + +def test_offset_walk_wraps_a_page_failure_with_recovery_guidance() -> None: + """A failed page fails the whole walk with the same actionable message the + sequential walk produces -- a partial frame silently returned would be + indistinguishable from a complete one.""" + + async def send(request: httpx.Request) -> httpx.Response: + if int(request.url.params["offset"]) == 10: + return httpx.Response(500, request=request) + return httpx.Response(200, json={"rows": list(range(10))}, request=request) + + client = mock.AsyncMock(spec=httpx.AsyncClient) + client.send.side_effect = send + + with pytest.raises(DataRetrievalError): + _walk(client, limit=10, width=4) + + +def test_offset_walk_cancels_sibling_page_after_failure() -> None: + """A retry must not overlap page requests abandoned by the failed attempt.""" + + async def scenario() -> None: + sibling_started = asyncio.Event() + sibling_cancelled = asyncio.Event() + + async def send(request: httpx.Request) -> httpx.Response: + offset = int(request.url.params["offset"]) + if offset == 0: + return httpx.Response( + 200, + json={"rows": list(range(10))}, + request=request, + ) + if offset == 10: + await sibling_started.wait() + return httpx.Response(500, request=request) + if offset == 20: + sibling_started.set() + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + sibling_cancelled.set() + raise + raise AssertionError(f"unexpected offset {offset}") + + client = mock.AsyncMock(spec=httpx.AsyncClient) + client.send.side_effect = send + + with pytest.raises(DataRetrievalError): + await paginate_by_offset( + build_page=lambda offset: httpx.Request( + "GET", + f"https://example.test/items?limit=10&offset={offset}", + ), + parse_page=_parse_rows, + raise_for_status=_raise_for_status, + client=client, + limit=10, + width=2, + ) + + assert sibling_cancelled.is_set(), ( + "the failed wave returned while a sibling request was still live" + ) + + from dataretrieval.transport.offsets import paginate_by_offset + + asyncio.run(scenario()) diff --git a/tests/utils_test.py b/tests/utils_test.py index b461b019..0368ac65 100644 --- a/tests/utils_test.py +++ b/tests/utils_test.py @@ -269,19 +269,19 @@ def test_chunk_interruptions_exported_at_top_level(self): dataretrieval.ChunkInterrupted, dataretrieval.DataRetrievalError ) - def test_parallel_chunks_exported_at_top_level_and_waterdata(self): - """The ``parallel_chunks`` context manager is reachable both from the top - level (``from dataretrieval import parallel_chunks``) and from the - user-facing ``dataretrieval.waterdata`` namespace, and both resolve to - the single object defined in ``dataretrieval.ogc.chunking``.""" + def test_parallel_chunks_is_gone(self): + """``parallel_chunks`` was removed: page parallelism now comes from the + offset walk (:mod:`dataretrieval.transport.offsets`), which overlaps a + single request's pages instead of splitting the query into more + sub-requests. Nothing should re-export the retired dial.""" import dataretrieval from dataretrieval import waterdata from dataretrieval.ogc import chunking - assert dataretrieval.parallel_chunks is chunking.parallel_chunks - assert waterdata.parallel_chunks is chunking.parallel_chunks - assert "parallel_chunks" in dataretrieval.__all__ - assert "parallel_chunks" in waterdata.__all__ + for module in (dataretrieval, waterdata, chunking): + assert not hasattr(module, "parallel_chunks") + assert "parallel_chunks" not in dataretrieval.__all__ + assert "parallel_chunks" not in waterdata.__all__ class Test_BaseMetadata: diff --git a/tests/waterdata_chunking_test.py b/tests/waterdata_chunking_test.py index 3454dc31..fc6227e4 100644 --- a/tests/waterdata_chunking_test.py +++ b/tests/waterdata_chunking_test.py @@ -50,10 +50,8 @@ from dataretrieval.ogc.chunking import ( ChunkedCall, _chunked_client, - _parallel_chunks, get_active_client, multi_value_chunked, - parallel_chunks, ) from dataretrieval.ogc.dates import _DATE_RANGE_PARAMS from dataretrieval.ogc.interruptions import ( @@ -1148,15 +1146,16 @@ def test_combine_chunk_frames_still_dedupes_overlapping_ids(): def test_list_axis_chunks_dedupe_repeated_feature_ids(): """Repeated list values can select the same feature in separate chunks.""" - @multi_value_chunked(build_request=_fake_build, url_limit=8000) + # Force the repeated values into separate chunks through the byte budget; + # fitting requests are no longer split merely to manufacture parallelism. + @multi_value_chunked(build_request=_fake_build, url_limit=201) async def fetch(args): return ( pd.DataFrame({"id": ["feature-1"], "site": [args["sites"][0]]}), _quota_response(500), ) - with parallel_chunks(2): - frame, _ = fetch({"sites": ["A", "A"]}) + frame, _ = fetch({"sites": ["A", "A"]}) assert frame.to_dict(orient="records") == [{"id": "feature-1", "site": "A"}] @@ -2177,293 +2176,81 @@ async def fetch(args): # --------------------------------------------------------------------------- -# Parallel chunks: the opt-in dial ``parallel_chunks(n)`` to fan a query out -# MORE finely than the byte limit alone requires (``ChunkPlan._refine`` + the -# ``parallel_chunks`` context manager). ``_fake_build``'s base is 200 bytes, so -# a handful of short atoms sits far under ``url_limit=8000`` — the byte pass -# passes it through untouched, and any splitting below is the ``n`` cap alone. -# ``ChunkPlan`` takes the integer cap (``max_chunks``) directly; -# ``parallel_chunks(n)`` publishes ``n`` onto it. The cap bounds the plan's -# *total* chunk count (the cartesian product across axes), not each axis -# independently — see ``test_cap_caps_the_total_across_axes``. +# The parallel-chunk dial is retired. ChunkPlan now divides only when required +# by the byte ceiling; page-wave parallelism is selected independently. # --------------------------------------------------------------------------- -def test_default_preserves_passthrough(): - """The default ``max_chunks`` (1 = off) must not perturb the existing - plan: a multi-value request that fits the byte limit is still the trivial - passthrough (no axes, ``total == 1``), byte-for-byte the pre-feature - behavior.""" +def test_fitting_request_is_always_a_passthrough(): args = {"monitoring_location_id": ["A", "B", "C", "D"]} - plan = ChunkPlan(args, _fake_build, url_limit=8000) # default max_chunks=1 - assert plan.axes == [] - assert plan.total == 1 - assert list(plan.iter_chunk_args()) == [args] - - -def test_unit_cap_preserves_passthrough(): - """``max_chunks=1`` means "no extra fan-out", so a fitting multi-value - request stays the trivial passthrough (no axes, ``total == 1``, - ``iter_chunk_args`` yields the original args verbatim) — identical to the - default (off), not a materialized one-chunk-per-axis plan.""" - args = {"monitoring_location_id": ["A", "B", "C", "D"]} - plan = ChunkPlan(args, _fake_build, url_limit=8000, max_chunks=1) + plan = ChunkPlan(args, _fake_build, url_limit=8000) assert plan.axes == [] assert plan.total == 1 assert list(plan.iter_chunk_args()) == [args] -@pytest.mark.parametrize("bad", [0, -1]) -def test_invalid_cap_raises(bad): - """``max_chunks`` is a chunk count, so a value below 1 (``0`` or - negative) is a caller bug, not a silent no-op: it raises ``ValueError`` at - construction. (The public ``parallel_chunks(n)`` already rejects ``n < 1``; - this pins the same guard on direct construction.)""" +def test_chunk_plan_rejects_the_retired_max_chunks_kwarg(): args = {"monitoring_location_id": ["A", "B", "C", "D"]} - with pytest.raises(ValueError, match="max_chunks must be >= 1"): - ChunkPlan(args, _fake_build, url_limit=8000, max_chunks=bad) + with pytest.raises(TypeError): + ChunkPlan(args, _fake_build, url_limit=8000, max_chunks=8) -@pytest.mark.parametrize( - ("max_chunks", "expected_pieces"), - [(1, 1), (2, 2), (8, 8), (16, 10), (32, 10)], -) -def test_cap_ramps_then_saturates(max_chunks, expected_pieces): - """A single 10-atom axis that fits the byte limit splits into - ``min(10, cap)`` pieces: 1 (off), 2, 8, then saturating at 10 (one atom per - chunk) once the cap overshoots the atom count. Monotonic and bounded, and - whenever it splits the partition is a cover — every atom exactly once. (The - cap-1 passthrough has no axis to cover; see the passthrough test.)""" - atoms = [f"S{i:02d}" for i in range(10)] - plan = ChunkPlan( - {"monitoring_location_id": atoms}, - _fake_build, - url_limit=8000, - max_chunks=max_chunks, - ) - assert plan.total == expected_pieces - if expected_pieces > 1: - flattened = [ - a for chunk in plan.chunks["monitoring_location_id"] for a in chunk - ] - assert sorted(flattened) == sorted(atoms) - - -def test_cap_bounds_fan_out_for_a_long_axis(): - """The cap holds fan-out to ``n``: at ``n=32`` a 100-atom axis fans into - ``n`` pieces — NOT 100 singletons — so ``parallel_chunks(32)`` on a huge - list can't detonate into hundreds of chunks. Every atom is still - covered exactly once.""" - high = 32 - atoms = [f"X{i:03d}" for i in range(100)] - plan = ChunkPlan( - {"monitoring_location_id": atoms}, - _fake_build, - url_limit=8000, - max_chunks=high, - ) - assert plan.total == high - flattened = [a for chunk in plan.chunks["monitoring_location_id"] for a in chunk] +def test_byte_driven_chunking_survives_the_removal(): + atoms = ["X" * 30, "Y" * 30, "Z" * 30, "W" * 30] + args = {"monitoring_location_id": atoms} + limit = 250 + plan = ChunkPlan(args, _fake_build, url_limit=limit) + assert plan.total > 1 + for chunk_args in plan.iter_chunk_args(): + assert _safe_request_bytes(_fake_build, chunk_args, limit) <= limit + flattened = [ + atom for chunk in plan.chunks["monitoring_location_id"] for atom in chunk + ] assert sorted(flattened) == sorted(atoms) -def test_cap_below_byte_split_does_not_reduce_fan_out(): - """The cap is purely additive — it can only split further, never coarsen. - A request the byte budget already fans into K>2 chunks is untouched by a - cap of 2 (below K), so the byte-driven plan is preserved.""" - # Heavy axis of four 30-char atoms; a limit tight enough that the byte pass - # must drive every atom into its own chunk (4 pieces > the cap of 2). - args = {"monitoring_location_id": ["X" * 30, "Y" * 30, "Z" * 30, "W" * 30]} - baseline = ChunkPlan(args, _fake_build, url_limit=250, max_chunks=1) - assert baseline.total > 2 # byte pass alone already fanned out past 2 - refined = ChunkPlan(args, _fake_build, url_limit=250, max_chunks=2) - # cap 2 < baseline pieces → refine is a no-op here. - assert refined.total == baseline.total - - -def test_cap_never_exceeds_the_byte_budget(): - """Refining on top of an over-budget request keeps the hard invariant: - every chunk still fits ``url_limit`` (splitting only ever shrinks - a chunk), and the fan-out is at least what the byte pass required.""" - args = {"monitoring_location_id": ["X" * 30, "Y" * 30, "Z" * 30, "W" * 30]} - limit = 310 - byte_only = ChunkPlan(args, _fake_build, url_limit=limit, max_chunks=1) - plan = ChunkPlan(args, _fake_build, url_limit=limit, max_chunks=32) - assert plan.total >= byte_only.total - for sub in plan.iter_chunk_args(): - assert _safe_request_bytes(_fake_build, sub, limit) <= limit - - -def test_cap_refines_the_filter_axis(): - """The dial treats the cql-text ``filter`` axis like any other: an - under-budget filter of N top-level OR-clauses is split along that axis - into ``min(N, cap)`` pieces.""" - clauses = [f"p='{i}'" for i in range(8)] - args = {"filter": " OR ".join(clauses)} - plan = ChunkPlan(args, _fake_build, url_limit=8000, max_chunks=4) - assert len(plan.chunks["filter"]) == 4 # min(8, 4) - assert plan.total == 4 - - -def test_cap_caps_the_total_across_axes(): - """With more than one multi-value axis the cap bounds the *total* - chunk count (the cartesian product), not each axis independently — - the blast-radius guardrail the dial exists for. Two 6-atom axes at a cap - of 4 top out at 4 chunks total, not 4x4=16; growth is distributed - round-robin across axes rather than one axis alone climbing to the cap.""" - args = { - "monitoring_location_id": [f"L{i}" for i in range(6)], - "parameter_code": [f"{i:05d}" for i in range(6)], - } - plan = ChunkPlan(args, _fake_build, url_limit=8000, max_chunks=4) - assert plan.total == 4 - # Every atom on every axis is still covered exactly once. - for key, atoms in ( - ("monitoring_location_id", args["monitoring_location_id"]), - ("parameter_code", args["parameter_code"]), - ): - flattened = [a for chunk in plan.chunks[key] for a in chunk] - assert sorted(flattened) == sorted(atoms) - - -def test_cap_bounds_fan_out_across_many_axes(): - """The guardrail holds regardless of axis count: three multi-value axes at - a cap of 30 fan out to *at most* 30 chunks total — never the - ``30 ** 3`` a per-axis cap would allow, and never *over* the cap either. - 30 is deliberately not evenly reachable by these axes: a single split - multiplies the plan by more than one, so the naive ``while total < cap`` - the first refine used stepped past 30 (to 32). The cap is a hard ceiling — - the property neither the single-axis-only cap nor that naive loop - guaranteed.""" - cap = 30 - # Three chunkable axes (two list axes + the filter OR-axis), each with 10 - # atoms — under the old per-axis cap this would have been cap**3. - args = { - "monitoring_location_id": [f"L{i}" for i in range(10)], - "parameter_code": [f"{i:05d}" for i in range(10)], - "filter": " OR ".join(f"p='{i}'" for i in range(10)), - } - plan = ChunkPlan(args, _fake_build, url_limit=8000, max_chunks=cap) - assert 1 < plan.total <= cap # fanned out, but never past the ceiling - - -@pytest.mark.parametrize( - "atoms_per_axis, cap", - [ - (4, 5), # pre-fix loop overshot 5 -> 6 - (8, 10), # pre-fix loop overshot 10 -> 12 - (10, 7), # pre-fix loop overshot 7 -> 8 - ], -) -def test_cap_is_a_hard_ceiling_never_overshoots(atoms_per_axis, cap): - """The cap is a hard ceiling, not a soft target. With two multi-value axes - a single split multiplies the plan by ``(k+1)/k`` for the split axis — - adding the product of the *other* axes, not one — so a naive - ``while total < cap`` loop steps *past* the cap. These are exactly the - (atoms, cap) combos that loop overshot (5->6, 10->12, 7->8). The plan must - fan out and cover every atom once, but never exceed the cap, landing below - it when no whole split lands on it exactly (two even axes reach 4, not 5).""" - args = { - "monitoring_location_id": [f"L{i:03d}" for i in range(atoms_per_axis)], - "parameter_code": [f"{i:05d}" for i in range(atoms_per_axis)], - } - plan = ChunkPlan(args, _fake_build, url_limit=8000, max_chunks=cap) - assert 1 < plan.total <= cap # fanned out, but never past the ceiling - # Every atom on every axis is still covered exactly once. - for key, atoms in args.items(): - flattened = [a for chunk in plan.chunks[key] for a in chunk] - assert sorted(flattened) == sorted(atoms) - - -def test_cap_does_not_mask_unchunkable(): - """A request with nothing to split that still busts the byte limit must - raise ``Unchunkable`` regardless of the cap — the soft pass has no axis to - act on and must not swallow the hard failure.""" +def test_unchunkable_still_raised_without_the_dial(): args = {"monitoring_location_id": "one-huge-scalar"} with pytest.raises(Unchunkable): - ChunkPlan(args, _fake_build, url_limit=10, max_chunks=32) + ChunkPlan(args, _fake_build, url_limit=10) -def test_parallel_chunks_publishes_n_on_the_ambient(): - """The context manager publishes ``n`` on the ambient for the block and - restores the previous value on exit — including proper nesting.""" - assert _parallel_chunks.get() == 1 # default (off, = no extra fan-out) - with parallel_chunks(32): - assert _parallel_chunks.get() == 32 - with parallel_chunks(2): - assert _parallel_chunks.get() == 2 - assert _parallel_chunks.get() == 32 # outer restored - assert _parallel_chunks.get() == 1 # default (off) outside any block +def test_page_concurrency_defaults_and_follows_the_env(monkeypatch): + monkeypatch.delenv("API_USGS_CONCURRENT", raising=False) + assert _chunking.page_concurrency() == 32 + monkeypatch.setenv("API_USGS_CONCURRENT", "4") + assert _chunking.page_concurrency() == 4 -@pytest.mark.parametrize( - "bad", - [ - 0, # not positive - -1, # negative - 1.5, # a float, not an int - "8", # a string, even a numeric one - "high", # the old level names are gone - None, # None not accepted - True, # bool is an int subclass but nonsensical here - ["8"], # a list - ], -) -def test_parallel_chunks_rejects_non_positive_int(bad): - """``n`` must be a positive integer; every other shape — zero, negative, a - float, a string (including a numeric one and the old level names), ``None``, - a ``bool``, a list — raises ``ValueError`` at ``with`` entry, before any - request, and leaves the ambient untouched.""" - with pytest.raises(ValueError, match="must be a positive integer"): - with parallel_chunks(bad): - pass - assert _parallel_chunks.get() == 1 # default (off) — unchanged by a rejected call + monkeypatch.setenv("API_USGS_CONCURRENT", "1") + assert _chunking.page_concurrency() == 1 -def test_parallel_chunks_drives_end_to_end_fan_out(): - """End-to-end: the same fitting request passes through as a single call by - default, but fans into ``n`` chunks inside a ``parallel_chunks(n)`` - block — and the combined result still recovers every atom exactly once.""" - sites = [f"S{i:02d}" for i in range(8)] +def test_page_concurrency_clamps_unbounded(monkeypatch): + monkeypatch.setenv("API_USGS_CONCURRENT", "unbounded") + assert _chunking.page_concurrency() == 32 - calls: list[tuple[str, ...]] = [] - @multi_value_chunked(build_request=_fake_build, url_limit=8000) - async def fetch(args): - chunk = tuple(args["monitoring_location_id"]) - calls.append(chunk) - return pd.DataFrame({"site": list(chunk)}), _ok_response() +def test_connection_pool_covers_chunks_times_pages(monkeypatch, httpx_mock): + import dataretrieval.transport.fanout as fanout + from dataretrieval.waterdata import get_daily - # Default: comfortably under the byte limit → one passthrough call. - df_plain, _ = fetch({"monitoring_location_id": sites}) - assert len(calls) == 1 - assert sorted(df_plain["site"]) == sorted(sites) - - calls.clear() - with parallel_chunks(8): - df_fine, _ = fetch({"monitoring_location_id": sites}) - # 8 atoms at n=8 → 8 singleton chunks. - assert len(calls) == 8 - assert all(len(chunk) == 1 for chunk in calls) - # Union across chunks recovers the original set, once each. - assert sorted(a for chunk in calls for a in chunk) == sorted(sites) - assert sorted(df_fine["site"]) == sorted(sites) - - -@pytest.mark.parametrize("n", [1, 2, 3, 8]) -def test_parallel_chunks_supports_arbitrary_n(n): - """An arbitrary ``n`` (not only 2/8/32) fans an under-limit request into - exactly ``n`` chunks, together covering every site once — including - ``n=1``, the explicit no-op that stays a single passthrough call.""" - sites = [f"S{i:02d}" for i in range(8)] - calls: list[int] = [] + monkeypatch.setenv("API_USGS_CONCURRENT", "8") + captured: dict[str, object] = {} + real_open = fanout.open_async_client - @multi_value_chunked(build_request=_fake_build, url_limit=8000) - async def fetch(args): - calls.append(len(args["monitoring_location_id"])) - return pd.DataFrame(), _ok_response() + def spy(**overrides): + captured["limits"] = overrides.get("limits") + return real_open(**overrides) + + monkeypatch.setattr(fanout, "open_async_client", spy) + httpx_mock.add_response( + json={"type": "FeatureCollection", "numberReturned": 0, "features": []}, + headers={"Content-Type": "application/geo+json"}, + ) + + get_daily(monitoring_location_id="USGS-01646500", limit=10) - with parallel_chunks(n): - fetch({"monitoring_location_id": sites}) - assert len(calls) == n - assert sum(calls) == 8 + limits = captured["limits"] + assert isinstance(limits, httpx.Limits) + assert limits.max_connections == 8 * _chunking.page_concurrency() diff --git a/tests/waterdata_offset_paging_test.py b/tests/waterdata_offset_paging_test.py new file mode 100644 index 00000000..7a886201 --- /dev/null +++ b/tests/waterdata_offset_paging_test.py @@ -0,0 +1,336 @@ +"""End-to-end tests for offset-parallel page fetching through a real getter. + +``tests/transport_test.py`` covers the service-neutral walk in isolation. What +this module pins is the *wiring*: that a Water Data getter actually dispatches +to the offset walk (rather than the cursor walk), that the offset requests carry +the parameters the API needs, and that the two documented escape hatches -- +a server ignoring ``offset``, and the API's 40000 offset ceiling -- produce a +complete, correct result rather than a silently wrong one. + +Every test here is fully mocked (``httpx_mock``); nothing touches the network. +The suite-wide conftest pins ``API_USGS_CONCURRENT=1``, which is the "page +sequentially" setting, so each test that wants the parallel path re-sets it. +""" + +from __future__ import annotations + +import dataclasses +import json + +import httpx +import pytest + +import dataretrieval.waterdata.utils as _wd_utils +from dataretrieval.waterdata import get_daily + +_ITEMS_URL = "https://api.waterdata.usgs.gov/ogcapi/v0/collections/daily/items" +_GEOJSON = {"Content-Type": "application/geo+json"} + + +def _page(rows, *, next_url: str | None = None) -> str: + """A GeoJSON FeatureCollection page. + + ``next_url`` adds the ``next`` link the *cursor* walk follows. The offset + walk ignores links entirely -- it computes its own offsets -- so a page can + carry one without affecting the offset path. + """ + rows = list(rows) + body = { + "type": "FeatureCollection", + "numberReturned": len(rows), + "features": [ + { + "type": "Feature", + "id": f"daily-{row}", + "geometry": None, + "properties": { + "monitoring_location_id": "USGS-01646500", + "value": str(row), + }, + } + for row in rows + ], + "links": [{"rel": "next", "href": next_url}] if next_url else [], + } + return json.dumps(body) + + +def _offset_of(request: httpx.Request) -> int | None: + raw = request.url.params.get("offset") + return None if raw is None else int(raw) + + +@pytest.fixture +def parallel_pages(monkeypatch): + """Undo the conftest's sequential pin so the offset walk fans out.""" + monkeypatch.setenv("API_USGS_CONCURRENT", "4") + + +def _serve(httpx_mock, total_rows: int, *, limit: int) -> list[int | None]: + """Register a callback serving ``total_rows`` rows in ``limit``-sized pages. + + Returns the list that records each request's ``offset``, so a test can + assert on the request count -- the quota cost, which is the whole reason + overlapping pages is preferable to splitting the query. A request with no + ``offset`` is the cursor walk's, and is answered with a linked page 1 so a + fallback can still complete. + """ + seen: list[int | None] = [] + + def respond(request: httpx.Request) -> httpx.Response: + offset = _offset_of(request) + seen.append(offset) + if offset is None: + return httpx.Response( + 200, + text=_page(range(min(limit, total_rows))), + headers=_GEOJSON, + ) + rows = range(offset, min(offset + limit, total_rows)) + return httpx.Response( + 200, + text=_page(rows), + headers={**_GEOJSON, "x-ratelimit-limit": "1000"}, + ) + + httpx_mock.add_callback(respond) + return seen + + +def test_getter_pages_by_offset_and_returns_every_row(httpx_mock, parallel_pages): + """The headline behavior: a multi-page result comes back complete, in + order, fetched via computed offsets rather than followed cursors.""" + seen = _serve(httpx_mock, total_rows=25, limit=10) + + df, _ = get_daily(monitoring_location_id="USGS-01646500", limit=10) + + # ``value`` is a Water Data numerical column, so ``convert_type`` (on by + # default) coerces it — hence ints, not the strings the mock served. + assert df["value"].tolist() == list(range(25)) + # Every request carried an offset, so no page came from a ``next`` link. + assert all(off is not None for off in seen) + assert 0 in seen + + +def test_offset_requests_preserve_the_query(httpx_mock, parallel_pages): + """Each page request is the planned request plus ``offset`` -- the filters + and page size must survive, or later pages would answer a different + question than the first.""" + _serve(httpx_mock, total_rows=25, limit=10) + + get_daily(monitoring_location_id="USGS-01646500", parameter_code="00060", limit=10) + + for request in httpx_mock.get_requests(): + params = request.url.params + assert params["monitoring_location_id"] == "USGS-01646500" + assert params["parameter_code"] == "00060" + assert params["limit"] == "10" + assert "offset" in params + + +def test_page_count_is_not_inflated_by_parallelism(httpx_mock, parallel_pages): + """Offsets are speculative -- a wave may overshoot the end -- but the walk + must not spend materially more quota than the sequential walk would. The + ramping wave width (1, 2, 4, ...) is what holds that line: 25 rows at limit + 10 needs 3 pages, and waves of 1 then 2 cover exactly those 3, so parallel + paging costs the *same* 3 requests the sequential walk would have spent.""" + seen = _serve(httpx_mock, total_rows=25, limit=10) + + get_daily(monitoring_location_id="USGS-01646500", limit=10) + + assert len(seen) == 3 + + +def test_sequential_setting_uses_the_cursor_walk(httpx_mock, monkeypatch): + """``API_USGS_CONCURRENT=1`` is the documented way back to strictly + sequential paging, and it must use *standard* OGC paging -- following + ``next`` links -- not offsets with a wave of one.""" + monkeypatch.setenv("API_USGS_CONCURRENT", "1") + seen = _serve(httpx_mock, total_rows=10, limit=10) + + df, _ = get_daily(monitoring_location_id="USGS-01646500", limit=10) + + assert len(df) == 10 + assert seen == [None] # no offset parameter was ever sent + + +def test_falls_back_to_cursors_when_the_server_ignores_offset( + httpx_mock, parallel_pages, caplog +): + """``offset`` is a Water Data extension, not part of OGC API - Features, and + an unrecognized query parameter is conventionally *ignored*, not rejected. A + server that ignores it answers every offset with page 1, so a naive walk + would concatenate the same rows N times and report success. The walk must + detect that and complete the query the standards-only way -- the result + stays correct, only the speed changes.""" + served: list[int | None] = [] + + def respond(request: httpx.Request) -> httpx.Response: + offset = _offset_of(request) + served.append(offset) + if request.url.params.get("cursor") == "c1": + return httpx.Response(200, text=_page(range(10, 15)), headers=_GEOJSON) + # Page 1 regardless of the offset asked for; a next link so the cursor + # fallback has somewhere to go. + return httpx.Response( + 200, + text=_page(range(10), next_url=f"{_ITEMS_URL}?cursor=c1"), + headers=_GEOJSON, + ) + + httpx_mock.add_callback(respond) + + with caplog.at_level("WARNING"): + df, _ = get_daily(monitoring_location_id="USGS-01646500", limit=10) + + # Correct, de-duplicated result via the cursor walk -- NOT page 1 repeated. + assert df["value"].tolist() == list(range(15)) + assert "sequential pagination" in caplog.text + # The offset attempt happened, then was abandoned in favor of a walk that + # sends no ``offset`` at all. + assert any(off is not None for off in served) + assert None in served + + +def test_ceiling_hands_off_to_a_cursor_walk_for_the_tail( + httpx_mock, parallel_pages, monkeypatch +): + """The API rejects ``offset > 40000``. That is a ceiling, not an + end-of-data signal, so the walk must continue with cursors (which have no + ceiling) rather than truncate. This is the invariant that keeps an + arbitrarily deep pull *complete*, not merely fast. The ceiling is lowered + here so the test stays small; the mechanism is identical at 40000.""" + limit, ceiling, total = 10, 25, 45 + cursor_rows = {"c1": range(30, 40), "c2": range(40, total)} + monkeypatch.setattr( + _wd_utils, + "WATERDATA_DIALECT", + dataclasses.replace(_wd_utils.WATERDATA_DIALECT, max_offset=ceiling), + ) + + def respond(request: httpx.Request) -> httpx.Response: + cursor = request.url.params.get("cursor") + if cursor is not None: + nxt = f"{_ITEMS_URL}?cursor=c2" if cursor == "c1" else None + return httpx.Response( + 200, text=_page(cursor_rows[cursor], next_url=nxt), headers=_GEOJSON + ) + offset = _offset_of(request) + assert offset is not None and offset <= ceiling, ( + f"walk issued offset={offset}, past the ceiling of {ceiling}" + ) + # Every offset page carries a next link. The offset walk ignores links + # entirely, so this only matters for the tail hand-off -- a cursor walk + # re-seeded at the last accepted offset, which follows it. + rows = range(offset, min(offset + limit, total)) + return httpx.Response( + 200, + text=_page(rows, next_url=f"{_ITEMS_URL}?cursor=c1"), + headers=_GEOJSON, + ) + + httpx_mock.add_callback(respond) + + df, _ = get_daily(monitoring_location_id="USGS-01646500", limit=limit) + + # Complete and seamless: offsets covered rows 0-19, the re-seeded cursor + # walk covered 20-44. No gap, and the one re-fetched page de-duplicates. + assert df["value"].tolist() == list(range(total)) + + +def test_small_result_costs_one_request_at_the_shipped_default(httpx_mock, monkeypatch): + """A one-page result must cost exactly one request -- at the *default* width, + which is what users actually get. + + This is the regression test for the bug that shipped: the wave width started + at the full ``API_USGS_CONCURRENT`` (32), so a single-page query fired 21 + requests (32 offsets clipped to the 40000 ceiling) to discover it was already + finished. Every other test in the suite pinned the width to 1 or 4 -- and the + conftest pins 1 -- so nothing exercised the shipped value. Deleting the env + var here, rather than setting a number, is the point of the test. + """ + monkeypatch.delenv("API_USGS_CONCURRENT", raising=False) + seen = _serve(httpx_mock, total_rows=5, limit=2000) + + df, _ = get_daily(monitoring_location_id="USGS-01646500", limit=2000) + + assert len(df) == 5 + assert len(seen) == 1, f"a one-page result cost {len(seen)} requests" + + +def test_request_count_stays_within_twice_the_pages_needed(httpx_mock, monkeypatch): + """The ramp's headline guarantee, at the default width: doubling means the + sum of all prior waves is less than the current one, so total requests stay + under 2x the pages that exist no matter how the result size falls between + wave boundaries. A flat wave would be ``width``x on every short result.""" + monkeypatch.delenv("API_USGS_CONCURRENT", raising=False) + limit, total = 2000, 19_000 # 10 pages: the deep-history case + seen = _serve(httpx_mock, total_rows=total, limit=limit) + + df, _ = get_daily(monitoring_location_id="USGS-01646500", limit=limit) + + pages_needed = -(-total // limit) + assert len(df) == total + assert len(seen) < 2 * pages_needed, ( + f"{len(seen)} requests for {pages_needed} pages exceeds the 2x bound" + ) + + +def test_offset_is_still_verified_when_the_first_wave_is_one_page( + httpx_mock, monkeypatch, caplog +): + """The ignore-detection guard compares two pages at different offsets. The + ramp makes the first wave a *single* page, so the comparison has to span + waves -- otherwise the guard silently never fires and a server that ignores + ``offset`` would yield page 1 concatenated N times.""" + monkeypatch.delenv("API_USGS_CONCURRENT", raising=False) + served: list[int | None] = [] + + def respond(request: httpx.Request) -> httpx.Response: + served.append(_offset_of(request)) + if request.url.params.get("cursor") == "c1": + return httpx.Response(200, text=_page(range(10, 15)), headers=_GEOJSON) + # Page 1 no matter which offset was asked for. + return httpx.Response( + 200, + text=_page(range(10), next_url=f"{_ITEMS_URL}?cursor=c1"), + headers=_GEOJSON, + ) + + httpx_mock.add_callback(respond) + + with caplog.at_level("WARNING"): + df, _ = get_daily(monitoring_location_id="USGS-01646500", limit=10) + + assert df["value"].tolist() == list(range(15)) + assert "sequential pagination" in caplog.text + + +def test_no_data_returns_an_empty_frame_not_an_error(httpx_mock, parallel_pages): + """A query matching nothing must return an empty DataFrame, exactly as the + sequential walk does. + + "A no-data result is *not* an error" is a documented, load-bearing promise + of the modern getters, and a query that matches nothing is ordinary -- a + typo'd site id, a parameter the site doesn't measure, a gap in the record. + The offset walk used to raise ``DataRetrievalError`` here: the single empty + page it fetched was *discarded* as past-the-end, which left it with no + response to report and it fell through to its "issued no requests" guard. + """ + _serve(httpx_mock, total_rows=0, limit=10) + + df, _ = get_daily(monitoring_location_id="USGS-99999999", limit=10) + + assert len(df) == 0 + + +def test_max_rows_is_exact_under_parallel_paging(httpx_mock, parallel_pages): + """A wave can fetch past the requested row count, so the cap has to be + applied to the combined frame -- otherwise ``max_rows`` would return + whatever a wave boundary happened to land on.""" + _serve(httpx_mock, total_rows=1000, limit=10) + + df, _ = get_daily(monitoring_location_id="USGS-01646500", limit=10, max_rows=25) + + assert len(df) == 25 + assert df["value"].tolist() == list(range(25))