diff --git a/NEWS.md b/NEWS.md index a9a63718..536f5bed 100644 --- a/NEWS.md +++ b/NEWS.md @@ -1,3 +1,5 @@ +**08/06/2026:** Fan-out execution is now shared across services. Chunking is how a query is divided structurally (a Water Data/NGWMN URL over the byte limit); fan-out is how the pieces are distributed operationally. Only the first is protocol-specific, so the executor moved to `dataretrieval.transport.fanout` (`FanOut`, over a three-member `FanOutPlan` protocol) while chunk planning stays in `dataretrieval.ogc`. Water Use no longer re-implements the fan-out gather and inherits resume, progress reporting, and `API_USGS_CONCURRENT`: a multi-location pull interrupted by a rate limit now raises a resumable interruption whose `.call.resume()` re-issues only the locations that did not finish, instead of discarding every completed one. The interruption taxonomy moved to the `dataretrieval.interruptions` leaf and its base class is now `FanOutInterrupted`; **`ChunkInterrupted` is a permanent alias of the same class**, so `except ChunkInterrupted` keeps working. **Breaking change:** a Water Use fan-out interrupted by a 5xx or 429 now raises `ServiceInterrupted`/`QuotaExhausted` rather than `ServiceUnavailable`/`RateLimited` — both are still `DataRetrievalError`, so broad handlers are unaffected, but a narrow `except ServiceUnavailable` around a Water Use call must widen. **Breaking change:** `wateruse.MAX_CONCURRENT_REQUESTS` is removed; set `API_USGS_CONCURRENT` (which now outranks any service default) or read `wateruse.DEFAULT_CONCURRENT_REQUESTS`. + **08/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/dataretrieval/__init__.py b/dataretrieval/__init__.py index 4226e247..edd964d1 100644 --- a/dataretrieval/__init__.py +++ b/dataretrieval/__init__.py @@ -45,23 +45,24 @@ 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, +# Resumable fan-out interruption exceptions. They are defined in +# ``dataretrieval.interruptions`` rather than ``dataretrieval.exceptions`` +# because they carry pandas/httpx state and a resumable ``FanOut`` handle, # which would pull heavy dependencies into the lightweight exceptions module. -# Surfaced here so callers get a stable public path: -# ``from dataretrieval import ChunkInterrupted``. -from dataretrieval.ogc.interruptions import ( +# They are not under ``ogc`` because Water Use raises them too. Surfaced here so +# callers get a stable public path: ``from dataretrieval import ChunkInterrupted``. +from dataretrieval.interruptions import ( ChunkInterrupted, + FanOutInterrupted, QuotaExhausted, ServiceInterrupted, ) +# Parallel-chunks control (a context manager). Defined with the chunker in +# ``dataretrieval.ogc.chunking``; surfaced here for a stable public path +# ``from dataretrieval import parallel_chunks``. +from dataretrieval.ogc.chunking import parallel_chunks + from . import ( exceptions, ngwmn, @@ -96,8 +97,9 @@ "TransientError", "URLTooLong", "Unchunkable", - # resumable chunk-interruption exceptions (defined in ogc.interruptions) + # resumable fan-out interruption exceptions (defined in interruptions) "ChunkInterrupted", + "FanOutInterrupted", "QuotaExhausted", "ServiceInterrupted", # parallel-chunks control (defined in ogc.chunking) diff --git a/dataretrieval/interruptions.py b/dataretrieval/interruptions.py new file mode 100644 index 00000000..424ac89b --- /dev/null +++ b/dataretrieval/interruptions.py @@ -0,0 +1,299 @@ +"""Resumable fan-out interruption exceptions — the public resume contract. + +When a fanned-out request fails mid-stream (a 429, a 5xx, or a bare transport +error), the work already completed is preserved and the call is resumable: the +raised exception carries a ``.call`` handle whose ``resume()`` re-issues only +the still-pending sub-requests. These exception types are that contract, +re-exported at the top level (``from dataretrieval import ChunkInterrupted``). +The execution machinery that raises and resumes them is +:class:`dataretrieval.transport.fanout.FanOut`. + +Vocabulary, consistently: a **fan-out** is one logical query the service forces +into several requests; a **sub-request** is one unit of a fan-out; a **chunk** +is specifically a *byte-driven* slice, which is OGC planning vocabulary and +belongs to :class:`~dataretrieval.ogc.planning.ChunkPlan`. Water Use fans out +without chunking anything — the NWDC simply accepts one location per request — +so the base class is :class:`FanOutInterrupted`. + +``ChunkInterrupted`` is retained as an alias of that same class, not a +deprecated shim to delete later: it is the name published in the user guide and +caught in user code, and aliasing costs nothing to keep. ``except +ChunkInterrupted`` and ``except FanOutInterrupted`` are the same handler. + +This is a top-level leaf rather than a member of ``ogc`` or ``transport``, +for the reason ADR 0006 gives for ``combining``, ``progress``, and +``credentials``: adapters need it whether or not they go through transport, and +an exception taxonomy is not HTTP execution policy. It stays out of +:mod:`dataretrieval.exceptions` because it carries pandas/httpx state, which +would pull heavy dependencies into that lightweight leaf. +""" + +from __future__ import annotations + +import socket +from typing import TYPE_CHECKING, Any, ClassVar + +import httpx +import pandas as pd + +from dataretrieval.exceptions import DataRetrievalError, RateLimited, TransientError + +if TYPE_CHECKING: + from dataretrieval.transport.fanout import FanOut + + +class FanOutInterrupted(DataRetrievalError): + """ + Base class for mid-stream sub-request failures whose completed work + is preserved and resumable. + + A ``FanOutInterrupted`` subclass means: a sub-request failed, but + ``FanOut`` still owns whatever completed successfully before + the failure. Call ``self.call.resume()`` to pick up where the + failure stopped you — only still-pending sub-requests are + re-issued. + + Subclasses describe *why* ``FanOut`` stopped so callers can + pick a retry policy: :class:`QuotaExhausted` for 429 (wait for the + rate-limit window), :class:`ServiceInterrupted` for 5xx (wait for + the upstream to recover). The ``.call`` handle is the same object + across every interruption of a single fanned-out call — frames + accumulate across retries. + + Attributes + ---------- + call : FanOut or None + Resumable handle into the ``FanOut`` that raised this + exception. ``None`` only on hand-constructed exceptions (test + fixtures), where ``.call``-derived accessors degrade to + empty/``None``. + retry_after : float or None + Seconds the server suggested waiting (``Retry-After`` header). + ``None`` when the server gave no hint. + completed_chunks : int + Number of sub-requests successfully completed before the failure. + total_chunks : int + Total sub-requests in the plan. + partial_frame : pandas.DataFrame + Combined frame of work completed by the moment this exception + was raised. Snapshot at raise time — does NOT advance on a + later ``call.resume()`` (use ``exc.call.partial_frame`` for + the live view). + partial_response : httpx.Response or None + Raw aggregate response covering the completed sub-requests at + raise time; ``None`` if nothing had completed yet. Same snapshot + semantics as ``partial_frame``. (Raw, not finalized — use + ``exc.call.resume()`` for the finalized ``(df, metadata)`` result.) + + Examples + -------- + Retry on any transient interruption, honoring the server's + ``Retry-After`` hint when present and falling back to a fixed wait + otherwise. Each new interruption keeps the already-completed work + intact — only the still-pending sub-requests are re-issued. + + .. code-block:: python + + import time + from dataretrieval import ChunkInterrupted + + # ``getter`` is any chunked OGC getter — e.g. + # ``waterdata.get_daily`` or ``ngwmn.get_water_level``. + try: + df, md = getter(monitoring_location_id=long_list_of_sites) + except ChunkInterrupted as exc: + while True: + time.sleep(exc.retry_after or 5 * 60) + try: + df, md = exc.call.resume() + break + except ChunkInterrupted as next_exc: + exc = next_exc + """ + + # Subclasses override with a ``str.format`` template; the format + # call sees ``completed_chunks`` and ``total_chunks`` as kwargs. + _MESSAGE_TEMPLATE: ClassVar[str] = ( + "Chunked request interrupted after {completed_chunks}/" + "{total_chunks} sub-requests; call .call.resume() to continue." + ) + + def __init__( + self, + *, + completed_chunks: int, + total_chunks: int, + call: FanOut | None = None, + retry_after: float | None = None, + cause: BaseException | None = None, + ) -> None: + message = self._MESSAGE_TEMPLATE.format( + completed_chunks=completed_chunks, total_chunks=total_chunks + ) + if cause is not None: + cause_msg = str(cause) or type(cause).__name__ + message = f"{message} Cause: {type(cause).__name__}: {cause_msg}" + super().__init__(message) + self.completed_chunks = completed_chunks + self.total_chunks = total_chunks + self.call = call + self.retry_after = retry_after + # Snapshot partial state at raise time so the exception stays a stable + # record of the failure moment: ``exc.partial_frame`` / + # ``.partial_response`` do NOT advance on a later ``call.resume()`` + # (that live view is on ``call.partial_frame`` / ``.partial_response``). + # This keeps each interruption in a resume loop a faithful record of + # what it saw, rather than every exception aliasing the shared call's + # advancing state. ``.copy()`` guards the single-chunk fast path, where + # the combined frame may be returned verbatim. + if call is None: + self.partial_frame: pd.DataFrame = pd.DataFrame() + self.partial_response: httpx.Response | None = None + else: + self.partial_frame = call.partial_frame.copy() + self.partial_response = call.partial_response + + def __getstate__(self) -> dict[str, Any]: + # Drop the live FanOut before pickling: its ``.fetch`` is an + # undecorated module function pickle can't reference by name, so the + # interruption can't cross a process boundary with ``.call`` attached. + # The degraded ``call=None`` form keeps the counts, retry hint, and the + # snapshotted partial frame / response — plain instance attributes the + # base ``__getstate__`` already pickles; only ``.resume()`` is lost + # (cross-process resume was never possible anyway). + return {**super().__getstate__(), "call": None} + + +class QuotaExhausted(FanOutInterrupted): + """ + A sub-request returned HTTP 429 — the per-key rate-limit window + is exhausted. Subclass of :class:`FanOutInterrupted`. + + The completed sub-requests are preserved on ``.call``; once the + rate-limit window resets, ``.call.resume()`` re-issues only the + still-pending work. ``partial_frame`` holds what completed + before the 429. + """ + + _MESSAGE_TEMPLATE = ( + "HTTP 429 after {completed_chunks}/{total_chunks} sub-requests; " + "catch QuotaExhausted (or FanOutInterrupted) to access " + ".partial_frame or .call.resume() once the rate-limit " + "window has rolled over." + ) + + +class ServiceInterrupted(FanOutInterrupted): + """ + A sub-request returned HTTP 5xx — the upstream service failed + transiently. Subclass of :class:`FanOutInterrupted`. + + The completed sub-requests are preserved on ``.call``; once the + upstream recovers, ``.call.resume()`` resumes only the + still-pending work. + """ + + _MESSAGE_TEMPLATE = ( + "Service error after {completed_chunks}/{total_chunks} " + "sub-requests; catch ServiceInterrupted (or FanOutInterrupted) " + "and call .call.resume() once the upstream service recovers." + ) + + +# Resolver failures that will not resolve differently on a later attempt. The +# temporary ones (notably EAI_AGAIN -- "try again", raised while a resolver is +# still coming up, on VPN reconnect, or after a laptop wakes) are deliberately +# absent: those are worth another try. Looked up defensively because the EAI_* +# constants are platform-dependent; an unrecognized code stays retryable, since +# spending a few seconds on a retry is cheaper than dropping a recoverable call. +_PERMANENT_DNS_ERRORS = frozenset( + code + for code in ( + getattr(socket, name, None) for name in ("EAI_NONAME", "EAI_FAIL", "EAI_NODATA") + ) + if code is not None +) + + +def _deterministic_failure(exc: BaseException) -> bool: + """Whether a transport failure would fail identically on every retry. + + An unsupported scheme or a request we built wrong is settled before a byte + goes out, and a hostname the resolver rejects outright won't be accepted on + the next attempt either -- so retrying only delays the error the caller + needs. A *temporary* resolver failure is not in that class and stays + retryable (see :data:`_PERMANENT_DNS_ERRORS`). + + The original failure is several layers down and not always an explicit + ``raise ... from``: a DNS failure reaches us as ``NetworkError`` -> + ``httpx.ConnectError`` -> ``httpcore.ConnectError`` -> ``socket.gaierror``, + linked by ``__context__`` (implicit chaining) rather than ``__cause__``. + + Both links of every frame are visited, not just the first one present. A + frame can carry an explicit ``__cause__`` *and* an unrelated ``__context__`` + (any ``raise X from Y`` inside an ``except`` block produces exactly that), so + following only the cause would walk off down the explicit branch and miss a + ``gaierror`` sitting on the implicit one -- spending the whole retry budget + on a hostname that will never resolve. The ``seen`` set keeps a chain that + rejoins itself, or points back at an ancestor, from looping. + """ + seen: set[int] = set() + pending: list[BaseException | None] = [exc] + while pending: + current = pending.pop() + if current is None or id(current) in seen: + continue + seen.add(id(current)) + if isinstance(current, (httpx.UnsupportedProtocol, httpx.LocalProtocolError)): + return True + if isinstance(current, socket.gaierror): + # Return, not continue: the first resolver code found settles the chain. + return current.errno in _PERMANENT_DNS_ERRORS + pending += [current.__cause__, current.__context__] + return False + + +def _classify_transient( + exc: BaseException, +) -> tuple[type[FanOutInterrupted], float | None] | None: + """Classify one failure as a resumable interruption.""" + if isinstance(exc, RateLimited): + return QuotaExhausted, exc.retry_after + if isinstance(exc, TransientError): + return ServiceInterrupted, exc.retry_after + if isinstance(exc, (httpx.HTTPError, httpx.InvalidURL)): + # Some failures will fail the same way every time -- a bad scheme, a + # hostname that doesn't resolve. Offering to resume one would just + # hide the real error behind a retry that can never work. + if _deterministic_failure(exc): + return None + return ServiceInterrupted, None + return None + + +def _classify_chunk_error( + exc: BaseException, +) -> tuple[type[FanOutInterrupted], float | None] | None: + """Walk a wrapped pagination failure for a resumable transport cause.""" + current: BaseException | None = exc + while current is not None: + result = _classify_transient(current) + if result is not None: + return result + current = current.__cause__ + return None + + +#: The name this taxonomy was published under, kept as a permanent alias so +#: ``except ChunkInterrupted`` keeps working. Same class object, not a subclass. +ChunkInterrupted = FanOutInterrupted + +__all__ = [ + "ChunkInterrupted", + "_deterministic_failure", + "FanOutInterrupted", + "QuotaExhausted", + "ServiceInterrupted", + "_classify_chunk_error", + "_classify_transient", +] diff --git a/dataretrieval/ogc/chunking.py b/dataretrieval/ogc/chunking.py index f15f226b..d2090afb 100644 --- a/dataretrieval/ogc/chunking.py +++ b/dataretrieval/ogc/chunking.py @@ -1,13 +1,21 @@ -"""Joint URL-byte chunking for the OGC getters. +"""URL-byte chunk planning and dispatch for the OGC getters. An OGC query has several chunkable axes: every multi-value list parameter (sites, parameter codes, …) plus the cql-text ``filter``, which splits along its top-level OR clauses. Any of them can fan the URL past the server's ~8 KB byte limit. ``ChunkPlan`` picks a fan-out for each axis that minimizes total sub-requests while keeping every -sub-request URL under the budget; ``ChunkedCall`` fetches the resulting -cartesian product of chunks. Requests that already fit get a trivial -single-step plan — ``ChunkedCall`` has one code path either way. +sub-request URL under the budget. Requests that already fit get a +trivial single-step plan — the executor has one code path either way. + +This module owns the OGC-specific half: the byte budget, the +``parallel_chunks`` dial, and the ``multi_value_chunked`` decorator that +ties a plan to a fetcher. Driving the resulting sub-requests to +completion — bounded concurrency, retry, failure precedence, resume — is +API-neutral and belongs to +:class:`dataretrieval.transport.fanout.FanOut`, which this module hands +its plan to. :class:`~dataretrieval.ogc.planning.ChunkPlan` satisfies +:class:`~dataretrieval.transport.fanout.FanOutPlan` structurally. Parallel chunks: the planner is conservative by default — it splits only as far as the byte limit forces. A caller who knows their result is large can opt into a @@ -15,50 +23,9 @@ out into ``n`` parallel sub-requests; ``n`` drives :meth:`ChunkPlan._refine`. See ``parallel_chunks`` for the why and the when. -This module owns the *execution* half — the event loop and bounded -concurrency that drive a plan to completion (``ChunkedCall``) plus the -public ``multi_value_chunked`` decorator. The neighboring concerns remain -separate: :mod:`~dataretrieval.ogc.planning` builds the -:class:`~dataretrieval.ogc.planning.ChunkPlan`; -:mod:`~dataretrieval.combining` assembles results; -:mod:`~dataretrieval.transport.retry` owns bounded retry policy; and -:mod:`~dataretrieval.ogc.interruptions` defines the resumable -:class:`~dataretrieval.ogc.interruptions.ChunkInterrupted` contract. - -Concurrency: ``multi_value_chunked`` fans every pending sub-request out -under one ``asyncio.gather`` sharing a single ``httpx.AsyncClient``. An -``asyncio.Semaphore`` — not the client's connection pool, which is -merely sized to match — caps the sub-requests in flight at ``N``; see -:meth:`ChunkedCall._run` for why the gate must be the semaphore rather -than the pool. ``API_USGS_CONCURRENT`` resolves ``N``: an integer N > 1 -allows N sub-requests in flight; ``1`` forces sequential dispatch (one -request at a time); the literal ``unbounded`` lifts the cap. ``N`` -bounds only how many of a chunked query's sub-requests are in flight at -once — a client-side trade-off between open connections and fan-out -latency. It does not affect the API rate limit: a chunked call issues -the same number of sub-requests regardless of ``N``, so ``N`` changes -their timing, not the total request volume. The USGS API rate-limits by -volume over time (HTTP 429), not by simultaneity; set ``API_USGS_PAT`` -to raise that quota. The default of 32 is a conservative cap that keeps -connection use modest. The fan-out runs in a short-lived worker thread -(an ``anyio`` blocking portal), so it works whether or not the caller is -already inside an event loop (Jupyter / IPython / async apps). - -Retries: each sub-request is retried on a transient failure (429, -5xx, connect/read timeout) with exponential backoff + full jitter, -honoring a server ``Retry-After`` when present. ``API_USGS_RETRIES`` -sets the cap (default 4; ``0`` disables). A ``Retry-After`` longer -than the per-call ceiling escalates to a resumable interruption. - -Interruption: any mid-stream transient failure — 429, 5xx, or a bare -transport error (connect/read timeout, oversize follow-up URL) — surfaces -as a ``ChunkInterrupted`` subclass: ``QuotaExhausted`` for 429, -``ServiceInterrupted`` for the rest. The exception carries ``.call``, a -``ChunkedCall`` handle that owns the already-completed sub-request -state (sparse-indexed, since gathered sub-requests complete out of -order). Call ``.call.resume()`` once the underlying condition clears; -only the still-pending sub-requests are re-issued. ``Retry-After`` (when -the server sets it) is surfaced on the exception as ``.retry_after``. +Concurrency, retries, and interruption semantics are documented on +:mod:`dataretrieval.transport.fanout`; ``API_USGS_CONCURRENT`` and +``API_USGS_RETRIES`` are read there. Dedup: list-axis chunks don't overlap; filter-axis chunks can, so ``_combine_chunk_frames`` dedupes by feature ``id``. ``properties``, @@ -69,32 +36,37 @@ from __future__ import annotations -import asyncio import functools -import os -from collections.abc import Awaitable, Callable, Iterator +from collections.abc import Callable, Iterator from contextlib import contextmanager -from contextvars import copy_context -from typing import Any, cast +from typing import Any import httpx import pandas as pd -from anyio.from_thread import start_blocking_portal -from dataretrieval import progress as _progress -from dataretrieval.combining import ( - _combine_chunk_frames, - _combine_chunk_responses, +from dataretrieval.transport.fanout import ( + FanOut, + _active_client, + _Fetch, + _Finalize, + _passthrough_result, + active_client, ) -from dataretrieval.exceptions import ConfigurationError -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.transport.retry import RetryPolicy from dataretrieval.utils import Ambient, _require_positive_int -from .interruptions import ChunkInterrupted from .planning import ChunkPlan -from .retry import _classify_chunk_error + +# Compatibility aliases. ``ChunkedCall`` was this module's executor before it +# moved down to transport as the API-neutral ``FanOut``; ``get_active_client`` +# and ``_chunked_client`` named its shared per-call client. Existing imports -- +# ``ogc.engine`` and the chunking/progress test modules -- still use these +# names, and the rename is not worth churning them over. They are aliases, not +# copies: the ambient in particular must be the *same* object transport +# publishes, or a test reading it here would never see the running client. +ChunkedCall = FanOut +get_active_client = active_client +_chunked_client = _active_client # Empirically the API replies HTTP 414 above ~8200 bytes of full URL — # matches nginx's default ``large_client_header_buffers`` of 8 KB. 8000 @@ -104,73 +76,6 @@ _OGC_URL_BYTE_LIMIT = 8000 -# Fan-out concurrency cap, read at call time (not import) so test -# ``monkeypatch.setenv`` applies. Value grammar in :func:`_read_concurrency_env`; -# the concurrency model is in the module docstring. -_CONCURRENCY_ENV = "API_USGS_CONCURRENT" -_CONCURRENCY_DEFAULT = 32 -_CONCURRENCY_UNBOUNDED = "unbounded" - - -def _read_concurrency_env() -> int | None: - """ - Resolve the ``API_USGS_CONCURRENT`` env var to a parallelism cap. - - Returns - ------- - int or None - ``1`` for sequential dispatch (one sub-request at a time); an - integer >1 for bounded concurrency; ``None`` to disable the - per-call cap entirely (``unbounded`` keyword). Unset → default - of ``_CONCURRENCY_DEFAULT``. - """ - raw = os.environ.get(_CONCURRENCY_ENV) - if raw is None: - return _CONCURRENCY_DEFAULT - raw = raw.strip() - if raw == "": - return _CONCURRENCY_DEFAULT - if raw.lower() == _CONCURRENCY_UNBOUNDED: - return None - try: - value = int(raw) - except ValueError as exc: - raise ConfigurationError( - f"{_CONCURRENCY_ENV} must be a positive integer or " - f"'{_CONCURRENCY_UNBOUNDED}'; got {raw!r}." - ) from exc - if value < 1: - raise ConfigurationError( - f"{_CONCURRENCY_ENV} must be >= 1 (got {value}); use " - f"'{_CONCURRENCY_UNBOUNDED}' to disable the cap." - ) - return value - - -# Shared per-call ``httpx.AsyncClient``, scoped via ``with _chunked_client(c):`` -# during ``ChunkedCall._run`` so paginated-loop helpers (``_walk_pages``) reuse -# the same connection pool across every sub-request. ``None`` outside a chunked -# call — paginated helpers then open their own short-lived client. -_chunked_client: Ambient[httpx.AsyncClient | None] = Ambient("_chunked_client", None) - - -def get_active_client() -> httpx.AsyncClient | None: - """ - Return the chunker's currently-published client, or ``None``. - - Used by the paginated-loop helpers (e.g. - :func:`dataretrieval.ogc.engine._client_for`) to reuse the - per-call connection pool. - - Returns - ------- - httpx.AsyncClient or None - The client scoped via ``with _chunked_client(...)`` if currently inside - a :class:`ChunkedCall` run; ``None`` otherwise. - """ - 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 @@ -283,425 +188,6 @@ def parallel_chunks(n: int) -> Iterator[None]: yield -# --------------------------------------------------------------------------- -# Type aliases for the ChunkedCall contract. -# --------------------------------------------------------------------------- - -# The per-sub-request fetcher the decorator wraps and ``ChunkedCall`` drives: -# an ``async def fetch(args) -> (df, response)``. -_Fetch = Callable[[dict[str, Any]], Awaitable[tuple[pd.DataFrame, httpx.Response]]] - -# Caller-supplied transform applied to the combined chunk result, so a -# resumed call returns the same shape as an un-interrupted one rather than -# the chunker's raw ``(frame, httpx.Response)``. This keeps the chunker -# generic: the OGC getters inject their post-processing (type coercion, -# column arrangement, ``BaseMetadata``) through ``_finalize_ogc``. -# The default is identity, so direct ``ChunkedCall`` use is unaffected. -_Finalize = Callable[[pd.DataFrame, httpx.Response], tuple[pd.DataFrame, Any]] - - -def _passthrough_result( - frame: pd.DataFrame, response: httpx.Response -) -> tuple[pd.DataFrame, Any]: - """Default :data:`_Finalize`: return the raw combined pair unchanged.""" - return frame, response - - -class ChunkedCall: - """ - Stateful handle for a chunked call. - - Holds the in-flight state (per-sub-request frames and responses) - and the async fetcher. A single :meth:`resume` entry point drives - the call from wherever it is to completion — used both for the - first invocation (from :meth:`ChunkPlan.execute`) and for subsequent - retries after a :class:`ChunkInterrupted`. - - :meth:`_run` gathers every pending sub-request over one shared - :class:`httpx.AsyncClient`, applies the failure-precedence rules, and - combines; :meth:`resume` drives it through an ``anyio`` blocking - portal so it works whether or not the caller is already inside an - event loop. Concurrency is bounded by a per-run ``asyncio.Semaphore`` - (see :meth:`_run`), so sequential dispatch - (``API_USGS_CONCURRENT=1``) is just a degenerate gather. - - A ``ChunkedCall`` is created internally when a :class:`ChunkPlan` - executes; callers reach it via :attr:`ChunkInterrupted.call` on - the exception raised by a mid-stream failure. - - :meth:`resume` is idempotent: :meth:`_run` iterates - :meth:`ChunkPlan.iter_sub_args` (deterministic order) and skips - any index whose result is already in ``self._chunks``. The - completion set is a sparse ``dict[int, (df, response)]`` so the - gather can record scattered completions (e.g. indices [0, 2, 5] - after siblings [1, 3, 4] failed) and a subsequent ``resume`` only - re-issues the missing indices. - - Parameters - ---------- - plan : ChunkPlan - The chunking plan to execute. - fetch : Callable - ``async def`` that issues a single sub-request, given the - substituted args dict, and returns ``(frame, response)``. - - Attributes - ---------- - plan : ChunkPlan - The plan being driven (read-only after construction). - fetch : Callable - The async per-sub-request fetch function. - finalize : Callable - Transform applied to the combined result (see :data:`_Finalize`) at - the terminal :meth:`_run` return, so a completed call yields the - caller's finished shape. The ``partial_*`` accessors deliberately - skip it and stay raw. - partial_frame : pandas.DataFrame - Raw combined frame of completed sub-requests (live; recomputed per - access). Not finalized — call :meth:`resume` for the finished shape. - partial_response : httpx.Response or None - Raw aggregate response (canonical URL restored), or ``None`` when - nothing has completed yet (live; recomputed per access). - """ - - def __init__( - self, - plan: ChunkPlan, - fetch: _Fetch, - retry_policy: RetryPolicy = _NO_RETRY, - finalize: _Finalize = _passthrough_result, - ) -> None: - self.plan = plan - self.fetch = fetch - self.retry_policy = retry_policy - self.finalize = finalize - # Snapshot the ambient context at construction time — i.e. inside the - # caller's ``with`` blocks (base URL, dialect, row cap, progress - # reporter). :meth:`resume` runs every drive inside this snapshot, so - # a *later* ``exc.call.resume()`` — which fires after those ``with`` - # blocks have exited and reset their ContextVars — still rebuilds - # sub-requests against the original API's base URL/dialect rather than - # the process defaults. ``build_request`` reads those ContextVars when - # it reconstructs each sub-request, so the snapshot must outlive them. - self._ctx = copy_context() - # Completed (frame, response) pairs keyed by sub-args index; sparse - # (gathered sub-requests complete out of order — see class docstring). - # ``_run``'s ``track`` closure is the only writer, so ``dict`` insertion - # order is completion order (relied on by :meth:`_combine_raw`). - self._chunks: dict[int, tuple[pd.DataFrame, httpx.Response]] = {} - - def wrap_failure(self, exc: BaseException) -> ChunkInterrupted | None: - """ - Build the matching :class:`ChunkInterrupted` carrying this - call when ``exc`` is a recognized transient transport failure; - return ``None`` for unrecognized failures so the caller can - re-raise. Encapsulates the - ``classify → instantiate-with-call-state`` recipe so - :class:`ChunkedCall`'s private fields stay private. - - Parameters - ---------- - exc : BaseException - The exception raised by a sub-request. - - Returns - ------- - ChunkInterrupted or None - The matching :class:`ChunkInterrupted` subclass carrying this - call for a recognized transient failure; ``None`` otherwise. - """ - classification = _classify_chunk_error(exc) - if classification is None: - return None - interrupted_class, retry_after = classification - return interrupted_class( - completed_chunks=self.completed_chunks, - total_chunks=self.plan.total, - call=self, - retry_after=retry_after, - cause=exc, - ) - - @property - def completed_chunks(self) -> int: - """Number of sub-requests completed so far.""" - return len(self._chunks) - - def _combine_raw(self) -> tuple[pd.DataFrame, httpx.Response]: - """Assemble the raw ``(frame, response)`` from completed sub-requests, - before :attr:`finalize` runs. - - Frames concatenate in sub-args *index* order (``sorted`` keys — - deterministic, independent of parallel completion order). The - aggregated response takes its headers from the response with the - lowest reported ``x-ratelimit-remaining`` value. If no response - reports that header, it falls back to the last completed response; - ``self._chunks`` preserves completion order because the ``track`` - closure in :meth:`_run` is its only writer. - - Returns - ------- - tuple of (pandas.DataFrame, httpx.Response) - The concatenated frame and the aggregated response, before - :attr:`finalize` is applied. - """ - frames = [self._chunks[i][0] for i in sorted(self._chunks)] - responses = [response for _, response in self._chunks.values()] - return ( - _combine_chunk_frames(frames), - _combine_chunk_responses(responses, self.plan.canonical_url), - ) - - @property - def partial_frame(self) -> pd.DataFrame: - """ - Raw combined frame of sub-requests that have completed so far. - - Live — recomputed on each access so it reflects current state - across resume attempts. Deliberately the *raw* combined frame - (``_combine_raw``), NOT the finalized result: this is a cheap, - side-effect-free snapshot for inspecting partial progress, so - reading it (or building a :class:`ChunkInterrupted` around it) - never triggers ``finalize`` work — which for OGC getters includes - a schema network fetch on an empty frame. Use ``call.resume()`` - for the finalized result. - - Returns - ------- - pandas.DataFrame - Combined frame of completed sub-requests, or an empty - ``DataFrame`` when nothing has completed. - """ - if not self._chunks: - return pd.DataFrame() - return self._combine_raw()[0] - - @property - def partial_response(self) -> httpx.Response | None: - """ - Raw aggregate response with the canonical URL restored to the - user's full original query. - - Live — recomputed on each access. Like :attr:`partial_frame`, this - is the *raw* aggregate (an :class:`httpx.Response`), not the - finalized result, so inspecting it is side-effect-free. - - Returns - ------- - httpx.Response or None - Aggregated response when at least one sub-request has - completed, ``None`` otherwise. - """ - if not self._chunks: - return None - return self._combine_raw()[1] - - def _pending(self) -> Iterator[tuple[int, dict[str, Any]]]: - """ - Yield ``(index, sub_args)`` for sub-requests not yet completed. - - Walks :meth:`ChunkPlan.iter_sub_args` in deterministic order - and skips any index already in ``self._chunks``. :meth:`_run` - uses this to pick up exactly the sub-requests it still owes — - first run and every resume alike. - - Yields - ------ - tuple of (int, dict) - The sub-args ``index`` and its ``sub_args`` dict for each - sub-request not yet completed. - """ - for index, sub_args in enumerate(self.plan.iter_sub_args()): - if index not in self._chunks: - yield index, sub_args - - def resume(self) -> tuple[pd.DataFrame, Any]: - """ - Drive the chunked call to completion and return the combined result. - - Runs :meth:`_run` through an ``anyio`` blocking portal (a - short-lived worker thread), so it works whether or not the caller - is already inside an event loop (Jupyter / IPython / async apps). - The portal copies the calling context, so the active progress - reporter still reaches the sub-requests. - - Idempotent: only sub-requests whose index isn't already in - ``self._chunks`` are re-issued. Sub-args order matches - :meth:`ChunkPlan.iter_sub_args` and is deterministic, so a - partial completion (sparse indices) resumes correctly. - - Returns - ------- - df : pandas.DataFrame - Combined data from every successful sub-request. - response - The finalized aggregate — a raw :class:`httpx.Response` - (canonical URL, headers from the response with the lowest reported - remaining quota, and summed response elapsed durations) by default, - or whatever - :attr:`finalize` produces (e.g. ``BaseMetadata`` for the OGC - getters). - - Raises - ------ - ChunkInterrupted - On a mid-stream transient failure — 429, 5xx, or a bare - transport error: :class:`QuotaExhausted` for 429, - :class:`ServiceInterrupted` for the rest. The resumable - handle is on ``exc.call`` — wait for the underlying - condition to clear and call ``exc.call.resume()`` again. - """ - # Drive inside the snapshot taken at construction (see ``__init__``). - # ``start_blocking_portal`` copies the *calling* context into its - # worker thread, and running here means that calling context is the - # snapshot — so the base URL / dialect / row cap / progress reporter - # active when the call was created reach the rebuilt sub-requests, - # even when this is a resume fired long after the original ``with`` - # blocks exited. - return self._ctx.run(self._resume_in_context) - - def _resume_in_context(self) -> tuple[pd.DataFrame, Any]: - """Body of :meth:`resume`, run inside the captured context.""" - concurrency = _read_concurrency_env() - with start_blocking_portal() as portal: - # ``portal.call`` returns ``Any`` because ``functools.partial`` - # erases ``_run``'s return type; restore the declared tuple. - return cast( - "tuple[pd.DataFrame, Any]", - portal.call(functools.partial(self._run, concurrency)), - ) - - async def _run(self, max_concurrent: int | None) -> tuple[pd.DataFrame, Any]: - """ - Gather every pending sub-request over one shared - :class:`httpx.AsyncClient` and return the combined, finalized result. - - Pending sub-requests (:meth:`_pending`) fan out under - ``asyncio.gather`` with ``return_exceptions=True`` so completed - sub-requests survive a sibling's transient failure. On a - recognized transient (:class:`RateLimited`, :class:`ServiceUnavailable`, - or a bare ``httpx.HTTPError`` / ``httpx.InvalidURL``) a - :class:`ChunkInterrupted` subclass is raised carrying ``self`` on - ``.call``; ``exc.call.resume()`` then re-issues only the unfinished - indices through this same runner. - - The gather dispatches *every* pending sub-request at once, but an - ``asyncio.Semaphore`` caps the number of concurrent fetches at - ``N = max_concurrent`` — ``None`` lifts the cap, ``N=1`` runs them - one at a time. The connection pool is sized to the same ``N`` - (``httpx.Limits(max_connections=N, max_keepalive_connections=N)``) - so the in-flight fetches reuse keepalive connections. - - 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 - against the pool-acquire timeout (60 s, from ``HTTPX_ASYNC_DEFAULTS``). - A batch of slow pages that keeps every connection busy past that - window would then trip ``httpx.PoolTimeout`` on the queued tail — - a purely client-side failure that consumes the retry budget and - surfaces as a spurious resumable ``ServiceInterrupted``. Holding - sub-requests at the semaphore keeps them out of the pool until a - slot frees, so the pool timeout only fires for a genuinely stuck - connection. - - The shared client is published on :data:`_chunked_client` so - the paginated-loop helpers reuse its connection pool. - - Parameters - ---------- - max_concurrent : int or None - Maximum sub-requests in flight (the semaphore value, and the - connection-pool size). ``None`` lifts the cap entirely. - - Returns - ------- - df : pandas.DataFrame - Combined data from every sub-request. - response - The finalized aggregate — a raw :class:`httpx.Response` - (canonical URL, headers from the response with the lowest reported - remaining quota, and summed response elapsed durations) by default, - or whatever - :attr:`finalize` produces (e.g. ``BaseMetadata`` for OGC getters). - - Raises - ------ - ChunkInterrupted - On a transient sub-request failure. ``.call`` is ``self``, - holding the sparse completed sub-requests; ``.call.resume()`` - re-issues the unfinished ones. - """ - # The semaphore is the throttle; the pool is merely sized to match - # it. Left at httpx's default client limits (``max_connections=100``, - # keepalive 20) the pool would bottleneck a wider cap or churn - # connections by keeping too few alive. See the method docstring for - # 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. - limits = httpx.Limits( - max_connections=max_concurrent, max_keepalive_connections=max_concurrent - ) - semaphore = asyncio.Semaphore( - self.plan.total if max_concurrent is None else max_concurrent - ) - - async with open_async_client(limits=limits) as client: - with _chunked_client(client): - reporter = _progress.current() - if reporter is not None: - reporter.set_chunks(self.plan.total) - - async def track( - index: int, args: dict[str, Any] - ) -> tuple[pd.DataFrame, httpx.Response]: - """One sub-request (with retry) + result-store + progress tick.""" - result = await _retry( - lambda: self.fetch(args), self.retry_policy, gate=semaphore - ) - self._chunks[index] = result - if reporter is not None: - # Chunks finish out of order under gather, so tick the - # completed *count* rather than a positional index. - reporter.start_chunk(self.completed_chunks) - return result - - # Dispatch every pending sub-request concurrently; the - # semaphore (held by ``_retry`` per attempt) is the only throttle. - # ``return_exceptions`` keeps completed pairs after a sibling - # fails, so partial state stays recoverable via :meth:`resume`. - # Failure precedence, in order: - # 1. Cancellation / interrupt signals (CancelledError, - # KeyboardInterrupt, SystemExit — non-Exception) propagate - # unmodified; wrapping them as a transient would swallow - # the user's stop signal. - # 2. A non-transient failure (a real bug — unrecognized by - # ``wrap_failure``) surfaces raw, so it isn't masked behind - # a resumable handle for a transient sibling that landed - # later. - # 3. Only when every failure is a recognized transient do we - # raise the first as a resumable ``ChunkInterrupted``. - results = await asyncio.gather( - *(track(index, args) for index, args in self._pending()), - return_exceptions=True, - ) - failures = [r for r in results if isinstance(r, BaseException)] - for exc in failures: - if not isinstance(exc, Exception): - raise exc - first_transient: tuple[ChunkInterrupted, BaseException] | None = None - for exc in failures: - interrupted = self.wrap_failure(exc) - if interrupted is None: - raise exc - if first_transient is None: - first_transient = (interrupted, exc) - if first_transient is not None: - interrupted, exc = first_transient - raise interrupted from exc - - return self.finalize(*self._combine_raw()) - - def multi_value_chunked( *, build_request: Callable[..., httpx.Request], diff --git a/dataretrieval/ogc/interruptions.py b/dataretrieval/ogc/interruptions.py index 8cb5723c..a9158fa7 100644 --- a/dataretrieval/ogc/interruptions.py +++ b/dataretrieval/ogc/interruptions.py @@ -1,180 +1,25 @@ -"""Resumable chunk-interruption exceptions — the public resume contract. - -When a transparently-chunked request fails mid-stream (a 429, a 5xx, or a -bare transport error), the work already completed is preserved and the call -is resumable: the raised exception carries a ``.call`` handle whose -``resume()`` re-issues only the still-pending sub-requests. These exception -types are that contract, re-exported at the top level -(``from dataretrieval import ChunkInterrupted``). The execution machinery -that raises and resumes them lives in :mod:`dataretrieval.ogc.chunking`. +"""Compatibility re-export: the interruption taxonomy moved to a top-level leaf. + +The resume contract is no longer OGC-specific — Water Use raises it too — so the +classes live in :mod:`dataretrieval.interruptions`, where the base class is +named :class:`~dataretrieval.interruptions.FanOutInterrupted`. This path is kept +because it is what existing code and tests import; new code should import from +the leaf, or the top level +(``from dataretrieval import FanOutInterrupted``). """ from __future__ import annotations -from typing import TYPE_CHECKING, Any, ClassVar - -import httpx -import pandas as pd - -from dataretrieval.exceptions import DataRetrievalError - -if TYPE_CHECKING: - from dataretrieval.ogc.chunking import ChunkedCall - - -class ChunkInterrupted(DataRetrievalError): - """ - Base class for mid-stream chunk failures whose completed work is - preserved and resumable. - - A ``ChunkInterrupted`` subclass means: a sub-request failed, but - ``ChunkedCall`` still owns whatever completed successfully before - the failure. Call ``self.call.resume()`` to pick up where the - failure stopped you — only still-pending sub-requests are - re-issued. - - Subclasses describe *why* ``ChunkedCall`` stopped so callers can - pick a retry policy: :class:`QuotaExhausted` for 429 (wait for the - rate-limit window), :class:`ServiceInterrupted` for 5xx (wait for - the upstream to recover). The ``.call`` handle is the same object - across every interruption of a single chunked call — frames - accumulate across retries. - - Attributes - ---------- - call : ChunkedCall or None - Resumable handle into the ``ChunkedCall`` that raised this - exception. ``None`` only on hand-constructed exceptions (test - fixtures), where ``.call``-derived accessors degrade to - empty/``None``. - retry_after : float or None - Seconds the server suggested waiting (``Retry-After`` header). - ``None`` when the server gave no hint. - completed_chunks : int - Number of sub-requests successfully completed before the failure. - total_chunks : int - Total sub-requests in the plan. - partial_frame : pandas.DataFrame - Combined frame of work completed by the moment this exception - was raised. Snapshot at raise time — does NOT advance on a - later ``call.resume()`` (use ``exc.call.partial_frame`` for - the live view). - partial_response : httpx.Response or None - Raw aggregate response covering the completed sub-requests at - raise time; ``None`` if nothing had completed yet. Same snapshot - semantics as ``partial_frame``. (Raw, not finalized — use - ``exc.call.resume()`` for the finalized ``(df, metadata)`` result.) - - Examples - -------- - Retry on any transient interruption, honoring the server's - ``Retry-After`` hint when present and falling back to a fixed wait - otherwise. Each new interruption keeps the already-completed work - intact — only the still-pending sub-requests are re-issued. - - .. code-block:: python - - import time - from dataretrieval import ChunkInterrupted - - # ``getter`` is any chunked OGC getter — e.g. - # ``waterdata.get_daily`` or ``ngwmn.get_water_level``. - try: - df, md = getter(monitoring_location_id=long_list_of_sites) - except ChunkInterrupted as exc: - while True: - time.sleep(exc.retry_after or 5 * 60) - try: - df, md = exc.call.resume() - break - except ChunkInterrupted as next_exc: - exc = next_exc - """ - - # Subclasses override with a ``str.format`` template; the format - # call sees ``completed_chunks`` and ``total_chunks`` as kwargs. - _MESSAGE_TEMPLATE: ClassVar[str] = ( - "Chunked request interrupted after {completed_chunks}/" - "{total_chunks} sub-requests; call .call.resume() to continue." - ) - - def __init__( - self, - *, - completed_chunks: int, - total_chunks: int, - call: ChunkedCall | None = None, - retry_after: float | None = None, - cause: BaseException | None = None, - ) -> None: - message = self._MESSAGE_TEMPLATE.format( - completed_chunks=completed_chunks, total_chunks=total_chunks - ) - if cause is not None: - cause_msg = str(cause) or type(cause).__name__ - message = f"{message} Cause: {type(cause).__name__}: {cause_msg}" - super().__init__(message) - self.completed_chunks = completed_chunks - self.total_chunks = total_chunks - self.call = call - self.retry_after = retry_after - # Snapshot partial state at raise time so the exception stays a stable - # record of the failure moment: ``exc.partial_frame`` / - # ``.partial_response`` do NOT advance on a later ``call.resume()`` - # (that live view is on ``call.partial_frame`` / ``.partial_response``). - # This keeps each interruption in a resume loop a faithful record of - # what it saw, rather than every exception aliasing the shared call's - # advancing state. ``.copy()`` guards the single-chunk fast path, where - # the combined frame may be returned verbatim. - if call is None: - self.partial_frame: pd.DataFrame = pd.DataFrame() - self.partial_response: httpx.Response | None = None - else: - self.partial_frame = call.partial_frame.copy() - self.partial_response = call.partial_response - - def __getstate__(self) -> dict[str, Any]: - # Drop the live ChunkedCall before pickling: its ``.fetch`` is an - # undecorated module function pickle can't reference by name, so the - # interruption can't cross a process boundary with ``.call`` attached. - # The degraded ``call=None`` form keeps the counts, retry hint, and the - # snapshotted partial frame / response — plain instance attributes the - # base ``__getstate__`` already pickles; only ``.resume()`` is lost - # (cross-process resume was never possible anyway). - return {**super().__getstate__(), "call": None} - - -class QuotaExhausted(ChunkInterrupted): - """ - A sub-request returned HTTP 429 — the per-key rate-limit window - is exhausted. Subclass of :class:`ChunkInterrupted`. - - The completed sub-requests are preserved on ``.call``; once the - rate-limit window resets, ``.call.resume()`` re-issues only the - still-pending work. ``partial_frame`` holds what completed - before the 429. - """ - - _MESSAGE_TEMPLATE = ( - "HTTP 429 after {completed_chunks}/{total_chunks} sub-requests; " - "catch QuotaExhausted (or ChunkInterrupted) to access " - ".partial_frame or .call.resume() once the rate-limit " - "window has rolled over." - ) - - -class ServiceInterrupted(ChunkInterrupted): - """ - A sub-request returned HTTP 5xx — the upstream service failed - transiently. Subclass of :class:`ChunkInterrupted`. - - The completed sub-requests are preserved on ``.call``; once the - upstream recovers, ``.call.resume()`` resumes only the - still-pending work. - """ - - _MESSAGE_TEMPLATE = ( - "Service error after {completed_chunks}/{total_chunks} " - "sub-requests; catch ServiceInterrupted (or ChunkInterrupted) " - "and call .call.resume() once the upstream service recovers." - ) +from dataretrieval.interruptions import ( + ChunkInterrupted, + FanOutInterrupted, + QuotaExhausted, + ServiceInterrupted, +) + +__all__ = [ + "ChunkInterrupted", + "FanOutInterrupted", + "QuotaExhausted", + "ServiceInterrupted", +] diff --git a/dataretrieval/ogc/retry.py b/dataretrieval/ogc/retry.py index 7eeafb44..894395b5 100644 --- a/dataretrieval/ogc/retry.py +++ b/dataretrieval/ogc/retry.py @@ -1,61 +1,20 @@ -"""OGC interruption classification over service-neutral transport retry policy. +"""Compatibility re-export: interruption classification moved to the taxonomy leaf. -Only the OGC-specific half of retry lives here: turning a transport failure into -the resumable :class:`~dataretrieval.ogc.interruptions.ChunkInterrupted` the -chunker reports. The policy itself -- backoff, bounds, classification of what is -transient -- belongs to :mod:`dataretrieval.transport.retry`, which callers -import directly; re-exporting its tunables here would hand out stale copies that -patching cannot reach. +Turning a transport failure into a resumable +:class:`~dataretrieval.interruptions.FanOutInterrupted` was never OGC-specific -- +it keys off the shared ``RateLimited``/``TransientError`` taxonomy and httpx -- +so it now lives beside the classes it produces, in +:mod:`dataretrieval.interruptions`. -"Should we retry this?" and "can the caller resume it?" are the same question -asked twice, so both answers come from one place in transport. Keeping a second -copy here is how they would end up disagreeing -- refusing to retry a failure -while still telling the caller it can be resumed. +The retry *policy* -- backoff, bounds, classification of what is transient -- +still belongs to :mod:`dataretrieval.transport.retry`, which callers import +directly; re-exporting its tunables here would hand out stale copies that +patching cannot reach. """ from __future__ import annotations -import httpx - -from dataretrieval.exceptions import RateLimited, TransientError -from dataretrieval.ogc.interruptions import ( - ChunkInterrupted, - QuotaExhausted, - ServiceInterrupted, -) -from dataretrieval.transport.retry import _deterministic_failure - - -def _classify_transient( - exc: BaseException, -) -> tuple[type[ChunkInterrupted], float | None] | None: - """Classify one failure as a resumable OGC interruption.""" - if isinstance(exc, RateLimited): - return QuotaExhausted, exc.retry_after - if isinstance(exc, TransientError): - return ServiceInterrupted, exc.retry_after - if isinstance(exc, (httpx.HTTPError, httpx.InvalidURL)): - # Some failures will fail the same way every time -- a bad scheme, a - # hostname that doesn't resolve. Offering to resume one would just - # hide the real error behind a retry that can never work. - if _deterministic_failure(exc): - return None - return ServiceInterrupted, None - return None - - -def _classify_chunk_error( - exc: BaseException, -) -> tuple[type[ChunkInterrupted], float | None] | None: - """Walk a wrapped pagination failure for a resumable transport cause.""" - current: BaseException | None = exc - while current is not None: - result = _classify_transient(current) - if result is not None: - return result - current = current.__cause__ - return None - +from dataretrieval.interruptions import _classify_chunk_error, _classify_transient __all__ = [ "_classify_chunk_error", diff --git a/dataretrieval/transport/fanout.py b/dataretrieval/transport/fanout.py new file mode 100644 index 00000000..476047e8 --- /dev/null +++ b/dataretrieval/transport/fanout.py @@ -0,0 +1,652 @@ +"""Bounded, resumable fan-out execution over a plan of sub-requests. + +A fan-out is one logical query the service forces into several requests. Two +unrelated reasons produce one: + +- a Water Data / NGWMN query whose URL exceeds the server's byte limit, split + along its multi-value axes by :class:`dataretrieval.ogc.planning.ChunkPlan`; +- a Water Use query naming several locations, which the NWDC accepts only one + at a time. + +Chunking is how you divide the data structurally; fan-out is how you distribute +the work operationally. The two are orthogonal, and only the first is protocol +knowledge: dividing a query needs the byte budget, the CQL2 grammar, and which +parameters are list-valued, while distributing the pieces needs none of it. Only +the Water Data / NGWMN case above involves chunking at all — Water Use fans out +without dividing anything, because the caller's locations were never one body to +split. + +So this module owns distribution and nothing else: concurrency bounded by a +semaphore, per-attempt retry, deterministic failure precedence, sparse +completion tracking, and resume. It names no protocol concept — an adapter +supplies a :class:`FanOutPlan` (whatever structure it divided into, if any) and +an ``async def fetch(args) -> (df, response)``. + +Concurrency: :meth:`FanOut._run` dispatches every pending sub-request under one +``asyncio.gather`` sharing a single ``httpx.AsyncClient``. An +``asyncio.Semaphore`` -- not the client's connection pool, which is merely sized +to match -- caps the sub-requests in flight at ``N``; see :meth:`FanOut._run` +for why the gate must be the semaphore rather than the pool. +``API_USGS_CONCURRENT`` resolves ``N``: an integer N > 1 allows N sub-requests +in flight; ``1`` forces sequential dispatch; the literal ``unbounded`` lifts the +cap. ``N`` bounds only how many of a query's sub-requests are in flight at once +-- a client-side trade-off between open connections and fan-out latency. It does +not affect the API rate limit: a fanned-out call issues the same number of +sub-requests regardless of ``N``, so ``N`` changes their timing, not the total +request volume. The USGS API rate-limits by volume over time (HTTP 429), not by +simultaneity; set ``API_USGS_PAT`` to raise that quota. The default of 32 is a +conservative cap that keeps connection use modest. The fan-out runs in a +short-lived worker thread (an ``anyio`` blocking portal), so it works whether or +not the caller is already inside an event loop (Jupyter / IPython / async apps). + +Retries: each sub-request is retried on a transient failure (429, 5xx, +connect/read timeout) with exponential backoff + full jitter, honoring a server +``Retry-After`` when present. ``API_USGS_RETRIES`` sets the cap (default 4; +``0`` disables). A ``Retry-After`` longer than the per-call ceiling escalates to +a resumable interruption. + +Interruption: any mid-stream transient failure surfaces as a +:class:`~dataretrieval.interruptions.FanOutInterrupted` subclass carrying +``.call``, a :class:`FanOut` handle owning the already-completed sub-request +state. Call ``.call.resume()`` once the underlying condition clears; only the +still-pending sub-requests are re-issued. +""" + +from __future__ import annotations + +import asyncio +import functools +import os +from collections.abc import Awaitable, Callable, Iterator +from contextvars import copy_context +from typing import Any, Protocol, cast + +import httpx +import pandas as pd +from anyio.from_thread import start_blocking_portal + +from dataretrieval import progress as _progress +from dataretrieval.combining import ( + _combine_chunk_frames, + _combine_chunk_responses, +) +from dataretrieval.exceptions import ConfigurationError +from dataretrieval.interruptions import FanOutInterrupted, _classify_chunk_error +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 + +# Fan-out concurrency cap, read at call time (not import) so test +# ``monkeypatch.setenv`` applies. Value grammar in :func:`_read_concurrency_env`; +# the concurrency model is in the module docstring. +_CONCURRENCY_ENV = "API_USGS_CONCURRENT" +_CONCURRENCY_DEFAULT = 32 +_CONCURRENCY_UNBOUNDED = "unbounded" + + +def _resolve_concurrency(default: int = _CONCURRENCY_DEFAULT) -> int | None: + """ + Resolve the parallelism cap: the general setting, or a module's default. + + ``API_USGS_CONCURRENT`` is the general knob and applies to every fanned-out + call in the package. A module may pass a different ``default`` when its + service warrants one — Water Use ships a lower figure than the OGC getters, + because the NWDC is only stress-tested to that level. + + The ordering is deliberate: an explicitly set environment variable wins over + a module's default, never the reverse. A module that could override the + general setting would make ``API_USGS_CONCURRENT=1`` a lie — the user + dialing concurrency down to be polite to the service would find one adapter + quietly ignoring them, which is precisely the defect this consolidates away. + Module defaults express "absent instruction, this service prefers N"; they + do not express "this service knows better than you". + + Parameters + ---------- + default : int + Cap to use when ``API_USGS_CONCURRENT`` is unset or empty. + + Returns + ------- + int or None + ``1`` for sequential dispatch (one sub-request at a time); an + integer >1 for bounded concurrency; ``None`` to disable the + per-call cap entirely (the ``unbounded`` keyword). + """ + raw = os.environ.get(_CONCURRENCY_ENV) + if raw is None: + return default + raw = raw.strip() + if raw == "": + return default + if raw.lower() == _CONCURRENCY_UNBOUNDED: + return None + try: + value = int(raw) + except ValueError as exc: + raise ConfigurationError( + f"{_CONCURRENCY_ENV} must be a positive integer or " + f"'{_CONCURRENCY_UNBOUNDED}'; got {raw!r}." + ) from exc + if value < 1: + raise ConfigurationError( + f"{_CONCURRENCY_ENV} must be >= 1 (got {value}); use " + f"'{_CONCURRENCY_UNBOUNDED}' to disable the cap." + ) + return value + + +# --------------------------------------------------------------------------- +# The plan contract +# --------------------------------------------------------------------------- + + +class FanOutPlan(Protocol): + """ + A fan-out's shape: how many sub-requests, their arguments, and the + identity of the whole query. + + Structural, not nominal: an implementation satisfies this by having the + three members, not by inheriting. That is the right relationship here + because the two implementations share an interface and no implementation + at all. :class:`~dataretrieval.ogc.planning.ChunkPlan` derives its + sub-requests from a URL byte budget over multi-value axes; a Water Use + plan simply lists the locations the caller named. Neither has anything + the other could inherit. + + Attributes + ---------- + total : int + Number of sub-requests in the plan. Bounds progress reporting and + sizes the degenerate semaphore when concurrency is unbounded. + canonical_url : str or None + URL identifying the query as a whole, restored onto the combined + response so the caller sees the request they made rather than + whichever sub-request happened to land last. + """ + + @property + def total(self) -> int: ... + + @property + def canonical_url(self) -> str | None: ... + + def iter_sub_args(self) -> Iterator[dict[str, Any]]: + """ + Yield each sub-request's arguments, in a deterministic order. + + Order is load-bearing: :meth:`FanOut.resume` keys completed work by + position, so a plan that yielded a different order on a second pass + would resume the wrong sub-requests. + """ + ... + + +# --------------------------------------------------------------------------- +# Shared per-call client +# --------------------------------------------------------------------------- + +# The per-call ``httpx.AsyncClient``, published for the duration of +# ``FanOut._run`` so paginated-loop helpers reuse the same connection pool +# across every sub-request. ``None`` outside a fan-out — paginated helpers then +# open their own short-lived client. Deliberately a plain ContextVar-backed +# ambient rather than a parameter: the fetch closure an adapter injects is often +# several frames below the client's owner. +_active_client: Ambient[httpx.AsyncClient | None] = Ambient("_fanout_client", None) + + +def active_client() -> httpx.AsyncClient | None: + """ + Return the fan-out's currently-published client, or ``None``. + + Used by paginated-loop helpers to reuse the per-call connection pool. + + Returns + ------- + httpx.AsyncClient or None + The client published for the duration of a :meth:`FanOut._run`; + ``None`` outside one. + """ + return _active_client.get() + + +# --------------------------------------------------------------------------- +# Type aliases for the FanOut contract +# --------------------------------------------------------------------------- + +# The per-sub-request fetcher an adapter injects and ``FanOut`` drives: +# an ``async def fetch(args) -> (df, response)``. +_Fetch = Callable[[dict[str, Any]], Awaitable[tuple[pd.DataFrame, httpx.Response]]] + +# Caller-supplied transform applied to the combined result, so a resumed call +# returns the same shape as an un-interrupted one rather than the executor's raw +# ``(frame, httpx.Response)``. This keeps the executor generic: the OGC getters +# inject their post-processing (type coercion, column arrangement, +# ``BaseMetadata``) through ``_finalize_ogc``. The default is identity. +_Finalize = Callable[[pd.DataFrame, httpx.Response], tuple[pd.DataFrame, Any]] + + +def _passthrough_result( + frame: pd.DataFrame, response: httpx.Response +) -> tuple[pd.DataFrame, Any]: + """Default :data:`_Finalize`: return the raw combined pair unchanged.""" + return frame, response + + +class FanOut: + """ + Stateful handle for a fanned-out call. + + Holds the in-flight state (per-sub-request frames and responses) + and the async fetcher. A single :meth:`resume` entry point drives + the call from wherever it is to completion — used both for the + first invocation and for subsequent retries after a + :class:`~dataretrieval.interruptions.FanOutInterrupted`. + + :meth:`_run` gathers every pending sub-request over one shared + :class:`httpx.AsyncClient`, applies the failure-precedence rules, and + combines; :meth:`resume` drives it through an ``anyio`` blocking + portal so it works whether or not the caller is already inside an + event loop. Concurrency is bounded by a per-run ``asyncio.Semaphore`` + (see :meth:`_run`), so sequential dispatch + (``API_USGS_CONCURRENT=1``) is just a degenerate gather. + + A ``FanOut`` is created internally when an adapter executes a plan; + callers reach it via ``FanOutInterrupted.call`` on the exception raised + by a mid-stream failure. + + :meth:`resume` is idempotent: :meth:`_run` iterates + :meth:`FanOutPlan.iter_sub_args` (deterministic order) and skips + any index whose result is already in ``self._chunks``. The + completion set is a sparse ``dict[int, (df, response)]`` so the + gather can record scattered completions (e.g. indices [0, 2, 5] + after siblings [1, 3, 4] failed) and a subsequent ``resume`` only + re-issues the missing indices. + + Parameters + ---------- + plan : FanOutPlan + The plan to execute. + fetch : Callable + ``async def`` that issues a single sub-request, given the + substituted args dict, and returns ``(frame, response)``. + client_options : dict, optional + Extra ``httpx.AsyncClient`` options for the shared client this run + opens (e.g. ``{"verify": False}``). + default_concurrent : int, optional + This service's preferred in-flight cap when ``API_USGS_CONCURRENT`` + is unset. Defaults to 32. + + Attributes + ---------- + plan : FanOutPlan + The plan being driven (read-only after construction). + fetch : Callable + The async per-sub-request fetch function. + finalize : Callable + Transform applied to the combined result (see :data:`_Finalize`) at + the terminal :meth:`_run` return, so a completed call yields the + caller's finished shape. The ``partial_*`` accessors deliberately + skip it and stay raw. + partial_frame : pandas.DataFrame + Raw combined frame of completed sub-requests (live; recomputed per + access). Not finalized — call :meth:`resume` for the finished shape. + partial_response : httpx.Response or None + Raw aggregate response (canonical URL restored), or ``None`` when + nothing has completed yet (live; recomputed per access). + """ + + def __init__( + self, + plan: FanOutPlan, + fetch: _Fetch, + retry_policy: RetryPolicy = _NO_RETRY, + finalize: _Finalize = _passthrough_result, + client_options: dict[str, Any] | None = None, + default_concurrent: int = _CONCURRENCY_DEFAULT, + ) -> None: + self.plan = plan + self.fetch = fetch + self.retry_policy = retry_policy + self.finalize = finalize + # This service's preferred cap when the user has not set + # ``API_USGS_CONCURRENT``. Resolved at resume time, not here, so a + # test's ``monkeypatch.setenv`` still applies. See + # :func:`_resolve_concurrency` for why the env var outranks it. + self.default_concurrent = default_concurrent + # Extra ``httpx.AsyncClient`` options merged into the shared client this + # run opens (``verify`` for the Water Use ``ssl_check`` flag, say). The + # executor owns client lifecycle, so an adapter with a per-call client + # requirement has to hand it down rather than open its own — opening its + # own would defeat the shared connection pool. Empty for OGC, which + # exposes no such flag. + self.client_options = client_options or {} + # Snapshot the ambient context at construction time — i.e. inside the + # caller's ``with`` blocks (base URL, dialect, row cap, progress + # reporter). :meth:`resume` runs every drive inside this snapshot, so + # a *later* ``exc.call.resume()`` — which fires after those ``with`` + # blocks have exited and reset their ContextVars — still rebuilds + # sub-requests against the original API's base URL/dialect rather than + # the process defaults. The adapter's request builder reads those + # ContextVars when it reconstructs each sub-request, so the snapshot + # must outlive them. The mechanism is generic; which ambients matter is + # the adapter's business. + self._ctx = copy_context() + # Completed (frame, response) pairs keyed by sub-args index; sparse + # (gathered sub-requests complete out of order — see class docstring). + # ``_run``'s ``track`` closure is the only writer, so ``dict`` insertion + # order is completion order (relied on by :meth:`_combine_raw`). + self._chunks: dict[int, tuple[pd.DataFrame, httpx.Response]] = {} + + def wrap_failure(self, exc: BaseException) -> FanOutInterrupted | None: + """ + Build the matching :class:`FanOutInterrupted` carrying this + call when ``exc`` is a recognized transient transport failure; + return ``None`` for unrecognized failures so the caller can + re-raise. Encapsulates the + ``classify → instantiate-with-call-state`` recipe so + :class:`FanOut`'s private fields stay private. + + Parameters + ---------- + exc : BaseException + The exception raised by a sub-request. + + Returns + ------- + FanOutInterrupted or None + The matching :class:`FanOutInterrupted` subclass carrying this + call for a recognized transient failure; ``None`` otherwise. + """ + classification = _classify_chunk_error(exc) + if classification is None: + return None + interrupted_class, retry_after = classification + return interrupted_class( + completed_chunks=self.completed_chunks, + total_chunks=self.plan.total, + call=self, + retry_after=retry_after, + cause=exc, + ) + + @property + def completed_chunks(self) -> int: + """Number of sub-requests completed so far.""" + return len(self._chunks) + + def _combine_raw(self) -> tuple[pd.DataFrame, httpx.Response]: + """Assemble the raw ``(frame, response)`` from completed sub-requests, + before :attr:`finalize` runs. + + Frames concatenate in sub-args *index* order (``sorted`` keys — + deterministic, independent of parallel completion order). The + aggregated response takes its headers from the response with the + lowest reported ``x-ratelimit-remaining`` value. If no response + reports that header, it falls back to the last completed response; + ``self._chunks`` preserves completion order because the ``track`` + closure in :meth:`_run` is its only writer. + + Returns + ------- + tuple of (pandas.DataFrame, httpx.Response) + The concatenated frame and the aggregated response, before + :attr:`finalize` is applied. + """ + frames = [self._chunks[i][0] for i in sorted(self._chunks)] + responses = [response for _, response in self._chunks.values()] + return ( + _combine_chunk_frames(frames), + _combine_chunk_responses(responses, self.plan.canonical_url), + ) + + @property + def partial_frame(self) -> pd.DataFrame: + """ + Raw combined frame of sub-requests that have completed so far. + + Live — recomputed on each access so it reflects current state + across resume attempts. Deliberately the *raw* combined frame + (``_combine_raw``), NOT the finalized result: this is a cheap, + side-effect-free snapshot for inspecting partial progress, so + reading it (or building a :class:`FanOutInterrupted` around it) + never triggers ``finalize`` work — which for OGC getters includes + a schema network fetch on an empty frame. Use ``call.resume()`` + for the finalized result. + + Returns + ------- + pandas.DataFrame + Combined frame of completed sub-requests, or an empty + ``DataFrame`` when nothing has completed. + """ + if not self._chunks: + return pd.DataFrame() + return self._combine_raw()[0] + + @property + def partial_response(self) -> httpx.Response | None: + """ + Raw aggregate response with the canonical URL restored to the + user's full original query. + + Live — recomputed on each access. Like :attr:`partial_frame`, this + is the *raw* aggregate (an :class:`httpx.Response`), not the + finalized result, so inspecting it is side-effect-free. + + Returns + ------- + httpx.Response or None + Aggregated response when at least one sub-request has + completed, ``None`` otherwise. + """ + if not self._chunks: + return None + return self._combine_raw()[1] + + def _pending(self) -> Iterator[tuple[int, dict[str, Any]]]: + """ + Yield ``(index, sub_args)`` for sub-requests not yet completed. + + Walks :meth:`FanOutPlan.iter_sub_args` in deterministic order + and skips any index already in ``self._chunks``. :meth:`_run` + uses this to pick up exactly the sub-requests it still owes — + the mechanism behind idempotent resume. + """ + for index, args in enumerate(self.plan.iter_sub_args()): + if index not in self._chunks: + yield index, args + + def resume(self) -> tuple[pd.DataFrame, Any]: + """ + Drive the call to completion and return the combined result. + + Runs :meth:`_run` through an ``anyio`` blocking portal (a + short-lived worker thread), so it works whether or not the caller + is already inside an event loop (Jupyter / IPython / async apps). + The portal copies the calling context, so the active progress + reporter still reaches the sub-requests. + + Idempotent: only sub-requests whose index isn't already in + ``self._chunks`` are re-issued. Sub-args order matches + :meth:`FanOutPlan.iter_sub_args` and is deterministic, so a + partial completion (sparse indices) resumes correctly. + + Returns + ------- + df : pandas.DataFrame + Combined data from every successful sub-request. + response + The finalized aggregate — a raw :class:`httpx.Response` + (canonical URL, headers from the response with the lowest reported + remaining quota, and summed response elapsed durations) by default, + or whatever :attr:`finalize` produces (e.g. ``BaseMetadata`` for + the OGC getters). + + Raises + ------ + FanOutInterrupted + On a mid-stream transient failure — 429, 5xx, or a bare + transport error: :class:`~dataretrieval.interruptions.QuotaExhausted` + for 429, :class:`~dataretrieval.interruptions.ServiceInterrupted` + for the rest. The resumable handle is on ``exc.call`` — wait for + the underlying condition to clear and call ``exc.call.resume()`` + again. + """ + # Drive inside the snapshot taken at construction (see ``__init__``). + # ``start_blocking_portal`` copies the *calling* context into its + # worker thread, and running here means that calling context is the + # snapshot — so the base URL / dialect / row cap / progress reporter + # active when the call was created reach the rebuilt sub-requests, + # even when this is a resume fired long after the original ``with`` + # blocks exited. + return self._ctx.run(self._resume_in_context) + + def _resume_in_context(self) -> tuple[pd.DataFrame, Any]: + """Body of :meth:`resume`, run inside the captured context.""" + concurrency = _resolve_concurrency(self.default_concurrent) + with start_blocking_portal() as portal: + # ``portal.call`` returns ``Any`` because ``functools.partial`` + # erases ``_run``'s return type; restore the declared tuple. + return cast( + "tuple[pd.DataFrame, Any]", + portal.call(functools.partial(self._run, concurrency)), + ) + + async def _run(self, max_concurrent: int | None) -> tuple[pd.DataFrame, Any]: + """ + Gather every pending sub-request over one shared + :class:`httpx.AsyncClient` and return the combined, finalized result. + + Pending sub-requests (:meth:`_pending`) fan out under + ``asyncio.gather`` with ``return_exceptions=True`` so completed + sub-requests survive a sibling's transient failure. On a + recognized transient (:class:`~dataretrieval.exceptions.RateLimited`, + :class:`~dataretrieval.exceptions.ServiceUnavailable`, or a bare + ``httpx.HTTPError`` / ``httpx.InvalidURL``) a + :class:`FanOutInterrupted` subclass is raised carrying ``self`` on + ``.call``; ``exc.call.resume()`` then re-issues only the unfinished + indices through this same runner. + + The gather dispatches *every* pending sub-request at once, but an + ``asyncio.Semaphore`` caps the number of concurrent fetches at + ``N = max_concurrent`` — ``None`` lifts the cap, ``N=1`` runs them + one at a time. The connection pool is sized to the same ``N`` + (``httpx.Limits(max_connections=N, max_keepalive_connections=N)``) + so the in-flight fetches reuse keepalive connections. + + 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 + against the pool-acquire timeout (60 s, from ``HTTPX_ASYNC_DEFAULTS``). + A batch of slow pages that keeps every connection busy past that + window would then trip ``httpx.PoolTimeout`` on the queued tail — + a purely client-side failure that consumes the retry budget and + surfaces as a spurious resumable ``ServiceInterrupted``. Holding + sub-requests at the semaphore keeps them out of the pool until a + slot frees, so the pool timeout only fires for a genuinely stuck + connection. + + The shared client is published on :data:`_active_client` so + the paginated-loop helpers reuse its connection pool. + + Parameters + ---------- + max_concurrent : int or None + Maximum sub-requests in flight (the semaphore value, and the + connection-pool size). ``None`` lifts the cap entirely. + + Returns + ------- + df : pandas.DataFrame + Combined data from every sub-request. + response + The finalized aggregate — a raw :class:`httpx.Response` + (canonical URL, headers from the response with the lowest reported + remaining quota, and summed response elapsed durations) by default, + or whatever :attr:`finalize` produces. + + Raises + ------ + FanOutInterrupted + On a transient sub-request failure. ``.call`` is ``self``, + holding the sparse completed sub-requests; ``.call.resume()`` + re-issues the unfinished ones. + """ + # The semaphore is the throttle; the pool is merely sized to match + # it. Left at httpx's default client limits (``max_connections=100``, + # keepalive 20) the pool would bottleneck a wider cap or churn + # connections by keeping too few alive. See the method docstring for + # 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. + limits = httpx.Limits( + max_connections=max_concurrent, max_keepalive_connections=max_concurrent + ) + semaphore = asyncio.Semaphore( + self.plan.total if max_concurrent is None else max_concurrent + ) + + async with open_async_client(limits=limits, **self.client_options) as client: + with _active_client(client): + reporter = _progress.current() + if reporter is not None: + reporter.set_chunks(self.plan.total) + + async def track( + index: int, args: dict[str, Any] + ) -> tuple[pd.DataFrame, httpx.Response]: + """One sub-request (with retry) + result-store + progress tick.""" + result = await _retry( + lambda: self.fetch(args), self.retry_policy, gate=semaphore + ) + self._chunks[index] = result + if reporter is not None: + # Chunks finish out of order under gather, so tick the + # completed *count* rather than a positional index. + reporter.start_chunk(self.completed_chunks) + return result + + # Dispatch every pending sub-request concurrently; the + # semaphore (held by ``_retry`` per attempt) is the only throttle. + # ``return_exceptions`` keeps completed pairs after a sibling + # fails, so partial state stays recoverable via :meth:`resume`. + # Failure precedence, in order: + # 1. Cancellation / interrupt signals (CancelledError, + # KeyboardInterrupt, SystemExit — non-Exception) propagate + # unmodified; wrapping them as a transient would swallow + # the user's stop signal. + # 2. A non-transient failure (a real bug — unrecognized by + # ``wrap_failure``) surfaces raw, so it isn't masked behind + # a resumable handle for a transient sibling that landed + # later. + # 3. Only when every failure is a recognized transient do we + # raise the first as a resumable ``FanOutInterrupted``. + results = await asyncio.gather( + *(track(index, args) for index, args in self._pending()), + return_exceptions=True, + ) + failures = [r for r in results if isinstance(r, BaseException)] + for exc in failures: + if not isinstance(exc, Exception): + raise exc + first_transient: tuple[FanOutInterrupted, BaseException] | None = None + for exc in failures: + interrupted = self.wrap_failure(exc) + if interrupted is None: + raise exc + if first_transient is None: + first_transient = (interrupted, exc) + if first_transient is not None: + interrupted, exc = first_transient + raise interrupted from exc + + return self.finalize(*self._combine_raw()) + + +__all__ = [ + "FanOut", + "FanOutPlan", + "active_client", +] diff --git a/dataretrieval/transport/retry.py b/dataretrieval/transport/retry.py index 577ed334..35a756ac 100644 --- a/dataretrieval/transport/retry.py +++ b/dataretrieval/transport/retry.py @@ -6,7 +6,6 @@ import math import os import random -import socket import time from collections.abc import Awaitable, Callable from dataclasses import dataclass @@ -22,6 +21,7 @@ NetworkError, TransientError, ) +from dataretrieval.interruptions import _deterministic_failure from dataretrieval.transport.liveness import ( credit_wait, elapsed_since_progress, @@ -50,19 +50,6 @@ # hint from waking together. Small on purpose: the server named the wait, so # jitter here decorrelates rather than extends it. _RETRY_AFTER_JITTER = 1.0 -# Resolver failures that will not resolve differently on a later attempt. The -# temporary ones (notably EAI_AGAIN -- "try again", raised while a resolver is -# still coming up, on VPN reconnect, or after a laptop wakes) are deliberately -# absent: those are worth another try. Looked up defensively because the EAI_* -# constants are platform-dependent; an unrecognized code stays retryable, since -# spending a few seconds on a retry is cheaper than dropping a recoverable call. -_PERMANENT_DNS_ERRORS = frozenset( - code - for code in ( - getattr(socket, name, None) for name in ("EAI_NONAME", "EAI_FAIL", "EAI_NODATA") - ) - if code is not None -) # Attempts the no-progress budget never withholds; see RetryPolicy.allows_wait. _STALL_EXEMPT_ATTEMPTS = 1 _STALL_TIMEOUT_ENV = "API_USGS_STALL_TIMEOUT" @@ -294,44 +281,6 @@ def backoff(self, attempt: int, retry_after: float | None) -> float: _NO_RETRY = RetryPolicy(max_retries=0) -def _deterministic_failure(exc: BaseException) -> bool: - """Whether a transport failure would fail identically on every retry. - - An unsupported scheme or a request we built wrong is settled before a byte - goes out, and a hostname the resolver rejects outright won't be accepted on - the next attempt either -- so retrying only delays the error the caller - needs. A *temporary* resolver failure is not in that class and stays - retryable (see :data:`_PERMANENT_DNS_ERRORS`). - - The original failure is several layers down and not always an explicit - ``raise ... from``: a DNS failure reaches us as ``NetworkError`` -> - ``httpx.ConnectError`` -> ``httpcore.ConnectError`` -> ``socket.gaierror``, - linked by ``__context__`` (implicit chaining) rather than ``__cause__``. - - Both links of every frame are visited, not just the first one present. A - frame can carry an explicit ``__cause__`` *and* an unrelated ``__context__`` - (any ``raise X from Y`` inside an ``except`` block produces exactly that), so - following only the cause would walk off down the explicit branch and miss a - ``gaierror`` sitting on the implicit one -- spending the whole retry budget - on a hostname that will never resolve. The ``seen`` set keeps a chain that - rejoins itself, or points back at an ancestor, from looping. - """ - seen: set[int] = set() - pending: list[BaseException | None] = [exc] - while pending: - current = pending.pop() - if current is None or id(current) in seen: - continue - seen.add(id(current)) - if isinstance(current, (httpx.UnsupportedProtocol, httpx.LocalProtocolError)): - return True - if isinstance(current, socket.gaierror): - # Return, not continue: the first resolver code found settles the chain. - return current.errno in _PERMANENT_DNS_ERRORS - pending += [current.__cause__, current.__context__] - return False - - def _retryable( exc: BaseException, statuses: frozenset[int] = _RETRYABLE_STATUSES ) -> tuple[bool, float | None]: diff --git a/dataretrieval/wateruse.py b/dataretrieval/wateruse.py index a3c788ee..f37efc96 100644 --- a/dataretrieval/wateruse.py +++ b/dataretrieval/wateruse.py @@ -41,24 +41,20 @@ from __future__ import annotations -import asyncio import io -from collections.abc import Callable, Iterable +from collections.abc import Callable, Iterable, Iterator from typing import Any import httpx import pandas as pd +from dataretrieval import progress as _progress from dataretrieval.codes.states import to_state -from dataretrieval.combining import ( - _combine_chunk_frames, - _combine_chunk_responses, -) from dataretrieval.exceptions import DataRetrievalError -from dataretrieval.transport.http import default_headers, open_async_client +from dataretrieval.transport.fanout import FanOut, active_client +from dataretrieval.transport.http import default_headers from dataretrieval.transport.pagination import paginate -from dataretrieval.transport.retry import RetryPolicy, retry_async -from dataretrieval.transport.sync import run_sync +from dataretrieval.transport.retry import RetryPolicy from dataretrieval.utils import BaseMetadata, _raise_for_status, to_str WATERUSE_URL = "https://api.water.usgs.gov/nwaa-data/data" @@ -80,13 +76,15 @@ #: Temporal resolutions: monthly, annual calendar year, annual water year. TIME_RESOLUTIONS = ("monthly", "annualcy", "annualwy") -#: Maximum locations fetched concurrently when a list of state/county/huc -#: selectors is fanned out (one request per location). Kept conservative -#: because every location retries independently, so the burst a rate-limit -#: episode produces is this number times the retry count; the NWDC tolerates -#: this level of concurrency without rate-limit errors (verified by stress -#: test). Set ``wateruse.MAX_CONCURRENT_REQUESTS = 1`` for serial. -MAX_CONCURRENT_REQUESTS = 4 +#: This service's preferred in-flight cap when ``API_USGS_CONCURRENT`` is +#: unset. Lower than the package default of 32 because every location retries +#: independently, so a rate-limit episode bursts this number times the retry +#: count; the NWDC tolerates this level without rate-limit errors (verified by +#: stress test) and higher has not been tested. Setting ``API_USGS_CONCURRENT`` +#: overrides it -- see :func:`dataretrieval.transport.fanout._resolve_concurrency` +#: for why the general setting outranks a module's default rather than the +#: reverse. +DEFAULT_CONCURRENT_REQUESTS = 4 # Page responses carry the HUC12 identifier in this column; it must stay a # string so leading zeros (e.g. "010900020502") survive the round trip. @@ -118,8 +116,12 @@ def get_wateruse( Each selector also accepts a list of values. The NWDC queries one area per request, so a list is fanned out into one request per value — up to - :data:`MAX_CONCURRENT_REQUESTS` in parallel — and the results are - concatenated in the order given. + ``API_USGS_CONCURRENT`` in parallel, defaulting to + :data:`DEFAULT_CONCURRENT_REQUESTS` for this service — and the results are + concatenated in the order given. A fan-out interrupted by a rate limit or an + upstream fault raises a resumable + :class:`~dataretrieval.interruptions.FanOutInterrupted`, whose + ``.call.resume()`` re-issues only the locations that did not complete. Parameters ---------- @@ -237,15 +239,7 @@ def get_wateruse( ) for location in _resolve_locations(state, county, huc) ] - # ``_run_sync`` drives the async fan-out via an anyio portal, so it is safe - # even inside an already-running event loop (e.g. a Jupyter notebook). - # ``error_url`` is the host reported in any connection-error message (this - # module builds its own requests, so it has no OGC request-builder base). - df, response = run_sync( - lambda: _fan_out(requests, headers, ssl_check), - service="wateruse", - error_url=WATERUSE_URL, - ) + df, response = _fan_out(requests, headers, ssl_check) return df, BaseMetadata(response) @@ -328,19 +322,59 @@ def _validate_huc(value: object) -> str: return code -async def _fan_out( +class _LocationPlan: + """The Water Use fan-out's shape: one pre-built request per location. + + Satisfies :class:`~dataretrieval.transport.fanout.FanOutPlan` structurally, + without inheriting from :class:`~dataretrieval.ogc.planning.ChunkPlan` -- + there is nothing to inherit. ``ChunkPlan`` divides one over-budget query + into byte-sized pieces; this divides nothing. The NWDC accepts one + ``location=`` per request, so the caller's locations arrive already + separate and the "plan" is just that list. Chunking is structural division; + this is only the operational distribution that follows. + """ + + def __init__(self, requests: list[httpx.Request]) -> None: + self._requests = requests + + @property + def total(self) -> int: + return len(self._requests) + + @property + def canonical_url(self) -> str | None: + """The first location's URL, standing for the query as a whole. + + There is no single URL expressing "all of these locations" -- the + service has no such request -- so the aggregate response reports the + first, matching what the un-fanned single-location call would show. + """ + return str(self._requests[0].url) if self._requests else None + + def iter_sub_args(self) -> Iterator[dict[str, Any]]: + for request in self._requests: + yield {"request": request} + + +def _fan_out( requests: list[httpx.Request], headers: dict[str, str], ssl_check: bool ) -> tuple[pd.DataFrame, httpx.Response]: - """Fetch every request (each paginated) concurrently over one shared client. + """Fetch every request (each paginated) over the shared fan-out executor. Each request is paginated by :func:`dataretrieval.transport.pagination.paginate` - with NWDC strategies: parse a CSV - page and read its ``Link`` header cursor (``parse``), follow that cursor - (``follow``), and raise the typed error carrying the NWDC ``detail`` - (``raise_for_status``). Concurrency is bounded by a semaphore at - :data:`MAX_CONCURRENT_REQUESTS`, and ``asyncio.gather`` preserves input - order, so the concatenation is deterministic. The shared - :class:`httpx.AsyncClient` keeps connections alive across pages and requests. + with NWDC strategies: parse a CSV page and read its ``Link`` header cursor + (``parse``), follow that cursor (``follow``), and raise the typed error + carrying the NWDC ``detail`` (``raise_for_status``). + + Everything else -- bounded concurrency, per-attempt retry, failure + precedence, progress, and resumable interruption -- belongs to + :class:`~dataretrieval.transport.fanout.FanOut`, which Water Data and NGWMN + drive too. This function is now only the NWDC-specific half: what a + sub-request is, and how to read one. + + The broad retry status set is on purpose: NWDC reports a bad query as a 400 + with a ``{"detail": ...}`` envelope, so unlike WQP and StreamStats its 5xx + really is an upstream fault worth re-sending. """ def parse(response: httpx.Response) -> tuple[pd.DataFrame, str | None]: @@ -352,67 +386,32 @@ async def follow(cursor: str, sess: httpx.AsyncClient) -> httpx.Response: def raise_for_status(response: httpx.Response) -> None: _raise_for_status(response, detail_from=_nwdc_error_detail) - # The broad status set on purpose: NWDC reports a bad query as a 400 with a - # ``{"detail": ...}`` envelope, so unlike WQP and StreamStats its 5xx really - # is an upstream fault worth re-sending. Note the cost is multiplied by the - # fan-out -- see MAX_CONCURRENT_REQUESTS. - policy = RetryPolicy.from_env() - async with open_async_client(verify=ssl_check) as client: - semaphore = asyncio.Semaphore(max(1, MAX_CONCURRENT_REQUESTS)) - - async def _one(request: httpx.Request) -> tuple[pd.DataFrame, httpx.Response]: - async def attempt() -> tuple[pd.DataFrame, httpx.Response]: - return await paginate( - request, - parse_response=parse, - follow_up=follow, - client=client, - raise_for_status=raise_for_status, - ) - - # ``retry_async`` owns the gate: the slot is acquired per attempt, - # so a location backing off isn't holding one. A later-page failure - # is intentionally wrapped by ``paginate`` and propagates instead of - # restarting a partially completed walk. - return await retry_async(attempt, policy, gate=semaphore) - - # ``return_exceptions`` so every location is joined before the client - # block exits. Letting the first failure propagate out of the gather - # closed the shared client from under its still-running siblings: a - # location mid-page-walk (or asleep on a ``Retry-After`` backoff) then - # failed with "Cannot send a request, as the client has been closed" on - # a task nobody was awaiting any more -- a spurious error, and an - # unretrieved-exception warning, both attributable to our own teardown. - # The cost is that a fatal error waits for the slowest sibling; that is - # the price of not abandoning in-flight work mid-request. - results = await asyncio.gather( - *(_one(req) for req in requests), return_exceptions=True + async def fetch(args: dict[str, Any]) -> tuple[pd.DataFrame, httpx.Response]: + """One location's full page walk, over the executor's shared client. + + ``active_client()`` is the client :meth:`FanOut._run` published for this + run; borrowing it keeps every location's pages on one connection pool + instead of opening a client per location. + """ + return await paginate( + args["request"], + parse_response=parse, + follow_up=follow, + client=active_client(), + raise_for_status=raise_for_status, ) - # A cancellation or interrupt signal (``CancelledError``, - # ``KeyboardInterrupt`` -- non-``Exception``) wins over any request failure: - # gathering with ``return_exceptions`` captures it like any other result, and - # reporting a sibling's HTTP error instead would swallow the user's stop - # signal. Otherwise raise in input order, so which failure a caller sees - # stays deterministic rather than depending on which location lost the race. - # (Same precedence the chunked fan-out applies -- see ``ChunkedCall._run``.) - failures = [result for result in results if isinstance(result, BaseException)] - for failure in failures: - if not isinstance(failure, Exception): - raise failure - if failures: - raise failures[0] - pairs = [result for result in results if not isinstance(result, BaseException)] - - # Reuse the transport combine helpers: drop empty frames and concat, and fold - # the per-location responses into one (headers from the response with the - # lowest reported remaining quota plus summed response durations), keeping - # the first request's URL as the query identity. - frames = [frame for frame, _ in pairs] - responses = [resp for _, resp in pairs] - return _combine_chunk_frames(frames), _combine_chunk_responses( - responses, str(requests[0].url) - ) + # ``progress_context`` activates the reporter ``FanOut`` ticks into; without + # it Water Use would run the shared executor but print nothing, which is + # what it did when it drove its own gather. + with _progress.progress_context(service="wateruse", target_url=WATERUSE_URL): + return FanOut( + _LocationPlan(requests), + fetch, + RetryPolicy.from_env(), + client_options={"verify": ssl_check}, + default_concurrent=DEFAULT_CONCURRENT_REQUESTS, + ).resume() def _read_csv_page(response: httpx.Response) -> pd.DataFrame: diff --git a/docs/source/architecture/decisions/0006-service-neutral-transport.rst b/docs/source/architecture/decisions/0006-service-neutral-transport.rst index f6d9fd52..acd30bf8 100644 --- a/docs/source/architecture/decisions/0006-service-neutral-transport.rst +++ b/docs/source/architecture/decisions/0006-service-neutral-transport.rst @@ -4,7 +4,9 @@ ADR 0006: Use a service-neutral transport layer Status ------ -Accepted +Accepted. The clause assigning resumable ``ChunkedCall`` state to OGC is +superseded by :doc:`0008-fan-out-execution`, which moves fan-out *execution* +into transport and leaves chunk *planning* in OGC. The rest stands. Context ------- diff --git a/docs/source/architecture/decisions/0008-fan-out-execution.rst b/docs/source/architecture/decisions/0008-fan-out-execution.rst new file mode 100644 index 00000000..effae8d3 --- /dev/null +++ b/docs/source/architecture/decisions/0008-fan-out-execution.rst @@ -0,0 +1,117 @@ +ADR 0008: Separate fan-out execution from chunk planning +======================================================== + +Status +------ + +Accepted. Supersedes the clause of :doc:`0006-service-neutral-transport` +assigning "resumable ``ChunkedCall`` state" to OGC's protocol concerns; the rest +of ADR 0006 stands. + +Context +------- + +Two services turn one logical query into several requests, for unrelated +reasons. A Water Data or NGWMN query whose URL exceeds the server's byte limit +is split along its multi-value axes. A Water Use query naming several locations +is split because the NWDC accepts one ``location=`` per request -- its URLs run +around 63 bytes against an 8000-byte budget, so the byte limit has nothing to do +with it. + +Chunking is how you divide the data structurally; fan-out is how you distribute +the work operationally. The two are orthogonal, and only the first is protocol +knowledge: dividing a query needs the byte budget, the CQL2 grammar, and which +parameters are list-valued, while distributing the pieces needs none of it. + +The package had not drawn that line. ``ChunkPlan`` (division) and +``ChunkedCall`` (distribution) sat side by side in ``dataretrieval.ogc`` as +siblings, and ADR 0006 grouped them together deliberately. That grouping was +correct while a byte plan was the only thing anyone fanned out over. It stopped +being correct once Water Use fanned out too: unable to reach an OGC-internal +executor, ``wateruse._fan_out`` re-implemented the semaphore, the +``asyncio.gather``, and the cancellation-beats-HTTP-error failure precedence, +with a comment naming ``ChunkedCall._run`` as the original. One subtle rule, +two copies, synchronized by prose. + +The duplicate was not merely redundant. It lacked resume, so a rate limit +partway through discarded every location that had already succeeded -- against +an hourly quota, on fan-outs that reach into the hundreds. It reported no +progress. And it read its own module-global concurrency cap, so a user setting +``API_USGS_CONCURRENT`` to be polite to the service found one adapter ignoring +them. + +Decision +-------- + +``dataretrieval.transport.fanout`` owns fan-out execution for every service: +bounded concurrency, per-attempt retry, deterministic failure precedence, sparse +completion tracking, and resume. It names no protocol concept. An adapter +supplies a ``FanOutPlan`` and an ``async def fetch(args) -> (df, response)``. + +``FanOutPlan`` is a ``Protocol`` of exactly three members -- ``total``, +``canonical_url``, and ``iter_sub_args()`` -- which is the whole surface the +executor ever touched. It is structural rather than nominal because its two +implementations share an interface and no implementation whatsoever: +``ChunkPlan`` derives sub-requests from a byte budget over multi-value axes, and +a Water Use plan lists locations the caller already named separately. Neither +has anything the other could inherit, so an abstract base would be ceremony. + +``dataretrieval.ogc`` keeps chunk planning: the byte budget, the axis +partitioning, the CQL2 filter split, the ``parallel_chunks`` dial. Those are +division, and division is protocol-specific. + +The interruption taxonomy moves to ``dataretrieval.interruptions``, a top-level +leaf, for the reason ADR 0006 gives for ``combining``, ``progress``, and +``credentials``: adapters need it whether or not they went through transport, +and an exception taxonomy is not HTTP execution policy. Its base class is +renamed ``FanOutInterrupted``, since Water Use raises it without chunking +anything. ``ChunkInterrupted`` is retained as a permanent alias of the same +class object -- not a shim scheduled for deletion -- because it is the name +published in the user guide and caught in user code. The subclasses +(``QuotaExhausted``, ``ServiceInterrupted``) were already neutral and are +unchanged. + +Concurrency is one general setting with per-service defaults. +``API_USGS_CONCURRENT`` applies to every fanned-out call; a service may declare +a different default for when it is unset. The precedence is deliberate: an +explicitly set environment variable outranks a service default, never the +reverse. A service that could override the general setting would make +``API_USGS_CONCURRENT=1`` a lie. Service defaults say "absent instruction, this +service prefers N"; they do not say "this service knows better than you". + +Consequences +------------ + +- Water Use gains resume, progress reporting, and the shared concurrency + setting, and sheds roughly 75 lines of duplicated orchestration. +- One implementation of failure precedence, so cancellation-beats-error and + deterministic failure ordering cannot drift between services. +- **Breaking:** a Water Use fan-out interrupted by a 5xx or 429 now raises + ``ServiceInterrupted`` / ``QuotaExhausted`` rather than ``ServiceUnavailable`` + / ``RateLimited``. Both remain ``DataRetrievalError``, so broad handlers are + unaffected, but a narrow ``except ServiceUnavailable`` around a Water Use call + must widen. This is convergence, not novelty -- it is what the OGC getters + have always done -- and it is what makes the failure resumable. +- **Breaking:** ``wateruse.MAX_CONCURRENT_REQUESTS`` is removed in favor of + ``API_USGS_CONCURRENT`` and ``wateruse.DEFAULT_CONCURRENT_REQUESTS``. +- Resume re-issues a failed location's entire page walk, so pages fetched before + the failure are fetched again. This already applied to OGC -- a partial walk + never enters the completion map -- and is a cost, not a correctness problem. +- Water Use frames carry ``huc12_id``, not ``id``, so ``_combine_chunk_frames`` + concatenates them without deduplicating. Correct, because locations partition + by construction, but the executor's dedup safety net does not apply there. +- ``transport`` is no longer purely leaf-shaped: ``fanout`` is a composite that + drives retry, pagination-borrowed clients, and combining. It remains HTTP + execution policy, which is the test the package applies. + +Compliance +---------- + +``tests/architecture_test.py`` asserts that ``wateruse`` contains no +``asyncio.gather``/``Semaphore``/``TaskGroup``, so the duplication cannot +return; that both ``ChunkPlan`` and the Water Use plan satisfy ``FanOutPlan``, +including that ``iter_sub_args()`` is stable across passes and agrees with +``total``, since resume keys completed work by position; that the Water Use plan +does not inherit ``ChunkPlan``; and that an interruption taxonomy does not +reappear inside ``transport``. Adapter tests cover Water Use resume re-issuing +only unfinished locations, progress ticks, and the concurrency precedence rule. diff --git a/docs/source/architecture/decisions/index.rst b/docs/source/architecture/decisions/index.rst index 92a9d3f1..006f513c 100644 --- a/docs/source/architecture/decisions/index.rst +++ b/docs/source/architecture/decisions/index.rst @@ -23,4 +23,5 @@ records sequentially. 0004-error-retry-resume 0005-legacy-nwis 0006-service-neutral-transport + 0008-fan-out-execution template diff --git a/docs/source/architecture/index.rst b/docs/source/architecture/index.rst index 09672a6e..6225c560 100644 --- a/docs/source/architecture/index.rst +++ b/docs/source/architecture/index.rst @@ -151,8 +151,10 @@ contracts; consistency alone is not sufficient reason for a breaking change. Failed requests derive from ``dataretrieval.DataRetrievalError``. Callers can inspect ``status_code``, ``retry_after``, and ``retryable`` without knowing the -concrete subtype. OGC calls may raise ``ChunkInterrupted`` subclasses carrying a -resumable call handle and completed partial state. +concrete subtype. A fanned-out call -- an over-large OGC request, or a Water Use +query naming several locations -- may raise ``FanOutInterrupted`` subclasses +(formerly, and still aliased as, ``ChunkInterrupted``) carrying a resumable call +handle and completed partial state. The public surface is defined by package/module exports and documentation. Underscore-prefixed symbols are implementation details even where existing diff --git a/docs/source/reference/exceptions.rst b/docs/source/reference/exceptions.rst index 1a963187..7b5c2909 100644 --- a/docs/source/reference/exceptions.rst +++ b/docs/source/reference/exceptions.rst @@ -7,16 +7,23 @@ dataretrieval.exceptions :members: :show-inheritance: -Resumable chunk interruptions +Resumable fan-out interruptions ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ -These are raised when a transparently-chunked request is interrupted -mid-stream; the completed work is preserved and ``exc.call.resume()`` continues -it. They are defined in ``dataretrieval.ogc.interruptions`` (they carry -pandas/httpx state) but are importable from the top level, e.g. -``from dataretrieval import ChunkInterrupted``. +These are raised when a fanned-out request is interrupted mid-stream; the +completed work is preserved and ``exc.call.resume()`` continues it. They are +defined in ``dataretrieval.interruptions`` (they carry pandas/httpx state) but +are importable from the top level, e.g. +``from dataretrieval import FanOutInterrupted``. -.. autoclass:: dataretrieval.ChunkInterrupted +``ChunkInterrupted`` is a permanent alias of ``FanOutInterrupted`` -- the same +class object under the name it was first published as -- so ``except +ChunkInterrupted`` and ``except FanOutInterrupted`` are the same handler. The +base class is named for the fan-out rather than for chunking because a Water Use +call fans out without dividing anything: the NWDC simply accepts one location +per request. + +.. autoclass:: dataretrieval.FanOutInterrupted :members: :show-inheritance: diff --git a/docs/source/userguide/errors.rst b/docs/source/userguide/errors.rst index 28da515f..f75c223d 100644 --- a/docs/source/userguide/errors.rst +++ b/docs/source/userguide/errors.rst @@ -71,31 +71,38 @@ honoring the server's ``Retry-After`` hint when present: raise time.sleep(e.retry_after or 2 ** attempt) -Resume a large Water Data request -================================= +Resume an interrupted request +============================= + +Some requests become several: the Water Data and NGWMN getters split an +over-large request into chunks, and a Water Use call with several locations +becomes one request per location. When a transient failure interrupts one +mid-stream, the work already completed is preserved: catch +``FanOutInterrupted`` and call ``exc.call.resume()`` once the condition clears +-- only the unfinished sub-requests are re-issued. -The Water Data getters transparently split an over-large request into chunks. -When a transient failure interrupts one mid-stream, the work already completed -is preserved: catch ``ChunkInterrupted`` and call ``exc.call.resume()`` once the -condition clears -- only the unfinished sub-requests are re-issued. +(``ChunkInterrupted`` is the same class under its original name; either works.) .. code-block:: python import time - from dataretrieval import ChunkInterrupted + from dataretrieval import FanOutInterrupted from dataretrieval.waterdata import get_daily try: df, md = get_daily(monitoring_location_id=long_list_of_sites) - except ChunkInterrupted as exc: + except FanOutInterrupted as exc: while True: time.sleep(exc.retry_after or 5 * 60) try: df, md = exc.call.resume() break - except ChunkInterrupted as again: + except FanOutInterrupted as again: exc = again +The same loop works for ``wateruse.get_wateruse`` with a list of states, +counties, or HUCs. + Chunk a large request more finely ================================= diff --git a/tests/architecture_test.py b/tests/architecture_test.py index c4d4c5b7..365178f3 100644 --- a/tests/architecture_test.py +++ b/tests/architecture_test.py @@ -372,6 +372,11 @@ def test_transport_is_execution_policy_only() -> None: misplaced = { "dataretrieval/transport/progress.py", "dataretrieval/transport/combining.py", + # An exception taxonomy is not HTTP execution policy either. ``fanout`` + # raises ``FanOutInterrupted`` and belongs here; defining it here would + # not, since adapters catch it whether or not they went through + # transport. + "dataretrieval/transport/interruptions.py", } present = { path @@ -489,3 +494,72 @@ def visit(module: str, path: tuple[str, ...]) -> None: for module in graph: visit(module, ()) + + +def test_wateruse_does_not_reimplement_fan_out_orchestration() -> None: + """Water Use must drive its locations through the shared fan-out executor. + + It previously ran its own ``asyncio.gather`` with a private semaphore and a + hand-copied failure-precedence rule, kept in sync with ``FanOut`` by a + comment. Two copies of that rule is how they drift, and the duplicate lost + resume, progress, and the shared concurrency setting. Assert the duplication + cannot quietly return. + """ + source = (PACKAGE_ROOT / "wateruse.py").read_text(encoding="utf-8") + tree = ast.parse(source) + offenders = { + f"{node.value.id}.{node.attr}" + for node in ast.walk(tree) + if isinstance(node, ast.Attribute) + and isinstance(node.value, ast.Name) + and node.value.id == "asyncio" + and node.attr in {"gather", "Semaphore", "wait", "TaskGroup"} + } + assert not offenders, ( + "Water Use re-implemented fan-out orchestration instead of using " + f"transport.fanout.FanOut: {sorted(offenders)}" + ) + + +def test_fan_out_plans_satisfy_the_plan_protocol() -> None: + """Every plan implementation must carry the three members ``FanOut`` drives. + + ``FanOutPlan`` is structural, so nothing forces an implementation to be + complete at definition time -- a missing ``canonical_url`` would surface as + an ``AttributeError`` mid-fan-out, after requests had already been issued. + Check both implementations up front instead. They are deliberately unrelated + by inheritance: chunking divides structurally, and a Water Use plan divides + nothing, so there is no shared base to inherit. + """ + import httpx + + from dataretrieval.ogc.planning import ChunkPlan + from dataretrieval.wateruse import _LocationPlan + + def _build(**args: object) -> httpx.Request: + return httpx.Request("GET", "https://example.invalid/items", params=args) + + plans = [ + ChunkPlan({"sites": ["a", "b"]}, _build, url_limit=8000), + _LocationPlan([httpx.Request("GET", "https://example.invalid/data")]), + ] + for plan in plans: + name = type(plan).__name__ + assert isinstance(plan.total, int), f"{name}.total is not an int" + assert plan.canonical_url is None or isinstance(plan.canonical_url, str), ( + f"{name}.canonical_url is neither str nor None" + ) + sub_args = list(plan.iter_sub_args()) + assert len(sub_args) == plan.total, ( + f"{name}.iter_sub_args() yielded {len(sub_args)}, total says {plan.total}" + ) + assert all(isinstance(item, dict) for item in sub_args), ( + f"{name}.iter_sub_args() must yield kwargs dicts" + ) + # Order is load-bearing for resume: a second pass must match the first. + assert [d.keys() for d in plan.iter_sub_args()] == [d.keys() for d in sub_args] + + assert not issubclass(_LocationPlan, ChunkPlan), ( + "A Water Use plan must satisfy FanOutPlan structurally, not by " + "inheriting ChunkPlan -- it has no byte budget or axes to inherit." + ) diff --git a/tests/waterdata_chunking_test.py b/tests/waterdata_chunking_test.py index 06657821..4ca74411 100644 --- a/tests/waterdata_chunking_test.py +++ b/tests/waterdata_chunking_test.py @@ -1913,7 +1913,7 @@ def test_retryable_skips_wrapped_midpagination_transient(): def test_retry_transient_then_recovers(monkeypatch): - monkeypatch.setattr(_chunking.asyncio, "sleep", _aiozero) + monkeypatch.setattr(_retry_mod.asyncio, "sleep", _aiozero) calls = {"n": 0} async def afn(): @@ -1928,7 +1928,7 @@ async def afn(): def test_retry_exhausted_reraises(monkeypatch): - monkeypatch.setattr(_chunking.asyncio, "sleep", _aiozero) + monkeypatch.setattr(_retry_mod.asyncio, "sleep", _aiozero) calls = {"n": 0} async def afn(): @@ -1943,7 +1943,7 @@ async def afn(): def test_retry_non_retryable_not_retried(monkeypatch): slept: list[float] = [] - monkeypatch.setattr(_chunking.asyncio, "sleep", _recording_sleep(slept)) + monkeypatch.setattr(_retry_mod.asyncio, "sleep", _recording_sleep(slept)) calls = {"n": 0} async def afn(): @@ -1958,7 +1958,7 @@ async def afn(): def test_retry_long_retry_after_escalates(monkeypatch): slept: list[float] = [] - monkeypatch.setattr(_chunking.asyncio, "sleep", _recording_sleep(slept)) + monkeypatch.setattr(_retry_mod.asyncio, "sleep", _recording_sleep(slept)) calls = {"n": 0} async def afn(): @@ -1974,7 +1974,7 @@ async def afn(): def test_retry_transient_then_success(monkeypatch): - monkeypatch.setattr(_chunking.asyncio, "sleep", _aiozero) + monkeypatch.setattr(_retry_mod.asyncio, "sleep", _aiozero) calls = {"n": 0} async def afn(): @@ -1994,7 +1994,7 @@ def test_chunker_retries_transient_then_completes(monkeypatch): """A transient on one sub-request is retried transparently; the decorated call completes with no ChunkInterrupted.""" monkeypatch.setenv("API_USGS_RETRIES", "3") - monkeypatch.setattr(_chunking.asyncio, "sleep", _aiozero) + monkeypatch.setattr(_retry_mod.asyncio, "sleep", _aiozero) state = {"failed": False} async def fetch(args): @@ -2033,7 +2033,7 @@ def test_chunker_exhausted_retries_still_resumable(monkeypatch): """When retries are exhausted the failure still surfaces as a resumable ChunkInterrupted — retries don't swallow the escape hatch.""" monkeypatch.setenv("API_USGS_RETRIES", "2") - monkeypatch.setattr(_chunking.asyncio, "sleep", _aiozero) + monkeypatch.setattr(_retry_mod.asyncio, "sleep", _aiozero) attempts = {"n": 0} async def fetch(args): @@ -2054,7 +2054,7 @@ def test_async_fan_out_retries_transient_then_completes(monkeypatch): """The parallel path retries a transient sub-request and completes.""" monkeypatch.setenv("API_USGS_RETRIES", "3") - monkeypatch.setattr(_chunking.asyncio, "sleep", _aiozero) + monkeypatch.setattr(_retry_mod.asyncio, "sleep", _aiozero) state = {"failed": False} async def fetch_async(args): @@ -2073,7 +2073,7 @@ def test_async_fan_out_surfaces_fatal_over_transient(monkeypatch): being masked behind a resumable interruption from a transient sibling.""" monkeypatch.setenv("API_USGS_RETRIES", "2") - monkeypatch.setattr(_chunking.asyncio, "sleep", _aiozero) + monkeypatch.setattr(_retry_mod.asyncio, "sleep", _aiozero) async def fetch_async(args): # One chunk carries a deterministic programmer error; the rest are diff --git a/tests/wateruse_test.py b/tests/wateruse_test.py index b6ee4e00..83cafd5d 100644 --- a/tests/wateruse_test.py +++ b/tests/wateruse_test.py @@ -12,6 +12,7 @@ import dataretrieval from dataretrieval import wateruse +from dataretrieval.transport import fanout as _fanout from dataretrieval.utils import BaseMetadata from dataretrieval.wateruse import _next_page_url, _resolve_locations, get_wateruse @@ -270,8 +271,8 @@ def test_multiple_states_fan_out_preserves_input_order(httpx_mock): def test_fan_out_is_serial_when_concurrency_is_one(httpx_mock, monkeypatch): - """``MAX_CONCURRENT_REQUESTS = 1`` still fans out correctly (serial path).""" - monkeypatch.setattr(wateruse, "MAX_CONCURRENT_REQUESTS", 1) + """``API_USGS_CONCURRENT=1`` still fans out correctly (serial path).""" + monkeypatch.setenv("API_USGS_CONCURRENT", "1") httpx_mock.add_response( method="GET", url=re.compile(r".*location=stateCd%3ARI.*"), text=_CSV_P1 ) @@ -311,7 +312,14 @@ def test_fan_out_surfaces_final_rate_limit_header(httpx_mock): def test_fan_out_failure_never_returns_partial_data(httpx_mock): - """A failed location aborts the call even when another location succeeded.""" + """A failed location aborts the call even when another location succeeded. + + The completed sibling is not returned as though the call had succeeded -- + it is carried on the raised interruption for ``resume()`` instead. Water Use + reports ``ServiceInterrupted`` rather than the bare ``ServiceUnavailable`` + it raised before sharing the fan-out executor: the same upstream 503, now + resumable. + """ httpx_mock.add_response( method="GET", url=re.compile(r".*location=stateCd%3ARI.*"), @@ -324,9 +332,16 @@ def test_fan_out_failure_never_returns_partial_data(httpx_mock): json={"detail": "temporarily unavailable"}, ) - with pytest.raises(dataretrieval.ServiceUnavailable): + with pytest.raises(dataretrieval.ServiceInterrupted) as excinfo: get_wateruse(model="wu-public-supply-wd", state=["RI", "WI"]) + # The 503 is still the reported cause, and the successful location survives + # on the exception rather than being passed off as the whole answer. + assert isinstance(excinfo.value.__cause__, dataretrieval.ServiceUnavailable) + assert excinfo.value.completed_chunks == 1 + assert excinfo.value.total_chunks == 2 + assert len(excinfo.value.partial_frame) == 2 + # --- _resolve_locations unit tests (no HTTP) ------------------------------- @@ -562,7 +577,7 @@ async def open_mock_client(**overrides): ) as client: yield client - monkeypatch.setattr(wateruse, "open_async_client", open_mock_client) + monkeypatch.setattr(_fanout, "open_async_client", open_mock_client) requests = [ httpx.Request("GET", wateruse.WATERUSE_URL, params={"location": location}) @@ -591,3 +606,129 @@ def test_next_page_url_rejects_cross_host_link(): # other failure rather than seeing a bare RuntimeError. with pytest.raises(dataretrieval.DataRetrievalError, match="outside.example"): _next_page_url(response) + + +# --- capabilities Water Use gained by sharing the fan-out executor ---------- + + +def test_interrupted_fan_out_resumes_only_the_unfinished_locations(httpx_mock): + """A rate-limited location is resumable; completed siblings are not re-fetched. + + Before Water Use shared the executor, a 429 anywhere in the fan-out + discarded every location that had already succeeded. That is the whole + reason a multi-location pull needed re-running from scratch against an + hourly quota. + """ + httpx_mock.add_response( + method="GET", url=re.compile(r".*location=stateCd%3ARI.*"), text=_CSV_P1 + ) + # WI is rate-limited on the first pass, then succeeds once resumed. + httpx_mock.add_response( + method="GET", + url=re.compile(r".*location=stateCd%3AWI.*"), + status_code=429, + json={"detail": "rate limited"}, + is_reusable=False, + ) + httpx_mock.add_response( + method="GET", url=re.compile(r".*location=stateCd%3AWI.*"), text=_CSV_P2 + ) + + with pytest.raises(dataretrieval.QuotaExhausted) as excinfo: + get_wateruse(model="wu-public-supply-wd", state=["RI", "WI"]) + + interrupted = excinfo.value + assert interrupted.completed_chunks == 1 + assert interrupted.total_chunks == 2 + requests_before = len(httpx_mock.get_requests()) + + df, _ = interrupted.call.resume() + + # Only WI was re-issued; RI's completed frame carried across the resume. + assert len(httpx_mock.get_requests()) == requests_before + 1 + assert len(df) == 3 + + +def test_fan_out_honors_the_general_concurrency_setting(monkeypatch): + """``API_USGS_CONCURRENT`` outranks this service's default. + + A user dialing concurrency down to be polite must not find Water Use + quietly ignoring them -- the defect that motivated consolidating the knob. + """ + monkeypatch.setenv("API_USGS_CONCURRENT", "7") + assert _fanout._resolve_concurrency(wateruse.DEFAULT_CONCURRENT_REQUESTS) == 7 + + monkeypatch.delenv("API_USGS_CONCURRENT", raising=False) + assert ( + _fanout._resolve_concurrency(wateruse.DEFAULT_CONCURRENT_REQUESTS) + == wateruse.DEFAULT_CONCURRENT_REQUESTS + ) + # The service default is deliberately below the package-wide 32. + assert wateruse.DEFAULT_CONCURRENT_REQUESTS < _fanout._CONCURRENCY_DEFAULT + + +def test_fan_out_reports_progress(httpx_mock, monkeypatch): + """The fan-out ticks the progress reporter, which it never did standalone.""" + seen = [] + + class _Recorder: + def set_chunks(self, total): + seen.append(("chunks", total)) + + def start_chunk(self, completed): + seen.append(("chunk", completed)) + + def set_rate_remaining(self, remaining, limit=None): + pass + + def add_page(self, rows): + seen.append(("page", rows)) + + monkeypatch.setattr(_fanout._progress, "current", lambda: _Recorder()) + httpx_mock.add_response( + method="GET", url=re.compile(r".*location=stateCd%3ARI.*"), text=_CSV_P1 + ) + httpx_mock.add_response( + method="GET", url=re.compile(r".*location=stateCd%3AWI.*"), text=_CSV_P2 + ) + + get_wateruse(model="wu-public-supply-wd", state=["RI", "WI"]) + + assert ("chunks", 2) in seen + assert ("chunk", 1) in seen and ("chunk", 2) in seen + + +def test_mid_page_walk_transient_is_still_resumable(httpx_mock): + """A 429 on page 2+ of a location must still be a resumable interruption. + + ``paginate`` re-wraps a later-page failure as a plain ``DataRetrievalError`` + (page 1's status check sits outside its ``try``), so the typed cause is only + reachable through ``__cause__``. ``_classify_chunk_error`` walks that chain + for exactly this reason; were it a single ``isinstance`` check, a mid-walk + rate limit would escape as a bare error and lose ``.call.resume()`` -- + inconsistently, since page 1 would still be resumable. + """ + httpx_mock.add_response( + method="GET", + url=re.compile(r".*location=stateCd%3ARI(?!.*cursor).*"), + text=_CSV_P1, + headers={ + "Link": '; rel="next"' + }, + ) + httpx_mock.add_response( + method="GET", + url=re.compile(r".*cursor=x.*"), + status_code=429, + json={"detail": "rate limited"}, + ) + httpx_mock.add_response( + method="GET", url=re.compile(r".*location=stateCd%3AWI.*"), text=_CSV_P2 + ) + + with pytest.raises(dataretrieval.QuotaExhausted) as excinfo: + get_wateruse(model="wu-public-supply-wd", state=["RI", "WI"]) + + assert excinfo.value.call is not None + assert excinfo.value.completed_chunks == 1 + assert excinfo.value.total_chunks == 2