From 31b014205d6d7c6200cb6beb1c26757d63d999fc Mon Sep 17 00:00:00 2001 From: Timothy Hodson <34148978+thodson-usgs@users.noreply.github.com> Date: Thu, 6 Aug 2026 13:39:36 -0500 Subject: [PATCH 1/2] feat(transport): bounded retry for active services, over an API-neutral layer (#350) * feat(transport): bounded retry for active services, over an API-neutral layer WQP, NLDI, StreamStats, and Water Use now retry transient failures instead of surfacing the first one. That is a resilience change, not a refactor, so it leads here; the layering change that made it tractable follows. Retry costs latency and quota on failing requests, so it is bounded on two independent axes and narrowed to failures a later attempt could survive. API_USGS_STALL_TIMEOUT (new; default 60 s, 0 disables) bounds how long a call may go without receiving any data. API_USGS_RETRIES counts attempts, not seconds, so on its own four retries of a request that times out after a minute is four silent minutes. Progress restarts the budget -- a page received, or a queued sub-request acquiring its concurrency slot, credited as the wait it was rather than restamped -- and an attempt already in flight is never interrupted. The first retry is never withheld, so one slow attempt cannot disable retry by itself. A dead connection costs about two read timeouts rather than five attempts. Which statuses are re-sent is per-adapter. WQP answers an over-large query with a 500 and StreamStats answers out-of-network coordinates with one, so those one-shot adapters re-send only for 429/502/503/504. The Water Data OGC API is a query interface where a 500 is an upstream fault, so the chunker keeps re-sending for every 5xx, as it always has. Failures already settled are not retried: an unsupported scheme, a request we built wrong, or a hostname the resolver rejects outright. A temporary resolver failure (EAI_AGAIN) stays retryable. Backoff always includes jitter, including on a server-named Retry-After, so sub-requests handed one hint do not wake in lockstep and a hint of 0 cannot become a zero-delay re-send. An unusable setting raises ConfigurationError -- both a DataRetrievalError and a ValueError -- rather than escaping a request path untyped. Measured against the live API: a 4-state, 30-year get_daily over 800 sites at parallel_chunks(1) runs 91.8 s and returns 581,070 rows with a worst inter-page silence of 12.1 s, so the budget does not threaten long successful queries. The layering half adds dataretrieval.transport, an internal API-neutral execution layer owning guarded client lifecycle and timeout defaults, host-scoped authentication, cursor pagination, bounded retry, response aggregation, progress, and sync-over-async dispatch. dataretrieval.ogc keeps its protocol concerns: dialects, CQL2, request construction, feature shaping, URL-byte chunk planning, resumable ChunkedCall state, and interruption types. Before this, generic execution behavior lived under OGC even where non-OGC services used it, so Water Use depended on private protocol modules and retry policy was uneven across services; there is now one policy to reason about. transport.liveness is a stdlib-only leaf recording when data last arrived, so the page loop that observes progress and the retry loop that acts on it depend on it rather than on each other. Architecture fitness functions enforce the dependency direction, an acyclic transport graph, and Water Use's isolation from OGC; ADR 0006 records the decision. Compatibility: public imports, service signatures, return shapes, metadata, deprecations, exception types, OGC chunking/resume behavior, and the four-symbol OGC facade are unchanged, and utils.query keeps its exact signature and performs no retry. Private compatibility aliases are kept where a consumer exists. Two modules were removed rather than aliased, since nothing imported them: ogc.progress and ogc.combining, now transport.progress and transport.combining. ogc.retry keeps only its OGC interruption classifiers. Also pins the CI test step to bash on every OS. Windows defaults to PowerShell, which does not halt on a failing native command and takes the step's exit status from the last one, so a coverage report following a failed pytest reported success -- every Windows test failure in this repository has been invisible. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_015UVSiGgUyDRSbqNWM5HBbL * Simplify pass * refactor: move presentation and credential policy out of transport Splits three concerns out of the shared transport layer and closes two credential-leak paths in server-supplied pagination links. The layer was named "API-neutral" but held `api.waterdata.usgs.gov`, read `API_USGS_PAT`, pinned `x-ratelimit-remaining`, and printed a USGS signup URL. It is neutral across USGS *services*, not across HTTP APIs, and the aspirational name invited generality nobody needs. ADR 0006 now says so plainly and is renamed to match. Two modules were in transport only because they had to leave `ogc/` during the earlier extraction: - `progress.py` is terminal presentation (Jupyter detection, status-line rewriting, broken-pipe handling), called *from* transport rather than part of it, and the sole reason a `progress -> http` edge existed. - `combining.py` is DataFrame assembly, consumed by `ogc/planning` and `wateruse` for reasons unrelated to HTTP. Both move to top-level leaves. Transport goes 1290 -> 766 lines and 7 -> 5 modules, and `http`/`liveness` become leaves. A new `credentials.py` leaf owns every answer about the API key. The code that attaches a credential and the code that strips it back off have to agree on which host is authorized, and the way they stop agreeing is a second copy of the host string. `waterdata/utils`, `ogc/policy`, and `ngwmn` each carried their own `BASE_URL` spelling of that same authority -- two of them with a comment documenting the duplication as deliberate -- so they now derive it from the one definition. Closes two ways a poisoned response body reached a credential: - `accepts_api_key` matched on host alone, so `http://` on the authorized host sent the key in cleartext. It now requires https. - `ogc/engine` checked the next-link host but not its userinfo, and `waterdata/ratings` followed STAC `next` hrefs with no check at all. httpx derives `Authorization: Basic` from userinfo, so a link carrying `user:pass@` minted a credential the caller never configured and sent it beside the real API key -- past the host check, which passes in exactly that case. The credential fitness function matched the quoted bare host, so the `https://`-prefixed form slipped past it and it reported success with three copies live. It now walks AST string values, excluding docstrings so prose naming the service is not mistaken for a second source of truth. Every new test was verified to fail against the unfixed source. Co-Authored-By: Claude Opus 5 * fix(ogc): don't offer a failure we refuse to retry as resumable "Should we retry this?" and "can the caller resume it?" are the same question asked twice, and the two answers disagreed. Retry already declines to re-send a failure no later attempt could survive -- a bad URL scheme, a malformed request, a hostname the resolver rejects outright. The interruption classifier mapped every httpx error to ServiceInterrupted regardless, so the caller got a .call.resume() whose every attempt fails identically, with the NetworkError that actually explained the problem buried underneath it. Both answers now come from one predicate in transport. A deterministic failure classifies as unrecognized, which is the existing "re-raise raw" path, so the caller sees the real error. The test asserts both answers on the same failures, so they cannot drift apart again. Note the chain shape matters: our wrapper raises with `from`, so the chunker's explicit-link walk reaches the httpx error, whose implicit links then lead to the resolver code -- a temporary resolver failure stays both retryable and resumable, decided only by the errno. Addresses finding 3 of the chunking review; finding 4 (a 5xx sibling masking a 429's Retry-After) remains open. Co-Authored-By: Claude Opus 5 * Update NEWS.md --------- Co-authored-by: Claude Opus 5 (1M context) --- .github/workflows/python-package.yml | 7 + NEWS.md | 2 + dataretrieval/__init__.py | 2 + dataretrieval/{ogc => }/combining.py | 14 +- dataretrieval/credentials.py | 97 +++ dataretrieval/exceptions.py | 18 +- dataretrieval/ngwmn.py | 6 +- dataretrieval/nldi.py | 4 +- dataretrieval/ogc/__init__.py | 7 +- dataretrieval/ogc/chunking.py | 63 +- dataretrieval/ogc/engine.py | 260 ++------ dataretrieval/ogc/errors.py | 77 +-- dataretrieval/ogc/planning.py | 11 +- dataretrieval/ogc/policy.py | 4 +- dataretrieval/ogc/requests.py | 11 +- dataretrieval/ogc/retry.py | 348 +---------- dataretrieval/{ogc => }/progress.py | 52 +- dataretrieval/streamstats.py | 11 +- dataretrieval/transport/__init__.py | 7 + dataretrieval/transport/http.py | 96 +++ dataretrieval/transport/liveness.py | 58 ++ dataretrieval/transport/pagination.py | 131 ++++ dataretrieval/transport/retry.py | 445 ++++++++++++++ dataretrieval/transport/sync.py | 29 + dataretrieval/utils.py | 233 ++++---- dataretrieval/waterdata/api.py | 14 +- dataretrieval/waterdata/ratings.py | 42 +- dataretrieval/waterdata/stats.py | 35 +- dataretrieval/waterdata/utils.py | 7 +- dataretrieval/wateruse.py | 138 +++-- dataretrieval/wqp.py | 8 +- .../decisions/0003-dependency-direction.rst | 12 +- .../decisions/0004-error-retry-resume.rst | 7 +- .../0006-service-neutral-transport.rst | 110 ++++ docs/source/architecture/decisions/index.rst | 1 + docs/source/architecture/index.rst | 86 ++- tests/architecture_test.py | 152 ++++- tests/conftest.py | 33 +- tests/headers_host_scoping_test.py | 36 ++ tests/nldi_test.py | 13 + tests/streamstats_test.py | 33 + tests/transport_test.py | 562 ++++++++++++++++++ tests/utils_test.py | 31 + tests/waterdata_chunking_test.py | 25 +- tests/waterdata_progress_test.py | 35 +- tests/waterdata_ratings_test.py | 52 ++ tests/waterdata_utils_test.py | 42 +- tests/wateruse_test.py | 179 +++++- tests/wqp_test.py | 18 + 49 files changed, 2680 insertions(+), 984 deletions(-) rename dataretrieval/{ogc => }/combining.py (94%) create mode 100644 dataretrieval/credentials.py rename dataretrieval/{ogc => }/progress.py (86%) create mode 100644 dataretrieval/transport/__init__.py create mode 100644 dataretrieval/transport/http.py create mode 100644 dataretrieval/transport/liveness.py create mode 100644 dataretrieval/transport/pagination.py create mode 100644 dataretrieval/transport/retry.py create mode 100644 dataretrieval/transport/sync.py create mode 100644 docs/source/architecture/decisions/0006-service-neutral-transport.rst create mode 100644 tests/transport_test.py diff --git a/.github/workflows/python-package.yml b/.github/workflows/python-package.yml index ade76452..c56bdf39 100644 --- a/.github/workflows/python-package.yml +++ b/.github/workflows/python-package.yml @@ -53,6 +53,7 @@ jobs: from pathlib import Path import dataretrieval + import dataretrieval.transport from dataretrieval import ngwmn, waterdata, wateruse from dataretrieval.ogc import engine @@ -105,6 +106,12 @@ jobs: python -m pip install --upgrade pip pip install .[test,nldi] - name: Test with pytest and report coverage + # Pinned to bash on every OS. The default Windows shell is PowerShell, + # which does not stop on a failing native command and takes the step's + # exit code from the last one -- so a pytest failure was masked by the + # coverage report that followed it, and the Windows matrix reported + # success while tests were red. + shell: bash run: | coverage run -m pytest tests/ coverage report -m diff --git a/NEWS.md b/NEWS.md index 861d2c5c..a9a63718 100644 --- a/NEWS.md +++ b/NEWS.md @@ -1,3 +1,5 @@ +**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. **08/02/2026:** Fixed source-distribution and wheel package discovery so the `dataretrieval.ogc` and `dataretrieval.waterdata` subpackages are included in installed artifacts. CI now builds and installs the wheel outside the source checkout before importing the core service modules. Added an architecture baseline, initial decision records, and executable dependency-direction guardrails for the existing modular-monolith boundaries. diff --git a/dataretrieval/__init__.py b/dataretrieval/__init__.py index 469fe0f5..4226e247 100644 --- a/dataretrieval/__init__.py +++ b/dataretrieval/__init__.py @@ -32,6 +32,7 @@ __version__ = "version-unknown" from dataretrieval.exceptions import ( + ConfigurationError, DataRetrievalError, HTTPError, NetworkError, @@ -84,6 +85,7 @@ # error taxonomy (canonical home: ``dataretrieval.exceptions``), re-exported # so callers can ``except dataretrieval.DataRetrievalError`` "exceptions", + "ConfigurationError", "DataRetrievalError", "HTTPError", "NetworkError", diff --git a/dataretrieval/ogc/combining.py b/dataretrieval/combining.py similarity index 94% rename from dataretrieval/ogc/combining.py rename to dataretrieval/combining.py index be4366c1..92e32f7a 100644 --- a/dataretrieval/ogc/combining.py +++ b/dataretrieval/combining.py @@ -2,12 +2,16 @@ These utilities assemble the output of a chunked/fan-out call from its individual per-sub-request results. They have no event-loop, retry, or -network state — they're pure data transforms imported by both the -chunked-call execution (:mod:`dataretrieval.ogc.chunking`) and the -per-page pagination (:mod:`dataretrieval.ogc.engine`). +network state — they're pure data transforms shared by protocol-specific +chunk execution, service fan-out, and +cursor-driven pagination. Separated from :mod:`dataretrieval.ogc.planning` so that module stays focused on *what* to split, while this module owns *how* to reassemble. + +A top-level leaf rather than part of :mod:`dataretrieval.transport`: these are +pandas transforms over already-fetched results, with no HTTP or event-loop +concern, consumed by chunk planning and service fan-out as well as by pagination. """ from __future__ import annotations @@ -104,8 +108,8 @@ def _merge_response( ``httpx.Headers`` means downstream mutations don't back-propagate into any underlying response — so callers may re-fold idempotently. This is the one low-level merge behind both pagination - (:func:`~dataretrieval.ogc.engine._paginate`) and the chunked / fan-out - aggregation (:func:`_combine_chunk_responses`).""" + (:func:`~dataretrieval.transport.pagination.paginate`) and the chunked / + fan-out aggregation (:func:`_combine_chunk_responses`).""" merged = copy.copy(base) merged.headers = httpx.Headers(headers_from.headers) merged.elapsed = elapsed diff --git a/dataretrieval/credentials.py b/dataretrieval/credentials.py new file mode 100644 index 00000000..82021b4e --- /dev/null +++ b/dataretrieval/credentials.py @@ -0,0 +1,97 @@ +"""Which host honors the USGS API key, and how it is attached and withheld. + +One leaf owns every answer about the ``API_USGS_PAT`` credential: the host that +accepts it, whether a given destination qualifies, and how it is stripped back +off a request bound somewhere else. Splitting those answers across the layers +that happen to need them is how a credential reaches a host nobody authorized: +the code that attaches a key and the code that removes it have to agree, and the +only way to guarantee they agree is to have them read the same predicate. + +This is deliberately a leaf. It sits below HTTP mechanics (which attaches the +header) and below progress reporting (which tells an unauthenticated caller where +to register), so neither has to depend on the other to learn the same fact. +""" + +from __future__ import annotations + +import os + +import httpx + +#: Environment variable holding the USGS Water Data personal access token. +API_KEY_ENV = "API_USGS_PAT" + +#: Where to register for a key. Surfaced once, by the progress reporter, when a +#: query against the authorized host runs without one -- unauthenticated callers +#: hit much lower rate limits (see the ``API_USGS_PAT`` note in the README). +SIGNUP_URL = "https://api.waterdata.usgs.gov/signup/" + +#: The only host that honors the key. Every other service this package talks to +#: ignores it, so sending it there would leak a credential for no benefit. +_AUTHORIZED_API_KEY_HOST = "api.waterdata.usgs.gov" + +#: Origin of the Water Data API, built from the authorized host rather than +#: spelled again. The host that serves these endpoints and the host that honors +#: the key are the same fact, and the failure mode of keeping two copies is +#: silent: the endpoint moves, the predicate does not follow, and either the key +#: quietly stops attaching or it is sent somewhere nobody authorized. The +#: adapters import this instead of restating the authority (enforced by +#: ``test_credential_policy_has_one_definition``). +WATERDATA_BASE_URL = f"https://{_AUTHORIZED_API_KEY_HOST}" + + +def accepts_api_key(target_url: str | httpx.URL | None) -> bool: + """Whether ``target_url`` names the host that honors :data:`API_KEY_ENV`. + + The single answer to "does this destination get the key" -- used when + attaching the credential, when stripping it back off at redirect time, and + when deciding whether "get an API key" is useful advice rather than noise, so + the three can't drift apart. + + The scheme has to be ``https``, not just the host. A bearer token sent over + cleartext is readable by anything on the path, and the destination that would + receive it is reachable through data we do not control: a redirect, or a + server-supplied next-page link naming ``http://`` on the very host that is + otherwise authorized. Matching on the host alone would hand the key over in + the clear on the strength of a hostname the attacker chose to keep. + """ + if target_url is None: + return False + try: + url = target_url if isinstance(target_url, httpx.URL) else httpx.URL(target_url) + except (httpx.InvalidURL, TypeError): + return False + return url.scheme == "https" and url.host == _AUTHORIZED_API_KEY_HOST + + +def without_embedded_credentials(url: httpx.URL) -> httpx.URL: + """Drop any ``user:pass@`` from a URL we were *handed* rather than built. + + A next-page link is data, not configuration. ``httpx`` derives an + ``Authorization: Basic`` header from userinfo in a URL, so a poisoned link + carrying ``user:pass@`` mints a credential the caller never configured and + sends it onward -- next to the real API key, when the host still checks out + and the host check therefore raises nothing. No USGS service authenticates + that way, so stripping it costs a legitimate caller nothing. + """ + return url.copy_with(userinfo=b"") if url.userinfo else url + + +def api_key() -> str | None: + """The configured token, or ``None``. + + Read through a function rather than captured at import so a caller that sets + the variable after import -- or a test that patches it -- is still honored. + """ + return os.getenv(API_KEY_ENV) + + +def strip_api_key_from_untrusted_host(request: httpx.Request) -> None: + """Remove Water Data credentials before sending to any other host.""" + if not accepts_api_key(request.url): + request.headers.pop("X-Api-Key", None) + + +async def strip_api_key_from_untrusted_host_async(request: httpx.Request) -> None: + """Async-client form of :func:`strip_api_key_from_untrusted_host`.""" + strip_api_key_from_untrusted_host(request) diff --git a/dataretrieval/exceptions.py b/dataretrieval/exceptions.py index fefb62c5..b40d62c4 100644 --- a/dataretrieval/exceptions.py +++ b/dataretrieval/exceptions.py @@ -11,7 +11,8 @@ of which :class:`TransientError` (429 / 5xx) is the retryable subset. The rest aren't a plain status: :class:`RequestTooLarge` (with :class:`URLTooLong` / :class:`Unchunkable`), :class:`NetworkError` (a failed connection, per above), -and :class:`NoSitesError`. :func:`error_for_status` maps a status to its type. +:class:`NoSitesError`, and :class:`ConfigurationError` for an unusable setting. +:func:`error_for_status` maps a status to its type. This module has no third-party runtime dependencies -- ``httpx`` is imported only for type checking -- so any module can import it without pulling in pandas / httpx @@ -36,6 +37,7 @@ "Unchunkable", "NetworkError", "NoSitesError", + "ConfigurationError", "error_for_status", ] @@ -240,6 +242,20 @@ class NetworkError(DataRetrievalError): retryable: ClassVar[bool] = True +# --- Bad configuration --------------------------------------------------- + + +class ConfigurationError(DataRetrievalError, ValueError): + """A ``dataretrieval`` setting -- an environment variable, a policy field -- + holds a value that can't be used, so no request was issued. + + It is a :class:`DataRetrievalError` so ``except`` around a retrieval catches + it rather than letting a bare ``ValueError`` escape a request path, and a + :class:`ValueError` so code that already treats a bad setting as one keeps + working. + """ + + # --- Empty result -------------------------------------------------------- diff --git a/dataretrieval/ngwmn.py b/dataretrieval/ngwmn.py index 83cb9726..c4522e08 100644 --- a/dataretrieval/ngwmn.py +++ b/dataretrieval/ngwmn.py @@ -24,11 +24,13 @@ import pandas as pd from dataretrieval.codes.states import apply_state +from dataretrieval.credentials import WATERDATA_BASE_URL from dataretrieval.ogc import OgcDialect, get_ogc_data, prepare_request_args from dataretrieval.utils import BaseMetadata -# The Water Data API base URL, defined locally to avoid importing policy internals. -BASE_URL = "https://api.waterdata.usgs.gov" +# The Water Data API base URL, from the credentials leaf rather than OGC policy +# internals: it names the same authority the API key is scoped to. +BASE_URL = WATERDATA_BASE_URL # The National Ground-Water Monitoring Network exposes its own OGC API at a # separate, unversioned base. diff --git a/dataretrieval/nldi.py b/dataretrieval/nldi.py index 9a169414..57d6048d 100644 --- a/dataretrieval/nldi.py +++ b/dataretrieval/nldi.py @@ -3,7 +3,7 @@ from json import JSONDecodeError from typing import Any, Literal, cast -from dataretrieval.utils import query +from dataretrieval.utils import _query_with_retry try: import geopandas as gpd @@ -23,7 +23,7 @@ def _query_nldi( # A helper function to query the NLDI API. ``query()`` already raises a # typed ``DataRetrievalError`` for any HTTP error response, so a returned # response is a success that we only need to parse. - response = query(url, payload=query_params) + response = _query_with_retry(url, payload=query_params) response_data: dict[str, Any] | list[Any] = {} try: response_data = response.json() diff --git a/dataretrieval/ogc/__init__.py b/dataretrieval/ogc/__init__.py index 46dd3918..ab3502cd 100644 --- a/dataretrieval/ogc/__init__.py +++ b/dataretrieval/ogc/__init__.py @@ -8,10 +8,9 @@ - :func:`fetch_ogc_request` — execute a pre-built request with pagination. Service adapters (NGWMN, Water Data's generic wrapper) import from this -facade rather than reaching into engine internals. The engine module remains -available for lower-level orchestration needs (e.g. ``_paginate``, -``_run_sync``) that sibling modules like ``wateruse`` use under the accepted -temporary variance. +facade rather than reaching into engine internals. Generic execution policy +lives in :mod:`dataretrieval.transport`; the engine retains compatibility +wrappers at previous private paths. """ from dataretrieval.ogc.engine import fetch_ogc_request, get_ogc_data diff --git a/dataretrieval/ogc/chunking.py b/dataretrieval/ogc/chunking.py index 79037a6b..f15f226b 100644 --- a/dataretrieval/ogc/chunking.py +++ b/dataretrieval/ogc/chunking.py @@ -17,15 +17,13 @@ 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 live in -sibling modules it imports, each with its own reason to change: -:mod:`~dataretrieval.ogc.planning` builds the -:class:`~dataretrieval.ogc.planning.ChunkPlan` and recombines per-chunk -frames and responses (pure, no I/O); :mod:`~dataretrieval.ogc.retry` holds -the transient-classification and exponential-backoff policy; and +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` exception -contract. +: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 @@ -83,23 +81,20 @@ import pandas as pd from anyio.from_thread import start_blocking_portal -from dataretrieval.utils import HTTPX_ASYNC_DEFAULTS, Ambient, _require_positive_int - -from . import progress as _progress -from .combining import ( +from dataretrieval import progress as _progress +from dataretrieval.combining import ( _combine_chunk_frames, _combine_chunk_responses, ) -from .interruptions import ( - ChunkInterrupted, -) +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.utils import Ambient, _require_positive_int + +from .interruptions import ChunkInterrupted from .planning import ChunkPlan -from .retry import ( - _NO_RETRY, - RetryPolicy, - _classify_chunk_error, - _retry, -) +from .retry import _classify_chunk_error # Empirically the API replies HTTP 414 above ~8200 bytes of full URL — # matches nginx's default ``large_client_header_buffers`` of 8 KB. 8000 @@ -140,12 +135,12 @@ def _read_concurrency_env() -> int | None: try: value = int(raw) except ValueError as exc: - raise ValueError( + raise ConfigurationError( f"{_CONCURRENCY_ENV} must be a positive integer or " f"'{_CONCURRENCY_UNBOUNDED}'; got {raw!r}." ) from exc if value < 1: - raise ValueError( + raise ConfigurationError( f"{_CONCURRENCY_ENV} must be >= 1 (got {value}); use " f"'{_CONCURRENCY_UNBOUNDED}' to disable the cap." ) @@ -650,31 +645,19 @@ async def _run(self, max_concurrent: int | None) -> tuple[pd.DataFrame, Any]: self.plan.total if max_concurrent is None else max_concurrent ) - async with httpx.AsyncClient(limits=limits, **HTTPX_ASYNC_DEFAULTS) as client: + 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 fetch_gated( - args: dict[str, Any], - ) -> tuple[pd.DataFrame, httpx.Response]: - """One fetch attempt under the concurrency gate. - - The slot is held for the attempt's full duration — - every page of a paginated sub-request — but acquired - per *attempt* (this is what ``_retry`` re-invokes), so - a sub-request sleeping off a retry backoff isn't - holding a slot while it isn't touching the server. - """ - async with semaphore: - return await self.fetch(args) - 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: fetch_gated(args), self.retry_policy) + 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 @@ -683,7 +666,7 @@ async def track( return result # Dispatch every pending sub-request concurrently; the - # semaphore (via ``fetch_gated``) is the only throttle. + # 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: diff --git a/dataretrieval/ogc/engine.py b/dataretrieval/ogc/engine.py index 31b81f5f..2ab300d7 100644 --- a/dataretrieval/ogc/engine.py +++ b/dataretrieval/ogc/engine.py @@ -1,8 +1,8 @@ """Generic OGC API engine shared by the Water Data and NGWMN getters. -This module holds the API-agnostic orchestration core for talking to an OGC -API Features service — async pagination, the sync bridge, and the chunked -fetch entry point :func:`get_ogc_data` that orchestrates them. Request +This module holds OGC API Features orchestration — OGC cursor/response +strategies and the chunked fetch entry point :func:`get_ogc_data`. Generic +pagination and sync dispatch live in :mod:`dataretrieval.transport`; request construction lives in :mod:`~dataretrieval.ogc.requests`. The surrounding concerns live in sibling modules it composes, each with its own reason to change: :mod:`~dataretrieval.ogc.dates` (time-parameter marshalling), @@ -27,23 +27,19 @@ import functools import logging from collections.abc import ( - AsyncIterator, Awaitable, Callable, ) -from contextlib import asynccontextmanager from typing import Any, TypeVar, cast import httpx import pandas as pd -from anyio.from_thread import start_blocking_portal import dataretrieval.ogc.chunking as chunking -import dataretrieval.ogc.progress as _progress -from dataretrieval.exceptions import DataRetrievalError +import dataretrieval.progress as _progress +from dataretrieval.credentials import without_embedded_credentials from dataretrieval.ogc.chunking import get_active_client -from dataretrieval.ogc.combining import _QUOTA_HEADER, _merge_response, _safe_elapsed -from dataretrieval.ogc.errors import _paginated_failure_message, _raise_for_non_200 +from dataretrieval.ogc.errors import _raise_for_non_200 from dataretrieval.ogc.policy import ( BASE_URL, # noqa: F401 — compatibility alias DEFAULT_DIALECT, @@ -71,11 +67,11 @@ prepare_request_args, ) from dataretrieval.ogc.shaping import GEOPANDAS, _finalize_ogc, _get_resp_data +from dataretrieval.transport.pagination import paginate +from dataretrieval.transport.sync import run_sync from dataretrieval.utils import ( - HTTPX_ASYNC_DEFAULTS, BaseMetadata, _default_headers, # noqa: F401 — compatibility re-export for tests - _network_error, _require_positive_int, ) @@ -137,8 +133,10 @@ def _next_req_url( # by falling open when host extraction isn't reliable. next_host: str | None cur_host: str | None + next_url: httpx.URL | None try: - next_host = httpx.URL(href).host + next_url = httpx.URL(href) + next_host = next_url.host resp_url = ( resp.url if isinstance(resp.url, httpx.URL) @@ -146,12 +144,19 @@ def _next_req_url( ) cur_host = resp_url.host except (httpx.InvalidURL, TypeError): + next_url = None next_host = cur_host = None if next_host and cur_host and next_host != cur_host: raise RuntimeError( f"Refusing to follow cross-host next-page URL: " f"{next_host} != {cur_host}" ) + # Matching hosts is not enough: a link may also carry ``user:pass@``, + # which httpx turns into an ``Authorization: Basic`` header on the + # follow-up request. The host check above passes in exactly that case, + # so strip it here rather than trusting the link we were handed. + if next_url is not None: + return str(without_embedded_credentials(next_url)) # ``href`` comes from the JSON ``links`` array (typed ``Any``); the # ``not href`` guard above already excluded empty/None, and it is a # URL string (passed to ``httpx.URL`` above). @@ -159,46 +164,6 @@ def _next_req_url( return None -@asynccontextmanager -async def _client_for( - client: httpx.AsyncClient | None, -) -> AsyncIterator[httpx.AsyncClient]: - """ - Yield a usable async client, picking the best available source. - - Resolution order: - - 1. ``client`` if the caller supplied one (borrowed; not closed - here — the caller owns its lifecycle). - 2. The chunker's shared async client if we're inside a - :class:`~dataretrieval.ogc.chunking.ChunkedCall` run (per - :func:`chunking.get_active_client`). Borrowed; the chunker - closes it on exit. - 3. A fresh short-lived ``httpx.AsyncClient`` opened here and closed - on context exit. - - Parameters - ---------- - client : httpx.AsyncClient or None - A caller-owned client to borrow, or ``None`` to defer to the - chunker's shared client or a temporary one. - - Yields - ------ - httpx.AsyncClient - The chosen client. - """ - if client is not None: - yield client - return - shared = get_active_client() - if shared is not None: - yield shared - return - async with httpx.AsyncClient(**HTTPX_ASYNC_DEFAULTS) as new: - yield new - - _Cursor = TypeVar("_Cursor") @@ -210,136 +175,16 @@ async def _paginate( client: httpx.AsyncClient | None = None, raise_for_status: Callable[[httpx.Response], None] = _raise_for_non_200, ) -> tuple[pd.DataFrame, httpx.Response]: - """ - Drive a paginated request to completion over an - :class:`httpx.AsyncClient`. - - The common shape behind the paginated fetch paths (e.g. - :func:`_walk_pages`): send the initial request, then loop calling - ``follow_up`` until ``parse_response`` reports a ``None`` cursor, - accumulating frames and elapsed time. Any mid-pagination failure - raises ``DataRetrievalError`` wrapping the cause — the API exposes no - resume cursor, so the caller's only recovery is to retry the whole - call. Issuing HTTP asynchronously lets the multiple sub-requests of a - chunked call run concurrently under - :meth:`~dataretrieval.ogc.chunking.ChunkedCall._run`. - - Parameters - ---------- - initial_req : httpx.Request - First-page request to send. - parse_response : callable - ``resp -> (df, next_cursor_or_None)``. Returns the page's - DataFrame and the cursor (URL, token, …) used to drive - ``follow_up`` for the next page; ``None`` terminates the loop. - follow_up : callable - ``(cursor, client) -> Awaitable[httpx.Response]``. Builds and - sends the next-page request. - client : httpx.AsyncClient, optional - Caller-borrowed client. ``None`` (default) means use the - chunker's shared client (if inside a chunked call) or open - a temporary one. - raise_for_status : callable, optional - ``resp -> None``; raises the typed error for a non-OK response. - Defaults to :func:`_raise_for_non_200` (the OGC ``{code, description}`` - envelope); wateruse passes its own to surface the NWDC ``detail``. - - Returns - ------- - df : pandas.DataFrame - Concatenation of every page's parsed frame. - response : httpx.Response - A shallow copy of the first-page response, with ``.headers`` - rebuilt as a fresh ``httpx.Headers`` reflecting the last page and - ``.elapsed`` set to the sum of the per-page response durations. The - canonical URL is preserved from the first page. The original first-page - response is not mutated. - - Raises - ------ - DataRetrievalError - On a non-200 initial response, the typed subclass for the status from - :func:`_raise_for_non_200` (a - :class:`~dataretrieval.exceptions.TransientError` for a retryable - 429 / 5xx, otherwise a fatal :class:`~dataretrieval.exceptions.HTTPError`); - or, on an initial-page parse failure or any subsequent-page failure, a - base ``DataRetrievalError`` wrapping the cause (built by - :func:`_paginated_failure_message`, original exception on ``__cause__``). - httpx.HTTPError - Network-level failures on the *initial* request (e.g. - ``ConnectError``, ``TimeoutException``) propagate unmodified - so callers can branch on the specific type; equivalent - failures on subsequent pages are wrapped per above. - """ - logger.debug("Requesting: %s", initial_req.url) - reporter = _progress.current() - - def report_page(page: httpx.Response, frame: pd.DataFrame) -> None: - """Tick the ambient progress reporter (a no-op when unset) for one page.""" - if reporter is not None: - reporter.set_rate_remaining( - page.headers.get(_QUOTA_HEADER), - limit=page.headers.get("x-ratelimit-limit"), - ) - reporter.add_page(rows=len(frame)) - - async with _client_for(client) as sess: - resp = await sess.send(initial_req) - raise_for_status(resp) - initial_response = resp - total_elapsed = _safe_elapsed(resp) - - try: - df, cursor = parse_response(resp) - except Exception as e: # noqa: BLE001 - # Initial-page parse failures (malformed JSON, missing - # ``features``, schema drift) get the same wrapped-message - # treatment as follow-up failures so callers see a consistent - # diagnostic regardless of which page broke. - logger.warning("Initial response parse failed.") - raise DataRetrievalError(_paginated_failure_message(0, e)) from e - dfs = [df] - # Stop following ``next`` links once the optional row cap is reached - # (see :func:`_row_cap`); ``None`` means uncapped. The concatenation - # is sliced to the cap below so a final over-budget page can't exceed it. - cap = _row_cap.get() - nrows = len(df) - # Guard a non-advancing or cyclic cursor (a server bug that would - # otherwise loop forever). OGC's next-URLs are unique, so this never - # fires for them; the Link-header pagers (e.g. wateruse) rely on it. - seen: set[Any] = set() - report_page(resp, df) - while ( - cursor is not None and cursor not in seen and (cap is None or nrows < cap) - ): - seen.add(cursor) - try: - resp = await follow_up(cursor, sess) - raise_for_status(resp) - df, cursor = parse_response(resp) - dfs.append(df) - nrows += len(df) - total_elapsed += _safe_elapsed(resp) - report_page(resp, df) - except Exception as e: # noqa: BLE001 - logger.warning( - "Request failed at cursor %r. Data download interrupted.", - cursor, - ) - raise DataRetrievalError(_paginated_failure_message(len(dfs), e)) from e - - # Fold the pages onto a COPY of the initial response so a caller that - # inspected it mid-pagination (a hook, a test fixture) never sees an - # in-place mutation. ``resp`` is the last page, whose headers carry the - # current ``x-ratelimit-remaining`` (monotonic, so the last page is the - # most depleted) — the same low-level merge the fan-out aggregation uses. - final_response = _merge_response( - initial_response, headers_from=resp, elapsed=total_elapsed - ) - result = pd.concat(dfs, ignore_index=True) - if cap is not None: - result = result.head(cap) - return result, final_response + """Compatibility wrapper around service-neutral cursor pagination.""" + active_client = client if client is not None else get_active_client() + return await paginate( + initial_req, + parse_response=parse_response, + follow_up=follow_up, + client=active_client, + raise_for_status=raise_for_status, + row_cap=_row_cap.get(), + ) def _ogc_parse_response( @@ -508,7 +353,10 @@ def get_ogc_data( extra_id_cols=extra_id_cols, dialect=dialect, ) - with _progress.progress_context(service=service), _row_cap(max_rows): + with ( + _progress.progress_context(service=service, target_url=base_url), + _row_cap(max_rows), + ): with _ogc_base_url(base_url), _dialect(dialect): return _fetch_once(args, finalize=finalize) @@ -541,44 +389,12 @@ def _run_sync( service: str, error_url: str | httpx.URL | None = None, ) -> tuple[pd.DataFrame, httpx.Response]: - """Drive an async OGC fetch to completion from synchronous code. - - Opens the service progress context and runs ``make_coro()`` through a - short-lived ``anyio`` blocking portal (a worker thread), so the - non-chunked getters work whether or not the caller is already inside an - event loop (Jupyter/async apps). The portal copies the calling context, - so the active progress reporter still reaches the sub-requests. - - Shared by the non-chunked fetch paths; the chunked OGC getters - drive their own portal - inside :meth:`chunking.ChunkedCall.resume`. - - A connection failure on the initial request is surfaced as a typed - ``NetworkError`` against ``error_url`` when given (callers that build their - own requests, e.g. ``wateruse``), else the request-builder base the caller - scoped via ``_ogc_base_url`` (the OGC / NGWMN getters). - """ - with _progress.progress_context(service=service): - with start_blocking_portal() as portal: - try: - # ``portal.call`` is ``Any`` (anyio is skipped by mypy — its - # source uses 3.10 syntax our 3.9 target can't parse), so cast - # to the declared return type, as ``ChunkedCall`` does too. - return cast( - "tuple[pd.DataFrame, httpx.Response]", portal.call(make_coro) - ) - except httpx.TransportError as exc: - # The initial-request connection failure ``_paginate`` lets - # through raw; mid-pagination failures are already typed. - # Report the base URL actually targeted: callers that build - # their own requests (``wateruse``) pass ``error_url``; the OGC - # getters leave it unset and fall back to the request-builder - # base they scoped via ``_ogc_base_url`` (NGWMN/sibling APIs set - # their own), not a hardcoded host. - raise _network_error( - error_url if error_url is not None else _ogc_base_url.get(), - exc, - ) from exc + """Compatibility wrapper around the service-neutral sync bridge.""" + return run_sync( + make_coro, + service=service, + error_url=error_url if error_url is not None else _ogc_base_url.get(), + ) def fetch_ogc_request( diff --git a/dataretrieval/ogc/errors.py b/dataretrieval/ogc/errors.py index 0bb39b0e..90b541de 100644 --- a/dataretrieval/ogc/errors.py +++ b/dataretrieval/ogc/errors.py @@ -10,7 +10,11 @@ import httpx -from dataretrieval.exceptions import RateLimited, error_for_status +from dataretrieval.exceptions import error_for_status +from dataretrieval.transport.pagination import ( + paginated_failure_message as _paginated_failure_message, # noqa: F401 +) +from dataretrieval.transport.retry import parse_retry_after as _parse_retry_after def _error_body(resp: httpx.Response) -> str: @@ -64,37 +68,6 @@ def _error_body(resp: httpx.Response) -> str: ) -def _parse_retry_after(value: str | None) -> float | None: - """ - Parse a USGS ``Retry-After`` header into seconds. - - Parameters - ---------- - value : str or None - The raw header value, or ``None`` if absent. - - Returns - ------- - float or None - Non-negative delta-seconds, clamped at zero. ``None`` when the - header is absent or unparseable; ``ChunkedCall`` treats - ``None`` as "fall back to my own retry policy". - - Notes - ----- - USGS sends ``Retry-After`` as integer delta-seconds (empirically - verified — e.g. ``Retry-After: 2619``). The HTTP spec also allows - HTTP-date form, but USGS doesn't use it, so this function doesn't - bother parsing it. - """ - if not value: - return None - try: - return max(0.0, float(value.strip())) - except ValueError: - return None - - def _raise_for_non_200(resp: httpx.Response) -> None: """ Raise a typed exception for any non-200 response. @@ -129,43 +102,3 @@ def _raise_for_non_200(resp: httpx.Response) -> None: _error_body(resp), retry_after=_parse_retry_after(resp.headers.get("Retry-After")), ) - - -def _paginated_failure_message(pages_collected: int, cause: BaseException) -> str: - """ - Build a user-facing message for a mid-pagination failure. - - The API exposes no resume cursor, so the caller's only recovery is - to retry the whole call — the message lists the practical knobs, - tailored to whether the failure was rate-limit (429) or something - else. - - Parameters - ---------- - pages_collected : int - Number of pages successfully fetched before the failure. - cause : BaseException - The underlying exception that interrupted pagination. - - Returns - ------- - str - A message suitable for the ``DataRetrievalError`` that the - paginated fetch paths raise from the original exception. - """ - cause_str = str(cause).removesuffix(".") - # Some ``httpx`` exceptions (e.g. ``TimeoutException()`` with no args) - # stringify to empty; fall back to the class name so the - # returned message is always informative. - if not cause_str.strip(): - cause_str = type(cause).__name__ - if isinstance(cause, RateLimited): - action = "wait for the rate-limit window to reset and retry" - else: - action = "retry the request (possibly after a short backoff)" - return ( - f"Paginated request failed after collecting {pages_collected} " - f"page(s): {cause_str}. To recover: {action}, reduce the " - f"request size (e.g. fewer locations, a shorter time range, or " - f"a smaller ``limit``), or obtain an API token." - ) diff --git a/dataretrieval/ogc/planning.py b/dataretrieval/ogc/planning.py index 15b397a0..76796df4 100644 --- a/dataretrieval/ogc/planning.py +++ b/dataretrieval/ogc/planning.py @@ -4,16 +4,15 @@ deciding how to split one over-budget OGC request into URL-fitting sub-requests (:class:`ChunkPlan` and the axis/byte-accounting helpers). It has no event loop, retry policy, or network state — those live in -:mod:`dataretrieval.ogc.chunking` (execution) and -:mod:`dataretrieval.ogc.retry` (retry policy), which import the plan and +:mod:`dataretrieval.ogc.chunking` (resumable execution) and +:mod:`dataretrieval.transport.retry` (retry policy), which import the plan and drive it. Result recombination — reassembling the per-chunk frames and responses back into one result -(:func:`~dataretrieval.ogc.combining._combine_chunk_frames`, -:func:`~dataretrieval.ogc.combining._combine_chunk_responses`, etc.) — lives in -the sibling :mod:`dataretrieval.ogc.combining` module, which callers import -directly. +(:func:`~dataretrieval.combining._combine_chunk_frames`, +:func:`~dataretrieval.combining._combine_chunk_responses`, etc.) — +lives in the top-level :mod:`dataretrieval.combining` module. """ from __future__ import annotations diff --git a/dataretrieval/ogc/policy.py b/dataretrieval/ogc/policy.py index 7aff5398..e8831b10 100644 --- a/dataretrieval/ogc/policy.py +++ b/dataretrieval/ogc/policy.py @@ -12,11 +12,13 @@ from dataclasses import dataclass, field +from dataretrieval.credentials import WATERDATA_BASE_URL + # --------------------------------------------------------------------------- # Endpoint constants # --------------------------------------------------------------------------- -BASE_URL = "https://api.waterdata.usgs.gov" +BASE_URL = WATERDATA_BASE_URL OGC_API_VERSION = "v0" OGC_API_URL = f"{BASE_URL}/ogcapi/{OGC_API_VERSION}" diff --git a/dataretrieval/ogc/requests.py b/dataretrieval/ogc/requests.py index eb5a201b..462cfe4b 100644 --- a/dataretrieval/ogc/requests.py +++ b/dataretrieval/ogc/requests.py @@ -23,7 +23,16 @@ from dataretrieval.ogc.dates import _DATE_RANGE_PARAMS, _format_api_dates from dataretrieval.ogc.errors import _raise_for_non_200 from dataretrieval.ogc.policy import DEFAULT_DIALECT, OGC_API_URL, OgcDialect -from dataretrieval.utils import HTTPX_DEFAULTS, Ambient, _default_headers, _get +from dataretrieval.transport.http import ( + HTTPX_DEFAULTS, +) +from dataretrieval.transport.http import ( + default_headers as _default_headers, +) +from dataretrieval.transport.http import ( + get as _get, +) +from dataretrieval.utils import Ambient logger = logging.getLogger(__name__) diff --git a/dataretrieval/ogc/retry.py b/dataretrieval/ogc/retry.py index bd45f275..7eeafb44 100644 --- a/dataretrieval/ogc/retry.py +++ b/dataretrieval/ogc/retry.py @@ -1,213 +1,45 @@ -"""Transient-failure retry policy for chunked sub-requests. - -Defines what counts as a retryable transient (:func:`_classify_chunk_error`, -:func:`_retryable`), the bounded exponential-backoff-with-jitter policy -(:class:`RetryPolicy`), and the driver that applies it (:func:`_retry`). Kept -separate from the execution engine in :mod:`dataretrieval.ogc.chunking` so the -retry/backoff behavior is one cohesive unit that changes independently of the -concurrency model. +"""OGC interruption classification over service-neutral transport retry policy. + +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. + +"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. """ from __future__ import annotations -import asyncio -import os -import random -from collections.abc import Awaitable, Callable -from dataclasses import dataclass - import httpx -import pandas as pd -import dataretrieval.ogc.progress as _progress from dataretrieval.exceptions import RateLimited, TransientError from dataretrieval.ogc.interruptions import ( ChunkInterrupted, QuotaExhausted, ServiceInterrupted, ) - -# Retry-with-backoff defaults for transient sub-request failures (429 / -# 5xx / connect-read timeouts): exponential backoff with full jitter, and -# honor a server ``Retry-After`` up to the cap below before escalating -# to a resumable interruption instead. -_RETRIES_ENV = "API_USGS_RETRIES" - - -_RETRIES_DEFAULT = 4 - - -_RETRY_BASE_BACKOFF = 0.5 - - -_RETRY_MAX_BACKOFF = 30.0 - - -_RETRY_AFTER_CAP = 60.0 - - -def _read_retries_env() -> int: - """ - Resolve the ``API_USGS_RETRIES`` env var to a max-retry count. - - Returns - ------- - int - Number of retries after the first attempt; ``0`` disables - retrying. Unset/blank → ``_RETRIES_DEFAULT``. - """ - raw = os.environ.get(_RETRIES_ENV) - if raw is None or raw.strip() == "": - return _RETRIES_DEFAULT - try: - value = int(raw.strip()) - except ValueError as exc: - raise ValueError( - f"{_RETRIES_ENV} must be a non-negative integer (got {raw!r})." - ) from exc - if value < 0: - raise ValueError(f"{_RETRIES_ENV} must be >= 0 (got {value}).") - return value - - -@dataclass(frozen=True) -class RetryPolicy: - """Bounded retry-with-backoff config for transient sub-request failures. - - An immutable value object that owns the *timing* decisions; the - exception taxonomy (which failures are retryable) lives in - :func:`_retryable`. Backoff is exponential with **full jitter** - (:func:`random.uniform` over ``[0, ceiling]``) so the concurrent - fan-out's retries don't re-burst in lockstep. A server ``Retry-After`` - hint, when present, overrides the computed backoff — unless it exceeds - :attr:`retry_after_cap`, in which case retrying stops and the failure - surfaces as a resumable :class:`ChunkInterrupted` (a multi-minute - quota-window reset shouldn't block the call inline). - - Attributes - ---------- - max_retries : int - Retries attempted after the first try; ``0`` disables retrying. - base_backoff : float - Seconds; the jitter ceiling for the first retry, doubled each - subsequent attempt. - max_backoff : float - Upper bound on any single attempt's backoff ceiling. - retry_after_cap : float - Largest ``Retry-After`` (seconds) honored inline; longer hints - escalate to a resumable interruption. - """ - - max_retries: int = _RETRIES_DEFAULT - base_backoff: float = _RETRY_BASE_BACKOFF - max_backoff: float = _RETRY_MAX_BACKOFF - retry_after_cap: float = _RETRY_AFTER_CAP - - def __post_init__(self) -> None: - # Catch invalid timing knobs here so a misconfiguration fails at - # construction, not deep in a later ``time.sleep`` (ValueError on - # a negative delay) or silently in ``asyncio.sleep`` (which - # treats negative as zero). - if self.max_retries < 0: - raise ValueError(f"max_retries must be >= 0 (got {self.max_retries}).") - if self.base_backoff < 0 or self.max_backoff < 0 or self.retry_after_cap < 0: - raise ValueError("retry backoff settings must be non-negative.") - - @classmethod - def from_env(cls) -> RetryPolicy: - """ - Build a policy from the module-level defaults, resolved now. - - Reads ``max_retries`` from ``API_USGS_RETRIES`` and the timing - knobs from the ``_RETRY_*`` module constants at call time — not - the dataclass field defaults (which freeze at class definition) - — so test ``monkeypatch.setattr`` on the constants takes effect. - - Returns - ------- - RetryPolicy - A policy built from the module-level defaults resolved at - call time. - """ - return cls( - max_retries=_read_retries_env(), - base_backoff=_RETRY_BASE_BACKOFF, - max_backoff=_RETRY_MAX_BACKOFF, - retry_after_cap=_RETRY_AFTER_CAP, - ) - - def should_retry(self, attempt: int, retry_after: float | None) -> bool: - """ - Whether a just-failed ``attempt`` (1-based) warrants another try. - - A ``Retry-After`` longer than ``retry_after_cap`` is *not* slept - off inline — it returns ``False`` so the failure escalates to a - resumable interruption instead of blocking the call for minutes. - - Parameters - ---------- - attempt : int - The just-failed attempt number (1-based). - retry_after : float or None - Seconds the server suggested waiting (``Retry-After`` hint), - or ``None`` when no hint was given. - - Returns - ------- - bool - ``True`` if another try is warranted, ``False`` otherwise. - """ - if attempt > self.max_retries: - return False - return retry_after is None or retry_after <= self.retry_after_cap - - def backoff(self, attempt: int, retry_after: float | None) -> float: - """ - Seconds to wait before retry ``attempt`` (1-based). - - Parameters - ---------- - attempt : int - The retry attempt number (1-based). - retry_after : float or None - Seconds the server suggested waiting (``Retry-After`` hint), - or ``None`` to use the computed exponential backoff instead. - - Returns - ------- - float - Seconds to wait before the retry. - """ - if retry_after is not None: - return retry_after - ceiling = min(self.max_backoff, self.base_backoff * 2 ** (attempt - 1)) - return random.uniform(0.0, ceiling) - - -# Default for direct ``ChunkedCall`` / ``ChunkPlan.execute`` construction -# (and tests): no retrying. The production decorator path explicitly passes -# ``RetryPolicy.from_env()`` so retries are on by default there. -_NO_RETRY = RetryPolicy(max_retries=0) +from dataretrieval.transport.retry import _deterministic_failure def _classify_transient( exc: BaseException, ) -> tuple[type[ChunkInterrupted], float | None] | None: - """Classify one exception as a transient, resumable failure. - - This function owns the shared exception taxonomy; it deliberately does not - walk ``__cause__``. :func:`_classify_chunk_error` walks wrapped pagination - failures, while :func:`_retryable` applies the narrower automatic-retry - policy to this classification. - """ + """Classify one failure as a resumable OGC interruption.""" if isinstance(exc, RateLimited): return QuotaExhausted, exc.retry_after if isinstance(exc, TransientError): - # Every typed transient other than a rate-limit error is a service - # interruption. This fallback keeps future TransientError subclasses - # resumable after their inline retries are exhausted. 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 @@ -215,139 +47,17 @@ def _classify_transient( def _classify_chunk_error( exc: BaseException, ) -> tuple[type[ChunkInterrupted], float | None] | None: - """ - Classify a fetch error as a known transient (resumable) failure. - - Walks the ``__cause__`` chain of ``exc`` looking for a known typed - transport failure. Returns the matching ``ChunkInterrupted`` - subclass and any ``Retry-After`` hint, or ``None`` if the error is - not a recognized transient — in which case ``ChunkedCall`` - re-raises rather than wrapping (programmer errors and unknown - failures shouldn't masquerade as resumable). - - Parameters - ---------- - exc : BaseException - The exception raised by a sub-request. - - Returns - ------- - tuple[type[ChunkInterrupted], float or None] or None - ``(interrupted_class, retry_after)`` for recognized transient - failures; ``None`` otherwise. - - Notes - ----- - ``_walk_pages`` re-wraps mid-pagination failures as a base - ``DataRetrievalError`` with the typed transport exception linked as - ``__cause__``, so this function must walk the chain rather than - just ``isinstance`` the top-level exception. - - Bare ``httpx.HTTPError`` (``ConnectError``, ``TimeoutException``, - etc.) and ``httpx.InvalidURL`` (server-supplied cursor URL too - long, oversize follow-up) are also treated as transport failures - and wrapped as :class:`ServiceInterrupted` — they aren't one of the - typed status errors above (and ``InvalidURL`` doesn't even inherit - from ``httpx.HTTPError``), so without explicit handling they would - escape classification with no resumable handle. - """ - cur: BaseException | None = exc - while cur is not None: - result = _classify_transient(cur) + """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 - cur = cur.__cause__ + current = current.__cause__ return None -def _retryable(exc: BaseException) -> tuple[bool, float | None]: - """Decide whether a top-level transient is worth an automatic retry. - - Wrapped mid-pagination failures are not retried from page one; they instead - escalate to a resumable :class:`ChunkInterrupted`. ``httpx.InvalidURL`` and - non-transport ``httpx.HTTPError`` instances are classified as resumable but - excluded from automatic retry by policy. - """ - classification = _classify_transient(exc) - if classification is None: - return False, None - - _, retry_after = classification - if isinstance(exc, (TransientError, httpx.TransportError)): - return True, retry_after - return False, None - - -def _retry_delay(exc: BaseException, attempt: int, policy: RetryPolicy) -> float | None: - """ - Decide the backoff for a just-failed ``attempt`` (1-based), or ``None`` - to give up and re-raise. - - Returns ``None`` in three cases — the error isn't a retryable - transient, the policy is exhausted, or the server's ``Retry-After`` - exceeds the cap (escalates to a resumable :class:`ChunkInterrupted` - instead). Otherwise returns the seconds to wait and emits the - progress-bar retry note. - - Parameters - ---------- - exc : BaseException - The exception raised by the just-failed attempt. - attempt : int - The just-failed attempt number (1-based). - policy : RetryPolicy - The retry-with-backoff policy governing the decision. - - Returns - ------- - float or None - Seconds to wait before retrying, or ``None`` to give up and - re-raise. - """ - retryable, retry_after = _retryable(exc) - if not retryable or not policy.should_retry(attempt, retry_after): - return None - delay = policy.backoff(attempt, retry_after) - # Surface the imminent retry on the active progress reporter, if any. - reporter = _progress.current() - if reporter is not None: - reporter.note_retry(attempt=attempt, wait=delay) - return delay - - -async def _retry( - afn: Callable[[], Awaitable[tuple[pd.DataFrame, httpx.Response]]], - policy: RetryPolicy, -) -> tuple[pd.DataFrame, httpx.Response]: - """ - Call ``afn`` with bounded retry-with-backoff on transient failures. - - A non-retryable or policy-exhausted failure (see :func:`_retry_delay`) - propagates unchanged so the caller's existing handling wraps it as a - resumable :class:`ChunkInterrupted`. The whole retry *decision* lives - in :func:`_retry_delay`; this driver only awaits the sleep between - attempts. - - Parameters - ---------- - afn : Callable - Zero-arg awaitable callable that issues a single sub-request and - returns ``(frame, response)``. - policy : RetryPolicy - The retry-with-backoff policy governing the retries. - - Returns - ------- - tuple of (pandas.DataFrame, httpx.Response) - The ``(frame, response)`` pair from the first successful call. - """ - attempt = 0 - while True: - try: - return await afn() - except Exception as exc: # noqa: BLE001 — re-raised unless retryable - attempt += 1 - delay = _retry_delay(exc, attempt, policy) - if delay is None: - raise - await asyncio.sleep(delay) +__all__ = [ + "_classify_chunk_error", + "_classify_transient", +] diff --git a/dataretrieval/ogc/progress.py b/dataretrieval/progress.py similarity index 86% rename from dataretrieval/ogc/progress.py rename to dataretrieval/progress.py index 6177c30f..132ea906 100644 --- a/dataretrieval/ogc/progress.py +++ b/dataretrieval/progress.py @@ -1,9 +1,10 @@ -"""A single self-updating status line for paginated / chunked OGC queries. +"""A single self-updating status line for paginated and chunked queries. -OGC getters fan out two ways the caller can't see: large multi-value +Retrieval adapters can fan out in ways the caller cannot see: large multi-value requests are split into URL-length-safe *chunks* (``chunking`` module), and each request follows ``next`` links across an unknown number of *pages* -(``engine._paginate``). This module surfaces that work as one line on stderr, +(``transport.pagination.paginate``). This module surfaces that work as one +line on stderr, rewritten in place as data arrives:: Retrieving: daily · 6 pages · 2,881 rows · 995/1,000 requests remaining @@ -14,6 +15,11 @@ page/row/rate-limit counts) both update without knowing about each other. Call :func:`progress_context` to activate one and :func:`current` to reach it. +This is a top-level leaf rather than part of :mod:`dataretrieval.transport`: it +is terminal presentation, not HTTP execution policy. Transport modules report +*into* it, so keeping it outside means the execution layer owns no rendering, and +every service adapter -- OGC or not -- reaches the same reporter. + By default the line is shown for interactive use — an interactive terminal or a Jupyter/IPython kernel, like ``tqdm`` — while redirected logs and CI stay clean. ``API_USGS_PROGRESS`` forces it on (``1``/``true``) or off (``0``/``false``). @@ -26,7 +32,12 @@ import sys from collections.abc import Iterator from contextlib import contextmanager -from typing import TextIO +from typing import TYPE_CHECKING, TextIO + +from dataretrieval.credentials import SIGNUP_URL, accepts_api_key, api_key + +if TYPE_CHECKING: + import httpx def _group_int(value: str) -> str: @@ -44,14 +55,9 @@ def _group_int(value: str) -> str: # state. (It does not give concurrent queries sharing one stderr separate # lines — they would still interleave.) _active: contextvars.ContextVar[ProgressReporter | None] = contextvars.ContextVar( - "ogc_progress", default=None + "dataretrieval_progress", default=None ) -# Where to register for an API key. Surfaced once when a query runs without an -# API key configured (no API_USGS_PAT), since unauthenticated callers hit much -# lower rate limits (see the API_USGS_PAT note in the README). -SIGNUP_URL = "https://api.waterdata.usgs.gov/signup/" - # Process-level latch so the "no API key" pointer is shown at most once. _api_key_hint_shown = False @@ -104,9 +110,12 @@ def __init__( service: str | None = None, stream: TextIO | None = None, enabled: bool | None = None, + target_url: str | httpx.URL | None = None, ) -> None: self._stream = stream if stream is not None else sys.stderr self.enabled = _enabled_default(self._stream) if enabled is None else enabled + # Whether the sign-up pointer on close() is advice this caller can act on. + self._key_helps = accepts_api_key(target_url) # The service/collection being retrieved (e.g. "daily", "peaks"), # shown as the line's leading label. self.service = service @@ -225,9 +234,9 @@ def _render(self) -> None: def close(self) -> None: """Finalize the line with a trailing newline so it persists on screen. - If no API key is configured (no ``API_USGS_PAT``), append a one-time - pointer to API-key registration, since unauthenticated callers hit much - lower rate limits. + If the query targeted the API-key host and no key is configured (no + ``API_USGS_PAT``), append a one-time pointer to API-key registration, + since unauthenticated callers hit much lower rate limits. """ if self._closed: return @@ -250,7 +259,7 @@ def close(self) -> None: def _maybe_hint_api_key(self) -> None: global _api_key_hint_shown - if _api_key_hint_shown or os.getenv("API_USGS_PAT"): + if not self._key_helps or _api_key_hint_shown or api_key(): return # Set the once-per-process latch only after a successful write, so a # failed write (broken pipe) doesn't silently burn the hint for every @@ -267,19 +276,24 @@ def progress_context( service: str | None = None, stream: TextIO | None = None, enabled: bool | None = None, + target_url: str | httpx.URL | None = None, ) -> Iterator[ProgressReporter]: """Activate a :class:`ProgressReporter` for the duration of a query. - ``service`` labels the line (e.g. ``"Retrieving: daily ..."``). If a reporter - is already active (a nested call), the existing one is yielded unchanged so - the outermost query owns the single line; only the outermost context closes - it (and ``service``/``stream``/``enabled`` of a nested call are ignored). + ``service`` labels the line (e.g. ``"Retrieving: daily ..."``), and + ``target_url`` is where the query is going -- it decides whether an + API-key pointer is worth showing when the line closes. If a reporter is + already active (a nested call), the existing one is yielded unchanged so the + outermost query owns the single line; only the outermost context closes it + (and every argument of a nested call is ignored). """ existing = _active.get() if existing is not None: yield existing return - reporter = ProgressReporter(service=service, stream=stream, enabled=enabled) + reporter = ProgressReporter( + service=service, stream=stream, enabled=enabled, target_url=target_url + ) token = _active.set(reporter) try: yield reporter diff --git a/dataretrieval/streamstats.py b/dataretrieval/streamstats.py index 1458494a..6727c2bd 100644 --- a/dataretrieval/streamstats.py +++ b/dataretrieval/streamstats.py @@ -12,7 +12,8 @@ import httpx -from dataretrieval.utils import HTTPX_DEFAULTS, _get, _raise_for_status +from dataretrieval.transport.http import HTTPX_DEFAULTS +from dataretrieval.utils import _get_with_retry def download_workspace(workspaceID: str, format: str = "") -> httpx.Response: @@ -37,9 +38,7 @@ def download_workspace(workspaceID: str, format: str = "") -> httpx.Response: payload = {"workspaceID": workspaceID, "format": format} url = "https://streamstats.usgs.gov/streamstatsservices/download" - r = _get(url, params=payload, **HTTPX_DEFAULTS) - - _raise_for_status(r) + r = _get_with_retry(url, params=payload, **HTTPX_DEFAULTS) return r # data = r.raw.read() @@ -144,9 +143,7 @@ def get_watershed( } url = "https://streamstats.usgs.gov/streamstatsservices/watershed.geojson" - r = _get(url, params=payload, **HTTPX_DEFAULTS) - - _raise_for_status(r) + r = _get_with_retry(url, params=payload, **HTTPX_DEFAULTS) if format == "geojson": return r diff --git a/dataretrieval/transport/__init__.py b/dataretrieval/transport/__init__.py new file mode 100644 index 00000000..e437aa13 --- /dev/null +++ b/dataretrieval/transport/__init__.py @@ -0,0 +1,7 @@ +"""Internal service-neutral HTTP transport and execution policy. + +The modules in this package own reusable client lifecycle, authentication, +pagination, retry, response aggregation, progress, and sync-dispatch behavior. +Service and protocol adapters consume these components; this package is not a +public framework API. +""" diff --git a/dataretrieval/transport/http.py b/dataretrieval/transport/http.py new file mode 100644 index 00000000..1b2e29b3 --- /dev/null +++ b/dataretrieval/transport/http.py @@ -0,0 +1,96 @@ +"""HTTP client lifecycle, timeout defaults, and host-scoped authentication.""" + +from __future__ import annotations + +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager +from importlib.metadata import PackageNotFoundError +from importlib.metadata import version as _pkg_version +from typing import Any + +import httpx + +from dataretrieval.credentials import ( + accepts_api_key, + api_key, + strip_api_key_from_untrusted_host, + strip_api_key_from_untrusted_host_async, +) +from dataretrieval.exceptions import NetworkError + +# Re-exported for the adapters that reach for credential policy through the +# transport surface they already import. ``dataretrieval.credentials`` is the +# single definition; these names are views on it, not copies of it. +__all__ = [ + "HTTPX_ASYNC_DEFAULTS", + "HTTPX_DEFAULTS", + "USER_AGENT", + "accepts_api_key", + "default_headers", + "get", + "network_error", + "open_async_client", + "strip_api_key_from_untrusted_host", + "strip_api_key_from_untrusted_host_async", +] + +try: + _PACKAGE_VERSION = _pkg_version("dataretrieval") +except PackageNotFoundError: + _PACKAGE_VERSION = "version-unknown" + +USER_AGENT = f"python-dataretrieval/{_PACKAGE_VERSION}" + +HTTPX_DEFAULTS: dict[str, Any] = { + "follow_redirects": True, + "timeout": httpx.Timeout(60.0, connect=10.0), +} + + +def default_headers(target_url: str | httpx.URL | None = None) -> dict[str, str]: + """Build standard headers, scoping the API key to its authorized host.""" + headers = { + "Accept-Encoding": "compress, gzip", + "Accept": "application/json", + "User-Agent": USER_AGENT, + "lang": "en-US", + } + token = api_key() + if token and accepts_api_key(target_url): + headers["X-Api-Key"] = token + return headers + + +HTTPX_ASYNC_DEFAULTS: dict[str, Any] = { + **HTTPX_DEFAULTS, + "event_hooks": {"request": [strip_api_key_from_untrusted_host_async]}, +} + + +def network_error(url: str | httpx.URL, exc: httpx.TransportError) -> NetworkError: + """Build a typed error for a failed round trip with no HTTP response.""" + detail = str(exc) or type(exc).__name__ + return NetworkError(f"Could not reach the service at {url}: {detail}") + + +def get(url: str | httpx.URL, **kwargs: Any) -> httpx.Response: + """Issue one guarded synchronous GET and map transport failures.""" + client_options: dict[str, Any] = { + key: kwargs.pop(key) + for key in ("follow_redirects", "timeout", "transport", "verify") + if key in kwargs + } + client_options["event_hooks"] = {"request": [strip_api_key_from_untrusted_host]} + try: + with httpx.Client(**client_options) as client: + return client.get(url, **kwargs) + except httpx.TransportError as exc: + raise network_error(url, exc) from exc + + +@asynccontextmanager +async def open_async_client(**overrides: Any) -> AsyncIterator[httpx.AsyncClient]: + """Open a short-lived async client with redirect-safe shared defaults.""" + options = {**HTTPX_ASYNC_DEFAULTS, **overrides} + async with httpx.AsyncClient(**options) as client: + yield client diff --git a/dataretrieval/transport/liveness.py b/dataretrieval/transport/liveness.py new file mode 100644 index 00000000..26d78936 --- /dev/null +++ b/dataretrieval/transport/liveness.py @@ -0,0 +1,58 @@ +"""When data last arrived, shared by the loops that produce and consume it. + +A retrieval can be slow for two very different reasons: it is downloading a lot +(fine, however long it takes) or it is receiving nothing at all (worth giving up +on). Telling those apart needs one fact -- when data last arrived -- that the +page-walking loop knows and the retry loop acts on. Keeping it in this leaf lets +both point *down* at it rather than at each other, and leaves any future producer +of liveness (a streaming body reader, a chunk-level fetch) somewhere to report. + +The stamp lives in a :class:`~contextvars.ContextVar` so concurrent retrievals -- +each sub-request of a chunked call, each location of a Water Use fan-out -- +measure their own silence instead of sharing one clock. +""" + +from __future__ import annotations + +import contextvars +import time + +_last_progress: contextvars.ContextVar[float | None] = contextvars.ContextVar( + "transport_last_progress", default=None +) + + +def note_progress() -> None: + """Restart the no-progress budget: data just arrived.""" + _last_progress.set(time.monotonic()) + + +def elapsed_since_progress() -> float | None: + """Seconds since data last arrived, or ``None`` if nothing has reported yet.""" + last = _last_progress.get() + return None if last is None else time.monotonic() - last + + +def credit_wait(seconds: float) -> None: + """Excuse ``seconds`` of sanctioned waiting from the no-progress budget. + + Two kinds of waiting are not silence: queueing behind a concurrency cap, and + sleeping off a delay the server itself named (see + :meth:`~dataretrieval.transport.retry.RetryPolicy.allows_wait` for why a + sanctioned delay costs the budget nothing). The deep tail of a wide fan-out + can wait past the whole budget and would otherwise start its first attempt + with nothing left to retry with. But neither is progress, and the difference + matters: crediting only the measured wait keeps the budget cumulative across + attempts, where restamping to "now" would also discard silence accumulated + by earlier attempts and quietly turn a bound on total silence into a + per-attempt latency bound. + + The credit never reaches past the present. A wait longer than the whole + budget would otherwise stamp the stamp into the *future*, making + :func:`elapsed_since_progress` negative -- and since nothing ever pulls it + back, that one long queue wait would disable the bound for the rest of the + call, which is precisely the silent-minutes case the budget exists to catch. + """ + last = _last_progress.get() + if last is not None: + _last_progress.set(min(time.monotonic(), last + seconds)) diff --git a/dataretrieval/transport/pagination.py b/dataretrieval/transport/pagination.py new file mode 100644 index 00000000..6199138c --- /dev/null +++ b/dataretrieval/transport/pagination.py @@ -0,0 +1,131 @@ +"""Callback-driven cursor pagination independent of any service protocol.""" + +from __future__ import annotations + +import logging +from collections.abc import AsyncIterator, Awaitable, Callable +from contextlib import asynccontextmanager +from typing import Any, TypeVar + +import httpx +import pandas as pd + +from dataretrieval import progress as _progress +from dataretrieval.combining import ( + _QUOTA_HEADER, + _merge_response, + _safe_elapsed, +) +from dataretrieval.exceptions import DataRetrievalError, RateLimited +from dataretrieval.transport.http import open_async_client +from dataretrieval.transport.liveness import note_progress + +logger = logging.getLogger(__name__) +_Cursor = TypeVar("_Cursor") + + +@asynccontextmanager +async def _client_for( + client: httpx.AsyncClient | None, +) -> AsyncIterator[httpx.AsyncClient]: + """Borrow a caller client or open a guarded short-lived client.""" + if client is not None: + yield client + return + async with open_async_client() as new: + yield new + + +def paginated_failure_message(pages_collected: int, cause: BaseException) -> str: + """Build a recovery-oriented message for an interrupted page walk.""" + cause_str = str(cause).removesuffix(".") + if not cause_str.strip(): + cause_str = type(cause).__name__ + if isinstance(cause, RateLimited): + action = "wait for the rate-limit window to reset and retry" + else: + action = "retry the request (possibly after a short backoff)" + return ( + f"Paginated request failed after collecting {pages_collected} " + f"page(s): {cause_str}. To recover: {action}, reduce the " + f"request size (e.g. fewer locations, a shorter time range, or " + f"a smaller ``limit``), or obtain an API token." + ) + + +async def paginate( + initial_req: httpx.Request, + *, + parse_response: Callable[[httpx.Response], tuple[pd.DataFrame, _Cursor | None]], + follow_up: Callable[[_Cursor, httpx.AsyncClient], Awaitable[httpx.Response]], + raise_for_status: Callable[[httpx.Response], None], + client: httpx.AsyncClient | None = None, + row_cap: int | None = None, +) -> tuple[pd.DataFrame, httpx.Response]: + """Fetch and combine pages until the injected parser returns no cursor. + + The service adapter supplies response parsing, cursor following, and status + mapping. This loop owns client lifecycle, repeated-cursor protection, + optional row capping, progress updates, failure wrapping, and response + metadata aggregation. + """ + logger.debug("Requesting: %s", initial_req.url) + reporter = _progress.current() + + def report_page(page: httpx.Response, frame: pd.DataFrame) -> None: + note_progress() # a walk still delivering pages is not stalled + if reporter is not None: + reporter.set_rate_remaining( + page.headers.get(_QUOTA_HEADER), + limit=page.headers.get("x-ratelimit-limit"), + ) + reporter.add_page(rows=len(frame)) + + async with _client_for(client) as session: + response = await session.send(initial_req) + raise_for_status(response) + initial_response = response + total_elapsed = _safe_elapsed(response) + + try: + frame, cursor = parse_response(response) + except Exception as exc: # noqa: BLE001 + logger.warning("Initial response parse failed.") + raise DataRetrievalError(paginated_failure_message(0, exc)) from exc + + frames = [frame] + nrows = len(frame) + seen: set[Any] = set() + report_page(response, frame) + + while ( + cursor is not None + and cursor not in seen + and (row_cap is None or nrows < row_cap) + ): + seen.add(cursor) + try: + response = await follow_up(cursor, session) + raise_for_status(response) + frame, cursor = parse_response(response) + frames.append(frame) + nrows += len(frame) + total_elapsed += _safe_elapsed(response) + report_page(response, frame) + except Exception as exc: # noqa: BLE001 + logger.warning( + "Request failed at cursor %r. Data download interrupted.", cursor + ) + raise DataRetrievalError( + paginated_failure_message(len(frames), exc) + ) from exc + + final_response = _merge_response( + initial_response, + headers_from=response, + elapsed=total_elapsed, + ) + result = pd.concat(frames, ignore_index=True) + if row_cap is not None: + result = result.head(row_cap) + return result, final_response diff --git a/dataretrieval/transport/retry.py b/dataretrieval/transport/retry.py new file mode 100644 index 00000000..577ed334 --- /dev/null +++ b/dataretrieval/transport/retry.py @@ -0,0 +1,445 @@ +"""Bounded retry policy and transient-failure classification.""" + +from __future__ import annotations + +import asyncio +import math +import os +import random +import socket +import time +from collections.abc import Awaitable, Callable +from dataclasses import dataclass +from datetime import datetime, timezone +from email.utils import parsedate_to_datetime +from typing import NamedTuple, TypeVar + +import httpx + +from dataretrieval import progress as _progress +from dataretrieval.exceptions import ( + ConfigurationError, + NetworkError, + TransientError, +) +from dataretrieval.transport.liveness import ( + credit_wait, + elapsed_since_progress, + note_progress, +) + +# Which error statuses a request may be re-sent for. Both are narrower than +# :attr:`~dataretrieval.exceptions.DataRetrievalError.retryable`, deliberately: +# that field tells a caller re-issuing *might* work, while spending someone's +# quota unasked needs a stricter bar. +# +# The default keeps every 5xx, because for a query interface like the Water Data +# OGC API a 500 is an upstream hiccup and re-sending is how a chunked call rides +# one out. The gateway-only set is for the single-shot adapters whose services +# answer a *bad query* with a 500 -- WQP does that for an over-large request, +# StreamStats for out-of-network coordinates -- where re-sending multiplies load +# on a request that can never succeed and delays the caller's error. +_RETRYABLE_STATUSES = frozenset({429, *range(500, 600)}) +_GATEWAY_STATUSES = frozenset({429, 502, 503, 504}) +_RETRIES_ENV = "API_USGS_RETRIES" +_RETRIES_DEFAULT = 4 +_RETRY_BASE_BACKOFF = 0.5 +_RETRY_MAX_BACKOFF = 30.0 +_RETRY_AFTER_CAP = 60.0 +# Most a server-named delay is nudged by, to keep sub-requests handed the same +# 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" +_STALL_TIMEOUT_DEFAULT = 60.0 + +_T = TypeVar("_T") +_Number = TypeVar("_Number", int, float) + + +def parse_retry_after(value: str | None) -> float | None: + """Parse a ``Retry-After`` header into seconds, or ``None`` for no usable hint. + + Both header forms mean the same thing and are treated the same way: the + seconds are returned as given, however large. A value past what a caller will + wait out inline stops the retry and surfaces a transient carrying the hint on + ``.retry_after``, so a long wait becomes the caller's decision (and, for a + chunked call, a resumable interruption) instead of being ignored. + + An over-long hint is honored rather than discarded. Dropping it would make + the client retry *harder* against a service that just asked for a long + pause, and would deny the caller the number it needs on ``.retry_after``. + Clock skew can inflate a date-form hint, but trusting one costs a + recoverable escalation while ignoring it costs hammering a service that is + already asking for room. + + A date that has *already* passed yields no hint at all rather than ``0.0``. + Read literally it says "retry now", but the likelier reading is that our + clock runs ahead of the server's -- and acting on it would re-send almost + immediately against a service that just asked for a pause. Falling back to + our own bounded backoff is right under either reading. (Delta-seconds is + clock-independent, so a literal ``Retry-After: 0`` is still honored as the + instruction it is, floored by :meth:`RetryPolicy.backoff`'s jitter.) + """ + if not value: + return None + raw = value.strip() + try: + seconds = float(raw) + except ValueError: + pass + else: + # ``inf``/``nan`` parse cleanly but poison every later comparison: an + # infinite hint would refuse retry forever and travel to the caller on + # ``.retry_after``. Treat them as no hint at all. + return max(0.0, seconds) if math.isfinite(seconds) else None + try: + retry_at = parsedate_to_datetime(raw) + except (TypeError, ValueError, OverflowError): + return None + if retry_at.tzinfo is None: + retry_at = retry_at.replace(tzinfo=timezone.utc) + delay = (retry_at - datetime.now(timezone.utc)).total_seconds() + return delay if delay > 0 else None + + +def _read_env_number( + name: str, default: _Number, cast: Callable[[str], _Number], expected: str +) -> _Number: + """Read a non-negative number from the environment, or ``default`` if unset. + + Raises :class:`~dataretrieval.exceptions.ConfigurationError` -- a + ``DataRetrievalError`` *and* a ``ValueError`` -- for an unusable value, so a + typo in the environment doesn't escape a request path as a bare + ``ValueError`` that ``except DataRetrievalError`` misses. + """ + raw = os.environ.get(name, "").strip() + if not raw: + return default + try: + value = cast(raw) + except ValueError as exc: + raise ConfigurationError(f"{name} must be {expected} (got {raw!r}).") from exc + # ``nan`` passes every ordering test, so a bare ``< 0`` guard lets it through + # and then silently makes each budget comparison false. + if not math.isfinite(value): + raise ConfigurationError(f"{name} must be {expected} (got {raw!r}).") + if value < 0: + raise ConfigurationError(f"{name} must be >= 0 (got {value}).") + return value + + +@dataclass(frozen=True) +class RetryPolicy: + """Immutable bounded exponential-backoff-with-full-jitter policy. + + Two independent bounds decide when to stop: :attr:`max_retries` caps *how + many* attempts a failure gets, and :attr:`stall_timeout` caps *how long* a + call may go on receiving nothing. + """ + + #: Attempts after the first. ``0`` disables retry entirely. + max_retries: int = _RETRIES_DEFAULT + #: First backoff ceiling; doubles per attempt up to :attr:`max_backoff`. + base_backoff: float = _RETRY_BASE_BACKOFF + #: Ceiling for our own exponential backoff. + max_backoff: float = _RETRY_MAX_BACKOFF + #: Longest server-named ``Retry-After`` we are willing to wait out inline. + #: A longer one stops the retry and surfaces a resumable transient, so the + #: caller decides whether to wait rather than blocking inside the request. + retry_after_cap: float = _RETRY_AFTER_CAP + #: Error statuses this policy will re-send for. Defaults to 429 and every + #: 5xx; single-shot adapters whose service reports a rejected query as a 500 + #: pass :data:`_GATEWAY_STATUSES` instead. + retryable_statuses: frozenset[int] = _RETRYABLE_STATUSES + #: Longest a call may go *without receiving any data* before retrying stops + #: and the failure surfaces -- the total of every silent attempt and every + #: unsanctioned wait since the last page arrived (a server-named + #: ``Retry-After`` and time queued behind the concurrency gate are excused; + #: see :meth:`allows_wait`). Bounds the wall-clock cost of a dead + #: connection or a service that keeps refusing, which :attr:`max_retries` + #: alone does not: it counts attempts, not seconds, so four retries of a + #: request that times out after a minute is four silent minutes. Progress + #: resets the clock (see + #: :func:`~dataretrieval.transport.liveness.note_progress`), so a slow but + #: productive download is never cut short, and an attempt already in flight + #: is never interrupted. ``0`` disables the bound. See :meth:`allows_wait` + #: for how it is applied. + stall_timeout: float = _STALL_TIMEOUT_DEFAULT + + def __post_init__(self) -> None: + if self.max_retries < 0: + raise ConfigurationError( + f"max_retries must be >= 0 (got {self.max_retries})." + ) + if ( + self.base_backoff < 0 + or self.max_backoff < 0 + or self.retry_after_cap < 0 + or self.stall_timeout < 0 + ): + raise ConfigurationError("retry backoff settings must be non-negative.") + + @classmethod + def from_env(cls, retryable_statuses: frozenset[int] | None = None) -> RetryPolicy: + """Build a policy from current environment and module defaults.""" + statuses = ( + _RETRYABLE_STATUSES if retryable_statuses is None else retryable_statuses + ) + return cls( + retryable_statuses=statuses, + max_retries=_read_env_number( + _RETRIES_ENV, _RETRIES_DEFAULT, int, "a non-negative integer" + ), + base_backoff=_RETRY_BASE_BACKOFF, + max_backoff=_RETRY_MAX_BACKOFF, + retry_after_cap=_RETRY_AFTER_CAP, + stall_timeout=_read_env_number( + _STALL_TIMEOUT_ENV, + _STALL_TIMEOUT_DEFAULT, + float, + "a non-negative number of seconds", + ), + ) + + def should_retry(self, attempt: int, retry_after: float | None) -> bool: + """Whether a just-failed 1-based attempt warrants another try.""" + if attempt > self.max_retries: + return False + return retry_after is None or retry_after <= self.retry_after_cap + + def allows_wait( + self, + attempt: int, + delay: float, + elapsed: float | None, + retry_after: float | None = None, + ) -> bool: + """Whether waiting ``delay`` more fits the no-progress budget. + + ``elapsed`` is the silence so far (see + :func:`~dataretrieval.transport.liveness.elapsed_since_progress`), passed + in rather than read here so the policy stays a pure value object. + + The first retry is always allowed. One slow attempt can spend the whole + budget on its own -- a heavy page against a loaded service, or any + attempt that runs to the read timeout -- and letting that suppress retry + entirely would turn a recoverable transient into an immediate failure + for exactly the large queries that most need retrying. So the budget + bounds *repeated* silence: with the defaults a dead connection costs + about two read timeouts rather than five attempts' worth. + + A delay the *server* named -- ``retry_after`` is not ``None``, the same + hint :meth:`should_retry` and :meth:`backoff` take -- costs the budget + nothing. Charging for it would mean a service that answers 429 with + ``Retry-After: 30`` gets fewer retries than one that says nothing at all + -- with the shipped defaults (a 60 s budget, a 60 s + :attr:`retry_after_cap`) any honored hint of half the budget or more + would allow exactly one retry no matter what + :attr:`max_retries` says. Waiting because we were told to is not the + service going quiet on us; it is the service telling us when to come + back. The driver credits the same wait back afterwards (see + :func:`~dataretrieval.transport.liveness.credit_wait`) so it doesn't + accumulate into the *next* attempt's silence either. + """ + if attempt <= _STALL_EXEMPT_ATTEMPTS: + return True + if self.stall_timeout <= 0 or elapsed is None: + return True + return ( + elapsed + (0.0 if retry_after is not None else delay) <= self.stall_timeout + ) + + def backoff(self, attempt: int, retry_after: float | None) -> float: + """Seconds to wait before a 1-based retry attempt. + + A jittered component is always included, even when the server named a + delay: a hint of ``0`` -- or a ``Retry-After`` date that has already + passed -- would otherwise become a zero-delay re-send against a service + that just asked us to slow down, and sub-requests handed the same hint + would all wake at the same instant and burst together. + + On a server hint that jitter is a small decorrelating nudge rather than + a second backoff, and the total is held to :attr:`retry_after_cap`: + full jitter on top of a hint already at the cap would sleep half again + as long as any bound this policy declares. It is bounded by + :attr:`max_backoff` rather than by this attempt's exponential ceiling, + so it survives a :attr:`base_backoff` of zero -- the case where the + ceiling collapses and a hint of ``0`` would otherwise become exactly the + zero-delay re-send this prevents. A policy that declares no backoff at + all still gets none. + """ + ceiling = min(self.max_backoff, self.base_backoff * 2 ** (attempt - 1)) + if retry_after is None: + return random.uniform(0.0, ceiling) + nudge = random.uniform(0.0, min(self.max_backoff, _RETRY_AFTER_JITTER)) + return min(retry_after + nudge, self.retry_after_cap) + + +_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]: + """Return whether ``exc`` is safe to retry and any server delay hint.""" + if isinstance(exc, TransientError): + if exc.status_code is not None and exc.status_code not in statuses: + return False, None + return True, exc.retry_after + if isinstance(exc, (NetworkError, httpx.TransportError)): + return not _deterministic_failure(exc), None + return False, None + + +class _Wait(NamedTuple): + """How long to hold off before a retry, and whether the server asked for it. + + ``sanctioned`` travels with the delay because only the driver knows when the + sleep finished, and a server-named wait has to be credited back to the + no-progress budget once it has been served (see + :meth:`RetryPolicy.allows_wait`). + """ + + delay: float + sanctioned: bool + + def settle(self) -> None: + """Credit a served server-named wait back to the no-progress budget. + + Paired with the sleep rather than left to each driver: a wait that + :meth:`RetryPolicy.allows_wait` excused going in has to be excused coming + out too, or it accumulates into the *next* attempt's silence and caps the + retries anyway. Both drivers sleep differently but settle identically. + """ + if self.sanctioned: + credit_wait(self.delay) + + +def _retry_delay(exc: BaseException, attempt: int, policy: RetryPolicy) -> _Wait | None: + """Return the bounded wait for a failed attempt, or ``None`` to stop.""" + retryable, retry_after = _retryable(exc, policy.retryable_statuses) + if not retryable or not policy.should_retry(attempt, retry_after): + return None + delay = policy.backoff(attempt, retry_after) + if not policy.allows_wait(attempt, delay, elapsed_since_progress(), retry_after): + return None + reporter = _progress.current() + if reporter is not None: + reporter.note_retry(attempt=attempt, wait=delay) + return _Wait(delay, retry_after is not None) + + +async def retry_async( + afn: Callable[[], Awaitable[_T]], + policy: RetryPolicy | None = None, + *, + gate: asyncio.Semaphore | None = None, +) -> _T: + """Call an awaitable with bounded retry on typed transient failures. + + ``gate`` bounds how many attempts run concurrently. Owning it here rather + than letting each caller wrap its own body keeps two rules in one place: the + slot is acquired per *attempt*, so a call sleeping off a backoff isn't + holding one while it isn't touching the server, and the time spent waiting + for it is credited back to the no-progress budget rather than counted as + silence. A caller that gated its own body would have to rediscover both, and + nothing would catch it getting them wrong. + """ + policy = RetryPolicy.from_env() if policy is None else policy + attempt = 0 + note_progress() + + async def attempt_once() -> _T: + if gate is None: + return await afn() + started = time.monotonic() + async with gate: + credit_wait(time.monotonic() - started) + return await afn() + + while True: + try: + return await attempt_once() + except Exception as exc: # noqa: BLE001 - re-raised unless retryable + attempt += 1 + wait = _retry_delay(exc, attempt, policy) + if wait is None: + raise + await asyncio.sleep(wait.delay) + wait.settle() + + +def retry_sync(fn: Callable[[], _T], policy: RetryPolicy | None = None) -> _T: + """Call a synchronous operation with bounded retry on typed transients. + + ``KeyboardInterrupt``, ``SystemExit``, and other cancellation signals are + not caught because the loop handles ``Exception`` rather than + ``BaseException``. + """ + policy = RetryPolicy.from_env() if policy is None else policy + attempt = 0 + note_progress() + while True: + try: + return fn() + except Exception as exc: # noqa: BLE001 - re-raised unless retryable + attempt += 1 + wait = _retry_delay(exc, attempt, policy) + if wait is None: + raise + time.sleep(wait.delay) + wait.settle() diff --git a/dataretrieval/transport/sync.py b/dataretrieval/transport/sync.py new file mode 100644 index 00000000..f9799d2a --- /dev/null +++ b/dataretrieval/transport/sync.py @@ -0,0 +1,29 @@ +"""Synchronous dispatch over asynchronous retrieval internals.""" + +from __future__ import annotations + +from collections.abc import Awaitable, Callable +from typing import TypeVar, cast + +import httpx +from anyio.from_thread import start_blocking_portal + +from dataretrieval import progress as _progress +from dataretrieval.transport.http import network_error + +_T = TypeVar("_T") + + +def run_sync( + make_coro: Callable[[], Awaitable[_T]], + *, + service: str, + error_url: str | httpx.URL, +) -> _T: + """Run an async retrieval from synchronous code in a blocking portal.""" + with _progress.progress_context(service=service, target_url=error_url): + with start_blocking_portal() as portal: + try: + return cast("_T", portal.call(make_coro)) + except httpx.TransportError as exc: + raise network_error(error_url, exc) from exc diff --git a/dataretrieval/utils.py b/dataretrieval/utils.py index 7506a469..bab5929e 100644 --- a/dataretrieval/utils.py +++ b/dataretrieval/utils.py @@ -5,38 +5,42 @@ from __future__ import annotations import numbers -import os import warnings from collections.abc import Callable, Iterable, Iterator from contextlib import contextmanager from contextvars import ContextVar -from importlib.metadata import PackageNotFoundError -from importlib.metadata import version as _pkg_version from typing import Any, Generic, TypeVar import httpx import pandas as pd +import dataretrieval.credentials as _credentials +import dataretrieval.transport.http as _transport_http from dataretrieval.codes import tz from dataretrieval.exceptions import ( - NetworkError, NoSitesError, URLTooLong, error_for_status, ) +from dataretrieval.transport.retry import ( + _GATEWAY_STATUSES, + RetryPolicy, + parse_retry_after, + retry_sync, +) -try: - _PACKAGE_VERSION = _pkg_version("dataretrieval") -except PackageNotFoundError: - _PACKAGE_VERSION = "version-unknown" - -# Typed as ``dict[str, Any]`` (not the inferred ``dict[str, object]``) so that -# splatting it as ``**HTTPX_DEFAULTS`` into ``httpx.get`` / ``httpx.AsyncClient`` -# type-checks: the values are a heterogeneous bag of httpx keyword arguments. -HTTPX_DEFAULTS: dict[str, Any] = { - "follow_redirects": True, - "timeout": httpx.Timeout(60.0, connect=10.0), -} +# Compatibility names retained at their historical utility paths. +_AUTHORIZED_API_KEY_HOST = _credentials._AUTHORIZED_API_KEY_HOST +HTTPX_ASYNC_DEFAULTS = _transport_http.HTTPX_ASYNC_DEFAULTS +HTTPX_DEFAULTS = _transport_http.HTTPX_DEFAULTS +USER_AGENT = _transport_http.USER_AGENT +_default_headers = _transport_http.default_headers +_get = _transport_http.get +_network_error = _transport_http.network_error +_strip_api_key_from_untrusted_host = _transport_http.strip_api_key_from_untrusted_host +_strip_api_key_from_untrusted_host_async = ( + _transport_http.strip_api_key_from_untrusted_host_async +) _T = TypeVar("_T") @@ -103,75 +107,6 @@ def _require_positive_int( raise ValueError(f"{name} must be a positive integer{eg} (got {value!r}).") -# The single authorized host for the API key. The key is a USGS personal -# access token issued for the Water Data API and must never be forwarded to -# non-USGS hosts, lookalikes, or external endpoints (including STAC rating -# asset downloads). -_AUTHORIZED_API_KEY_HOST = "api.waterdata.usgs.gov" - - -def _default_headers(target_url: str | httpx.URL | None = None) -> dict[str, str]: - """Build the default HTTP headers for a USGS web-API request. - - Always sets a descriptive ``User-Agent`` plus ``Accept`` / - ``Accept-Encoding`` and ``lang``. If the ``API_USGS_PAT`` environment - variable is set AND ``target_url`` points to the explicitly authorized - Water Data host (``api.waterdata.usgs.gov``), its value is added as the - ``X-Api-Key`` header. The key is never sent to other hosts. - - Parameters - ---------- - target_url : str or httpx.URL or None - The URL the request will be sent to. When provided, the API key is - included only if the host matches the authorized Water Data host. - When ``None`` (legacy callers), the key is NOT included — callers - must pass the concrete URL. - - Returns - ------- - dict[str, str] - Headers suitable for an ``httpx`` request. - """ - headers = { - "Accept-Encoding": "compress, gzip", - "Accept": "application/json", - "User-Agent": f"python-dataretrieval/{_PACKAGE_VERSION}", - "lang": "en-US", - } - token = os.getenv("API_USGS_PAT") - if token and target_url is not None: - try: - host = httpx.URL(str(target_url)).host - except (httpx.InvalidURL, TypeError): - host = None - if host == _AUTHORIZED_API_KEY_HOST: - headers["X-Api-Key"] = token - return headers - - -def _strip_api_key_from_untrusted_host(request: httpx.Request) -> None: - """Remove Water Data credentials from any request to another host. - - HTTPX retains arbitrary custom headers across cross-origin redirects. This - hook runs for the initial request and every redirect, making host scoping an - execution-time invariant rather than relying only on initial header - construction. - """ - if request.url.host != _AUTHORIZED_API_KEY_HOST: - request.headers.pop("X-Api-Key", None) - - -async def _strip_api_key_from_untrusted_host_async(request: httpx.Request) -> None: - """Async-client form of :func:`_strip_api_key_from_untrusted_host`.""" - _strip_api_key_from_untrusted_host(request) - - -HTTPX_ASYNC_DEFAULTS: dict[str, Any] = { - **HTTPX_DEFAULTS, - "event_hooks": {"request": [_strip_api_key_from_untrusted_host_async]}, -} - - def to_str(listlike: object, delimiter: str = ",") -> str | None: """Translates list-like objects into strings. @@ -427,35 +362,6 @@ def _url_too_long_error(detail: str) -> URLTooLong: ) -def _network_error(url: str | httpx.URL, exc: httpx.TransportError) -> NetworkError: - """Build the :class:`~dataretrieval.exceptions.NetworkError` for a failed - round-trip ``exc`` (no HTTP response: timeout, DNS, refused connection).""" - # Some httpx transport errors stringify empty (e.g. ``ConnectTimeout()``); - # fall back to the class name so the message is always informative. - detail = str(exc) or type(exc).__name__ - return NetworkError(f"Could not reach the service at {url}: {detail}") - - -def _get(url: str | httpx.URL, **kwargs: Any) -> httpx.Response: - """Issue one guarded synchronous GET and map transport failures. - - A short-lived client supplies a request hook for both the initial request - and every redirect. The hook removes ``X-Api-Key`` unless the destination - host is the authorized Water Data API host. - """ - client_options: dict[str, Any] = { - key: kwargs.pop(key) - for key in ("follow_redirects", "timeout", "transport", "verify") - if key in kwargs - } - client_options["event_hooks"] = {"request": [_strip_api_key_from_untrusted_host]} - try: - with httpx.Client(**client_options) as client: - return client.get(url, **kwargs) - except httpx.TransportError as exc: - raise _network_error(url, exc) from exc - - def _raise_for_status( response: httpx.Response, *, @@ -487,14 +393,53 @@ def _raise_for_status( if detail: message += f": {detail}" message += f" (URL: {response.url})" - raise error_for_status(status, message) + raise error_for_status( + status, + message, + retry_after=parse_retry_after(response.headers.get("Retry-After")), + ) -def query( +def _single_request_policy() -> RetryPolicy: + """Retry policy for the one-shot adapters (WQP, NLDI, StreamStats). + + These services answer a rejected query with a 500, so only the gateway + statuses are worth re-sending; the Water Data chunker keeps the broader + default, where a 5xx is an upstream hiccup worth riding out. + """ + return RetryPolicy.from_env(retryable_statuses=_GATEWAY_STATUSES) + + +def _get_with_retry( + url: str | httpx.URL, + *, + detail_from: Callable[[httpx.Response], str | None] | None = None, + retry_policy: RetryPolicy | None = None, + **kwargs: Any, +) -> httpx.Response: + """GET with status mapping and bounded retry on typed transients.""" + + def attempt() -> httpx.Response: + response = _get(url, **kwargs) + _raise_for_status(response, detail_from=detail_from) + return response + + try: + return retry_sync( + attempt, + _single_request_policy() if retry_policy is None else retry_policy, + ) + except httpx.InvalidURL as exc: + raise _url_too_long_error(f"httpx rejected the URL client-side: {exc}") from exc + + +def _query_impl( url: str, payload: dict[str, Any], delimiter: str = ",", ssl_check: bool = True, + *, + retry_policy: RetryPolicy, ) -> httpx.Response: """Send a query. @@ -535,20 +480,16 @@ def query( # Drop them. (``to_str`` returns None for non-iterable scalars like bools.) payload = {k: v for k, v in payload.items() if v is not None} - user_agent = {"user-agent": f"python-dataretrieval/{_PACKAGE_VERSION}"} + user_agent = {"user-agent": USER_AGENT} - try: - response = _get( - url, - params=payload, - headers=user_agent, - verify=ssl_check, - **HTTPX_DEFAULTS, - ) - except httpx.InvalidURL as exc: - raise _url_too_long_error(f"httpx rejected the URL client-side: {exc}") from exc - - _raise_for_status(response) + response = _get_with_retry( + url, + params=payload, + headers=user_agent, + verify=ssl_check, + retry_policy=retry_policy, + **HTTPX_DEFAULTS, + ) # USGS waterservices signals an empty result with a 200 whose body starts # "No sites/data ..." (its legacy wording); surface it as NoSitesError. @@ -556,3 +497,37 @@ def query( raise NoSitesError(response.url) return response + + +def query( + url: str, + payload: dict[str, Any], + delimiter: str = ",", + ssl_check: bool = True, +) -> httpx.Response: + return _query_impl( + url, + payload, + delimiter, + ssl_check, + retry_policy=RetryPolicy(max_retries=0), + ) + + +query.__doc__ = _query_impl.__doc__ + + +def _query_with_retry( + url: str, + payload: dict[str, Any], + delimiter: str = ",", + ssl_check: bool = True, +) -> httpx.Response: + """Active-service form of :func:`query` with bounded transient retry.""" + return _query_impl( + url, + payload, + delimiter, + ssl_check, + retry_policy=_single_request_policy(), + ) diff --git a/dataretrieval/waterdata/api.py b/dataretrieval/waterdata/api.py index 7073d183..3f2e8dcb 100644 --- a/dataretrieval/waterdata/api.py +++ b/dataretrieval/waterdata/api.py @@ -25,14 +25,16 @@ _construct_cql_request, _switch_properties_id, ) -from dataretrieval.utils import ( +from dataretrieval.transport.http import ( HTTPX_DEFAULTS, - BaseMetadata, - _attach_datetime_columns, - _default_headers, - _get, - to_str, ) +from dataretrieval.transport.http import ( + default_headers as _default_headers, +) +from dataretrieval.transport.http import ( + get as _get, +) +from dataretrieval.utils import BaseMetadata, _attach_datetime_columns, to_str from dataretrieval.waterdata import stats from dataretrieval.waterdata.types import ( CODE_SERVICES, diff --git a/dataretrieval/waterdata/ratings.py b/dataretrieval/waterdata/ratings.py index cbaab057..6cdc3896 100644 --- a/dataretrieval/waterdata/ratings.py +++ b/dataretrieval/waterdata/ratings.py @@ -17,13 +17,22 @@ import httpx import pandas as pd +from dataretrieval.credentials import without_embedded_credentials from dataretrieval.exceptions import DataRetrievalError from dataretrieval.ogc.dates import _DURATION_RE, _format_api_dates from dataretrieval.ogc.errors import _raise_for_non_200 from dataretrieval.ogc.filters import _quote_cql_str from dataretrieval.ogc.requests import _check_monitoring_location_id from dataretrieval.rdb import extract_rdb_comment, read_rdb -from dataretrieval.utils import HTTPX_DEFAULTS, _default_headers, _get +from dataretrieval.transport.http import ( + HTTPX_DEFAULTS, +) +from dataretrieval.transport.http import ( + default_headers as _default_headers, +) +from dataretrieval.transport.http import ( + get as _get, +) from .utils import BASE_URL @@ -234,6 +243,29 @@ def _search( STAC ``next`` link is followed until exhausted so a result set larger than one page isn't silently truncated. """ + + def _checked_next_url(href: str, current: httpx.URL) -> str: + """Resolve a server-supplied ``next`` href into a URL safe to request.""" + try: + target = httpx.URL(href) + except (httpx.InvalidURL, TypeError) as exc: + raise DataRetrievalError( + f"The ratings service returned an unusable next-page link: " + f"{href!r}. The page walk cannot continue; report this if it " + f"persists." + ) from exc + if not target.is_absolute_url: + target = current.join(target) + if target.host != current.host: + raise DataRetrievalError( + f"Refusing to follow a ratings next-page link pointing at " + f"{target.host} rather than {current.host}. Following it would " + f"send this request, and any credentials on it, to a host you " + f"did not ask for. Retrying will not help; report this if it " + f"persists." + ) + return str(without_embedded_credentials(target)) + query_params: dict[str, Any] = {"limit": min(limit, 10000)} if filter_str is not None: query_params["filter"] = filter_str @@ -261,10 +293,16 @@ def _search( # The STAC ``next`` link is a fully-formed GET href carrying the # limit/filter/bbox and a continuation token, so follow it verbatim # (dropping our own params) until the server stops emitting one. - url = next( + href = next( (lnk["href"] for lnk in body.get("links", []) if lnk.get("rel") == "next"), None, ) + # Verbatim except for the credentials: the href is response data, so it + # is checked before it becomes a request. A link to another host would + # carry this request's API key off the authorized host, and one carrying + # ``user:pass@`` would mint an ``Authorization: Basic`` header the caller + # never configured. + url = None if href is None else _checked_next_url(href, response.url) params = None return features diff --git a/dataretrieval/waterdata/stats.py b/dataretrieval/waterdata/stats.py index abba5deb..ba4097fc 100644 --- a/dataretrieval/waterdata/stats.py +++ b/dataretrieval/waterdata/stats.py @@ -3,8 +3,9 @@ Wraps ``https://api.waterdata.usgs.gov/statistics/v0`` — the daily-statistics service (period-of-record and date-range normals/intervals). This is a *separate*, non-OGC API: it has no chunkable multi-value axes, so it drives -:func:`engine._paginate` directly through a blocking portal rather than going -through ``multi_value_chunked``. The typed getters ``get_stats_por`` and +:func:`dataretrieval.transport.pagination.paginate` through the shared sync +bridge rather than going through ``multi_value_chunked``. The typed getters +``get_stats_por`` and ``get_stats_date_range`` in :mod:`dataretrieval.waterdata.api` call :func:`get_data` here. """ @@ -16,17 +17,17 @@ import httpx import pandas as pd -from dataretrieval.ogc.engine import ( - _paginate, - _run_sync, -) +from dataretrieval.ogc.errors import _raise_for_non_200 from dataretrieval.ogc.shaping import ( _CRS, GEOPANDAS, _attach_coordinates, _empty_feature_frame, ) -from dataretrieval.utils import BaseMetadata, _default_headers +from dataretrieval.transport.http import default_headers +from dataretrieval.transport.pagination import paginate +from dataretrieval.transport.sync import run_sync +from dataretrieval.utils import BaseMetadata from dataretrieval.waterdata.utils import BASE_URL # ``_handle_nesting``'s geopandas branch calls ``gpd.GeoDataFrame.from_features`` @@ -220,7 +221,7 @@ def get_data( to the specified parameters. The stats path doesn't go through ``multi_value_chunked`` (its query - shape has no chunkable list axes), so it drives :func:`engine._paginate` + shape has no chunkable list axes), so it drives transport pagination directly through an ``anyio`` blocking portal. The portal runs the pagination loop in a short-lived worker thread, so this works whether or not the caller is already inside an event loop. @@ -238,8 +239,12 @@ def get_data( True and the user requests a computation_type other than percentiles, a percentile column is still returned. client : httpx.AsyncClient, optional - Caller-borrowed async client. ``None`` (default) opens a - temporary one inside the portal. Primarily a test seam. + Caller-borrowed async client. ``None`` (default) opens a temporary one + inside the portal. Primarily a test seam. Deliberately does *not* fall + back to the chunker's shared client: that client belongs to the + chunker's event loop, and this runs in its own portal loop, so driving + it from here would corrupt the connection pool. Statistics is a + standalone API and never runs nested inside a chunked call anyway. Returns ------- @@ -251,7 +256,8 @@ def get_data( Raises ------ DataRetrievalError - The typed subclass for an HTTP error response (see :func:`engine._paginate`); + The typed subclass for an HTTP error response (see + :func:`transport.pagination.paginate`); or :class:`~dataretrieval.exceptions.NetworkError` if the initial request can't reach the service (timeout / DNS), the ``httpx`` exception chained on ``__cause__``. @@ -261,7 +267,7 @@ def get_data( req = httpx.Request( method="GET", url=url, - headers=_default_headers(url), + headers=default_headers(url), params=args, ) method = req.method @@ -282,14 +288,15 @@ async def follow_up(cursor: str, sess: httpx.AsyncClient) -> httpx.Response: ) async def _run() -> tuple[pd.DataFrame, httpx.Response]: - return await _paginate( + return await paginate( req, parse_response=parse_response, follow_up=follow_up, client=client, + raise_for_status=_raise_for_non_200, ) - df, response = _run_sync(_run, service=service) + df, response = run_sync(_run, service=service, error_url=url) if expand_percentiles: df = _expand_percentiles(df) diff --git a/dataretrieval/waterdata/utils.py b/dataretrieval/waterdata/utils.py index a6d91272..aff35a7a 100644 --- a/dataretrieval/waterdata/utils.py +++ b/dataretrieval/waterdata/utils.py @@ -27,6 +27,7 @@ import dataretrieval.ogc.dates as _ogc_dates import dataretrieval.ogc.shaping as _ogc_shaping from dataretrieval.codes.states import apply_state +from dataretrieval.credentials import WATERDATA_BASE_URL from dataretrieval.ogc import OgcDialect, prepare_request_args from dataretrieval.ogc import get_ogc_data as _facade_get_ogc_data from dataretrieval.utils import BaseMetadata @@ -37,10 +38,12 @@ ) # --------------------------------------------------------------------------- -# Water Data endpoint constants (defined locally, not imported from OGC policy) +# Water Data endpoint constants. The authority comes from the credentials leaf +# -- the host that serves these endpoints is the host that honors the API key -- +# while the paths below stay local rather than importing OGC policy internals. # --------------------------------------------------------------------------- -BASE_URL = "https://api.waterdata.usgs.gov" +BASE_URL = WATERDATA_BASE_URL OGC_API_URL = f"{BASE_URL}/ogcapi/v0" SAMPLES_URL = f"{BASE_URL}/samples-data" diff --git a/dataretrieval/wateruse.py b/dataretrieval/wateruse.py index 53de7597..a3c788ee 100644 --- a/dataretrieval/wateruse.py +++ b/dataretrieval/wateruse.py @@ -12,11 +12,9 @@ (:mod:`dataretrieval.ngwmn`), the NWDC is a plain CSV REST service rather than an OGC API Features collection. This module supplies the NWDC-specific bits — request building, CSV parsing, the ``Link``-header cursor, and the ``{detail}`` -error envelope — but reuses the OGC engine's generic, API-agnostic pagination -and sync-from-async plumbing (:func:`~dataretrieval.ogc.engine._paginate` and -:func:`~dataretrieval.ogc.engine._run_sync`) rather than re-implementing it. It -follows the same conventions: shared request headers -(:func:`~dataretrieval.utils._default_headers`), the typed +error envelope — and uses the service-neutral transport layer for cursor pagination, +response aggregation, client lifecycle, and sync-from-async dispatch. It follows +the same conventions: host-scoped request headers, the typed :class:`~dataretrieval.exceptions.DataRetrievalError` taxonomy, and a ``(DataFrame, BaseMetadata)`` return. @@ -52,18 +50,22 @@ import pandas as pd from dataretrieval.codes.states import to_state -from dataretrieval.exceptions import DataRetrievalError -from dataretrieval.ogc.combining import _combine_chunk_frames, _combine_chunk_responses -from dataretrieval.ogc.engine import _paginate, _run_sync -from dataretrieval.utils import ( - HTTPX_ASYNC_DEFAULTS, - BaseMetadata, - _default_headers, - _raise_for_status, - to_str, +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.pagination import paginate +from dataretrieval.transport.retry import RetryPolicy, retry_async +from dataretrieval.transport.sync import run_sync +from dataretrieval.utils import BaseMetadata, _raise_for_status, to_str WATERUSE_URL = "https://api.water.usgs.gov/nwaa-data/data" +_WATERUSE_HOST = httpx.URL(WATERUSE_URL).host +# Hosts a ``rel="next"`` cursor may name for this same service; each is +# rewritten to :data:`_WATERUSE_HOST` rather than followed as given. +_WATERUSE_HOST_ALIASES = frozenset({_WATERUSE_HOST, "water.usgs.gov"}) #: Water-use models (categories) served by the NWDC. The catalog at #: https://water.usgs.gov/nwaa-data/ lists the variables available within each. @@ -80,9 +82,10 @@ #: Maximum locations fetched concurrently when a list of state/county/huc #: selectors is fanned out (one request per location). Kept conservative -#: because this module intentionally carries no request backoff/retry; the -#: NWDC tolerates this level of concurrency without rate-limit errors (verified -#: by stress test). Set ``wateruse.MAX_CONCURRENT_REQUESTS = 1`` for serial. +#: 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 # Page responses carry the HUC12 identifier in this column; it must stay a @@ -222,9 +225,9 @@ def get_wateruse( base_params = {k: v for k, v in base_params.items() if v is not None} # The NWDC queries one location per request, so fan a multi-value selector - # out into one request per location, each paginated by the OGC engine's - # shared pager (``_paginate``), and concatenate the results. - headers = _default_headers(WATERUSE_URL) + # out into one request per location, each handled by shared transport + # pagination, and concatenate the results. + headers = default_headers(WATERUSE_URL) requests = [ httpx.Request( "GET", @@ -238,7 +241,7 @@ def get_wateruse( # 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( + df, response = run_sync( lambda: _fan_out(requests, headers, ssl_check), service="wateruse", error_url=WATERUSE_URL, @@ -330,8 +333,8 @@ async def _fan_out( ) -> tuple[pd.DataFrame, httpx.Response]: """Fetch every request (each paginated) concurrently over one shared client. - Each request is paginated by the engine's - :func:`~dataretrieval.ogc.engine._paginate` with NWDC strategies: parse a CSV + 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 @@ -349,12 +352,17 @@ async def follow(cursor: str, sess: httpx.AsyncClient) -> httpx.Response: def raise_for_status(response: httpx.Response) -> None: _raise_for_status(response, detail_from=_nwdc_error_detail) - async with httpx.AsyncClient(verify=ssl_check, **HTTPX_ASYNC_DEFAULTS) as client: + # 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 with semaphore: - return await _paginate( + async def attempt() -> tuple[pd.DataFrame, httpx.Response]: + return await paginate( request, parse_response=parse, follow_up=follow, @@ -362,14 +370,46 @@ async def _one(request: httpx.Request) -> tuple[pd.DataFrame, httpx.Response]: raise_for_status=raise_for_status, ) - results = await asyncio.gather(*(_one(req) for req in requests)) + # ``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 + ) - # Reuse the engine's combine helpers: drop empty frames and concat, and fold + # 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 results] - responses = [resp for _, resp in results] + 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) ) @@ -392,14 +432,44 @@ def _next_page_url(response: httpx.Response) -> str | None: """Return the absolute URL of the next page, or None if this is the last. Reads the standard ``Link: <...>; rel="next"`` header (parsed by httpx into - ``response.links``). A next link served against the bare ``water.usgs.gov`` - host is normalized to the public ``api.water.usgs.gov`` gateway so the - follow-up request reaches the API. + ``response.links``). The cursor is normalized before it is trusted, because + the service spells it inconsistently: a relative reference is resolved + against the page it came from, and the bare ``water.usgs.gov`` host is + rewritten to the public ``api.water.usgs.gov`` gateway (over https, whatever + scheme the link used) so the follow-up request reaches the API. Only a + cursor that still points somewhere else after that is refused -- following + it would send Water Use requests, and any credentials on them, to a host the + caller never asked for. """ url = response.links.get("next", {}).get("url") if not url: return None - return str(url).replace("https://water.usgs.gov", "https://api.water.usgs.gov", 1) + try: + target = httpx.URL(url) + except (httpx.InvalidURL, TypeError) as exc: + raise DataRetrievalError( + f"Water Use returned an unusable next-page link: {url!r}. The page " + f"walk cannot continue; report this if it persists." + ) from exc + if not target.is_absolute_url: + target = response.url.join(target) + if target.host not in _WATERUSE_HOST_ALIASES: + raise DataRetrievalError( + f"Refusing to follow a Water Use next-page link pointing at " + f"{target.host} rather than {_WATERUSE_HOST}. Following it would " + f"send this request, and any credentials on it, to a host you did " + f"not ask for. Retrying will not help; report this if it persists." + ) + # Drop any explicit port and any embedded userinfo along with the + # scheme/host rewrite. A port that went with the link's original scheme + # (``http://…:8080``) would otherwise survive into an https request and be + # dialed under TLS; userinfo (``http://user:pass@…``) would survive into an + # ``Authorization: Basic`` header that httpx derives from it and send a + # credential the caller never configured to the rewritten host -- the very + # thing the host check above exists to prevent. + return str( + target.copy_with(scheme="https", host=_WATERUSE_HOST, port=None, userinfo=b"") + ) def _nwdc_error_detail(response: httpx.Response) -> str | None: diff --git a/dataretrieval/wqp.py b/dataretrieval/wqp.py index ffbaee91..a0b5e642 100644 --- a/dataretrieval/wqp.py +++ b/dataretrieval/wqp.py @@ -17,7 +17,7 @@ import pandas as pd -from .utils import BaseMetadata, _attach_datetime_columns, query +from .utils import BaseMetadata, _attach_datetime_columns, _query_with_retry if TYPE_CHECKING: import httpx @@ -179,7 +179,7 @@ def get_results( if legacy is not True and profile is None: kwargs["dataProfile"] = "fullPhysChem" - response = query(url, kwargs, delimiter=";", ssl_check=ssl_check) + response = _query_with_retry(url, kwargs, delimiter=";", ssl_check=ssl_check) df = _read_wqp_csv(response.text) df = _attach_datetime_columns(df) @@ -208,7 +208,9 @@ def _what( else: url = _legacy_only_url(service, legacy=legacy) - response = query(url, payload=kwargs, delimiter=";", ssl_check=ssl_check) + response = _query_with_retry( + url, payload=kwargs, delimiter=";", ssl_check=ssl_check + ) df = _read_wqp_csv(response.text) return df, WQP_Metadata(response, **kwargs) diff --git a/docs/source/architecture/decisions/0003-dependency-direction.rst b/docs/source/architecture/decisions/0003-dependency-direction.rst index 3f7d8ec9..8e24b229 100644 --- a/docs/source/architecture/decisions/0003-dependency-direction.rst +++ b/docs/source/architecture/decisions/0003-dependency-direction.rst @@ -18,13 +18,15 @@ unintended cross-package contracts. Decision -------- -Dependencies point from public facades to service/protocol adapters and then to -stable shared policy and third-party infrastructure. In particular: +Dependencies point from public facades to service/protocol adapters, then to +service-neutral transport and stable policy, and finally to third-party +infrastructure. In particular: - ``dataretrieval.exceptions`` is a runtime-dependency-light leaf. - ``dataretrieval.ogc`` must not import Water Data, NGWMN, Water Use, or NWIS. - Modern modules must not import deprecated NWIS. -- New non-OGC services must not obtain generic transport behavior by importing +- Service-neutral transport must not import OGC modules or service adapters. +- Non-OGC services must obtain generic execution behavior from transport, not private OGC implementation symbols. Underscore-prefixed symbols remain implementation details even when existing @@ -49,8 +51,8 @@ second copy of that mutable inventory. Focused fitness functions verify the current boundaries: NGWMN's only OGC dependency is the facade, ``waterdata.utils`` does not bulk re-export private -OGC helpers, ``ogc.shaping`` does not depend on ``ogc.engine``, and the full OGC -runtime graph is acyclic. +OGC helpers, ``ogc.shaping`` does not depend on ``ogc.engine``, Water Use has +no OGC dependency, and both the OGC and transport runtime graphs are acyclic. The exact allowlist should shrink as private seams move. Any growth requires explicit architecture review, and a change to the dependency policy requires diff --git a/docs/source/architecture/decisions/0004-error-retry-resume.rst b/docs/source/architecture/decisions/0004-error-retry-resume.rst index 50c9e3d3..d05d7522 100644 --- a/docs/source/architecture/decisions/0004-error-retry-resume.rst +++ b/docs/source/architecture/decisions/0004-error-retry-resume.rst @@ -28,9 +28,10 @@ cancellation. OGC fan-out retains completed subrequests and raises a typed ``ChunkInterrupted`` with a handle that resumes only missing work. Fatal or unknown failures are not disguised as resumable transients. -This decision does not assert that every upstream API supports pagination, -chunking, or resume. Those capabilities remain explicit per service until a -shared API-neutral transport contract is introduced. +The shared transport layer supplies bounded retry and callback-driven cursor +pagination, but each adapter opts in only where its requests are idempotent and +its protocol exposes a cursor. Chunk planning and resumable partial state remain +OGC-specific capabilities rather than assumptions imposed on every service. Consequences ------------ diff --git a/docs/source/architecture/decisions/0006-service-neutral-transport.rst b/docs/source/architecture/decisions/0006-service-neutral-transport.rst new file mode 100644 index 00000000..f6d9fd52 --- /dev/null +++ b/docs/source/architecture/decisions/0006-service-neutral-transport.rst @@ -0,0 +1,110 @@ +ADR 0006: Use a service-neutral transport layer +=============================================== + +Status +------ + +Accepted + +Context +------- + +Several service adapters need the same low-level capabilities: guarded HTTP +clients, cursor pagination, bounded retry, response aggregation, progress, and a +sync-over-async bridge. Locating those capabilities inside a protocol package +would make unrelated services depend on protocol-specific implementation +details -- Water Use previously imported its page walker and sync bridge from +``ogc.engine``, a dependency with no conceptual basis. Duplicating them would +allow authentication, timeout, retry, and failure behavior to drift. + +"Neutral" here means neutral across the USGS services this package talks to, not +across HTTP APIs in general. The layer knows the ``API_USGS_*`` environment +variables and the quota header USGS returns. Claiming broader neutrality than +that invites generality no caller needs. + +Decision +-------- + +``dataretrieval.transport`` is the internal service-neutral execution layer -- +neutral across the USGS services this package talks to, not across HTTP APIs in +general. It owns: + +- synchronous and asynchronous HTTP client lifecycle and timeout defaults; +- attaching the API key and stripping it at redirect time, over the predicate + ``dataretrieval.credentials`` defines; +- callback-driven cursor pagination; +- bounded retry with exponential backoff, full jitter, capped ``Retry-After`` + handling, and a no-progress budget bounding how long a call may receive + nothing at all; and +- the sync-over-async blocking-portal bridge. + +Three concerns are deliberately *outside* it, as top-level leaves, because they +are not HTTP execution policy and every adapter needs them whether or not it goes +through transport: + +- ``dataretrieval.credentials`` -- which host honors the key, whether a + destination qualifies, and how the key is withheld. One definition, so the code + that attaches a credential and the code that removes it cannot disagree. +- ``dataretrieval.progress`` -- terminal rendering. Transport reports *into* it. +- ``dataretrieval.combining`` -- pandas frame and response assembly. Transport + returns results *through* it. + +Transport depends only on stable package leaves and third-party infrastructure. +It must not import OGC modules or service adapters. Service adapters inject +request construction, response parsing, cursor extraction, and API-specific +error details. + +OGC retains its protocol concerns: dialects, CQL2, request construction, feature +shaping, URL-byte chunk planning, resumable ``ChunkedCall`` state, and typed +interruption handles. Thin imports at previous private OGC and utility paths +preserve compatibility where a consumer still uses them; a path no consumer +imports is deleted rather than kept as a module that exists to satisfy its own +test. Tunables are never re-exported by value: a copy taken at import time is +one a caller can patch without reaching the policy that reads it, so +``transport.retry`` is the single place they are read from. + +Automatic retry is enabled only on active, idempotent request paths, and only +for failures a later attempt could survive -- rate limiting, gateway 5xx, and +transport failures that are not settled before the request leaves. A server +error reporting that *this* request was rejected is surfaced on the first +attempt rather than multiplied against an already-failing service. Deprecated +NWIS calls retain their compatibility behavior. A failed pagination or fan-out +operation raises rather than returning successful siblings as an apparently +complete result. + +Two independent bounds limit retry: an attempt count and a no-progress budget +measured in seconds since data last arrived. Attempts alone leave elapsed time +unbounded, since each attempt may itself block until its timeout; the budget +alone would cut short a slow but productive download. Receiving a page restarts +the budget, and an attempt already in flight is never interrupted. + +Consequences +------------ + +- Water Use has no dependency on OGC implementation modules. +- OGC and non-OGC adapters share authentication, timeout, retry, pagination, + aggregation, progress, and sync-dispatch policy where their semantics match. +- Service-specific request and result contracts remain explicit instead of + being forced into a universal adapter abstraction. +- Retry can increase latency and quota consumption, so attempt counts, waits, + and total silent time remain bounded and cancellation signals are never + wrapped. +- Guidance the progress reporter prints is gated on the host it applies to, so a + service that cannot use an API key is not told to obtain one. +- The transport package is internal infrastructure, not a new public API + promise. +- Keeping presentation and frame assembly out means transport is roughly 570 + lines across five modules, each recognizably HTTP execution policy. Retry is + the one intricate module, and it is intricate because two independent bounds + are what make retry safe against a slow service. + +Compliance +---------- + +``tests/architecture_test.py`` enforces transport dependency direction, an +acyclic transport graph, Water Use isolation from OGC, that presentation and +frame-assembly modules do not reappear inside transport, and that only +``dataretrieval.credentials`` names the API-key host. Component and adapter +tests cover cursor termination, row caps, response aggregation, retry +exhaustion, ``Retry-After`` limits, the no-progress budget, which failures are +re-sent, cancellation, no-partial fan-out behavior, and credential host scoping. diff --git a/docs/source/architecture/decisions/index.rst b/docs/source/architecture/decisions/index.rst index f11aa4ba..92a9d3f1 100644 --- a/docs/source/architecture/decisions/index.rst +++ b/docs/source/architecture/decisions/index.rst @@ -22,4 +22,5 @@ records sequentially. 0003-dependency-direction 0004-error-retry-resume 0005-legacy-nwis + 0006-service-neutral-transport template diff --git a/docs/source/architecture/index.rst b/docs/source/architecture/index.rst index 906c07a4..09672a6e 100644 --- a/docs/source/architecture/index.rst +++ b/docs/source/architecture/index.rst @@ -79,14 +79,13 @@ Public service facades translation, and :class:`OgcDialect`. ``dataretrieval.wateruse`` - NWDC Water Use facade. Builds CSV requests and follows ``Link`` headers. - It currently reuses generic pagination and response-combining helpers from - private OGC modules; this is an explicitly recorded variance. + NWDC Water Use facade. Builds CSV requests, follows ``Link`` headers, and + uses service-neutral transport for bounded fan-out, retry, pagination, response + aggregation, and synchronous dispatch. It does not depend on OGC modules. ``dataretrieval.wqp``, ``dataretrieval.nldi``, and ``dataretrieval.streamstats`` - Service-specific adapters over the synchronous request infrastructure in - ``dataretrieval.utils``. Their return types intentionally reflect their - upstream data models. + Service-specific adapters over shared synchronous HTTP and bounded retry + policy. Their return types intentionally reflect their upstream data models. ``dataretrieval.nwis`` Deprecated legacy NWIS facade, scheduled for removal on or after @@ -101,30 +100,41 @@ Shared components ``prepare_request_args``, ``get_ogc_data``, and ``fetch_ogc_request``. Internally, ``policy`` defines the dialect type and endpoint constants (depends only on stdlib); ``requests`` owns request construction, argument - normalization, and queryables/schema lookup; ``engine`` orchestrates - pagination and sync-from-async; ``planning`` determines chunk boundaries; - ``chunking`` executes plans and retains resumable state; ``interruptions`` - defines the resumable failure contract; ``retry`` owns the bounded retry - policy; ``combining`` assembles results; and ``shaping``, ``dates``, - ``filters``, ``errors``, and ``progress`` isolate their named concerns. The - full runtime OGC graph, including the facade, is acyclic — - enforced by ``tests/architecture_test.py``. + normalization, and queryables/schema lookup; ``engine`` supplies OGC cursor + and response strategies to transport pagination; ``planning`` determines + chunk boundaries; ``chunking`` executes plans and retains resumable state; + ``interruptions`` defines the resumable failure contract; ``retry`` + classifies failures into OGC interruption types; and ``shaping``, ``dates``, + ``filters``, and ``errors`` isolate their named protocol concerns. The full + runtime OGC graph, including the facade, is acyclic — enforced by + ``tests/architecture_test.py``. + +``dataretrieval.transport`` + Internal service-neutral execution layer. Owns guarded client lifecycle and + timeouts, host-scoped authentication, cursor pagination, bounded retry, + response aggregation, progress, and sync-over-async dispatch. Internally, + ``liveness`` is a stdlib-only leaf recording when data last arrived, so the + page loop that observes progress and the retry loop that acts on it both + depend on it rather than on each other. It imports no service adapter or OGC + protocol module, and it is not exposed as a public framework API. ``dataretrieval.exceptions`` Stable error-policy leaf. It has no runtime third-party dependency and may be imported by every service without creating an infrastructure cycle. ``dataretrieval.utils`` - Shared metadata, data-shaping helpers, ambient context support, and the - legacy synchronous request path. Its broad responsibility is known debt; - new service-specific behavior should not be added there by default. + Shared metadata, data-shaping helpers, ambient context support, legacy + request composition, and compatibility imports for transport names that + historically lived here. New service-specific behavior should not be added + there by default. ``dataretrieval.codes`` and ``dataretrieval.rdb`` State/time-zone code conversion and RDB parsing leaves. The intended direction is:: - public facade -> service/protocol adapter -> shared policy/infrastructure + public facade -> service/protocol adapter -> service-neutral transport + -> stable policy/infrastructure -> third-party library / network Dependencies must not point from shared infrastructure back to a public service @@ -170,9 +180,10 @@ The retained ``ChunkedCall`` reissues only missing chunks and applies the same finalization path when resumed. Cancellation and non-transient programming errors take precedence over retry/resume wrapping. -Non-OGC services use simpler request paths where their protocols do not provide -the same paging or resume semantics. Later transport consolidation must preserve -those public contracts and must not invent unsupported upstream capabilities. +Non-OGC services use the same transport policy only where their protocols have +matching semantics. Retry and cursor pagination remain explicit adapter choices; +chunk planning and resumable interruptions remain OGC capabilities rather than +invented features of upstream APIs that do not provide them. Resource and configuration view ------------------------------- @@ -189,16 +200,34 @@ Resource and configuration view the execution throttle. ``API_USGS_RETRIES`` - Number of OGC retries after the first attempt; defaults to four. Backoff is - exponential with full jitter and honors bounded ``Retry-After`` values. + Number of retries after the first attempt on supported active request paths; + defaults to four. Backoff is exponential with full jitter and honors bounded + ``Retry-After`` values. Only failures a later attempt could survive are + re-sent: 429 and gateway 5xx, not a 500 rejecting the query itself, and not a + transport failure that is settled before the request leaves (unresolvable + host, unsupported scheme). Deprecated NWIS compatibility paths do not opt in. + +``API_USGS_STALL_TIMEOUT`` + Seconds a call may go without receiving any data before retrying stops and + the failure surfaces; defaults to 60, and ``0`` disables the bound. It + complements ``API_USGS_RETRIES``, which caps attempts rather than elapsed + time: without it, four retries of a request that times out after a minute is + four silent minutes. Progress restarts the budget — a page received, or a + queued sub-request acquiring its concurrency slot — so neither a slow but + productive download nor the tail of a wide fan-out is cut short, and an + attempt already in flight is never interrupted. The first retry is never + withheld by this bound, so one slow attempt cannot disable retry by itself; + the budget decides whether to continue after that. A dead connection + therefore costs about two read timeouts rather than five attempts' worth. ``API_USGS_PROGRESS`` Controls best-effort progress display. Reporting failures must never change retrieval results. -HTTP timeouts and connection limits are centralized for existing paths. -``wateruse`` currently has its own smaller fan-out cap. These differences must -remain visible until a shared transport policy replaces them deliberately. +HTTP timeout, redirect, and authentication policy is centralized in +``dataretrieval.transport``. OGC subrequest fan-out and Water Use location +fan-out retain separate explicit concurrency caps because their upstream costs +and request shapes differ. Known architectural debt ------------------------ @@ -207,12 +236,7 @@ This view records categories and representative locations of debt. The fitness functions in ``tests/architecture_test.py`` are authoritative for exact current dependency allowlists. -- ``wateruse`` depends on private generic helpers located under ``ogc`` even - though NWDC is not an OGC service. - ``waterdata/api.py`` and ``ogc/engine.py`` contain multiple reasons to change. -- Active non-OGC services do not yet share OGC's retry/resume capabilities. -- ``utils.py`` combines metadata, shaping, configuration, and transport duties. - These are documented so guardrails distinguish accepted current dependencies from new erosion. They should be removed through small, test-protected changes, not a rewrite. diff --git a/tests/architecture_test.py b/tests/architecture_test.py index 8ba2d864..c4d4c5b7 100644 --- a/tests/architecture_test.py +++ b/tests/architecture_test.py @@ -18,19 +18,11 @@ "dataretrieval.wqp", ) -# These top-level modules currently reach into OGC. NGWMN is an OGC adapter -# that uses the small facade (``dataretrieval.ogc``) exclusively. Water Use's -# imports are an accepted temporary variance under ADR 0003. This allowlist is -# the authoritative exact inventory; the ADR owns the policy and rationale. -# Exact equality makes either growth or removal intentional. +# NGWMN is the only top-level OGC consumer and uses the small facade +# (``dataretrieval.ogc``) exclusively. Exact equality makes growth or removal +# an intentional architecture change. _ALLOWED_TOP_LEVEL_OGC_IMPORTS = { - "dataretrieval.ngwmn": { - "dataretrieval.ogc", - }, - "dataretrieval.wateruse": { - "dataretrieval.ogc.combining", - "dataretrieval.ogc.engine", - }, + "dataretrieval.ngwmn": {"dataretrieval.ogc"}, } _ENGINE_REQUEST_IMPORTS = { @@ -349,7 +341,7 @@ def test_default_header_calls_are_target_scoped() -> None: if isinstance(node.func, ast.Attribute) else None ) - if function_name != "_default_headers": + if function_name not in {"_default_headers", "default_headers"}: continue has_target = bool(node.args) or any( keyword.arg == "target_url" for keyword in node.keywords @@ -363,3 +355,137 @@ def test_default_header_calls_are_target_scoped() -> None: "_default_headers calls without destination URL context:\n" + "\n".join(violations) ) + + +# --- Shared execution-layer boundaries --- + + +def test_transport_is_execution_policy_only() -> None: + """Transport owns HTTP execution, not presentation or result assembly. + + Terminal rendering (``progress``) and pandas result assembly (``combining``) + are top-level leaves that transport reports *into* and returns *through*. + They lived here only because they had to leave ``ogc`` and this was the + nearest home; keeping them out is what makes "transport is HTTP execution + policy" a checkable claim rather than a description of a grab bag. + """ + misplaced = { + "dataretrieval/transport/progress.py", + "dataretrieval/transport/combining.py", + } + present = { + path + for path in misplaced + if (PACKAGE_ROOT.parent / path).exists() # repo-root-relative + } + assert not present, ( + "Presentation or frame-assembly code reappeared inside transport: " + f"{sorted(present)}" + ) + + +def test_credential_policy_has_one_definition() -> None: + """Only ``dataretrieval.credentials`` may name the API-key host. + + Attaching the key and stripping it back off have to agree about which host + is authorized; the way they stop agreeing is a second copy of the host + string. ``transport.http`` re-exports the predicate, it does not restate it. + """ + host = "api.waterdata.usgs.gov" + # Walked as AST string *values*, not as source text. A line-substring match + # is wrong in both directions: it missed the ``"https://…"`` form three + # modules use to spell the same authority, and it flagged docstring prose + # that merely names the service. Docstrings are excluded here (they are + # documentation, not a second source of truth) while every other literal -- + # bare host or full base URL -- counts. + offenders: list[str] = [] + for path in sorted(PACKAGE_ROOT.rglob("*.py")): + if path.name == "credentials.py": + continue + tree = ast.parse(path.read_text(encoding="utf-8")) + docstrings = { + text + for node in ast.walk(tree) + if isinstance( + node, (ast.Module, ast.ClassDef, ast.FunctionDef, ast.AsyncFunctionDef) + ) + for text in [ast.get_docstring(node, clean=False)] + if text is not None + } + for node in ast.walk(tree): + if not isinstance(node, ast.Constant) or not isinstance(node.value, str): + continue + if host in node.value and node.value not in docstrings: + offenders.append(f"{_module_name(path)}:{node.lineno}") + assert not offenders, ( + "The API-key host must come from dataretrieval.credentials, " + f"not be restated at: {offenders}" + ) + + +def test_transport_does_not_depend_on_ogc_or_services() -> None: + """Transport policy must point inward, never back to protocol adapters.""" + violations: list[str] = [] + transport_root = PACKAGE_ROOT / "transport" + for path in sorted(transport_root.rglob("*.py")): + module = _module_name(path) + for dependency in _runtime_imports(path): + if ( + dependency == "dataretrieval.ogc" + or dependency.startswith("dataretrieval.ogc.") + or dependency.startswith(_SERVICE_PREFIXES) + ): + violations.append(f"{module} -> {dependency}") + assert not violations, "Transport crossed an adapter boundary:\n" + "\n".join( + violations + ) + + +def test_wateruse_has_no_ogc_dependency() -> None: + """The non-OGC Water Use adapter must consume transport directly.""" + imports = _runtime_imports(PACKAGE_ROOT / "wateruse.py") + ogc_dependencies = { + dependency + for dependency in imports + if dependency == "dataretrieval.ogc" + or dependency.startswith("dataretrieval.ogc.") + } + assert not ogc_dependencies, ( + f"Water Use imported OGC implementation modules: {sorted(ogc_dependencies)}" + ) + + +def test_transport_runtime_graph_is_acyclic() -> None: + """The service-neutral transport package must remain a directed acyclic graph.""" + graph = { + module: { + dependency + for dependency in imports + if dependency == "dataretrieval.transport" + or dependency.startswith("dataretrieval.transport.") + } + for module, imports in _package_import_graph().items() + if module == "dataretrieval.transport" + or module.startswith("dataretrieval.transport.") + } + visiting: set[str] = set() + visited: set[str] = set() + + def visit(module: str, path: tuple[str, ...]) -> None: + if module in visiting: + start = path.index(module) + cycle = (*path[start:], module) + raise AssertionError( + f"Cycle in transport runtime graph: {' -> '.join(cycle)}" + ) + if module in visited: + return + visiting.add(module) + for dependency in graph.get(module, set()): + if dependency in graph: + visit(dependency, (*path, module)) + visiting.remove(module) + visited.add(module) + + for module in graph: + visit(module, ()) diff --git a/tests/conftest.py b/tests/conftest.py index 85f7739e..fcc6df6c 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -3,11 +3,10 @@ * Relaxes ``pytest-httpx``'s strict-mode flags so unconsumed mocks and unmatched requests don't fail the suite (keeps mocked-URL setup terse). -* Pins ``API_USGS_CONCURRENT=1`` and ``API_USGS_RETRIES=0`` for every - test by default, so sub-request dispatch is deterministic and a single - transient surfaces immediately (no backoff). Concurrency and retry - tests opt in by re-setting the env vars inside their body via - ``monkeypatch.setenv``. +* Pins the chunker env for every test (see ``_pin_chunker_env``), so + sub-request dispatch is deterministic and mocked retries measure attempt + counts rather than wall clock. Concurrency and retry tests opt in by + re-setting the env vars inside their body via ``monkeypatch.setenv``. """ from __future__ import annotations @@ -65,14 +64,24 @@ def non_mocked_hosts() -> list[str]: @pytest.fixture(autouse=True) def _pin_chunker_env(monkeypatch): - """Pin every test to one connection and no retries. + """Pin every test to one connection, no retries, and no stall budget. - Production defaults ``API_USGS_CONCURRENT`` to 32 and - ``API_USGS_RETRIES`` to 4. Pinning ``API_USGS_CONCURRENT=1`` keeps - sub-request dispatch deterministic for the mocked suite, and - ``API_USGS_RETRIES=0`` makes a single transient surface immediately - rather than be retried. Concurrency and retry tests opt in by - overriding the env inside their body. + Production defaults ``API_USGS_CONCURRENT`` to 32, + ``API_USGS_RETRIES`` to 4, and ``API_USGS_STALL_TIMEOUT`` to 60 s. + Pinning ``API_USGS_CONCURRENT=1`` keeps sub-request dispatch + deterministic for the mocked suite, and ``API_USGS_RETRIES=0`` makes + a single transient surface immediately rather than be retried. + Concurrency and retry tests opt in by overriding the env inside + their body. + + ``API_USGS_STALL_TIMEOUT=0`` is pinned too so that an opting-in retry + test measures the thing it names -- attempt counts -- and not the wall + clock of the machine running it. Left at the production 60 s, a test + that sets ``API_USGS_RETRIES`` would have its retries silently capped + by whatever real time its mocked attempts consumed, which is both flaky + on a loaded CI box and a way for a stall-budget bug to hide behind a + passing retry test. Tests of the budget itself set it explicitly. """ monkeypatch.setenv("API_USGS_CONCURRENT", "1") monkeypatch.setenv("API_USGS_RETRIES", "0") + monkeypatch.setenv("API_USGS_STALL_TIMEOUT", "0") diff --git a/tests/headers_host_scoping_test.py b/tests/headers_host_scoping_test.py index 85dfc270..82dbbb22 100644 --- a/tests/headers_host_scoping_test.py +++ b/tests/headers_host_scoping_test.py @@ -79,6 +79,42 @@ def test_non_auth_headers_always_present(self): # Key should NOT be sent to example.com assert "X-Api-Key" not in headers + def test_key_excluded_over_cleartext_on_the_authorized_host(self): + """The right host over plain http is still the wrong destination. + + Matching on the host alone would send a bearer token in the clear on + the strength of a hostname an attacker chose to keep -- reachable via a + redirect or a server-supplied ``http://`` next-page link. + """ + headers = _default_headers("http://api.waterdata.usgs.gov/ogcapi/v0/daily") + assert "X-Api-Key" not in headers + + def test_sync_transport_withholds_key_on_downgrade_to_cleartext(self): + """The guard runs at send time, not only where headers are built.""" + seen: list[httpx.Request] = [] + + def handler(request: httpx.Request) -> httpx.Response: + seen.append(request) + if len(seen) == 1: + return httpx.Response( + 302, + headers={"Location": "http://api.waterdata.usgs.gov/next"}, + request=request, + ) + return httpx.Response(200, request=request) + + url = "https://api.waterdata.usgs.gov/start" + _get( + url, + headers=_default_headers(url), + follow_redirects=True, + transport=httpx.MockTransport(handler), + ) + + assert seen[0].headers.get("X-Api-Key") == self.FAKE_TOKEN + assert seen[1].url.scheme == "http", "the redirect under test must downgrade" + assert "X-Api-Key" not in seen[1].headers + def test_generic_ogc_request_excludes_key_for_custom_host(self): """A caller-supplied OGC base URL never inherits Water Data auth.""" from dataretrieval.ogc.requests import _construct_api_requests, _ogc_base_url diff --git a/tests/nldi_test.py b/tests/nldi_test.py index 9092bbf3..1b7a5957 100644 --- a/tests/nldi_test.py +++ b/tests/nldi_test.py @@ -1,3 +1,5 @@ +from unittest import mock + import pytest from geopandas import GeoDataFrame @@ -46,6 +48,17 @@ def mock_request_data_sources(httpx_mock): ) +def test_query_nldi_opts_into_retry(monkeypatch): + """NLDI explicitly enables shared retry while NWIS remains unchanged.""" + response = mock.Mock() + response.json.return_value = {} + query = mock.Mock(return_value=response) + monkeypatch.setattr(nldi, "_query_with_retry", query) + + assert nldi._query_nldi("https://example.test", {}) == {} + query.assert_called_once_with("https://example.test", payload={}) + + def mock_request(httpx_mock, request_url, file_path): with open(file_path) as text: httpx_mock.add_response( diff --git a/tests/streamstats_test.py b/tests/streamstats_test.py index ee528693..4c481e06 100644 --- a/tests/streamstats_test.py +++ b/tests/streamstats_test.py @@ -4,6 +4,7 @@ import pytest +import dataretrieval from dataretrieval.streamstats import Watershed, get_watershed # Minimal StreamStats watershed payload shaped like the service response @@ -60,3 +61,35 @@ def test_get_watershed_shape_raises_not_implemented(httpx_mock): httpx_mock.add_response(text=json.dumps(_SAMPLE)) with pytest.raises(NotImplementedError): get_watershed("NY", -74.524, 43.939, format="shape") + + +def test_get_watershed_does_not_retry_a_rejected_query(httpx_mock, monkeypatch): + """A 500 means the service rejected *this* request, so re-sending it only + multiplies load on a failing service and delays the caller's error.""" + import dataretrieval.transport.retry as retry + + httpx_mock.add_response(status_code=500) + monkeypatch.setenv("API_USGS_RETRIES", "4") + monkeypatch.setattr(retry, "_RETRY_BASE_BACKOFF", 0.0) + monkeypatch.setattr(retry, "_RETRY_MAX_BACKOFF", 0.0) + + with pytest.raises(dataretrieval.ServiceUnavailable): + get_watershed("XX", -74.524, 43.939) + + assert len(httpx_mock.get_requests()) == 1 + + +def test_get_watershed_retries_transient_failure(httpx_mock, monkeypatch): + """StreamStats retries a bounded transient before returning normally.""" + import dataretrieval.transport.retry as retry + + httpx_mock.add_response(status_code=503) + httpx_mock.add_response(text=json.dumps(_SAMPLE)) + monkeypatch.setenv("API_USGS_RETRIES", "1") + monkeypatch.setattr(retry, "_RETRY_BASE_BACKOFF", 0.0) + monkeypatch.setattr(retry, "_RETRY_MAX_BACKOFF", 0.0) + + response = get_watershed("NY", -74.524, 43.939) + + assert response.status_code == 200 + assert len(httpx_mock.get_requests()) == 2 diff --git a/tests/transport_test.py b/tests/transport_test.py new file mode 100644 index 00000000..dd5e9278 --- /dev/null +++ b/tests/transport_test.py @@ -0,0 +1,562 @@ +"""Component tests for the internal service-neutral transport layer.""" + +from __future__ import annotations + +import asyncio +import datetime +import itertools +import socket +from unittest import mock + +import httpx +import pandas as pd +import pytest + +import dataretrieval.transport.liveness as liveness +import dataretrieval.transport.retry as retry +from dataretrieval.exceptions import ( + ConfigurationError, + DataRetrievalError, + HTTPError, + NetworkError, + RateLimited, + ServiceUnavailable, +) +from dataretrieval.transport.pagination import paginate +from dataretrieval.transport.sync import run_sync +from dataretrieval.utils import _raise_for_status + + +def _response( + status: int = 200, *, url: str = "https://example.test/page" +) -> httpx.Response: + return httpx.Response(status, request=httpx.Request("GET", url)) + + +def test_paginate_follows_cursor_and_aggregates_response() -> None: + first = _response(url="https://example.test/page/1") + second = _response(url="https://example.test/page/2") + first.headers["x-ratelimit-remaining"] = "9" + second.headers["x-ratelimit-remaining"] = "8" + client = mock.AsyncMock(spec=httpx.AsyncClient) + client.send.return_value = first + client.get.return_value = second + + cursors = {str(first.url): "next", str(second.url): None} + + def parse(response: httpx.Response) -> tuple[pd.DataFrame, str | None]: + return pd.DataFrame({"value": [str(response.url)]}), cursors[str(response.url)] + + async def follow(cursor: str, session: httpx.AsyncClient) -> httpx.Response: + assert cursor == "next" + return await session.get("https://example.test/page/2") + + frame, response = asyncio.run( + paginate( + httpx.Request("GET", first.url), + parse_response=parse, + follow_up=follow, + raise_for_status=_raise_for_status, + client=client, + ) + ) + + assert frame["value"].tolist() == [str(first.url), str(second.url)] + assert response.url == first.url + assert response.headers["x-ratelimit-remaining"] == "8" + + +def test_paginate_stops_on_repeated_cursor_and_respects_row_cap() -> None: + first = _response(url="https://example.test/page/1") + second = _response(url="https://example.test/page/2") + client = mock.AsyncMock(spec=httpx.AsyncClient) + client.send.return_value = first + client.get.return_value = second + + def parse(response: httpx.Response) -> tuple[pd.DataFrame, str]: + return pd.DataFrame({"value": [1, 2]}), "same-cursor" + + async def follow(cursor: str, session: httpx.AsyncClient) -> httpx.Response: + return await session.get(str(second.url)) + + frame, _ = asyncio.run( + paginate( + httpx.Request("GET", first.url), + parse_response=parse, + follow_up=follow, + raise_for_status=_raise_for_status, + client=client, + row_cap=3, + ) + ) + + assert frame["value"].tolist() == [1, 2, 1] + assert client.get.await_count == 1 + + +def test_retry_sync_retries_transient_then_succeeds(monkeypatch) -> None: + attempts = 0 + slept: list[float] = [] + monkeypatch.setattr(retry.time, "sleep", slept.append) + + def operation() -> str: + nonlocal attempts + attempts += 1 + if attempts == 1: + raise ServiceUnavailable("temporary") + return "ok" + + result = retry.retry_sync( + operation, + retry.RetryPolicy(max_retries=1, base_backoff=0, max_backoff=0), + ) + + assert result == "ok" + assert attempts == 2 + assert slept == [0] + + +def test_retry_sync_honors_cap_and_does_not_catch_cancellation(monkeypatch) -> None: + sleep = mock.Mock() + monkeypatch.setattr(retry.time, "sleep", sleep) + policy = retry.RetryPolicy(max_retries=2, retry_after_cap=60) + + with pytest.raises(RateLimited): + retry.retry_sync( + lambda: (_ for _ in ()).throw(RateLimited("later", retry_after=61)), + policy, + ) + sleep.assert_not_called() + + with pytest.raises(KeyboardInterrupt): + retry.retry_sync( + lambda: (_ for _ in ()).throw(KeyboardInterrupt()), + policy, + ) + + +def test_shared_status_mapping_preserves_retry_after() -> None: + response = httpx.Response( + 429, + headers={"Retry-After": "2.5"}, + request=httpx.Request("GET", "https://example.test"), + ) + with pytest.raises(RateLimited) as exc_info: + _raise_for_status(response) + assert exc_info.value.retry_after == 2.5 + + +def test_sync_bridge_runs_async_operation() -> None: + async def operation() -> str: + return "ok" + + assert run_sync(operation, service="test", error_url="https://example.test") == "ok" + + +def test_retry_tunables_have_a_single_home() -> None: + """Patching the tunables must reach the policy that reads them. + + Re-exporting them from ``ogc.retry`` would hand out copies taken at import + time, so patching that path would change a value nothing consults. That + module owns OGC classification only. + """ + import dataretrieval.ogc.retry as ogc_retry + + assert not [name for name in vars(ogc_retry) if name.startswith("_RETRY")] + assert set(ogc_retry.__all__) == {"_classify_chunk_error", "_classify_transient"} + + +def test_parse_retry_after_accepts_http_date() -> None: + """A date in the future is honored; one already past is not a hint. + + Read literally an elapsed date says "retry now", but the likelier cause is + our clock running ahead of the server's, and acting on it would re-send + almost immediately against a service that just asked for a pause. + """ + soon = datetime.datetime.now(datetime.timezone.utc) + datetime.timedelta(seconds=30) + parsed = retry.parse_retry_after(soon.strftime("%a, %d %b %Y %H:%M:%S GMT")) + assert parsed is not None and 0 < parsed <= 30 + + assert retry.parse_retry_after("Wed, 21 Oct 2015 07:28:00 GMT") is None + assert retry.parse_retry_after("not-a-date") is None + # Delta-seconds is clock-independent, so a literal 0 stays an instruction. + assert retry.parse_retry_after("0") == 0.0 + + +def test_both_retry_after_forms_are_honored_alike() -> None: + """The two header spellings mean the same thing and must behave the same. + + Discarding an over-long date hint (returning ``None``) made the client retry + *harder* against a service asking for a long pause, and dropped the number + the caller needs from ``.retry_after``. + """ + far_future = datetime.datetime.now(datetime.timezone.utc) + datetime.timedelta( + minutes=30 + ) + header = far_future.strftime("%a, %d %b %Y %H:%M:%S GMT") + + parsed = retry.parse_retry_after(header) + assert parsed is not None and 1750 < parsed <= 1800 + assert retry.parse_retry_after("1800") == 1800.0 + # Either spelling, over the cap, stops the retry rather than being ignored. + policy = retry.RetryPolicy(max_retries=4) + assert not policy.should_retry(attempt=1, retry_after=parsed) + assert not policy.should_retry(attempt=1, retry_after=1800.0) + + +def test_elapsed_retry_after_still_backs_off() -> None: + """A ``Retry-After`` of zero must not become a zero-delay re-send.""" + policy = retry.RetryPolicy(base_backoff=0.5, max_backoff=30.0) + + assert policy.backoff(attempt=1, retry_after=0.0) > 0.0 + # The nudge is bounded by max_backoff, not this attempt's exponential + # ceiling: keying it to the ceiling made it vanish whenever base_backoff was + # zero -- exactly when a hint of 0 would become a zero-delay re-send. + assert retry.RetryPolicy(base_backoff=0.0).backoff(attempt=1, retry_after=0.0) > 0.0 + # A server-named delay is honored, plus a small decorrelating nudge so + # concurrent sub-requests handed the same hint do not all wake together -- + # and never enough to push the wait past the policy's own bounds. + assert 5.0 < policy.backoff(attempt=1, retry_after=5.0) <= 6.0 + # A hint already at the cap is never nudged past it -- the jitter would + # otherwise sleep longer than any bound the policy declares. + at_cap = policy.backoff(attempt=8, retry_after=policy.retry_after_cap) + assert at_cap == policy.retry_after_cap + + +def _dns_failure(errno: int) -> NetworkError: + """A DNS failure shaped the way one actually reaches the retry loop. + + httpx and httpcore link their wrappers with ``__context__`` (implicit + chaining), not ``__cause__``, so a walker following only explicit causes + never reaches the ``gaierror``. + """ + resolution_failed = socket.gaierror(errno, "name resolution failed") + transport_failed = httpx.ConnectError("name resolution failed") + transport_failed.__context__ = resolution_failed + wrapped = NetworkError("could not reach host") + wrapped.__context__ = transport_failed + return wrapped + + +def test_deterministic_failures_are_not_retried() -> None: + """Only failures a later attempt could survive are worth re-sending. + + The ``EAI_*`` values are platform-specific -- ``EAI_NONAME`` is 8 on + macOS and -2 on Linux -- so these must come from :mod:`socket` rather + than being written out, or the test only holds on the platform it was + written on. + """ + assert retry._retryable(_dns_failure(socket.EAI_NONAME)) == (False, None) + assert retry._retryable(httpx.UnsupportedProtocol("no scheme")) == (False, None) + assert retry._retryable(httpx.ConnectTimeout("timed out")) == (True, None) + + +def test_temporary_name_resolution_is_still_retried() -> None: + """``gaierror`` is not one condition: ``EAI_AGAIN`` means "try again". + + A resolver still coming up, a VPN reconnect, or a laptop waking all + surface this way, and they are exactly the failures retry exists for. + """ + assert retry._retryable(_dns_failure(socket.EAI_AGAIN)) == (True, None) + # An unrecognized code is retried too: a wasted attempt is cheaper than + # dropping a call we could have recovered. + assert retry._retryable(_dns_failure(0)) == (True, None) + + +def test_resolver_failure_found_past_an_unrelated_explicit_cause() -> None: + """Both chain links are walked, not just the first one present. + + ``raise X from Y`` inside an ``except`` block leaves an explicit + ``__cause__`` *and* an unrelated ``__context__`` on the same frame. Following + only the cause walks off down the explicit branch and never reaches the + ``gaierror``, so an unresolvable hostname spends the whole retry budget + instead of failing fast. + """ + failure = _dns_failure(socket.EAI_NONAME) + failure.__cause__ = ValueError("an unrelated explicit cause") + + assert retry._retryable(failure) == (False, None) + + # The walk still distinguishes the temporary code on the same shape. + temporary = _dns_failure(socket.EAI_AGAIN) + temporary.__cause__ = ValueError("an unrelated explicit cause") + assert retry._retryable(temporary) == (True, None) + + +def test_chain_walk_terminates_on_a_self_referential_cause() -> None: + """A chain pointing back at itself must not hang the classifier.""" + looped = NetworkError("could not reach host") + looped.__context__ = looped + + assert retry._retryable(looped) == (True, None) + + +def test_retryable_statuses_are_per_adapter() -> None: + """A 500 means different things to different services, so the set differs. + + WQP answers an over-large query with a 500 and StreamStats answers + out-of-network coordinates with one, so re-sending can never help there. The + Water Data OGC API is a query interface where a 500 is an upstream hiccup, so + the chunker keeps riding those out — applying WQP's rationale to it would + quietly drop retries the chunked getters have always had. + """ + rejected_query = ServiceUnavailable("bad query", status_code=500) + gateway = ServiceUnavailable("bad gateway", status_code=502) + + # Default (Water Data chunker): every 5xx is worth another try. + assert retry._retryable(rejected_query)[0] + assert retry._retryable(gateway)[0] + + # One-shot adapters: only the gateway family. + strict = retry._GATEWAY_STATUSES + assert not retry._retryable(rejected_query, strict)[0] + assert retry._retryable(gateway, strict)[0] + assert retry._retryable(RateLimited("slow down", retry_after=1.0), strict) == ( + True, + 1.0, + ) + # Never a plain client error, under either set. + assert retry._retryable(HTTPError("not found", status_code=404)) == (False, None) + + +def test_stall_timeout_stops_a_silent_call(monkeypatch) -> None: + """Retrying stops once a call has gone quiet for the whole budget. + + Without this, a request that times out is retried until the attempts run + out, turning one 60 s timeout into minutes of apparent hang. The first + retry is exempt (see below), so a silent call costs two attempts, not five. + """ + attempts = 0 + + def operation() -> str: + nonlocal attempts + attempts += 1 + raise ServiceUnavailable("busy", status_code=503) + + # Every attempt appears to consume 100 s against a 60 s budget. + clock = itertools.count(0.0, 100.0) + monkeypatch.setattr(liveness.time, "monotonic", lambda: next(clock)) + monkeypatch.setattr(retry.time, "sleep", mock.Mock()) + + with pytest.raises(ServiceUnavailable): + retry.retry_sync( + operation, retry.RetryPolicy(max_retries=4, stall_timeout=60.0) + ) + + assert attempts == 2, "first retry is exempt; the budget stops the rest" + + +def test_server_named_delay_does_not_consume_the_stall_budget(monkeypatch) -> None: + """Honoring ``Retry-After`` must not cost a call its retries. + + The budget bounds *silence*; a delay the service named is the opposite of + going quiet. Charging for it meant the more politely a service asked for + room, the fewer retries it got: with the shipped defaults a + ``Retry-After: 30`` against a 60 s budget allowed exactly one retry no + matter what ``API_USGS_RETRIES`` said, silently capping the feature this + layer exists to provide. + """ + attempts = 0 + + def operation() -> str: + nonlocal attempts + attempts += 1 + raise RateLimited("slow down", status_code=429, retry_after=30.0) + + # A clock that advances by exactly what we sleep, so the only thing that can + # exhaust the budget is the server-named wait itself. + now = 0.0 + + def sleep(seconds: float) -> None: + nonlocal now + now += seconds + + monkeypatch.setattr(liveness.time, "monotonic", lambda: now) + monkeypatch.setattr(retry.time, "sleep", sleep) + + with pytest.raises(RateLimited): + retry.retry_sync( + operation, + retry.RetryPolicy(max_retries=4, stall_timeout=60.0, retry_after_cap=60.0), + ) + + assert attempts == 5, "a sanctioned wait costs the no-progress budget nothing" + + +def test_credited_wait_never_credits_past_the_present(monkeypatch) -> None: + """A wait longer than the budget must not disable the budget. + + ``credit_wait`` moves the progress stamp forward; without a ceiling at + "now", one long queue wait pushed it into the future, made + ``elapsed_since_progress`` negative, and -- since nothing ever pulls it back + -- left that call exempt from the stall bound for the rest of its life. + """ + now = 0.0 + monkeypatch.setattr(liveness.time, "monotonic", lambda: now) + policy = retry.RetryPolicy(stall_timeout=60.0) + + liveness.note_progress() + liveness.credit_wait(300.0) # a deep-tail task queued past the whole budget + assert liveness.elapsed_since_progress() == 0.0, "clamped to now, not negative" + + # The budget is spent again by real silence, not permanently disabled. + now = 200.0 + assert not policy.allows_wait(5, 30.0, liveness.elapsed_since_progress()) + + +def test_arriving_pages_restart_the_stall_budget(monkeypatch) -> None: + """A slow but productive download keeps earning more time.""" + now = 0.0 + monkeypatch.setattr(liveness.time, "monotonic", lambda: now) + policy = retry.RetryPolicy(stall_timeout=60.0) + + liveness.note_progress() + now = 100.0 + assert not policy.allows_wait(2, 0.5, liveness.elapsed_since_progress()) + + liveness.note_progress() # a page arrived + assert policy.allows_wait(2, 0.5, liveness.elapsed_since_progress()) + + +def test_bad_retry_environment_raises_a_catchable_error(monkeypatch) -> None: + """A typo in the environment must not escape as a bare ValueError. + + Every retrieval path builds its policy from the environment, so an + unparseable value would otherwise bypass ``except DataRetrievalError`` in + caller code and abort the run with an unrelated-looking error. + """ + monkeypatch.setenv("API_USGS_RETRIES", "off") + with pytest.raises(DataRetrievalError): + retry.RetryPolicy.from_env() + + monkeypatch.setenv("API_USGS_RETRIES", "2") + monkeypatch.setenv("API_USGS_STALL_TIMEOUT", "none") + with pytest.raises(ConfigurationError): + retry.RetryPolicy.from_env() + + monkeypatch.setenv("API_USGS_STALL_TIMEOUT", "10") + assert retry.RetryPolicy.from_env().stall_timeout == 10.0 + # Still a ValueError, so existing handling of a bad setting keeps working. + assert issubclass(ConfigurationError, ValueError) + + +def test_queued_work_keeps_its_retries() -> None: + """Time spent waiting for a concurrency slot is not silence. + + The no-progress budget starts when a retry loop is entered, but a fan-out + task may sit behind a full semaphore long after that. Without excusing the + wait, the tail of a wide fan-out enters its first attempt with the budget + already spent, while the tasks dispatched ahead of it get the full + allowance. + """ + + async def drive() -> dict[int, int]: + gate = asyncio.Semaphore(1) + attempts: dict[int, int] = {} + policy = retry.RetryPolicy( + max_retries=2, stall_timeout=1.0, base_backoff=0.001, max_backoff=0.001 + ) + + async def one(index: int) -> None: + attempts[index] = 0 + + async def attempt() -> str: + attempts[index] += 1 + if index == 0: + await asyncio.sleep(1.2) # hold the gate past the budget + return "ok" + raise ServiceUnavailable("busy", status_code=503) + + try: + await retry.retry_async(attempt, policy, gate=gate) + except ServiceUnavailable: + pass + + await asyncio.gather(*(one(i) for i in range(3))) + return attempts + + attempts = asyncio.run(drive()) + assert attempts[0] == 1 + # Queued behind a 1.2 s hold with a 1.0 s budget, these still get retried. + assert attempts[1] == 3, attempts + assert attempts[2] == 3, attempts + + +def test_gate_does_not_reset_silence_from_earlier_attempts() -> None: + """Excusing the queue wait must not also forgive accumulated silence. + + The gated body is what the retry loop re-invokes, so stamping "now" on every + slot acquisition would restart the clock each attempt and quietly turn a + bound on *total* silence into a per-attempt latency bound -- five slow + failures would each look brief while the call sat silent for their sum. + """ + + async def drive() -> int: + gate = asyncio.Semaphore(4) # never contended: no waiting to excuse + attempts = 0 + policy = retry.RetryPolicy( + max_retries=4, stall_timeout=1.0, base_backoff=0.001, max_backoff=0.001 + ) + + async def attempt() -> str: + nonlocal attempts + attempts += 1 + await asyncio.sleep(0.4) # each attempt is silent for 0.4 s + raise ServiceUnavailable("gateway", status_code=504) + + try: + await retry.retry_async(attempt, policy, gate=gate) + except ServiceUnavailable: + pass + return attempts + + # 0.4 s per attempt against a 1.0 s budget: attempt 1 is exempt, attempt 2 + # accumulates past the budget. Five attempts would mean the budget stopped + # counting across attempts. + assert asyncio.run(drive()) == 3 + + +def _wrapped_dns_failure(errno: int) -> NetworkError: + """A DNS failure shaped the way one reaches the chunker. + + Our own wrapper uses ``raise ... from``, so the outer error links to the + httpx one explicitly; only the layers beneath it chain implicitly. The + chunker follows explicit links only, so this shape -- not + :func:`_dns_failure`'s -- is the one that decides whether an unresolvable + host is offered as resumable. + """ + resolution_failed = socket.gaierror(errno, "name resolution failed") + transport_failed = httpx.ConnectError("name resolution failed") + transport_failed.__context__ = resolution_failed + wrapped = NetworkError("Could not reach the service at https://nope.invalid") + wrapped.__cause__ = transport_failed + return wrapped + + +def test_deterministic_failures_are_not_offered_as_resumable() -> None: + """ "Retryable" and "resumable" are one judgement and must agree. + + ``_retryable`` already refuses to re-send a hostname the resolver rejects + outright. If the interruption classifier still mapped it to + ``ServiceInterrupted``, the caller would be handed a ``.call.resume()`` + whose every attempt fails identically -- a resumable handle for something + that cannot be resumed, hiding the ``NetworkError`` that actually explains + the failure. + """ + from dataretrieval.ogc.retry import _classify_chunk_error + + permanent = _wrapped_dns_failure(socket.EAI_NONAME) + assert retry._retryable(permanent) == (False, None) + assert _classify_chunk_error(permanent) is None + + unsupported = httpx.UnsupportedProtocol("no scheme") + assert retry._retryable(unsupported) == (False, None) + assert _classify_chunk_error(unsupported) is None + + # The converse still holds: a failure a later attempt could survive stays + # both retryable and resumable. A temporary resolver failure is the sharp + # case -- same exception type, same chain shape, opposite verdict, decided + # only by the errno. + temporary = _wrapped_dns_failure(socket.EAI_AGAIN) + assert retry._retryable(temporary) == (True, None) + assert _classify_chunk_error(temporary) is not None diff --git a/tests/utils_test.py b/tests/utils_test.py index 30950294..2a743cb6 100644 --- a/tests/utils_test.py +++ b/tests/utils_test.py @@ -33,6 +33,23 @@ def test_header(self, httpx_mock): assert response.status_code == 200 # GET was successful assert "user-agent" in response.request.headers + def test_opt_in_retry_recovers_from_transient(self, httpx_mock, monkeypatch): + """Active adapters can opt into bounded retry without changing NWIS.""" + import dataretrieval.transport.retry as retry + + url = "https://example.invalid/x" + request_url = f"{url}?a=1" + httpx_mock.add_response(method="GET", url=request_url, status_code=503) + httpx_mock.add_response(method="GET", url=request_url, text="ok") + monkeypatch.setenv("API_USGS_RETRIES", "1") + monkeypatch.setattr(retry, "_RETRY_BASE_BACKOFF", 0.0) + monkeypatch.setattr(retry, "_RETRY_MAX_BACKOFF", 0.0) + + response = utils._query_with_retry(url, {"a": "1"}) + + assert response.text == "ok" + assert len(httpx_mock.get_requests()) == 2 + class Test_error_taxonomy: """The unified request-error hierarchy. @@ -386,3 +403,17 @@ def test_resolves_an_iterable_element_wise(self): # A bad element fails the whole call (fail-fast). with pytest.raises(ValueError, match="not a recognized US state"): to_state(["WI", "XX"]) + + +def test_retrying_get_maps_invalid_url(monkeypatch): + """Direct active-service GETs do not leak raw httpx InvalidURL errors.""" + import httpx + + monkeypatch.setattr( + utils, + "_get", + mock.Mock(side_effect=httpx.InvalidURL("invalid URL")), + ) + + with pytest.raises(exceptions.URLTooLong): + utils._get_with_retry("https://example.invalid") diff --git a/tests/waterdata_chunking_test.py b/tests/waterdata_chunking_test.py index ddefafcc..06657821 100644 --- a/tests/waterdata_chunking_test.py +++ b/tests/waterdata_chunking_test.py @@ -31,6 +31,11 @@ import pandas as pd import pytest +from dataretrieval.combining import ( + _QUOTA_HEADER, + _combine_chunk_frames, + _combine_chunk_responses, +) from dataretrieval.exceptions import ( DataRetrievalError, RateLimited, @@ -40,7 +45,6 @@ ) from dataretrieval.ogc import chunking as _chunking from dataretrieval.ogc import engine as _engine -from dataretrieval.ogc import retry as _retry_mod from dataretrieval.ogc.chunking import ( ChunkedCall, _chunked_client, @@ -49,11 +53,6 @@ multi_value_chunked, parallel_chunks, ) -from dataretrieval.ogc.combining import ( - _QUOTA_HEADER, - _combine_chunk_frames, - _combine_chunk_responses, -) from dataretrieval.ogc.dates import _DATE_RANGE_PARAMS from dataretrieval.ogc.interruptions import ( ChunkInterrupted, @@ -70,12 +69,15 @@ _safe_request_bytes, ) from dataretrieval.ogc.requests import _construct_api_requests -from dataretrieval.ogc.retry import ( +from dataretrieval.transport import retry as _retry_mod +from dataretrieval.transport.retry import ( _RETRIES_DEFAULT, RetryPolicy, - _retry, _retryable, ) +from dataretrieval.transport.retry import ( + retry_async as _retry, +) from dataretrieval.utils import HTTPX_DEFAULTS @@ -1818,13 +1820,6 @@ def _wrap_cause(transport_exc): # -- RetryPolicy (pure value object) ---------------------------------------- -def test_retry_policy_backoff_honors_retry_after(): - policy = RetryPolicy() - # A server Retry-After overrides the computed backoff verbatim. - assert policy.backoff(attempt=1, retry_after=7.5) == 7.5 - assert policy.backoff(attempt=4, retry_after=2.0) == 2.0 - - def test_retry_policy_backoff_full_jitter_within_ceiling(): policy = RetryPolicy(base_backoff=2.0, max_backoff=30.0) for attempt, ceiling in [(1, 2.0), (2, 4.0), (3, 8.0), (5, 30.0)]: diff --git a/tests/waterdata_progress_test.py b/tests/waterdata_progress_test.py index dc752591..33e246f1 100644 --- a/tests/waterdata_progress_test.py +++ b/tests/waterdata_progress_test.py @@ -17,11 +17,11 @@ import pandas as pd import pytest -from dataretrieval.ogc import progress as _progress +from dataretrieval import progress as _progress from dataretrieval.ogc.chunking import ChunkedCall from dataretrieval.ogc.engine import _paginate, _walk_pages from dataretrieval.ogc.planning import ChunkPlan -from dataretrieval.ogc.progress import ( +from dataretrieval.progress import ( ProgressReporter, current, progress_context, @@ -40,6 +40,11 @@ def _run_walk_pages(*, geopd, req, client): return asyncio.run(_walk_pages(geopd=geopd, req=req, client=client)) +# The Water Data host is the only one that honors ``API_USGS_PAT``, and so the +# only one where pointing the user at API-key registration is useful advice. +_KEYED_URL = "https://api.waterdata.usgs.gov/ogcapi/v0/" + + @pytest.fixture(autouse=True) def _reset_api_key_hint_latch(monkeypatch): """The 'no API key' pointer is latched once per process; reset it so each @@ -220,7 +225,7 @@ def test_reporter_swallows_stream_errors_and_disables(monkeypatch): def test_hints_api_key_when_no_key_configured(monkeypatch): monkeypatch.delenv("API_USGS_PAT", raising=False) stream = io.StringIO() - reporter = ProgressReporter(stream=stream, enabled=True) + reporter = ProgressReporter(stream=stream, enabled=True, target_url=_KEYED_URL) reporter.add_page(rows=5) reporter.close() assert _progress.SIGNUP_URL in stream.getvalue() @@ -231,7 +236,7 @@ def test_hint_fires_even_when_rate_limit_was_seen(monkeypatch): # — not absence of the header — is what drives the pointer. monkeypatch.delenv("API_USGS_PAT", raising=False) stream = io.StringIO() - reporter = ProgressReporter(stream=stream, enabled=True) + reporter = ProgressReporter(stream=stream, enabled=True, target_url=_KEYED_URL) reporter.set_rate_remaining("891") reporter.add_page(rows=5) reporter.close() @@ -241,12 +246,28 @@ def test_hint_fires_even_when_rate_limit_was_seen(monkeypatch): def test_no_hint_when_api_key_present(monkeypatch): monkeypatch.setenv("API_USGS_PAT", "secret") stream = io.StringIO() - reporter = ProgressReporter(stream=stream, enabled=True) + reporter = ProgressReporter(stream=stream, enabled=True, target_url=_KEYED_URL) reporter.add_page(rows=5) # no rate-limit, but a key is configured reporter.close() assert _progress.SIGNUP_URL not in stream.getvalue() +def test_no_hint_for_a_service_the_key_does_not_cover(monkeypatch): + """Only the host that honors ``API_USGS_PAT`` gets the sign-up pointer. + + Water Use is on a different host and never receives the key, so telling its + users to register sends them after a fix that changes nothing. + """ + monkeypatch.delenv("API_USGS_PAT", raising=False) + stream = io.StringIO() + reporter = ProgressReporter( + stream=stream, enabled=True, target_url="https://api.water.usgs.gov/nwaa-data/" + ) + reporter.add_page(rows=5) + reporter.close() + assert _progress.SIGNUP_URL not in stream.getvalue() + + def test_no_hint_when_disabled(monkeypatch): monkeypatch.delenv("API_USGS_PAT", raising=False) stream = io.StringIO() @@ -260,13 +281,13 @@ def test_api_key_hint_shown_at_most_once(monkeypatch): monkeypatch.delenv("API_USGS_PAT", raising=False) first = io.StringIO() - r1 = ProgressReporter(stream=first, enabled=True) + r1 = ProgressReporter(stream=first, enabled=True, target_url=_KEYED_URL) r1.add_page(rows=5) r1.close() assert _progress.SIGNUP_URL in first.getvalue() second = io.StringIO() - r2 = ProgressReporter(stream=second, enabled=True) + r2 = ProgressReporter(stream=second, enabled=True, target_url=_KEYED_URL) r2.add_page(rows=5) r2.close() assert _progress.SIGNUP_URL not in second.getvalue() diff --git a/tests/waterdata_ratings_test.py b/tests/waterdata_ratings_test.py index 54558f42..bd1d8e43 100644 --- a/tests/waterdata_ratings_test.py +++ b/tests/waterdata_ratings_test.py @@ -4,6 +4,7 @@ import pandas as pd import pytest +from dataretrieval.exceptions import DataRetrievalError from dataretrieval.waterdata import get_ratings from dataretrieval.waterdata.ratings import _build_filter @@ -198,3 +199,54 @@ def test_get_ratings_multi_type_filters_via_property(httpx_mock, tmp_path): search_req = httpx_mock.get_requests()[0] qs = parse_qs(urlsplit(str(search_req.url)).query) assert "file_type" not in qs["filter"][0] + + +def test_stac_next_link_refuses_another_host(httpx_mock): + """The STAC page walk must not follow a link off the ratings host. + + The search request carries the Water Data API key; a ``next`` href naming + another host would take it somewhere the caller never asked for. Unlike the + OGC engine, this walk had no host check at all. + """ + httpx_mock.add_response( + method="GET", + url=re.compile(r".*/stac/v0/search.*"), + json={ + "features": [{"id": "a", "properties": {}, "assets": {}}], + "links": [{"rel": "next", "href": "https://evil.example/page2"}], + }, + ) + with pytest.raises(DataRetrievalError, match="rather than"): + get_ratings(monitoring_location_id="USGS-X") + + +def test_stac_next_link_strips_embedded_credentials(httpx_mock): + """A same-host ``next`` href must not smuggle in ``user:pass@``. + + The host check passes by construction here, so only the strip catches it. + """ + httpx_mock.add_response( + method="GET", + url=re.compile(r"^https://api\.waterdata\.usgs\.gov/stac/v0/search\?.*"), + json={ + "features": [], + "links": [ + { + "rel": "next", + "href": "https://u:p@api.waterdata.usgs.gov/stac/v0/search?page=2", + } + ], + }, + ) + # Second page: no ``next``, so the walk terminates. + httpx_mock.add_response( + method="GET", + url="https://api.waterdata.usgs.gov/stac/v0/search?page=2", + json={"features": [], "links": []}, + ) + + assert get_ratings(monitoring_location_id="USGS-X") == {} + + followed = httpx_mock.get_requests()[1] + assert followed.url.userinfo == b"" + assert followed.headers.get("Authorization") is None diff --git a/tests/waterdata_utils_test.py b/tests/waterdata_utils_test.py index c39a8b19..453a3de2 100644 --- a/tests/waterdata_utils_test.py +++ b/tests/waterdata_utils_test.py @@ -878,13 +878,15 @@ def test_parse_retry_after_clamps_negative_delta_to_zero(): assert _parse_retry_after("-0.5") == 0.0 -def test_parse_retry_after_returns_none_for_unparseable(): - """Garbage values (including the RFC 1123 HTTP-date form that the - HTTP spec allows but USGS doesn't actually send) surface as - ``None``, letting the chunker fall back to its own retry policy - instead of guessing a delay.""" +def test_parse_retry_after_supports_http_date_and_rejects_garbage(): + """Both standard header forms are accepted; malformed values use backoff. + + A date is converted to seconds exactly like the delta-seconds form, however + far out it lands: an over-long wait stops the retry and travels to the + caller on ``.retry_after`` rather than being silently ignored. + """ assert _parse_retry_after("not-a-date") is None - assert _parse_retry_after("Wed, 21 Oct 2099 07:28:00 GMT") is None + assert _parse_retry_after("Wed, 21 Oct 2099 07:28:00 GMT") > 0 def test_raise_for_non_200_raises_service_unavailable_for_5xx(): @@ -944,6 +946,34 @@ def test_next_req_url_rejects_cross_host(): _next_req_url(resp, body=body) +def test_next_req_url_strips_embedded_credentials(): + """A same-host next link carrying ``user:pass@`` must not survive. + + The cross-host guard passes here by construction -- the host matches -- so + nothing else would catch it. httpx derives ``Authorization: Basic`` from + userinfo, so following the link verbatim would mint a credential the caller + never configured and send it alongside the real API key. + """ + resp = mock.MagicMock() + resp.url = httpx.URL("https://api.waterdata.usgs.gov/page1") + body = { + "numberReturned": 1, + "features": [{"id": "1"}], + "links": [ + {"rel": "next", "href": "https://attacker:pw@api.waterdata.usgs.gov/page2"} + ], + } + following = _next_req_url(resp, body=body) + assert following == "https://api.waterdata.usgs.gov/page2" + # Asserted through a real client: ``httpx.Request`` alone never derives the + # header from userinfo, so checking it there would pass for any URL. + with httpx.Client( + transport=httpx.MockTransport(lambda request: httpx.Response(200)) + ) as client: + sent = client.build_request("GET", following) + assert sent.headers.get("Authorization") is None + + def test_check_ogc_requests_raises_typed_on_5xx(httpx_mock): """``_check_ogc_requests`` routes a non-200 through ``_raise_for_non_200``, so a 5xx surfaces as the typed ``ServiceUnavailable`` — the same typed diff --git a/tests/wateruse_test.py b/tests/wateruse_test.py index 00a843c8..b6ee4e00 100644 --- a/tests/wateruse_test.py +++ b/tests/wateruse_test.py @@ -306,10 +306,28 @@ def test_fan_out_surfaces_final_rate_limit_header(httpx_mock): assert md.header["x-ratelimit-remaining"] == "850" -# (response aggregation now reuses ogc.combining._combine_chunk_responses; the +# (response aggregation uses combining._combine_chunk_responses; the # integration test above pins the rate-limit-header behavior end-to-end.) +def test_fan_out_failure_never_returns_partial_data(httpx_mock): + """A failed location aborts the call even when another location succeeded.""" + 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.*"), + status_code=503, + json={"detail": "temporarily unavailable"}, + ) + + with pytest.raises(dataretrieval.ServiceUnavailable): + get_wateruse(model="wu-public-supply-wd", state=["RI", "WI"]) + + # --- _resolve_locations unit tests (no HTTP) ------------------------------- @@ -411,6 +429,165 @@ def test_next_page_url_leaves_api_host_untouched(): assert _next_page_url(resp) == url +def test_next_page_url_normalizes_other_spellings_of_the_same_service(): + """The cursor is normalized by host, not by one literal prefix. + + A plain-http or relative ``next`` link is the same service; refusing it + would throw away every page already collected for that location. + """ + plain_http = httpx.Response( + 200, + text="", + headers={"link": '; rel="next"'}, + ) + assert _next_page_url(plain_http) == ( + "https://api.water.usgs.gov/nwaa-data/data?skip=600" + ) + + relative = httpx.Response( + 200, + text="", + headers={"link": '; rel="next"'}, + request=httpx.Request("GET", "https://api.water.usgs.gov/nwaa-data/data"), + ) + assert _next_page_url(relative) == ( + "https://api.water.usgs.gov/nwaa-data/data?skip=600" + ) + + +def test_next_page_url_strips_credentials_from_the_cursor(): + """Userinfo on a cursor must not become an Authorization header. + + httpx derives ``Authorization: Basic ...`` from a URL's userinfo, so a + cursor spelled ``http://user:pass@water.usgs.gov/...`` would send a + credential the caller never configured to the rewritten host -- exactly what + the host check exists to prevent, arriving through the host check's own + normalization. The port is dropped for the same reason. + """ + response = httpx.Response( + 200, + text="", + headers={ + "link": ( + "; rel="next"' + ) + }, + ) + + cursor = _next_page_url(response) + + assert cursor == "https://api.water.usgs.gov/nwaa-data/data?skip=600" + assert "s3cret" not in cursor + assert httpx.URL(cursor).userinfo == b"" + + # Assert at the layer that actually synthesizes the header: ``httpx.Request`` + # never derives Basic auth from userinfo (so asserting there would pass for + # any URL) -- the ``Client`` does it at send time. + sent: dict[str, str | None] = {} + + def capture(request: httpx.Request) -> httpx.Response: + sent["auth"] = request.headers.get("Authorization") + return httpx.Response(200, text="") + + with httpx.Client(transport=httpx.MockTransport(capture)) as client: + client.get(cursor) + assert sent["auth"] is None + + def test_module_exposes_catalog_constants(): assert "wu-public-supply-wd" in wateruse.MODELS assert set(wateruse.TIME_RESOLUTIONS) == {"monthly", "annualcy", "annualwy"} + + +def test_initial_transient_is_retried(httpx_mock, monkeypatch): + """Water Use retries an initial transient without holding its semaphore.""" + import dataretrieval.transport.retry as retry + + url = re.compile(r".*location=stateCd%3ARI.*") + httpx_mock.add_response(method="GET", url=url, status_code=503) + httpx_mock.add_response(method="GET", url=url, text=_CSV_P1) + monkeypatch.setenv("API_USGS_RETRIES", "1") + monkeypatch.setattr(retry, "_RETRY_BASE_BACKOFF", 0.0) + monkeypatch.setattr(retry, "_RETRY_MAX_BACKOFF", 0.0) + + df, _ = get_wateruse(model="wu-public-supply-wd", state="RI") + + assert len(df) == 2 + assert len(httpx_mock.get_requests()) == 2 + + +def test_fatal_failure_waits_for_siblings_before_closing_the_client(monkeypatch): + """A fan-out failure must not close the client under its own siblings. + + Every location shares one ``httpx.AsyncClient`` scoped to the fan-out. When + the first failure propagated straight out of the ``gather``, that block + exited while siblings were still walking pages, and the next page they asked + for failed with "Cannot send a request, as the client has been closed" -- on + a task nobody was awaiting any more, so it also surfaced as an unretrieved + exception. Both are artifacts of our own teardown, not of the service. + """ + import asyncio + from contextlib import asynccontextmanager + + pages = {"n": 0} + + async def handler(request: httpx.Request) -> httpx.Response: + if request.url.params.get("location") == "stateCd:AA": + return httpx.Response(400, json={"detail": "Invalid model name: bad"}) + pages["n"] += 1 + # A real suspension window, deliberately: the sibling has to still be + # mid-walk when the failure propagates, and an ``Event`` set by the + # failing branch is already set by the time this runs -- it returns + # without suspending, and the test then passes against the old code too. + await asyncio.sleep(0.05) + if pages["n"] == 1: + return httpx.Response( + 200, + text=_CSV_P1, + headers={ + "link": ( + "; rel="next"' + ) + }, + ) + return httpx.Response(200, text=_CSV_P2) + + @asynccontextmanager + async def open_mock_client(**overrides): + overrides.pop("verify", None) + async with httpx.AsyncClient( + transport=httpx.MockTransport(handler), **overrides + ) as client: + yield client + + monkeypatch.setattr(wateruse, "open_async_client", open_mock_client) + + requests = [ + httpx.Request("GET", wateruse.WATERUSE_URL, params={"location": location}) + for location in ("stateCd:AA", "stateCd:BB") + ] + + async def drive() -> int: + with pytest.raises(dataretrieval.DataRetrievalError, match="Invalid model"): + await wateruse._fan_out(requests, {}, True) + # Nothing left running: the sibling was joined before teardown, so no + # task can later fail against the closed client. + return len( + [task for task in asyncio.all_tasks() if task is not asyncio.current_task()] + ) + + assert asyncio.run(drive()) == 0 + assert pages["n"] == 2, "the sibling finished its walk rather than being abandoned" + + +def test_next_page_url_rejects_cross_host_link(): + response = httpx.Response( + 200, + headers={"link": '; rel="next"'}, + ) + # Typed, so a caller's ``except DataRetrievalError`` catches it like any + # other failure rather than seeing a bare RuntimeError. + with pytest.raises(dataretrieval.DataRetrievalError, match="outside.example"): + _next_page_url(response) diff --git a/tests/wqp_test.py b/tests/wqp_test.py index e4d0dba0..e7f34e67 100644 --- a/tests/wqp_test.py +++ b/tests/wqp_test.py @@ -4,6 +4,7 @@ import pytest from pandas import DataFrame +import dataretrieval.wqp as wqp from dataretrieval.wqp import ( WQP_Metadata, _check_kwargs, @@ -37,6 +38,23 @@ def _assert_wqp_metadata(md, request_url): assert md.comment is None +def test_get_results_opts_into_retry(monkeypatch): + """WQP explicitly enables retry at its shared query boundary.""" + response = mock.Mock( + text="ResultIdentifier,ResultMeasureValue\nA,1.0\n", + url="https://example.test", + elapsed=datetime.timedelta(), + headers={}, + ) + query = mock.Mock(return_value=response) + monkeypatch.setattr(wqp, "_query_with_retry", query) + + df, _ = wqp.get_results(legacy=True) + + assert len(df) == 1 + assert query.call_count == 1 + + def test_read_wqp_csv_preserves_leading_zero_codes(): """Regression: WQP code columns (HUCs, parameter codes, FIPS) carry significant leading zeros; a bare ``read_csv`` inferred them as int/float From 518db33acc475c6183db76437e6e8b70c0e85c17 Mon Sep 17 00:00:00 2001 From: thodson-usgs Date: Thu, 6 Aug 2026 12:51:25 -0500 Subject: [PATCH 2/2] refactor(transport)!: share fan-out execution across services 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. Unable to reach an OGC-internal executor, wateruse._fan_out re-implemented the semaphore, the gather, and the failure-precedence rule, with a comment naming ChunkedCall._run as the original -- one subtle rule, two copies, synchronized by prose. The duplicate lacked resume (a 429 partway through discarded every completed location), reported no progress, and ignored API_USGS_CONCURRENT. Move execution down; leave planning up. transport.fanout.FanOut drives any FanOutPlan -- a Protocol of the three members the executor already used (total, canonical_url, iter_sub_args). It is structural because its two implementations share an interface and no implementation: ChunkPlan derives sub-requests from a byte budget, a Water Use plan lists locations the caller already named separately. Water Use sheds ~75 lines and gains resume, progress, and the shared concurrency setting. Concurrency is now one general knob with per-service defaults, and an explicitly set API_USGS_CONCURRENT outranks a service default -- never the reverse, or the setting would be a lie. The interruption taxonomy moves to the dataretrieval.interruptions leaf, since adapters need it whether or not they went through transport. Its base is renamed FanOutInterrupted, because Water Use raises it without chunking anything; ChunkInterrupted stays as a permanent alias of the same class object, so `except ChunkInterrupted` keeps working. _deterministic_failure moves to that leaf too, and transport.retry imports it back. Whether a failure is worth retrying and whether it can be resumed are one judgement about what the exception means, not two -- and the leaf is where meaning lives. Leaving it in transport would have forced the leaf to import transport to ask. BREAKING CHANGE: a Water Use fan-out interrupted by 5xx/429 now raises ServiceInterrupted/QuotaExhausted rather than ServiceUnavailable/ RateLimited. Both remain DataRetrievalError, so broad handlers are unaffected, but a narrow `except ServiceUnavailable` must widen. This is convergence with the OGC getters, and it is what makes the failure resumable. wateruse.MAX_CONCURRENT_REQUESTS is removed in favor of API_USGS_CONCURRENT / wateruse.DEFAULT_CONCURRENT_REQUESTS. Supersedes the ADR 0006 clause assigning resumable ChunkedCall state to OGC; see ADR 0008. Co-Authored-By: Claude Opus 5 --- NEWS.md | 2 + dataretrieval/__init__.py | 26 +- dataretrieval/interruptions.py | 299 ++++++++ dataretrieval/ogc/chunking.py | 586 +--------------- dataretrieval/ogc/interruptions.py | 197 +----- dataretrieval/ogc/retry.py | 63 +- dataretrieval/transport/fanout.py | 652 ++++++++++++++++++ dataretrieval/transport/retry.py | 53 +- dataretrieval/wateruse.py | 189 +++-- .../0006-service-neutral-transport.rst | 4 +- .../decisions/0008-fan-out-execution.rst | 117 ++++ docs/source/architecture/decisions/index.rst | 1 + docs/source/architecture/index.rst | 6 +- docs/source/reference/exceptions.rst | 21 +- docs/source/userguide/errors.rst | 25 +- tests/architecture_test.py | 74 ++ tests/waterdata_chunking_test.py | 18 +- tests/wateruse_test.py | 151 +++- 18 files changed, 1514 insertions(+), 970 deletions(-) create mode 100644 dataretrieval/interruptions.py create mode 100644 dataretrieval/transport/fanout.py create mode 100644 docs/source/architecture/decisions/0008-fan-out-execution.rst 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