diff --git a/NEWS.md b/NEWS.md index a9a63718..ba21d320 100644 --- a/NEWS.md +++ b/NEWS.md @@ -1,3 +1,5 @@ +**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/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. **08/02/2026:** Phase 1 OGC boundary stabilization: the `dataretrieval.ogc` package now exposes a deliberate, small facade (`OgcDialect`, `prepare_request_args`, `get_ogc_data`, `fetch_ogc_request`) that service adapters (NGWMN, Water Data's generic wrapper) import. Internal request-construction helpers moved to a new `ogc.requests` module; the dialect type and endpoint constants live in the leaf `ogc.policy` module. `ogc.shaping` no longer depends on `ogc.engine` at all, and the complete runtime OGC import graph is now acyclic. `_default_headers` now accepts a target URL and adds `X-Api-Key` only for `api.waterdata.usgs.gov`; shared sync and async clients also strip the key before following any cross-host redirect, including external rating-asset downloads. `waterdata.utils` no longer bulk-re-exports private OGC symbols, and consumers import implementation helpers from their canonical modules. No public API, return-value, or deprecation changes. diff --git a/README.md b/README.md index 88b076ab..ffb96e29 100644 --- a/README.md +++ b/README.md @@ -106,57 +106,50 @@ 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 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. It 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 just adds requests, so it is a deliberate, -scoped `with` block, never the default. +#### Large downloads are paged in parallel automatically + +A large result arrives one page at a time, and cursor pagination is inherently +sequential: page *N+1*'s URL is only revealed by page *N*, so a 10-page result +costs 10 round trips end to end. The Water Data API also accepts an `offset` +parameter, which means every page's URL is computable up front +(`offset = i * limit`) — so `dataretrieval` fetches them **concurrently** +instead. Nothing to opt into; it is the default path for every getter. ```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", - ) +# 20 years of daily discharge for one gage: ~7,300 rows over several pages, +# fetched concurrently rather than one after another. +df, md = waterdata.get_daily( + monitoring_location_id="USGS-01646500", + parameter_code="00060", # discharge + time="2004-01-01/2023-12-31", +) ``` -`n` is the number of sub-requests to fan the call out into. It is capped by how -many values there are to split, and each sub-request costs a request against -your hourly [rate limit](https://api.waterdata.usgs.gov/signup/); since how many -run *at once* is capped separately by `API_USGS_CONCURRENT` (default 32), 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` was run 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, so the exact multiplier 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. +The **request count is unchanged** — only their timing is. That matters because +the hourly [rate limit](https://api.waterdata.usgs.gov/signup/) counts requests, +so the speedup is free of quota. How many pages are in flight at once is capped +by `API_USGS_CONCURRENT` (default 32); set it to `1` to page sequentially. + +Measured, against the live API with `limit=2000`: + +| pull | pages | sequential | offset-parallel | speedup | +| ------------------------------------------ | ----- | ---------- | --------------- | ----------- | +| `get_daily`, 8 sites, 2 years | 8 | 4.6 s | 2.2 s | 2.1× | +| `get_daily`, 1 site, full history (~19k rows) | 10 | 6.3 s | 1.5–2.0 s | 3.1–4.2× | + +The second row is the case that previously had no answer at all: a single-site +deep history has no multi-value argument to divide, so splitting the *query* to +gain parallelism was impossible. Paging by offset parallelizes the pages +themselves, so it applies whether or not the query can be split. + +Because `offset` is a Water Data extension rather than part of OGC API - +Features, the walk is defensive: it verifies the server is honoring the +parameter before trusting any rows, and falls back to standard cursor +pagination if not. Past the API's `offset` ceiling of 40,000 rows it continues +with a sequential cursor walk, so arbitrarily deep results still come back +complete. Visit the [API Reference](https://doi-usgs.github.io/dataretrieval-python/reference/waterdata.html) diff --git a/dataretrieval/__init__.py b/dataretrieval/__init__.py index 4226e247..48b5c865 100644 --- a/dataretrieval/__init__.py +++ b/dataretrieval/__init__.py @@ -45,11 +45,6 @@ URLTooLong, ) -# 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 - # Resumable chunk-interruption exceptions. They are defined in # ``dataretrieval.ogc.interruptions`` rather than ``dataretrieval.exceptions`` # because they carry pandas/httpx state and a resumable ``ChunkedCall`` handle, @@ -101,6 +96,5 @@ "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 f15f226b..c85b1c08 100644 --- a/dataretrieval/ogc/chunking.py +++ b/dataretrieval/ogc/chunking.py @@ -9,11 +9,11 @@ cartesian product of chunks. Requests that already fit get a trivial single-step plan — ``ChunkedCall`` has one code path either way. -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 sub-requests; ``n`` drives :meth:`ChunkPlan._refine`. See -``parallel_chunks`` for the why and the when. +Chunking splits only as far as the byte limit forces. Parallelism *within* a +sub-request is the offset-parallel page walk's job +(:func:`dataretrieval.transport.offsets.paginate_by_offset`), which overlaps a +result's pages without issuing extra requests — so the planner no longer has a +fan-out dial. See :func:`page_concurrency` for the wave width. This module owns the *execution* half — the event loop and bounded concurrency that drive a plan to completion (``ChunkedCall``) plus the @@ -73,7 +73,6 @@ import functools import os from collections.abc import Awaitable, Callable, Iterator -from contextlib import contextmanager from contextvars import copy_context from typing import Any, cast @@ -90,7 +89,7 @@ from dataretrieval.transport.http import open_async_client from dataretrieval.transport.retry import _NO_RETRY, RetryPolicy from dataretrieval.transport.retry import retry_async as _retry -from dataretrieval.utils import Ambient, _require_positive_int +from dataretrieval.utils import Ambient from .interruptions import ChunkInterrupted from .planning import ChunkPlan @@ -171,116 +170,23 @@ def get_active_client() -> httpx.AsyncClient | None: return _chunked_client.get() -# 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 -# sub-request 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) +# Page-fetch concurrency: how many pages of one sub-request are fetched at +# once by the offset-parallel walk. Resolved from ``API_USGS_CONCURRENT`` (the +# same dial that bounds sub-request fan-out) so one env var governs all +# outstanding requests, and ``1`` still means strictly sequential. +def page_concurrency() -> int: + """Pages to fetch per wave in the offset-parallel walk. - -@contextmanager -def parallel_chunks(n: int) -> Iterator[None]: + Reuses the ``API_USGS_CONCURRENT`` cap rather than adding a second dial: + both bound how many requests this library has in flight, and the quota they + spend is volume-based, not simultaneity-based. ``unbounded`` is clamped to a + finite wave width because a wave is *speculative* — an unbounded wave would + issue arbitrarily many past-the-end requests to discover one short page. """ - Fan the OGC getters' multi-value requests out into ``n`` parallel sub-requests. - - 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 sub-requests that - fit. That is the safe default, but it can be *needlessly* conservative: - because every sub-request paginates, splitting a large result further costs - little or no extra quota *as long as each sub-request still spans many - pages* — rows-per-chunk far exceeding the page size (ten states pulled as - one request then page nearly as many times as ten per-state requests - would). When a split leaves each sub-request 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 — trading roughly the same pages for more, smaller sub-requests, which - gives smoother progress, more even concurrency, and a smaller unit of - retry/resume. - - 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), this is a *deliberate* per-call knob rather than an - automatic behavior or a process-wide environment variable — 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. - - Parameters - ---------- - n : int - The number of sub-requests to fan the whole call out into — a positive - integer such as ``2``, ``8``, or ``32``. It caps the plan's *total* - sub-request 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 sub-request 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 sub-request fetches at least one page, so it costs at least one - request against your hourly rate limit — a larger ``n`` spends more - quota. And because how many sub-requests run *at once* is capped - separately by ``API_USGS_CONCURRENT`` (default 32), 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 sub-request 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 sub-requests, 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.exceptions.ChunkInterrupted` - (or ``QuotaExhausted``) carrying the completed sub-requests, which you - finish with ``exc.call.resume()``. - - Cross-sub-request 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``. - """ - # 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 = _read_concurrency_env() + if resolved is None: + return _CONCURRENCY_DEFAULT + return resolved # --------------------------------------------------------------------------- @@ -592,6 +498,11 @@ async def _run(self, max_concurrent: int | None) -> tuple[pd.DataFrame, Any]: (``httpx.Limits(max_connections=N, max_keepalive_connections=N)``) so the in-flight fetches reuse keepalive connections. + Because each sub-request additionally fans its own pages out via + ``offset``, peak in-flight requests reach ``max_concurrent x + page_concurrency()``, so the pool is sized to that product rather than + to ``max_concurrent``. + The semaphore, not the pool, is deliberately the throttle. If the pool throttled instead, the excess sub-requests would queue *inside* httpx waiting for a connection, and that wait counts @@ -638,8 +549,21 @@ 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. + # Size the pool for the *product*, not the sub-request cap. Each + # sub-request now runs its own offset-parallel page wave, and those + # pages are not held by this semaphore (they can't be — a sub-request + # holds a permit for its whole attempt, so pages waiting on the same + # semaphore would deadlock against their own parent). Peak in-flight is + # therefore up to ``max_concurrent * page_width``. Sizing the pool to + # ``max_concurrent`` alone leaves the excess queued *inside* httpx + # against the 60 s pool-acquire timeout — the spurious-``PoolTimeout`` + # failure this method's docstring warns about, which measurably fired + # as a mid-walk ``ReadError`` before this was widened. + pool_size = ( + None if max_concurrent is None else max_concurrent * page_concurrency() + ) limits = httpx.Limits( - max_connections=max_concurrent, max_keepalive_connections=max_concurrent + max_connections=pool_size, max_keepalive_connections=pool_size ) semaphore = asyncio.Semaphore( self.plan.total if max_concurrent is None else max_concurrent @@ -715,9 +639,10 @@ 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 sub-request 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 + sub-request then fetches its own pages concurrently via ``offset`` (see + :func:`page_concurrency`). See the module docstring for the concurrency + model. Parameters ---------- @@ -761,14 +686,10 @@ 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 sub-request cap). It only affects *planning*, done - # here up front, so a later resume — which re-issues the - # already-planned sub-requests — needs no snapshot. - plan = ChunkPlan( - args, build_request, limit, max_chunks=_parallel_chunks.get() - ) + # Chunk only as far as the byte budget forces. Parallelism within a + # sub-request now comes from the offset-parallel page walk, so the + # planner no longer splits beyond what the URL limit requires. + 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, diff --git a/dataretrieval/ogc/engine.py b/dataretrieval/ogc/engine.py index 2ab300d7..82f91202 100644 --- a/dataretrieval/ogc/engine.py +++ b/dataretrieval/ogc/engine.py @@ -66,7 +66,17 @@ _switch_properties_id, prepare_request_args, ) + +# Genuine dependencies of the offset-parallel page walk (not compatibility +# re-exports): the engine derives each page's request from the planned one. +from dataretrieval.ogc.requests import ( + page_limit as _page_limit, +) +from dataretrieval.ogc.requests import ( + with_offset as _with_offset, +) from dataretrieval.ogc.shaping import GEOPANDAS, _finalize_ogc, _get_resp_data +from dataretrieval.transport.offsets import OffsetUnsupported, paginate_by_offset from dataretrieval.transport.pagination import paginate from dataretrieval.transport.sync import run_sync from dataretrieval.utils import ( @@ -187,6 +197,15 @@ async def _paginate( ) +def _ogc_parse_page(resp: httpx.Response, *, geopd: bool) -> pd.DataFrame: + """Parse one OGC API page to its frame, ignoring links. + + The offset walk's parser. It needs no cursor — offsets are computed, not + discovered — so this is :func:`_ogc_parse_response` without the link hop. + """ + return _get_resp_data(resp, geopd=geopd) + + def _ogc_parse_response( resp: httpx.Response, *, geopd: bool ) -> tuple[pd.DataFrame, str | None]: @@ -204,6 +223,70 @@ def _ogc_parse_response( ) +async def _walk_pages_by_offset( + geopd: bool, + req: httpx.Request, + client: httpx.AsyncClient | None = None, + *, + width: int, + max_offset: int | None, +) -> tuple[pd.DataFrame, httpx.Response]: + """Fetch this request's pages concurrently via ``offset``, then finish + sequentially if the offset ceiling is reached before the data runs out. + + The fast path for a deep result. Where :func:`_walk_pages` must parse page + ``N`` to learn page ``N+1``'s cursor, every offset is known up front, so the + pages overlap. Same request count, same rows, lower wall clock. + + Falls back to :func:`_walk_pages` for the remainder past ``max_offset`` + (the API's ceiling, 40000 for Water Data), which keeps the result complete + for arbitrarily large pulls: offsets cover the parallelizable prefix, + cursors cover the unbounded tail. + """ + limit = _page_limit(req) + if limit is None: + # No usable ``limit`` on the request, so no stride and no short-page + # signal. Fall back rather than guess. + return await _walk_pages(geopd, req, client) + + async def tail_walk( + resume_offset: int, rows_so_far: int, session: httpx.AsyncClient + ) -> tuple[pd.DataFrame, httpx.Response]: + """Cursor-walk the remainder, starting at the ceiling offset.""" + # Rebase the row cap onto what's left: the ambient cap counts the whole + # result, but this continuation only sees the tail, so handing it the + # original cap would over-fetch by the rows already held. + cap = _row_cap.get() + remaining = None if cap is None else max(cap - rows_so_far, 0) + with _row_cap(remaining): + return await _walk_pages(geopd, _with_offset(req, resume_offset), session) + + 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.get(), + tail_walk=tail_walk, + ) + except OffsetUnsupported as exc: + # The API stopped honoring ``offset`` (or never did for this + # collection). ``offset`` is a non-standard extension, so this is a + # supported outcome, not a bug: warn once and complete the query the + # standards-only way. Raised before any rows were returned, so the + # re-walk can't duplicate data. + 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, req, client) + + async def _walk_pages( geopd: bool, req: httpx.Request, @@ -317,7 +400,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. if max_rows is not None: _require_positive_int(max_rows, "max_rows") @@ -380,7 +463,20 @@ async def _fetch_once( ``(frame, response)``. """ req = _construct_api_requests(**args) - return await _walk_pages(geopd=GEOPANDAS, req=req) + width = chunking.page_concurrency() + max_offset = _dialect.get().max_offset + # Offset-parallel only when the API declares an offset ceiling (i.e. it + # supports the non-standard ``offset`` parameter at all) and the caller + # hasn't pinned the walk to one page at a time. + if max_offset is None or width <= 1: + return await _walk_pages(geopd=GEOPANDAS, req=req) + return await _walk_pages_by_offset( + GEOPANDAS, + req, + get_active_client(), + width=width, + max_offset=max_offset, + ) def _run_sync( diff --git a/dataretrieval/ogc/planning.py b/dataretrieval/ogc/planning.py index 76796df4..61e923d8 100644 --- a/dataretrieval/ogc/planning.py +++ b/dataretrieval/ogc/planning.py @@ -245,12 +245,10 @@ 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 - preserves the partition invariants every consumer relies on: *coverage* - (each atom survives, exactly once) and *contiguous, deterministic order* - (resume and :meth:`ChunkPlan.iter_sub_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). + The primitive :meth:`ChunkPlan._plan` 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_sub_args` depend on it). """ chunk = chunks[idx] mid = len(chunk) // 2 @@ -283,21 +281,6 @@ class ChunkPlan: url_limit : int Byte budget for the request (URL + body) — a hard ceiling every sub-request must fit. - max_chunks : int, optional - Hard cap on the plan's total sub-request count (default ``1`` = off). - ``1`` chunks only as much as ``url_limit`` requires — the most - conservative plan, fewest sub-requests — so a fitting request is a - passthrough. A cap of ``2`` or more fans the plan out to up to - ``max_chunks`` sub-requests overall (the cartesian product across axes, - never fewer than the byte budget already forces) — capped as a whole, - not per axis, so several multi-value axes can't multiply past the cap. - The plan never exceeds the cap and may land below it when no whole - split lands on it exactly. ``max_chunks`` is a sub-request 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 ---------- args : dict @@ -321,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__( @@ -330,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 sub-request *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]]] = {} @@ -350,9 +319,9 @@ 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 - # verbatim (the common passthrough). ``_safe_request_bytes`` treats + # No chunkable axis: nothing to split. If the single request fits, + # run it verbatim (the common passthrough) — its pages are then + # overlapped by the offset walk. ``_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: return @@ -391,25 +360,19 @@ 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_sub_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: + # A request that already fits is the common passthrough: leave + # ``axes``/``chunks`` empty so ``total == 1`` and ``iter_sub_args`` + # yields the original args verbatim. Splitting a fitting request for + # parallelism is no longer the planner's job — the offset-parallel page + # walk overlaps that request's pages instead, without extra requests. + 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 sub-request 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) + # Greedy-halve until every worst-case sub-request fits the byte budget + # (may raise ``Unchunkable``). + self._plan(build_request, url_limit) if self.canonical_url is None: # Original URL was un-constructable (httpx.InvalidURL); fall @@ -463,68 +426,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 — - 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`` sub-requests, 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`` sub-requests (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 sub-requests 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 dict representing the largest sub-request the current diff --git a/dataretrieval/ogc/policy.py b/dataretrieval/ogc/policy.py index e8831b10..8f8e67d4 100644 --- a/dataretrieval/ogc/policy.py +++ b/dataretrieval/ogc/policy.py @@ -51,6 +51,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) @@ -58,6 +65,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 462cfe4b..b30930cf 100644 --- a/dataretrieval/ogc/requests.py +++ b/dataretrieval/ogc/requests.py @@ -128,6 +128,57 @@ def _ogc_query_params( return params +def with_offset(request: httpx.Request, offset: int) -> httpx.Request: + """Rebuild ``request`` with ``offset`` applied, preserving everything else. + + Used by the offset-parallel page walk to derive each page's request from the + single request the chunker planned, so the page URLs stay byte-for-byte + identical to the sequential path apart from the added parameter. Rebuilding + (rather than mutating) keeps the original request reusable across waves. + + ``offset=0`` is still written explicitly: it makes the first page's URL + match the shape of its siblings, so a cached or logged URL set is uniform. + + Parameters + ---------- + request : httpx.Request + The planned page request (GET, or POST with a CQL2 body). + offset : int + Absolute row offset to request. + + Returns + ------- + httpx.Request + A new request with ``offset`` set in the query string. The method, + headers, and body (for POST/CQL2) are carried over unchanged. + """ + 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: + """The ``limit`` on ``request``'s URL, or ``None`` when absent/unparseable. + + The offset walk needs the page size to compute its stride and to recognize + a short (final) page. Reading it back off the request keeps a single source + of truth: whatever ``_ogc_query_params`` decided, including its 50000 clamp. + """ + 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 _construct_api_requests( service: str, properties: list[str] | None = None, diff --git a/dataretrieval/transport/offsets.py b/dataretrieval/transport/offsets.py new file mode 100644 index 00000000..dfeb2a81 --- /dev/null +++ b/dataretrieval/transport/offsets.py @@ -0,0 +1,473 @@ +"""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 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 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 waves until a stop condition is met. + + The offset-parallel counterpart to + :func:`dataretrieval.transport.pagination.paginate`. Returns the same + ``(frame, response)`` shape so either strategy can back the same getter. + + Parameters + ---------- + build_page : callable + Maps an absolute row offset to the :class:`httpx.Request` for that page. + parse_page : callable + Maps a response to its frame. + raise_for_status : callable + Maps a non-success response to a typed exception. + client : httpx.AsyncClient, optional + Client to send on. When supplied the caller owns its lifecycle, so a + chunked call can share one connection pool across sub-requests; + ``None`` opens a short-lived guarded client for this walk. + limit : int + Page size, and the offset stride. + width : int + Pages to fetch per wave. + max_offset : int or None, optional + Largest offset the service accepts (Water Data: 40000). The walk never + issues a request beyond it; instead it hands off to ``tail_walk``. + row_cap : int or None, optional + Stop once this many rows are collected. + tail_walk : callable, optional + Sequential continuation invoked when the offset ceiling is reached + before the data runs out. Called as + ``tail_walk(resume_offset, rows_so_far, client)`` and must return the + same ``(frame, response)`` shape for the *remainder* of the result. + Without it, hitting the ceiling logs a warning and returns a partial + frame — so callers that need completeness must supply it. + + Returns + ------- + pandas.DataFrame + Concatenated pages, in offset order. + httpx.Response + Aggregated metadata: the first page's URL, the last-completed page's + headers, and ``elapsed`` summed across pages. + + Raises + ------ + DataRetrievalError + If any page fails, wrapped with the recovery guidance the sequential + walk uses. Sibling requests in the wave are cancelled. + """ + async with _client_for(client) as session: + reporter = _progress.current() + frames: list[pd.DataFrame] = [] + first_response: httpx.Response | None = None + last_response: httpx.Response | None = None + total_elapsed = timedelta(0) + offset = 0 + ceiling_reached = False + offset_verified = False + # Ramp the wave width from 1 up to ``width`` (see the module docstring). + # Starting at the full width would make a one-page result cost ``width`` + # requests; starting at 1 and doubling keeps the small-query cost + # identical to the sequential walk while still reaching full parallelism + # within a few waves for a genuinely large result. + wave_width = 1 + + while True: + offsets = plan_offsets( + limit=limit, width=wave_width, start=offset, max_offset=max_offset + ) + if not offsets: + # The offset ceiling — NOT end of data. Hand off to the sequential + # continuation so the result stays complete; without one, fall + # through to the truncation warning below. + ceiling_reached = True + break + + try: + responses = await asyncio.gather( + *[ + _fetch_page(build_page, off, session, raise_for_status, None) + for off in offsets + ] + ) + except Exception as exc: # noqa: BLE001 + logger.warning( + "Offset-parallel page fetch failed at offsets %r.", offsets + ) + raise DataRetrievalError( + paginated_failure_message(len(frames), exc) + ) from exc + + try: + wave = [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(len(frames), exc) + ) from exc + + if not offset_verified: + # Verify the server honors ``offset`` as soon as two pages at + # *different* offsets are in hand, and before this function + # returns anything. Because the wave width ramps from 1, the + # first wave holds a single page, so the comparison usually + # spans waves: page 0 (kept from the previous wave) against the + # first page of this one. + probe = wave if len(wave) >= 2 else [*frames[-1:], *wave] + if _offset_ignored(probe): + raise OffsetUnsupported( + "The service returned identical pages for different " + "`offset` values, so it appears to ignore `offset`." + ) + # Two distinct offsets returned distinct pages: settled. + offset_verified = len(probe) >= 2 + + if first_response is None: + first_response = responses[0] + if last_response is None: + # Seed the metadata response before the keep/discard split below. + # A query that matches nothing returns one empty page, keeps zero + # pages, and would otherwise leave this ``None`` and raise the + # "issued no requests" error — turning a legitimate empty result + # into a failure. An empty result is not an error here (the + # sequential walk returns an empty frame), so the response that + # *reported* the emptiness is the right metadata to carry. + last_response = responses[0] + + stop_at = _stop_index( + wave, + limit=limit, + rows_before=sum(len(f) for f in frames), + row_cap=row_cap, + ) + keep = wave if stop_at is None else wave[: stop_at + 1] + + # Attribute progress and metadata only to pages we keep, so a + # speculative overshoot past the end doesn't inflate the page count. + for response, frame in zip(responses[: len(keep)], keep, strict=False): + total_elapsed += _safe_elapsed(response) + last_response = response + 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)) + + frames.extend(keep) + + if stop_at is not None: + break + offset = offsets[-1] + limit + # Every page in this wave was full, so more data is likely: widen. + wave_width = min(wave_width * 2, width) + + if ceiling_reached and tail_walk is not None and frames: + # The next offset the walk would need is *past* the ceiling, so it + # can't seed the continuation either — that request would earn the + # same rejection the offset walk just avoided. Rewind one page: drop + # the last page fetched and let the continuation re-fetch it from + # ``offsets[-1]``, the largest offset the service still accepts. + # Costs one duplicate request per deep query and buys a seam with no + # gap (missing rows) and no overlap (duplicate rows). + offset -= limit + frames.pop() + + rows_so_far = sum(len(f) for f in frames) + + if ceiling_reached: + if tail_walk is not None: + logger.debug( + "Offset ceiling (%s) reached after %d row(s); continuing " + "sequentially from offset %d.", + max_offset, + rows_so_far, + offset, + ) + tail_frame, tail_response = await tail_walk( + offset, rows_so_far, session + ) + if len(tail_frame): + frames.append(tail_frame) + total_elapsed += _safe_elapsed(tail_response) + last_response = tail_response + else: + 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, + rows_so_far, + ) + + if first_response is None or last_response is None: + # No wave completed — only reachable when the very first plan was empty + # (``max_offset`` below the first offset), which the caller prevents. + raise DataRetrievalError( + "Offset-parallel pagination issued no requests; " + f"max_offset={max_offset!r} leaves no valid page offset." + ) + + result = pd.concat(frames, ignore_index=True) if frames else pd.DataFrame() + if row_cap is not None: + result = result.head(row_cap) + return ( + result, + _merge_response( + first_response, headers_from=last_response, elapsed=total_elapsed + ), + ) diff --git a/dataretrieval/utils.py b/dataretrieval/utils.py index bab5929e..62b3be2b 100644 --- a/dataretrieval/utils.py +++ b/dataretrieval/utils.py @@ -86,11 +86,7 @@ def _require_positive_int( not only ``int``) but rejects ``bool`` — an ``Integral`` subtype that is nonsensical as a count. A non-integer (float, str, ``None``) or a value ``< 1`` raises before any I/O, rather than crashing later (e.g. deep in - ``pd.DataFrame.head``). Shared by the user-facing count knobs ``max_rows`` - and ``parallel_chunks(n)`` so their boundary validation can't drift. - (``ChunkPlan.max_chunks`` is an internal, already-``int`` precondition with - its own domain-specific message, so it keeps a lighter ``< 1`` guard rather - than routing through here.) + ``pd.DataFrame.head``). Used by the user-facing count knob ``max_rows``. Parameters ---------- diff --git a/dataretrieval/waterdata/__init__.py b/dataretrieval/waterdata/__init__.py index eb231469..99b6e178 100644 --- a/dataretrieval/waterdata/__init__.py +++ b/dataretrieval/waterdata/__init__.py @@ -9,7 +9,6 @@ from __future__ import annotations -from dataretrieval.ogc.chunking import parallel_chunks from dataretrieval.ogc.filters import FILTER_LANG # Public API exports @@ -51,7 +50,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 aff35a7a..f141004b 100644 --- a/dataretrieval/waterdata/utils.py +++ b/dataretrieval/waterdata/utils.py @@ -106,6 +106,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, ) # Iterable-shaped params that ``_get_args`` must NOT push through diff --git a/docs/source/architecture/decisions/0006-service-neutral-transport.rst b/docs/source/architecture/decisions/0006-service-neutral-transport.rst index f6d9fd52..8616e052 100644 --- a/docs/source/architecture/decisions/0006-service-neutral-transport.rst +++ b/docs/source/architecture/decisions/0006-service-neutral-transport.rst @@ -10,7 +10,7 @@ Context ------- Several service adapters need the same low-level capabilities: guarded HTTP -clients, cursor pagination, bounded retry, response aggregation, progress, and a +clients, page walking, bounded retry, response aggregation, progress, and a sync-over-async bridge. Locating those capabilities inside a protocol package would make unrelated services depend on protocol-specific implementation details -- Water Use previously imported its page walker and sync bridge from @@ -32,7 +32,9 @@ 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 page fetching for services that honor + ``offset``; - 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 @@ -54,6 +56,20 @@ 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, in separate modules, because both +are HTTP execution policy and neither is protocol-specific. Cursor pagination +(``transport.pagination``) is the standards-only walk: page ``N+1``'s URL is +discovered by parsing page ``N``, so it is inherently sequential and always +available. Offset-parallel fetching (``transport.offsets``) computes every +page's URL up front, so pages overlap; it is opt-in per service because +``offset`` is a server extension, not part of OGC API - Features Part 1. Which +strategy runs is a *dialect* decision — an adapter declares the offset ceiling +its API accepts — so the choice stays with the code that knows the service, not +with transport. Offsets also carry a ceiling where cursors do not, so the offset +walk takes a sequential continuation callback and hands the tail back to the +cursor walk; the two strategies compose rather than compete, and the fallback +path is the one that needs no extensions. + 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 @@ -84,6 +100,13 @@ 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 be made faster by declaring an offset ceiling rather than by + changing a page walker, and a service that later stops honoring ``offset`` + degrades to the sequential walk instead of returning duplicated rows. +- Parallelism belongs to the page walk, not the chunk planner: splitting a + request that already fits the byte budget would spend extra quota, whereas + overlapping the pages that request was going to fetch anyway spends none. The + planner therefore splits only as far as the URL byte limit forces. - 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, @@ -93,10 +116,12 @@ 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 means transport is roughly 1,200 + lines across six modules, each recognizably HTTP execution policy. Retry and + the offset walk are the two intricate modules, and each is intricate for a + bounded reason: retry because two independent bounds are what make it safe + against a slow service, the offset walk because a result's page count cannot + be known before it is fetched. Compliance ---------- @@ -108,3 +133,8 @@ frame-assembly modules do not reappear inside transport, and that only 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. +``tests/transport_test.py`` covers the offset walk's stop conditions in +isolation, and ``tests/waterdata_offset_paging_test.py`` pins the wiring through +a real getter: that offset requests preserve the query, that a server ignoring +``offset`` falls back to cursors with a de-duplicated result, and that the +offset ceiling hands off to the cursor walk with neither a gap nor an overlap. diff --git a/docs/source/userguide/errors.rst b/docs/source/userguide/errors.rst index 28da515f..40418e94 100644 --- a/docs/source/userguide/errors.rst +++ b/docs/source/userguide/errors.rst @@ -96,44 +96,44 @@ condition clears -- only the unfinished sub-requests are re-issued. except ChunkInterrupted as again: exc = again -Chunk a large request more finely -================================= +Large pulls are already paged in parallel +========================================= + +A large result is delivered one page at a time, and standard OGC pagination is +sequential: page *N+1*'s URL is only revealed by page *N*. The Water Data API +also honors an ``offset`` parameter, so every page's URL is computable up front +(``offset = i * limit``) and ``dataretrieval`` fetches the pages concurrently. +This is the default for every getter -- there is nothing to enable, and no +sub-request fan-out to configure. -By default the getters split an over-large request only as much as the -server's ~8 KB URL limit forces -- the fewest sub-requests. Because each -sub-request paginates, splitting a large result further costs little or no -extra quota *as long as each sub-request 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 sub-request only a page or two adds its partial -final page). So if you *know* your pull is large you can ask for a finer split -with ``parallel_chunks(n)`` -- trading roughly the same pages for more, smaller -sub-requests, which gives smoother progress, more even concurrency, and a -smaller unit of retry/resume. It is a scoped ``with`` -block, so an aggressive setting can't leak into unrelated calls and -accidentally spend quota: +Crucially the **request count is unchanged**: the same pages are fetched, just +overlapped. Your hourly rate limit counts requests, so the speedup costs no +quota. How many pages are in flight at once is bounded by +``API_USGS_CONCURRENT`` (default 32): .. code-block:: python - from dataretrieval import waterdata + import os - with waterdata.parallel_chunks(32): - df, md = waterdata.get_daily( - monitoring_location_id=many_sites, parameter_code="00060" - ) + os.environ["API_USGS_CONCURRENT"] = "1" # page strictly sequentially + +Two behaviors are worth knowing about, because ``offset`` is a Water Data +extension rather than part of OGC API - Features: + +* If a service silently ignores ``offset`` (the conventional handling of an + unrecognized query parameter), every offset would answer with page 1. The walk + detects that on the first wave -- before returning any rows -- and re-runs the + query using standard cursor pagination, logging a warning. The result is + correct either way; only the speed changes. +* The API rejects ``offset`` beyond 40000. That is a ceiling, not an + end-of-data signal, so the walk hands off to a sequential cursor walk for the + remainder. A pull deeper than 40000 rows is fast over its first 40000 and + complete over the rest. -``n`` is a positive integer (e.g. ``2``, ``8``, ``32``) -- the number of -sub-requests to fan the call out into; a non-integer or non-positive value -raises ``ValueError`` at the ``with``. It caps the *total* sub-request 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, so several multi-value arguments can't multiply past it and ``n=1`` asks -for no extra fan-out. Each sub-request costs a request against your hourly rate -limit, and because how many run *at once* is capped separately by -``API_USGS_CONCURRENT`` (default 32) 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: simply 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. +The chunker still splits an over-large request along its multi-value arguments +whenever the URL would exceed the server's ~8 KB limit; that is a correctness +requirement (see :class:`~dataretrieval.Unchunkable` below) and is unrelated to +paging speed. The full taxonomy ================= diff --git a/tests/architecture_test.py b/tests/architecture_test.py index c4d4c5b7..21ae0f00 100644 --- a/tests/architecture_test.py +++ b/tests/architecture_test.py @@ -25,7 +25,11 @@ "dataretrieval.ngwmn": {"dataretrieval.ogc"}, } -_ENGINE_REQUEST_IMPORTS = { +# Names the engine re-exports purely for backward compatibility. This set is +# frozen: it may shrink as legacy callers migrate to the canonical +# ``dataretrieval.ogc.requests``, but it must never grow — that would rebuild +# the request hub the module split dissolved. +_ENGINE_COMPAT_REQUEST_IMPORTS = { "_NO_NORMALIZE_PARAMS", "_as_str_list", "_check_monitoring_location_id", @@ -44,6 +48,17 @@ "prepare_request_args", } +# Names the engine actually *calls*. Unlike the compat set these are real +# dependencies of engine logic, so this set legitimately changes when engine +# behavior does; it is enumerated (rather than unconstrained) so a new entry is +# a reviewed decision instead of drift. +_ENGINE_LIVE_REQUEST_IMPORTS = { + "page_limit", # page size -> the offset walk's stride + "with_offset", # derive each page's request from the planned one +} + +_ENGINE_REQUEST_IMPORTS = _ENGINE_COMPAT_REQUEST_IMPORTS | _ENGINE_LIVE_REQUEST_IMPORTS + def _module_name(path: Path) -> str: """Return the import name for one Python file below ``PACKAGE_ROOT``.""" diff --git a/tests/transport_test.py b/tests/transport_test.py index dd5e9278..0f5c9c11 100644 --- a/tests/transport_test.py +++ b/tests/transport_test.py @@ -560,3 +560,209 @@ def test_deterministic_failures_are_not_offered_as_resumable() -> None: temporary = _wrapped_dns_failure(socket.EAI_AGAIN) assert retry._retryable(temporary) == (True, None) assert _classify_chunk_error(temporary) is not 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) diff --git a/tests/utils_test.py b/tests/utils_test.py index 2a743cb6..33c03b61 100644 --- a/tests/utils_test.py +++ b/tests/utils_test.py @@ -225,19 +225,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 06657821..f27bdb37 100644 --- a/tests/waterdata_chunking_test.py +++ b/tests/waterdata_chunking_test.py @@ -48,10 +48,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 ( @@ -1140,15 +1138,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) + # 202 bytes forces ``["A", "A"]`` (203 bytes) to split into two singleton + # chunks (201 bytes each), so both sub-requests select the same feature. + @multi_value_chunked(build_request=_fake_build, url_limit=202) 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"}] @@ -2142,293 +2141,115 @@ 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* sub-request count (the cartesian product across axes), not each axis -# independently — see ``test_cap_caps_the_total_across_axes``. +# Page concurrency: ``parallel_chunks(n)`` and ``ChunkPlan.max_chunks`` are gone. +# Splitting a *fitting* request to buy parallelism was the only thing they did, +# and the offset-parallel page walk +# (:mod:`dataretrieval.transport.offsets`) now overlaps that one request's pages +# instead — same request count, no extra quota, and it works on single-value +# queries the old dial could not split at all. What remains here pins the +# invariants that replaced it: byte-driven chunking is untouched, a fitting +# request stays a passthrough at any concurrency, and the wave width comes from +# the existing ``API_USGS_CONCURRENT`` dial rather than a second knob. # --------------------------------------------------------------------------- -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.""" - 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_sub_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_sub_args`` yields the original args verbatim) — identical to the - default (off), not a materialized one-chunk-per-axis plan.""" +def test_fitting_request_is_always_a_passthrough(): + """A multi-value request that fits the byte limit is the trivial passthrough + (no axes, ``total == 1``, original args verbatim). There is no longer any + dial that splits it further — parallelism comes from paging that single + request by offset, so the planner's only job is the byte budget.""" 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_sub_args()) == [args] -@pytest.mark.parametrize("bad", [0, -1]) -def test_invalid_cap_raises(bad): - """``max_chunks`` is a sub-request 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(): + """``max_chunks`` was the parallelism dial's entry point into planning. It + is gone, so passing it is a ``TypeError`` — not silently ignored, which + would let stale callers believe they were still fanning out.""" 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) - - -@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 sub-requests. 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] - 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 sub-request (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 sub-request 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 + with pytest.raises(TypeError): + ChunkPlan(args, _fake_build, url_limit=8000, max_chunks=8) + + +def test_byte_driven_chunking_survives_the_removal(): + """The half of chunking that is a *correctness* requirement is untouched: + an over-budget request still fans out until every sub-request fits, and the + partition still covers every atom exactly once.""" + 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 sub in plan.iter_sub_args(): assert _safe_request_bytes(_fake_build, sub, limit) <= limit + flattened = [a for chunk in plan.chunks["monitoring_location_id"] for a in chunk] + assert sorted(flattened) == sorted(atoms) -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_unchunkable_still_raised_without_the_dial(): + """A request with nothing to split that busts the byte limit still raises + ``Unchunkable`` rather than shipping for an opaque HTTP 414.""" + args = {"monitoring_location_id": "one-huge-scalar"} + with pytest.raises(Unchunkable): + ChunkPlan(args, _fake_build, url_limit=10) -def test_cap_caps_the_total_across_axes(): - """With more than one multi-value axis the cap bounds the *total* - sub-request 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 sub-requests 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 sub-requests 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 +def test_page_concurrency_defaults_and_follows_the_env(monkeypatch): + """The wave width reuses ``API_USGS_CONCURRENT`` instead of introducing a + second dial: both bound how many requests are in flight, and the quota they + spend is volume-based, not simultaneity-based. ``1`` means strictly + sequential paging.""" + monkeypatch.delenv("API_USGS_CONCURRENT", raising=False) + assert _chunking.page_concurrency() == _chunking._CONCURRENCY_DEFAULT -@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.""" - args = {"monitoring_location_id": "one-huge-scalar"} - with pytest.raises(Unchunkable): - ChunkPlan(args, _fake_build, url_limit=10, max_chunks=32) + monkeypatch.setenv("API_USGS_CONCURRENT", "4") + assert _chunking.page_concurrency() == 4 + monkeypatch.setenv("API_USGS_CONCURRENT", "1") + assert _chunking.page_concurrency() == 1 -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_clamps_unbounded(monkeypatch): + """``unbounded`` disables the *sub-request* cap, but a wave is speculative — + an unbounded wave would issue arbitrarily many past-the-end requests to + discover one short page — so the wave width is clamped to a finite default.""" + monkeypatch.setenv("API_USGS_CONCURRENT", "unbounded") + assert _chunking._read_concurrency_env() is None + assert _chunking.page_concurrency() == _chunking._CONCURRENCY_DEFAULT -@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 +def test_connection_pool_covers_subrequests_times_pages(monkeypatch, httpx_mock): + """The pool must be sized for the *product* of the two fan-outs. -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`` sub-requests 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)] + Sub-requests are gated by a semaphore, but the pages inside each one are + not (they can't be — a sub-request holds its permit for the whole attempt, + so its own pages would deadlock waiting on it). Peak in-flight is therefore + ``API_USGS_CONCURRENT * page_concurrency()``. Sizing the pool to the + sub-request cap alone left the excess queued inside httpx against the 60 s + pool-acquire timeout, which surfaced as a mid-walk ``ReadError`` and a + spurious resumable ``ServiceInterrupted``. + """ + monkeypatch.setenv("API_USGS_CONCURRENT", "8") + captured: dict[str, object] = {} + real_open = _chunking.open_async_client - calls: list[tuple[str, ...]] = [] + def spy(**overrides): + captured["limits"] = overrides.get("limits") + return real_open(**overrides) - @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() + 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 sub-requests. - 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`` sub-requests, 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.setattr(_chunking, "open_async_client", spy) + httpx_mock.add_response( + json={"type": "FeatureCollection", "numberReturned": 0, "features": []}, + headers={"Content-Type": "application/geo+json"}, + ) - @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() + 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 limits is not None, "the chunker must configure explicit pool 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))