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/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/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..71ea5516 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.transport.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.exceptions import ConfigurationError +from dataretrieval.transport import progress as _progress +from dataretrieval.transport.combining import ( _combine_chunk_frames, _combine_chunk_responses, ) -from .interruptions import ( - ChunkInterrupted, -) +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..060c7a56 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,18 @@ 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.transport.progress as _progress 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 +66,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, ) @@ -159,46 +154,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 +165,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 API-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 +343,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 +379,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 API-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..a4a09eee 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.transport.combining._combine_chunk_frames`, +:func:`~dataretrieval.transport.combining._combine_chunk_responses`, etc.) — +lives in the API-neutral :mod:`dataretrieval.transport.combining` module. """ from __future__ import annotations 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..e669e02d 100644 --- a/dataretrieval/ogc/retry.py +++ b/dataretrieval/ogc/retry.py @@ -1,25 +1,17 @@ -"""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 API-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. """ 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, @@ -27,185 +19,14 @@ 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) - 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)): return ServiceInterrupted, None @@ -215,139 +36,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/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..9405b7f4 --- /dev/null +++ b/dataretrieval/transport/__init__.py @@ -0,0 +1,7 @@ +"""Internal API-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/ogc/combining.py b/dataretrieval/transport/combining.py similarity index 96% rename from dataretrieval/ogc/combining.py rename to dataretrieval/transport/combining.py index be4366c1..0975da2f 100644 --- a/dataretrieval/ogc/combining.py +++ b/dataretrieval/transport/combining.py @@ -2,9 +2,9 @@ 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. @@ -104,7 +104,7 @@ 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 + (: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) diff --git a/dataretrieval/transport/http.py b/dataretrieval/transport/http.py new file mode 100644 index 00000000..7cd5a4f5 --- /dev/null +++ b/dataretrieval/transport/http.py @@ -0,0 +1,106 @@ +"""HTTP client lifecycle, timeout defaults, and host-scoped authentication.""" + +from __future__ import annotations + +import os +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.exceptions import NetworkError + +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), +} + +_AUTHORIZED_API_KEY_HOST = "api.waterdata.usgs.gov" + + +def accepts_api_key(target_url: str | httpx.URL | None) -> bool: + """Whether ``target_url`` names the host that honors ``API_USGS_PAT``. + + The single answer to "does this destination get the key" -- used both when + attaching the credential and when stripping it back off at redirect time, so + the two can't drift apart. It is also what makes "get an API key" useful + advice rather than noise: every other service this package talks to is on a + different host and ignores the key entirely. + """ + 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.host == _AUTHORIZED_API_KEY_HOST + + +def default_headers(target_url: str | httpx.URL | None = None) -> dict[str, str]: + """Build standard headers, scoping ``API_USGS_PAT`` to its authorized host.""" + headers = { + "Accept-Encoding": "compress, gzip", + "Accept": "application/json", + "User-Agent": USER_AGENT, + "lang": "en-US", + } + token = os.getenv("API_USGS_PAT") + if token and accepts_api_key(target_url): + headers["X-Api-Key"] = token + return headers + + +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) + + +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..6cd8921c --- /dev/null +++ b/dataretrieval/transport/liveness.py @@ -0,0 +1,49 @@ +"""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 waiting-for-a-turn from the no-progress budget. + + Queueing behind a concurrency cap is not silence -- 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 it 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. + """ + last = _last_progress.get() + if last is not None: + _last_progress.set(last + seconds) diff --git a/dataretrieval/transport/pagination.py b/dataretrieval/transport/pagination.py new file mode 100644 index 00000000..a14eae90 --- /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.exceptions import DataRetrievalError, RateLimited +from dataretrieval.transport import progress as _progress +from dataretrieval.transport.combining import ( + _QUOTA_HEADER, + _merge_response, + _safe_elapsed, +) +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/ogc/progress.py b/dataretrieval/transport/progress.py similarity index 86% rename from dataretrieval/ogc/progress.py rename to dataretrieval/transport/progress.py index 6177c30f..3520da7e 100644 --- a/dataretrieval/ogc/progress.py +++ b/dataretrieval/transport/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 @@ -26,7 +27,12 @@ import sys from collections.abc import Iterator from contextlib import contextmanager -from typing import TextIO +from typing import TYPE_CHECKING, TextIO + +from dataretrieval.transport.http import accepts_api_key + +if TYPE_CHECKING: + import httpx def _group_int(value: str) -> str: @@ -44,12 +50,13 @@ 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 + "transport_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). +# Where to register for an API key. Surfaced once when a query against the host +# that accepts one (see ``transport.http.accepts_api_key``) 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. @@ -104,9 +111,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 +235,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 +260,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 os.getenv("API_USGS_PAT"): 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 +277,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/transport/retry.py b/dataretrieval/transport/retry.py new file mode 100644 index 00000000..4a8cc2f5 --- /dev/null +++ b/dataretrieval/transport/retry.py @@ -0,0 +1,386 @@ +"""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 TypeVar + +import httpx + +from dataretrieval.exceptions import ( + ConfigurationError, + NetworkError, + TransientError, +) +from dataretrieval.transport import progress as _progress +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. + + Discarding an over-long hint was tried and is worse: it makes the client + retry *harder* against a service that just asked for a long pause, and drops + the number the caller needs from ``.retry_after``. Clock skew can inflate a + date-form hint, but the cost of trusting it is a recoverable escalation, + while the cost of ignoring it is 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 wait and every silent + #: attempt since the last page arrived. 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.""" + return cls( + retryable_statuses=( + _RETRYABLE_STATUSES + if retryable_statuses is None + else retryable_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) -> 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. + """ + if attempt <= _STALL_EXEMPT_ATTEMPTS: + return True + if self.stall_timeout <= 0 or elapsed is None: + return True + return elapsed + 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. The nudge is bounded by + :attr:`max_backoff`, not by this attempt's exponential ceiling -- keying + it to the ceiling made it vanish whenever :attr:`base_backoff` was zero, + which is exactly when a hint of ``0`` would become the zero-delay + re-send this is here to prevent. 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__``. So + both links are followed, preferring an explicit cause where one exists. + """ + seen: set[int] = set() + current: BaseException | None = exc + while current is not None and id(current) not in seen: + seen.add(id(current)) + if isinstance(current, (httpx.UnsupportedProtocol, httpx.LocalProtocolError)): + return True + if isinstance(current, socket.gaierror): + return current.errno in _PERMANENT_DNS_ERRORS + current = current.__cause__ or 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 + + +def _retry_delay(exc: BaseException, attempt: int, policy: RetryPolicy) -> float | None: + """Return the bounded delay 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()): + return None + reporter = _progress.current() + if reporter is not None: + reporter.note_retry(attempt=attempt, wait=delay) + return delay + + +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 + delay = _retry_delay(exc, attempt, policy) + if delay is None: + raise + await asyncio.sleep(delay) + + +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 + delay = _retry_delay(exc, attempt, policy) + if delay is None: + raise + time.sleep(delay) diff --git a/dataretrieval/transport/sync.py b/dataretrieval/transport/sync.py new file mode 100644 index 00000000..523d0faf --- /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.transport 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..d9c99140 100644 --- a/dataretrieval/utils.py +++ b/dataretrieval/utils.py @@ -5,38 +5,41 @@ 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.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 = _transport_http._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 +106,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 +361,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 +392,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 +479,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 +496,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..95249574 100644 --- a/dataretrieval/waterdata/ratings.py +++ b/dataretrieval/waterdata/ratings.py @@ -23,7 +23,15 @@ 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 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/wateruse.py b/dataretrieval/wateruse.py index 53de7597..18b44e7a 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 API-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. @@ -53,17 +51,21 @@ 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.transport.combining import ( + _combine_chunk_frames, + _combine_chunk_responses, ) +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,9 +370,15 @@ async def _one(request: httpx.Request) -> tuple[pd.DataFrame, httpx.Response]: 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) + results = await asyncio.gather(*(_one(req) for req in requests)) - # Reuse the engine's combine helpers: drop empty frames and concat, and fold + # 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. @@ -392,14 +406,38 @@ 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 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. + return str(target.copy_with(scheme="https", host=_WATERUSE_HOST, port=None)) 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..32c14131 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 +API-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 +- API-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-api-neutral-transport.rst b/docs/source/architecture/decisions/0006-api-neutral-transport.rst new file mode 100644 index 00000000..f5fb3eab --- /dev/null +++ b/docs/source/architecture/decisions/0006-api-neutral-transport.rst @@ -0,0 +1,86 @@ +ADR 0006: Use an API-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. Duplicating them would allow authentication, timeout, retry, and +failure behavior to drift. + +Decision +-------- + +``dataretrieval.transport`` is the internal API-neutral execution layer. It owns: + +- synchronous and asynchronous HTTP client lifecycle and timeout defaults; +- host-scoped API-key construction and redirect-time credential stripping; +- 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; +- DataFrame and HTTP-response aggregation; +- best-effort progress reporting; and +- the sync-over-async blocking-portal bridge. + +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 transport layer 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. + +Compliance +---------- + +``tests/architecture_test.py`` enforces transport dependency direction, an +acyclic transport graph, and Water Use isolation from OGC. 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..ef4c05e8 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-api-neutral-transport template diff --git a/docs/source/architecture/index.rst b/docs/source/architecture/index.rst index 906c07a4..9ed2a0e9 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 API-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 API-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 -> API-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..c89e42ed 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,74 @@ def test_default_header_calls_are_target_scoped() -> None: "_default_headers calls without destination URL context:\n" + "\n".join(violations) ) + + +# --- API-neutral transport boundaries --- + + +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 API-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/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..d895efef --- /dev/null +++ b/tests/transport_test.py @@ -0,0 +1,429 @@ +"""Component tests for the internal API-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. + + ``ogc.retry`` used to re-export these by value, so patching them there + changed a copy nothing consulted. It now owns only OGC classification. + """ + 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_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_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 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..d4331854 100644 --- a/tests/waterdata_chunking_test.py +++ b/tests/waterdata_chunking_test.py @@ -40,7 +40,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 +48,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 +64,20 @@ _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.combining import ( + _QUOTA_HEADER, + _combine_chunk_frames, + _combine_chunk_responses, +) +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..a75b4f86 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.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.transport import progress as _progress +from dataretrieval.transport.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_utils_test.py b/tests/waterdata_utils_test.py index c39a8b19..49d87365 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(): diff --git a/tests/wateruse_test.py b/tests/wateruse_test.py index 00a843c8..65572998 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 transport.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,60 @@ 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_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_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