diff --git a/.importlinter b/.importlinter index 2f3091bdd..f6b951381 100644 --- a/.importlinter +++ b/.importlinter @@ -10,10 +10,9 @@ [importlinter] root_package = dataretrieval -; Contracts describe what runs, matching the AST suite. ``ogc.interruptions`` -; and ``ogc.chunking`` reference each other's types under ``TYPE_CHECKING``; -; that is a documentation edge, not a runtime one, and no cycle exists at -; import time. +; Contracts describe what runs, matching the AST suite. Type-checking imports +; document structural protocols and callback types without creating runtime +; dependency edges. exclude_type_checking_imports = True [importlinter:contract:layers] @@ -25,9 +24,13 @@ layers = ngwmn | nldi | nwis | streamstats | waterdata | wateruse | wqp ogc utils + _querying transport progress - _ambient | _response_metadata | codes | combining | rdb +; Response-format conventions sit above the pure leaves because they read the +; code tables, and below every adapter that shapes a response with them. + _wqx + _ambient | _response_metadata | codes | combining | interruptions | rdb credentials exceptions ; Every top-level module must be placed in the stack deliberately. A new @@ -55,10 +58,9 @@ allowed_importers = dataretrieval.ngwmn dataretrieval.waterdata ignore_imports = -; The package __init__ re-exports the resumable-call and interruption types; -; they are part of the documented public surface, not a service reaching in. +; The package __init__ re-exports the parallel-chunks context manager; it is +; part of the documented public surface, not a service reaching into OGC. dataretrieval -> dataretrieval.ogc.chunking - dataretrieval -> dataretrieval.ogc.interruptions [importlinter:contract:ogc-facade] name = NGWMN consumes the OGC facade only, never its internals (ADR 0007) @@ -91,16 +93,19 @@ type = forbidden source_modules = dataretrieval.ogc forbidden_modules = + dataretrieval._querying dataretrieval.utils [importlinter:contract:nwis-quarantine] name = Deprecated NWIS has no dependents (ADR 0005) type = forbidden source_modules = + dataretrieval._querying dataretrieval.codes dataretrieval.combining dataretrieval.credentials dataretrieval.exceptions + dataretrieval.interruptions dataretrieval.ngwmn dataretrieval.nldi dataretrieval.ogc diff --git a/.pyscn-known-clones.json b/.pyscn-known-clones.json new file mode 100644 index 000000000..41448e734 --- /dev/null +++ b/.pyscn-known-clones.json @@ -0,0 +1,108 @@ +{ + "_comment": [ + "Accepted pyscn clone groups. Purpose: make a NEW clone group a visible event.", + "The duplication sub-score is dominated by these five accepted families, so it", + "barely moves when real duplication is added -- this file is what the weekly", + "Code Health workflow diffs against instead.", + "", + "Groups are keyed by the SET OF FUNCTION NAMES they contain, not by line numbers,", + "so ordinary edits above a getter do not churn this file. Regenerate only when", + "adding or removing a getter from an accepted family, and say why in the commit.", + "", + "These are NOT a backlog. Collapsing them means **kwargs and the loss of the typed", + "public surface tests/contracts/public_api_test.py freezes.", + "Verified 2026-08-09: at similarity_threshold 0.55 (default 0.65) the package still", + "reports exactly these five groups, so no latent clone sits below the bar.", + "Verified 2026-08-09: stripping docstrings makes duplication far WORSE (60 -> 0),", + "so suppressing them is not an option -- they are what keeps similarity below 0.86." + ], + "pyscn_version": "1.29.0", + "default_thresholds": { + "similarity_threshold": 0.65, + "min_lines": 10, + "min_nodes": 20 + }, + "observed": { + "total_fragments": 283, + "cloned_fragments": 25, + "groups": 5 + }, + "groups": [ + { + "members": [ + "dataretrieval/ngwmn.py::get_lithology", + "dataretrieval/ngwmn.py::get_providers", + "dataretrieval/ngwmn.py::get_well_construction" + ], + "similarity": 0.831, + "clone_type": 2, + "rationale": "NGWMN collection getters over the OGC facade, one per collection." + }, + { + "members": [ + "dataretrieval/nwis.py::get_dv", + "dataretrieval/nwis.py::get_iv" + ], + "similarity": 0.85, + "clone_type": 2, + "rationale": "Deprecated NWIS getters sharing the waterservices call shape; bodies already delegate to _get_json_values." + }, + { + "members": [ + "dataretrieval/waterdata/measurements.py::get_channel", + "dataretrieval/waterdata/measurements.py::get_field_measurements", + "dataretrieval/waterdata/measurements.py::get_peaks", + "dataretrieval/waterdata/metadata.py::get_field_measurements_metadata", + "dataretrieval/waterdata/metadata.py::get_monitoring_locations", + "dataretrieval/waterdata/metadata.py::get_time_series_metadata", + "dataretrieval/waterdata/time_series.py::get_continuous", + "dataretrieval/waterdata/time_series.py::get_daily", + "dataretrieval/waterdata/time_series.py::get_latest_continuous", + "dataretrieval/waterdata/time_series.py::get_latest_daily" + ], + "similarity": 0.805, + "clone_type": 2, + "rationale": "Water Data collection-family getters: ~30 explicitly typed parameters and a worked-example docstring around a thin body." + }, + { + "members": [ + "dataretrieval/waterdata/time_series.py::get_stats_date_range", + "dataretrieval/waterdata/time_series.py::get_stats_por" + ], + "similarity": 0.85, + "clone_type": 2, + "rationale": "Water Data collection-family getters: ~30 explicitly typed parameters and a worked-example docstring around a thin body." + }, + { + "members": [ + "dataretrieval/wqp.py::what_activities", + "dataretrieval/wqp.py::what_activity_metrics", + "dataretrieval/wqp.py::what_detection_limits", + "dataretrieval/wqp.py::what_habitat_metrics", + "dataretrieval/wqp.py::what_organizations", + "dataretrieval/wqp.py::what_project_weights", + "dataretrieval/wqp.py::what_projects", + "dataretrieval/wqp.py::what_sites" + ], + "similarity": 0.85, + "clone_type": 2, + "rationale": "WQP what_* wrappers: one shared query per service endpoint, differing only by the service path segment and the column each documents." + } + ], + "sub_threshold_accepted": { + "_comment": "Families that are NOT reported at default thresholds (they sit under min_lines=10 / min_nodes=20) but were found by a deliberate scan at min_nodes=8 on 2026-08-09. Recorded so they are not rediscovered and re-litigated. They are NOT counted in `observed` above.", + "groups": [ + { + "members": [ + "dataretrieval/nwis.py::get_discharge_measurements", + "dataretrieval/nwis.py::get_gwlevels", + "dataretrieval/nwis.py::get_pmcodes", + "dataretrieval/nwis.py::get_qwdata", + "dataretrieval/nwis.py::get_water_use" + ], + "similarity": 0.83, + "rationale": "Deprecated NWIS getters (ADR 0005). The surface is frozen pending removal, so collapsing it buys nothing and churns a public API users are being migrated off." + } + ] + } +} diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 8a58c41df..c6ce3b8ff 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -113,6 +113,17 @@ about the upstream service rather than about this package. ### Coding Standards and Style +**Before adding a small helper, check whether a leaf already generalizes it.** +This package keeps its general mechanisms in dependency-free leaves -- +`_ambient.Ambient` for scoped context values, `transport.retry._read_env_number` +for `API_USGS_*` settings, `transport.links.resolve_next_url` for pagination +cursors. Each of those has been re-implemented at least once by someone who +did not know it was there, and the copies drift: the same question gets a +different cycle guard, a different error message, a different edge case. None +of the automated checks catch it, because two eight-line helpers are below the +clone detector's floor and neither one couples or complicates anything. A grep +for the mechanism you are about to write is the only thing that does. + The continuous integration and pre-commit configurations enforce formatting, linting, and strict type checking. Run the relevant checks before opening a PR: diff --git a/NEWS.md b/NEWS.md index 7c9264769..bb5ae8cd3 100644 --- a/NEWS.md +++ b/NEWS.md @@ -1,3 +1,9 @@ +**08/09/2026:** Every retrieval path now runs through one executor. `waterdata.get_cql` (via the OGC `fetch_ogc_request`) and `waterdata.get_stats_por` / `get_stats_date_range` (via the Statistics page walk) previously bypassed `dataretrieval.transport.fanout.FanOut` through a private sync bridge, which meant they were the only getters in the package with **no retry**: a mid-page-walk 429 or 503 failed the whole call while every typed getter and Water Use rode it out. Both now run as a one-item fan-out and the 25-line `transport/sync.py` is gone. **Behavior change:** those three getters now retry transient failures (`API_USGS_RETRIES`, default 4) and, when the retries are exhausted, raise the resumable `ServiceInterrupted` / `QuotaExhausted` rather than `ServiceUnavailable` / `RateLimited` / `NetworkError` — all remain `DataRetrievalError`, so broad handlers are unaffected, but narrow handlers around those calls must widen, and `.call.resume()` is now available on the interruption. A failure that retrying cannot fix (bad scheme, a hostname that does not resolve) still surfaces as `NetworkError` immediately. The progress line moved with it: `FanOut.resume()` opens the reporter it ticks into, so a driver can no longer run the shared executor and silently print nothing, and a `.call.resume()` fired long after the interruption now reports progress instead of running mute. Internal tidying with no public effect: the WQX3 / legacy-WQP CSV datetime shaping moved out of `dataretrieval.utils` (whose docstring reserves it for non-service-specific shaping) into the `dataretrieval._wqx` leaf; the five Water Data endpoint URLs are declared once in `dataretrieval.waterdata.endpoints` instead of being derived in three modules; the OGC queryables document is parsed by `dataretrieval.ogc.schema` so every OGC adapter can offer the table, with `waterdata.get_queryables` unchanged as its documented wrapper; and `ogc/engine.py` imports each symbol from the module that defines it. + +**08/09/2026:** Internal structure cleanup, no public API change. Validating a server-supplied next-page link is now one policy in `dataretrieval.transport.links` instead of three divergent copies (the OGC engine, the ratings STAC walk, and Water Use). Two of those copies were fixed by the merge: the OGC page walk now resolves a *relative* `next` href against the page it came from (it previously handed the unresolved reference back as the pagination cursor) and refuses an unparseable one rather than following it unchecked. Cross-host refusal, credential stripping, and Water Use's host-alias rewrite are unchanged, as is the error type each walk raises. `parse_retry_after` moved to `dataretrieval.exceptions`, next to the `DataRetrievalError.retry_after` field it exists to produce. The one-shot HTTP query path (`query`, `to_str`, and their helpers) moved out of `dataretrieval.utils` into the private `dataretrieval._querying`; `dataretrieval.utils.query` and `dataretrieval.utils.to_str` remain the documented public paths, as `Ambient` and `BaseMetadata` already do. `waterdata` profile validation moved next to the tables it validates in `waterdata.types`, and `nwis.get_dv`/`get_iv` now share one body. + +**08/06/2026:** Fan-out execution is now shared across services. Chunking is how a query is divided structurally (a Water Data/NGWMN URL over the byte limit); fan-out is how the pieces are distributed operationally. Only the first is protocol-specific, so the executor moved to `dataretrieval.transport.fanout` (`FanOut`, over a three-member `FanOutPlan` protocol) while chunk planning stays in `dataretrieval.ogc`. Water Use no longer re-implements the fan-out gather and inherits resume, progress reporting, and `API_USGS_CONCURRENT`: a multi-location pull interrupted by a rate limit now raises a resumable interruption whose `.call.resume()` re-issues only the locations that did not finish, instead of discarding every completed one. The interruption taxonomy moved to the `dataretrieval.interruptions` leaf and its base class is now `FanOutInterrupted`; **`ChunkInterrupted` is a permanent alias of the same class**, so `except ChunkInterrupted` keeps working. **Breaking change:** a Water Use fan-out interrupted by a 5xx, 429, or recoverable connection failure now raises `ServiceInterrupted`/`QuotaExhausted` rather than `ServiceUnavailable`/`RateLimited`/`NetworkError` — all remain `DataRetrievalError`, so broad handlers are unaffected, but narrow handlers around a Water Use call must widen. **Breaking change:** `wateruse.MAX_CONCURRENT_REQUESTS` is removed; set `API_USGS_CONCURRENT` (which now outranks any service default) or read `wateruse.DEFAULT_CONCURRENT_REQUESTS`. + **08/03/2026:** Split the Water Data implementation into focused time-series, metadata, measurements, reference, samples, and CQL collection-family modules behind the unchanged `waterdata.api` facade. Active service modules now declare explicit exports; public Water Data imports, signatures, function identities, deprecations, and return contracts are protected by executable contract snapshots. OGC ambient context and schema/queryables execution are separated from request construction, adapter-to-adapter reach-through is prohibited by architecture tests, and service-specific output shapes are documented rather than forced into one model. **08/02/2026:** Added an internal API-neutral transport layer for guarded HTTP clients, host-scoped authentication, cursor pagination, bounded retry, response aggregation, progress, and sync-over-async dispatch. Water Use and the non-OGC Statistics API now consume transport directly instead of private OGC execution helpers; WQP, NLDI, and StreamStats opt into bounded transient retry while deprecated NWIS behavior remains unchanged. OGC retains CQL2, request construction, feature shaping, chunk planning, resumable calls, and interruption types, with compatibility imports at previous private paths. Failed pagination and fan-out still raise rather than returning partial data, and no public signatures or return shapes changed. diff --git a/dataretrieval/__init__.py b/dataretrieval/__init__.py index 4226e2473..29e288d7b 100644 --- a/dataretrieval/__init__.py +++ b/dataretrieval/__init__.py @@ -19,9 +19,10 @@ A failed request raises a subclass of :class:`dataretrieval.DataRetrievalError` (the taxonomy lives in ``dataretrieval.exceptions``); connection-level failures -(timeouts, DNS) are wrapped as :class:`dataretrieval.NetworkError`. A large -request interrupted mid-stream raises :class:`dataretrieval.ChunkInterrupted`, -whose ``.call.resume()`` continues from the work already completed. +(timeouts, DNS) are wrapped as :class:`dataretrieval.NetworkError`. A fanned-out +request interrupted mid-stream raises :class:`dataretrieval.FanOutInterrupted` +(also available under its original ``ChunkInterrupted`` name), whose +``.call.resume()`` continues from the work already completed. """ from importlib.metadata import PackageNotFoundError, version @@ -45,23 +46,24 @@ URLTooLong, ) -# Parallel-chunks control (a context manager). Defined with the chunker in -# ``dataretrieval.ogc.chunking``; surfaced here for a stable public path -# ``from dataretrieval import parallel_chunks``. -from dataretrieval.ogc.chunking import parallel_chunks - -# Resumable chunk-interruption exceptions. They are defined in -# ``dataretrieval.ogc.interruptions`` rather than ``dataretrieval.exceptions`` -# because they carry pandas/httpx state and a resumable ``ChunkedCall`` handle, +# Resumable fan-out interruption exceptions. They are defined in +# ``dataretrieval.interruptions`` rather than ``dataretrieval.exceptions`` +# because they carry pandas/httpx state and a resumable ``FanOut`` handle, # which would pull heavy dependencies into the lightweight exceptions module. -# Surfaced here so callers get a stable public path: -# ``from dataretrieval import ChunkInterrupted``. -from dataretrieval.ogc.interruptions import ( +# They are not under ``ogc`` because Water Use raises them too. Surfaced here so +# callers get a stable public path: ``from dataretrieval import ChunkInterrupted``. +from dataretrieval.interruptions import ( ChunkInterrupted, + FanOutInterrupted, QuotaExhausted, ServiceInterrupted, ) +# Parallel-chunks control (a context manager). Defined with the chunker in +# ``dataretrieval.ogc.chunking``; surfaced here for a stable public path +# ``from dataretrieval import parallel_chunks``. +from dataretrieval.ogc.chunking import parallel_chunks + from . import ( exceptions, ngwmn, @@ -96,8 +98,9 @@ "TransientError", "URLTooLong", "Unchunkable", - # resumable chunk-interruption exceptions (defined in ogc.interruptions) + # resumable fan-out interruption exceptions (defined in interruptions) "ChunkInterrupted", + "FanOutInterrupted", "QuotaExhausted", "ServiceInterrupted", # parallel-chunks control (defined in ogc.chunking) diff --git a/dataretrieval/_querying.py b/dataretrieval/_querying.py new file mode 100644 index 000000000..9fcde81aa --- /dev/null +++ b/dataretrieval/_querying.py @@ -0,0 +1,250 @@ +"""The one-shot HTTP query path behind the legacy service adapters. + +"Compose a USGS query URL, send it, map the status, retry a transient" -- the +half of the old ``utils`` module that talks to the network, as used by ``nwis``, +``wqp``, ``nldi``, ``streamstats`` and ``wateruse``. Its other half (pandas +column munging) shared nothing with this but a filename: no caller wanted both, +and the two have disjoint dependencies -- this one needs ``exceptions`` and +``transport``, that one needs ``codes`` and pandas. + +The module is private because the *names* are not: ``query`` and ``to_str`` keep +their documented ``dataretrieval.utils`` path, the way ``Ambient`` and +``BaseMetadata`` do from their own implementation leaves. This is legacy +machinery for the deprecated single-request adapters; new service code belongs +on the chunked transport instead. +""" + +from __future__ import annotations + +from collections.abc import Callable, Iterable +from typing import Any + +import httpx + +from dataretrieval.exceptions import ( + NoSitesError, + URLTooLong, + error_for_status, + parse_retry_after, +) +from dataretrieval.transport.http import HTTPX_DEFAULTS, USER_AGENT +from dataretrieval.transport.http import get as _get +from dataretrieval.transport.retry import ( + _GATEWAY_STATUSES, + RetryPolicy, + retry_sync, +) + +__all__ = ["query", "to_str"] + + +def to_str(listlike: object, delimiter: str = ",") -> str | None: + """Translate a list-like object into a delimited string. + + Parameters + ---------- + listlike: list-like object + A list, or a list-like object + (e.g. ``pandas.core.series.Series``). + delimiter: string, optional + String placed between entries of ``listlike`` when it is turned into a + string. Default value is a comma. + + Returns + ------- + listlike: string + The listlike object as a string separated by the delimiter. + + Examples + -------- + .. doctest:: + + >>> dataretrieval.utils.to_str([1, "a", 2]) + '1,a,2' + + >>> dataretrieval.utils.to_str([0, 10, 42], delimiter="+") + '0+10+42' + + """ + if isinstance(listlike, str): + return listlike + + if isinstance(listlike, Iterable): + return delimiter.join(map(str, listlike)) + + return None + + +_URL_TOO_LONG_EXAMPLE = """ + # n is the number of chunks to divide the query into \n + split_list = np.array_split(site_list, n) + data_list = [] # list to store chunk results in \n + # loop through chunks and make requests \n + for site_list in split_list: \n + data = nwis.get_record(sites=site_list, service='dv', \n + start=start, end=end) \n + data_list.append(data) # append results to list""" + + +def _url_too_long_error(detail: str) -> URLTooLong: + return URLTooLong( + "Request URL too long. Modify your query to use fewer sites. " + f"{detail}. Pseudo-code example of how to split your query: " + f"\n {_URL_TOO_LONG_EXAMPLE}" + ) + + +def _raise_for_status( + response: httpx.Response, + *, + detail_from: Callable[[httpx.Response], str | None] | None = None, +) -> None: + """Raise the typed :class:`DataRetrievalError` for an HTTP error response. + + A success status returns ``None``. Shared by the legacy :func:`query` path + (and ``streamstats`` / ``wateruse``). Delegates the status-to-type mapping to + :func:`dataretrieval.exceptions.error_for_status`, except a too-long-URL + status (413 / 414): that gets the same actionable "split your query" + remediation as the client-side over-long-URL case below, rather than a bare + ``HTTP 414`` (both still raise :class:`~dataretrieval.exceptions.URLTooLong`). + + ``detail_from``, when given, is called *only on an error response* to pull an + API-specific detail string (e.g. a JSON error envelope's message) out of the + body; a truthy result is appended to the raised message. This lets callers + surface their API's error wording without re-implementing the status-to-type + mapping and message format. + """ + status = response.status_code + if status < 400: + return + if status in (413, 414): + raise _url_too_long_error(f"API response reason: {response.reason_phrase}") + message = f"HTTP {status} {response.reason_phrase}".rstrip() + detail = detail_from(response) if detail_from is not None else None + if detail: + message += f": {detail}" + message += f" (URL: {response.url})" + raise error_for_status( + status, + message, + retry_after=parse_retry_after(response.headers.get("Retry-After")), + ) + + +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_with_retry( + url: str, + payload: dict[str, Any], + delimiter: str = ",", + ssl_check: bool = True, + *, + retry_policy: RetryPolicy | None = None, +) -> httpx.Response: + """Send an active-service query with bounded transient retry by default.""" + + for key, value in payload.items(): + payload[key] = to_str(value, delimiter) + # httpx serializes None params as ``foo=``; USGS rejects with 400. + # 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": USER_AGENT} + + 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. + if response.text.startswith("No sites/data"): + raise NoSitesError(response.url) + + return response + + +def query( + url: str, + payload: dict[str, Any], + delimiter: str = ",", + ssl_check: bool = True, +) -> httpx.Response: + """Send a query. + + Wrapper for ``httpx.get`` that handles errors, converts listed query + parameters to comma-separated strings, and returns the response. + + Parameters + ---------- + url: string + URL to query. + payload: dict + Query parameters passed to ``httpx.get``. + delimiter: string + Delimiter to use with lists. + ssl_check: bool + Whether to check SSL certificates. Default is True. + + Returns + ------- + response: ``httpx.Response`` + The response from the API query ``httpx.get`` function call. + + Raises + ------ + DataRetrievalError + On an HTTP error response, the typed subclass for the status (see + :func:`dataretrieval.exceptions.error_for_status` for the mapping); or + :class:`~dataretrieval.exceptions.NoSitesError` when a 200 response + reports no data matched; or :class:`~dataretrieval.exceptions.NetworkError` + on a connection-level failure (timeout, DNS), with the underlying + ``httpx`` exception on ``__cause__``. + """ + return _query_with_retry( + url, + payload, + delimiter, + ssl_check, + retry_policy=RetryPolicy(max_retries=0), + ) + + +# Preserve the documented function paths from the v1.2.0 utility API. +to_str.__module__ = "dataretrieval.utils" +query.__module__ = "dataretrieval.utils" diff --git a/dataretrieval/_wqx.py b/dataretrieval/_wqx.py new file mode 100644 index 000000000..8089c68d0 --- /dev/null +++ b/dataretrieval/_wqx.py @@ -0,0 +1,121 @@ +"""WQX3 / legacy-WQP CSV column conventions. + +The Samples database and the Water Quality Portal both split an instant across +three columns -- a date, a time, and a time-zone abbreviation -- and they spell +the trio two different ways. Recognizing either spelling and folding it into one +UTC column is knowledge about those response formats, so it lives in its own +leaf rather than in :mod:`dataretrieval.utils`, whose docstring reserves that +module for shaping that is not service-specific. + +Depends on pandas and the time-zone table only; nothing here issues a request. +""" + +from __future__ import annotations + +import pandas as pd + +from dataretrieval.codes import tz + +# (time-suffix, tz-suffix) pairs that follow a "Date" column. +_TIME_TZ_SUFFIXES = ( + # WQX3 / Samples, e.g. + # Activity_StartDate / Activity_StartTime / Activity_StartTimeZone + ("Time", "TimeZone"), + # Legacy WQP (slash-separated), e.g. + # ActivityStartDate / ActivityStartTime/Time / ActivityStartTime/TimeZoneCode + ("Time/Time", "Time/TimeZoneCode"), +) + + +def _build_utc_datetime( + date_series: pd.Series, time_series: pd.Series, tz_series: pd.Series +) -> pd.Series: + """Combine date + time + tz-abbreviation columns into a UTC pandas Series. + + Unknown timezone codes (and rows missing any of the three values) yield + ``NaT``. The input columns are not mutated. + """ + offsets = tz_series.map(tz) + combined = ( + date_series.astype("string") + + " " + + time_series.astype("string") + + " " + + offsets.astype("string") + ) + return pd.to_datetime( + combined, format="%Y-%m-%d %H:%M:%S %z", utc=True, errors="coerce" + ) + + +def _attach_datetime_columns(df: pd.DataFrame) -> pd.DataFrame: + """Append a UTC ``DateTime`` column per Date/Time/TimeZone triplet. + + Detects two naming patterns that appear in USGS Samples and Water Quality + Portal CSV responses: + + * **WQX3** — ``Date``, ``Time``, ``TimeZone`` + * **Legacy WQP** — ``Date``, ``Time/Time``, + ``Time/TimeZoneCode`` + + For every triplet present, a new ``DateTime`` column is appended + holding a UTC ``Timestamp`` (offsets resolved via + :data:`dataretrieval.codes.tz`). The original Date/Time/TimeZone columns + are left intact, and an existing ``DateTime`` column is never + overwritten. + + Rows are sorted (and the index reset) by the canonical activity-start + datetime when present — ``Activity_StartDateTime`` (WQX3) or + ``ActivityStartDateTime`` (legacy WQP) — falling back to the first + detected ``*Date`` column. Mirrors R ``dataRetrieval``'s + end-of-pipeline sort in ``importWQP.R``. + + Parameters + ---------- + df : ``pandas.DataFrame`` + DataFrame returned from a Samples or WQP CSV endpoint. + + Returns + ------- + df : ``pandas.DataFrame`` + A new DataFrame with derivable ``DateTime`` columns appended + and rows sorted by the activity-start datetime (if any date column + was detected). + """ + columns = set(df.columns) + new_columns = {} + first_date_col = None + for col in df.columns: + if not col.endswith("Date"): + continue + if first_date_col is None: + first_date_col = col + prefix = col.removesuffix("Date") + target = prefix + "DateTime" + if target in columns or target in new_columns: + continue + for time_suffix, tz_suffix in _TIME_TZ_SUFFIXES: + time_col = prefix + time_suffix + tz_col = prefix + tz_suffix + if time_col in columns and tz_col in columns: + new_columns[target] = _build_utc_datetime( + df[col], df[time_col], df[tz_col] + ) + break + if new_columns: + # Concat in one shot — per-column assignment on a wide CSV-derived + # frame triggers pandas' fragmentation PerformanceWarning. + df = pd.concat([df, pd.DataFrame(new_columns, index=df.index)], axis=1) + sort_key: str | None + if "Activity_StartDateTime" in df.columns: + sort_key = "Activity_StartDateTime" + elif "ActivityStartDateTime" in df.columns: + sort_key = "ActivityStartDateTime" + else: + sort_key = first_date_col + if sort_key is not None: + df = df.sort_values(by=sort_key, ignore_index=True) + return df + + +__all__ = ["_attach_datetime_columns", "_build_utc_datetime"] diff --git a/dataretrieval/exceptions.py b/dataretrieval/exceptions.py index 372171a09..99a77be3d 100644 --- a/dataretrieval/exceptions.py +++ b/dataretrieval/exceptions.py @@ -4,8 +4,9 @@ ``streamstats``) raises a subclass of :class:`DataRetrievalError` when a request fails, so one ``except dataretrieval.DataRetrievalError`` catches them all. That includes connection-level failures (timeouts, DNS, refused connections), which -are wrapped as :class:`NetworkError` with the underlying ``httpx`` exception on -``__cause__``. +remain inside this taxonomy rather than leaking ``httpx`` exceptions. A +deterministic failure is :class:`NetworkError`; a recoverable failure that +exhausts retries during fan-out is a resumable ``ServiceInterrupted``. Most failures are an :class:`HTTPError` carrying the response ``.status_code``, of which :class:`TransientError` (429 / 5xx) is the retryable subset. The rest @@ -21,6 +22,9 @@ from __future__ import annotations +import math +from datetime import datetime, timezone +from email.utils import parsedate_to_datetime from typing import TYPE_CHECKING, Any, ClassVar if TYPE_CHECKING: @@ -39,6 +43,7 @@ "NoSitesError", "ConfigurationError", "error_for_status", + "parse_retry_after", ] @@ -59,8 +64,9 @@ class DataRetrievalError(Exception): else: raise - Connection-level failures (timeouts, DNS) are wrapped as - :class:`NetworkError`, so this single clause covers them too. + Connection-level failures (timeouts, DNS) remain subclasses of this base: + :class:`NetworkError` when deterministic, or a resumable + ``ServiceInterrupted`` when recoverable fan-out retries are exhausted. """ #: HTTP status that triggered the error, or ``None`` for errors without one @@ -315,3 +321,50 @@ def error_for_status( if 500 <= status < 600: return ServiceUnavailable(message, status_code=status, retry_after=retry_after) return HTTPError(message, status_code=status) + + +def parse_retry_after(value: str | None) -> float | None: + """Parse a ``Retry-After`` header into seconds, or ``None`` for no usable hint. + + Both header forms mean the same thing and are treated the same way: the + seconds are returned as given, however large. A value past what a caller will + wait out inline stops the retry and surfaces a transient carrying the hint on + ``.retry_after``, so a long wait becomes the caller's decision (and, for a + chunked call, a resumable interruption) instead of being ignored. + + An over-long hint is honored rather than discarded. Dropping it would make + the client retry *harder* against a service that just asked for a long + pause, and would deny the caller the number it needs on ``.retry_after``. + Clock skew can inflate a date-form hint, but trusting one costs a + recoverable escalation while ignoring it costs hammering a service that is + already asking for room. + + A date that has *already* passed yields no hint at all rather than ``0.0``. + Read literally it says "retry now", but the likelier reading is that our + clock runs ahead of the server's -- and acting on it would re-send almost + immediately against a service that just asked for a pause. Falling back to + our own bounded backoff is right under either reading. (Delta-seconds is + clock-independent, so a literal ``Retry-After: 0`` is still honored as the + instruction it is, floored by + :meth:`~dataretrieval.transport.retry.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 diff --git a/dataretrieval/interruptions.py b/dataretrieval/interruptions.py new file mode 100644 index 000000000..2d408ab70 --- /dev/null +++ b/dataretrieval/interruptions.py @@ -0,0 +1,330 @@ +"""Resumable fan-out interruption exceptions — the public resume contract. + +When a fanned-out request fails mid-stream (a 429, a 5xx, or a bare transport +error), the work already completed is preserved and the call is resumable: the +raised exception carries a ``.call`` handle whose ``resume()`` re-issues only +the still-pending sub-requests. These exception types are that contract, +re-exported at the top level (``from dataretrieval import ChunkInterrupted``). +The execution machinery that raises and resumes them is +:class:`dataretrieval.transport.fanout.FanOut`. + +Vocabulary, consistently: a **fan-out** is one logical query the service forces +into several requests; a **sub-request** is one unit of a fan-out; a **chunk** +is specifically a *byte-driven* slice, which is OGC planning vocabulary and +belongs to :class:`~dataretrieval.ogc.planning.ChunkPlan`. Water Use fans out +without chunking anything — the NWDC simply accepts one location per request — +so the base class is :class:`FanOutInterrupted`. + +``ChunkInterrupted`` is retained as an alias of that same class, not a +deprecated shim to delete later: it is the name published in the user guide and +caught in user code, and aliasing costs nothing to keep. ``except +ChunkInterrupted`` and ``except FanOutInterrupted`` are the same handler. + +This is a top-level leaf rather than a member of ``ogc`` or ``transport``, +for the reason ADR 0006 gives for ``combining``, ``progress``, and +``credentials``: adapters need it whether or not they go through transport, and +an exception taxonomy is not HTTP execution policy. It stays out of +:mod:`dataretrieval.exceptions` because it carries pandas/httpx state, which +would pull heavy dependencies into that lightweight leaf. +""" + +from __future__ import annotations + +import socket +from collections.abc import Iterator +from typing import TYPE_CHECKING, Any, ClassVar + +import httpx +import pandas as pd + +from dataretrieval.exceptions import DataRetrievalError, RateLimited, TransientError + +if TYPE_CHECKING: + from dataretrieval.transport.fanout import FanOut + + +class FanOutInterrupted(DataRetrievalError): + """ + Base class for mid-stream sub-request failures whose completed work + is preserved and resumable. + + A ``FanOutInterrupted`` subclass means: a sub-request failed, but + ``FanOut`` still owns whatever completed successfully before + the failure. Call ``self.call.resume()`` to pick up where the + failure stopped you — only still-pending sub-requests are + re-issued. + + Subclasses describe *why* ``FanOut`` stopped so callers can + pick a retry policy: :class:`QuotaExhausted` for 429 (wait for the + rate-limit window), :class:`ServiceInterrupted` for 5xx (wait for + the upstream to recover). The ``.call`` handle is the same object + across every interruption of a single fanned-out call — frames + accumulate across retries. + + Attributes + ---------- + call : FanOut or None + Resumable handle into the ``FanOut`` that raised this + exception. ``None`` only on hand-constructed exceptions (test + fixtures), where ``.call``-derived accessors degrade to + empty/``None``. + retry_after : float or None + Seconds the server suggested waiting (``Retry-After`` header). + ``None`` when the server gave no hint. + completed_chunks : int + Number of sub-requests successfully completed before the failure. + total_chunks : int + Total sub-requests in the plan. + partial_frame : pandas.DataFrame + Combined frame of work completed by the moment this exception + was raised. Snapshot at raise time — does NOT advance on a + later ``call.resume()`` (use ``exc.call.partial_frame`` for + the live view). + partial_response : httpx.Response or None + Raw aggregate response covering the completed sub-requests at + raise time; ``None`` if nothing had completed yet. Same snapshot + semantics as ``partial_frame``. (Raw, not finalized — use + ``exc.call.resume()`` for the finalized ``(df, metadata)`` result.) + + Examples + -------- + Retry on any transient interruption, honoring the server's + ``Retry-After`` hint when present and falling back to a fixed wait + otherwise. Each new interruption keeps the already-completed work + intact — only the still-pending sub-requests are re-issued. + + .. code-block:: python + + import time + from dataretrieval import ChunkInterrupted + + # ``getter`` is any chunked OGC getter — e.g. + # ``waterdata.get_daily`` or ``ngwmn.get_water_level``. + try: + df, md = getter(monitoring_location_id=long_list_of_sites) + except ChunkInterrupted as exc: + while True: + time.sleep(exc.retry_after or 5 * 60) + try: + df, md = exc.call.resume() + break + except ChunkInterrupted as next_exc: + exc = next_exc + """ + + # Subclasses override with a ``str.format`` template; the format + # call sees ``completed_chunks`` and ``total_chunks`` as kwargs. + _MESSAGE_TEMPLATE: ClassVar[str] = ( + "Fan-out interrupted after {completed_chunks}/" + "{total_chunks} sub-requests; call .call.resume() to continue." + ) + retryable: ClassVar[bool] = True + + def __init__( + self, + *, + completed_chunks: int, + total_chunks: int, + call: FanOut[Any] | None = None, + retry_after: float | None = None, + cause: BaseException | None = None, + ) -> None: + message = self._MESSAGE_TEMPLATE.format( + completed_chunks=completed_chunks, total_chunks=total_chunks + ) + if cause is not None: + cause_msg = str(cause) or type(cause).__name__ + message = f"{message} Cause: {type(cause).__name__}: {cause_msg}" + super().__init__(message) + self.completed_chunks = completed_chunks + self.total_chunks = total_chunks + self.call = call + self.retry_after = retry_after + self.status_code = getattr(type(self), "_DEFAULT_STATUS", None) + if self.status_code is None and cause is not None: + # The status is usually a few frames down: a typed error raised + # ``from`` the httpx failure that carried it. + for current in _walk_causes(cause): + status = getattr(current, "status_code", None) + if status is not None: + self.status_code = status + break + # Snapshot partial state at raise time so the exception stays a stable + # record of the failure moment: ``exc.partial_frame`` / + # ``.partial_response`` do NOT advance on a later ``call.resume()`` + # (that live view is on ``call.partial_frame`` / ``.partial_response``). + # This keeps each interruption in a resume loop a faithful record of + # what it saw, rather than every exception aliasing the shared call's + # advancing state. ``.copy()`` guards the single-chunk fast path, where + # the combined frame may be returned verbatim. + if call is None: + self.partial_frame: pd.DataFrame = pd.DataFrame() + self.partial_response: httpx.Response | None = None + else: + self.partial_frame = call.partial_frame.copy() + self.partial_response = call.partial_response + + def __getstate__(self) -> dict[str, Any]: + # Drop the live FanOut before pickling: its ``.fetch`` is an + # undecorated module function pickle can't reference by name, so the + # interruption can't cross a process boundary with ``.call`` attached. + # The degraded ``call=None`` form keeps the counts, retry hint, and the + # snapshotted partial frame / response — plain instance attributes the + # base ``__getstate__`` already pickles; only ``.resume()`` is lost + # (cross-process resume was never possible anyway). + return {**super().__getstate__(), "call": None} + + +class QuotaExhausted(FanOutInterrupted): + """ + A sub-request returned HTTP 429 — the per-key rate-limit window + is exhausted. Subclass of :class:`FanOutInterrupted`. + + The completed sub-requests are preserved on ``.call``; once the + rate-limit window resets, ``.call.resume()`` re-issues only the + still-pending work. ``partial_frame`` holds what completed + before the 429. + """ + + _MESSAGE_TEMPLATE = ( + "HTTP 429 after {completed_chunks}/{total_chunks} sub-requests; " + "catch QuotaExhausted (or FanOutInterrupted) to access " + ".partial_frame or .call.resume() once the rate-limit " + "window has rolled over." + ) + _DEFAULT_STATUS = 429 + + +class ServiceInterrupted(FanOutInterrupted): + """ + A sub-request returned HTTP 5xx — the upstream service failed + transiently. Subclass of :class:`FanOutInterrupted`. + + The completed sub-requests are preserved on ``.call``; once the + upstream recovers, ``.call.resume()`` resumes only the + still-pending work. + """ + + _MESSAGE_TEMPLATE = ( + "Service error after {completed_chunks}/{total_chunks} " + "sub-requests; catch ServiceInterrupted (or FanOutInterrupted) " + "and call .call.resume() once the upstream service recovers." + ) + + +# Resolver failures that will not resolve differently on a later attempt. The +# temporary ones (notably EAI_AGAIN -- "try again", raised while a resolver is +# still coming up, on VPN reconnect, or after a laptop wakes) are deliberately +# absent: those are worth another try. Looked up defensively because the EAI_* +# constants are platform-dependent; an unrecognized code stays retryable, since +# spending a few seconds on a retry is cheaper than dropping a recoverable call. +_PERMANENT_DNS_ERRORS = frozenset( + code + for code in ( + getattr(socket, name, None) for name in ("EAI_NONAME", "EAI_FAIL", "EAI_NODATA") + ) + if code is not None +) + + +def _walk_causes( + exc: BaseException, *, follow_context: bool = False +) -> Iterator[BaseException]: + """Yield ``exc`` and the exceptions it chains to, each at most once. + + Every question this module asks about a failure -- is it transient, is it + deterministic, what status did it carry -- is "find the first exception in + this chain that satisfies P". One traversal answers all of them, so the + cycle guard and the choice of links cannot drift between callers. + + ``__cause__`` (explicit ``raise ... from``) is always followed. + ``__context__`` (implicit chaining, from raising inside an ``except`` + block) is followed only when ``follow_context`` is set, because it can + lead away from the failure being classified into whatever unrelated error + happened to be in flight. + + The ``seen`` set keeps a chain that rejoins itself, or points back at an + ancestor, from looping. + """ + seen: set[int] = set() + pending: list[BaseException | None] = [exc] + while pending: + current = pending.pop() + if current is None or id(current) in seen: + continue + seen.add(id(current)) + yield current + if follow_context: + pending += [current.__cause__, current.__context__] + else: + pending.append(current.__cause__) + + +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`). + + Walks ``__context__`` as well as ``__cause__``, because 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 implicit chaining. Following only the cause would walk off down + the explicit branch and miss a ``gaierror`` sitting on the implicit one -- + spending the whole retry budget on a hostname that will never resolve. + """ + for current in _walk_causes(exc, follow_context=True): + if isinstance(current, (httpx.UnsupportedProtocol, httpx.LocalProtocolError)): + return True + if isinstance(current, socket.gaierror): + # Return, not continue: the first resolver code found settles the chain. + return current.errno in _PERMANENT_DNS_ERRORS + return False + + +def _classify_transient( + exc: BaseException, +) -> tuple[type[FanOutInterrupted], float | None] | None: + """Classify one failure as a resumable interruption.""" + if isinstance(exc, RateLimited): + return QuotaExhausted, exc.retry_after + if isinstance(exc, TransientError): + return ServiceInterrupted, exc.retry_after + if isinstance(exc, (httpx.HTTPError, httpx.InvalidURL)): + # Some failures will fail the same way every time -- a bad scheme, a + # hostname that doesn't resolve. Offering to resume one would just + # hide the real error behind a retry that can never work. + if _deterministic_failure(exc): + return None + return ServiceInterrupted, None + return None + + +def _classify_chunk_error( + exc: BaseException, +) -> tuple[type[FanOutInterrupted], float | None] | None: + """Walk a wrapped pagination failure for a resumable transport cause.""" + return next( + ( + result + for current in _walk_causes(exc) + if (result := _classify_transient(current)) is not None + ), + None, + ) + + +#: The name this taxonomy was published under, kept as a permanent alias so +#: ``except ChunkInterrupted`` keeps working. Same class object, not a subclass. +ChunkInterrupted = FanOutInterrupted + +__all__ = [ + "ChunkInterrupted", + "FanOutInterrupted", + "QuotaExhausted", + "ServiceInterrupted", +] diff --git a/dataretrieval/ngwmn.py b/dataretrieval/ngwmn.py index a71d2b47b..154899faf 100644 --- a/dataretrieval/ngwmn.py +++ b/dataretrieval/ngwmn.py @@ -18,15 +18,17 @@ from __future__ import annotations from collections.abc import Iterable -from typing import Any +from typing import TYPE_CHECKING, Any import pandas as pd -from dataretrieval._response_metadata import BaseMetadata from dataretrieval.codes.states import apply_state from dataretrieval.credentials import WATERDATA_BASE_URL from dataretrieval.ogc import OgcDialect, get_ogc_data, prepare_request_args +if TYPE_CHECKING: + from dataretrieval._response_metadata import BaseMetadata + __all__ = [ "get_sites", "get_water_level", diff --git a/dataretrieval/nldi.py b/dataretrieval/nldi.py index b211bfefc..c32efd10e 100644 --- a/dataretrieval/nldi.py +++ b/dataretrieval/nldi.py @@ -13,7 +13,7 @@ from json import JSONDecodeError from typing import Any, Literal, cast -from dataretrieval.utils import _query_with_retry +from dataretrieval._querying import _query_with_retry __all__ = [ "get_flowlines", @@ -68,6 +68,14 @@ def _features_to_gdf(feature_collection: dict[str, Any]) -> gpd.GeoDataFrame: return gpd.GeoDataFrame.from_features(feature_collection, crs=_CRS) +def _query_features( + url: str, query_params: dict[str, str], as_json: bool +) -> gpd.GeoDataFrame | dict[str, Any]: + """Run an NLDI query and return the raw FeatureCollection or a GeoDataFrame.""" + feature_collection = cast("dict[str, Any]", _query_nldi(url, query_params)) + return feature_collection if as_json else _features_to_gdf(feature_collection) + + def get_flowlines( navigation_mode: str, distance: int = 5, @@ -122,20 +130,19 @@ def get_flowlines( _validate_feature_source_comid(feature_source, feature_id, comid) if feature_source: _validate_data_source(feature_source) - url = f"{NLDI_API_BASE_URL}/{feature_source}/{feature_id}/navigation" - else: - url = f"{NLDI_API_BASE_URL}/comid/{comid}/navigation" - query_params = {"distance": str(distance), "trimStart": str(trim_start).lower()} - - url += f"/{navigation_mode}/flowlines" + url, query_params = _navigation_request( + feature_source=feature_source, + feature_id=feature_id, + comid=comid, + navigation_mode=navigation_mode, + distance=distance, + tail="flowlines", + ) + query_params["trimStart"] = str(trim_start).lower() if stop_comid is not None: query_params["stopComid"] = str(stop_comid) - feature_collection = cast("dict[str, Any]", _query_nldi(url, query_params)) - if as_json: - return feature_collection - gdf = _features_to_gdf(feature_collection) - return gdf + return _query_features(url, query_params, as_json) def get_basin( @@ -185,11 +192,7 @@ def get_basin( "simplified": simplified_str, "splitCatchment": split_catchment_str, } - feature_collection = cast("dict[str, Any]", _query_nldi(url, query_params)) - if as_json: - return feature_collection - gdf = _features_to_gdf(feature_collection) - return gdf + return _query_features(url, query_params, as_json) def get_features( @@ -272,11 +275,28 @@ def get_features( stop_comid=stop_comid, ) - feature_collection = cast("dict[str, Any]", _query_nldi(url, query_params)) - if as_json: - return feature_collection - gdf = _features_to_gdf(feature_collection) - return gdf + return _query_features(url, query_params, as_json) + + +def _navigation_request( + *, + feature_source: str | None, + feature_id: str | None, + comid: int | None, + navigation_mode: str, + distance: int, + tail: str, +) -> tuple[str, dict[str, str]]: + """URL and query params for an NLDI navigation from a validated origin. + + The single home for the navigation path grammar — ``{origin}/navigation/ + {mode}/{tail}`` — and its ``distance`` knob. Callers add the knobs specific + to their endpoint (``trimStart``, ``stopComid``) afterwards, so the query + string keeps its documented parameter order. + """ + origin = f"{feature_source}/{feature_id}" if feature_source else f"comid/{comid}" + url = f"{NLDI_API_BASE_URL}/{origin}/navigation/{navigation_mode}/{tail}" + return url, {"distance": str(distance)} def _get_features_request( @@ -323,9 +343,14 @@ def _get_features_request( return f"{NLDI_API_BASE_URL}/{feature_source}/{feature_id}", {} navigation_mode = _validate_navigation_mode(navigation_mode) - origin = f"{feature_source}/{feature_id}" if feature_source else f"comid/{comid}" - url = f"{NLDI_API_BASE_URL}/{origin}/navigation/{navigation_mode}/{data_source}" - query_params = {"distance": str(distance)} + url, query_params = _navigation_request( + feature_source=feature_source, + feature_id=feature_id, + comid=comid, + navigation_mode=navigation_mode, + distance=distance, + tail=f"{data_source}", + ) if stop_comid is not None: query_params["stopComid"] = str(stop_comid) return url, query_params @@ -447,7 +472,7 @@ def search( if (lat is None) != (long is None): raise ValueError("Both lat and long are required") - find = cast(Literal["basin", "flowlines", "features"], find.lower()) + find = cast("Literal['basin', 'flowlines', 'features']", find.lower()) if find not in ("basin", "flowlines", "features"): raise ValueError( f"Invalid value for find: {find} - allowed values are:" diff --git a/dataretrieval/nwis.py b/dataretrieval/nwis.py index 82121c9f2..e8354d719 100644 --- a/dataretrieval/nwis.py +++ b/dataretrieval/nwis.py @@ -19,7 +19,7 @@ from dataretrieval._response_metadata import BaseMetadata from dataretrieval.rdb import read_rdb -from .utils import query +from ._querying import query try: import geopandas as gpd @@ -103,7 +103,7 @@ def wrapper(*args: Any, **kwargs: Any) -> Any: finally: _deprecation_state.active = False - return cast(F, wrapper) + return cast("F", wrapper) def _parse_json_or_raise(response: httpx.Response) -> pd.DataFrame: @@ -460,6 +460,37 @@ def query_waterservices( return query(url, payload=kwargs, ssl_check=ssl_check) +def _get_json_values( + service: str, + sites: list[str] | str | None, + start: str | None, + end: str | None, + multi_index: bool, + ssl_check: bool, + kwargs: dict[str, Any], +) -> tuple[pd.DataFrame, NWIS_Metadata]: + """Shared body of the JSON waterservices time-series getters (dv / iv). + + The caller-facing ``sites`` / ``start`` / ``end`` arguments are aliases: an + explicit waterservices keyword of the same meaning wins over them. Note that + ``multi_index`` travels through ``kwargs`` so that :func:`format_response` + sees it. + """ + _check_sites_value_types(sites) + + kwargs["startDT"] = kwargs.pop("startDT", start) + kwargs["endDT"] = kwargs.pop("endDT", end) + kwargs["sites"] = kwargs.pop("sites", sites) + kwargs["multi_index"] = multi_index + + response = query_waterservices( + service, format="json", ssl_check=ssl_check, **kwargs + ) + df = _parse_json_or_raise(response) + + return format_response(df, **kwargs), NWIS_Metadata(response, **kwargs) + + @_deprecated def get_dv( sites: list[str] | str | None = None, @@ -517,17 +548,7 @@ def get_dv( >>> df, md = dataretrieval.nwis.get_dv(sites="01646500") """ - _check_sites_value_types(sites) - - kwargs["startDT"] = kwargs.pop("startDT", start) - kwargs["endDT"] = kwargs.pop("endDT", end) - kwargs["sites"] = kwargs.pop("sites", sites) - kwargs["multi_index"] = multi_index - - response = query_waterservices("dv", format="json", ssl_check=ssl_check, **kwargs) - df = _parse_json_or_raise(response) - - return format_response(df, **kwargs), NWIS_Metadata(response, **kwargs) + return _get_json_values("dv", sites, start, end, multi_index, ssl_check, kwargs) @_deprecated @@ -701,19 +722,7 @@ def get_iv( ... ) """ - _check_sites_value_types(sites) - - kwargs["startDT"] = kwargs.pop("startDT", start) - kwargs["endDT"] = kwargs.pop("endDT", end) - kwargs["sites"] = kwargs.pop("sites", sites) - kwargs["multi_index"] = multi_index - - response = query_waterservices( - service="iv", format="json", ssl_check=ssl_check, **kwargs - ) - - df = _parse_json_or_raise(response) - return format_response(df, **kwargs), NWIS_Metadata(response, **kwargs) + return _get_json_values("iv", sites, start, end, multi_index, ssl_check, kwargs) def get_pmcodes(**kwargs: Any) -> NoReturn: diff --git a/dataretrieval/ogc/chunking.py b/dataretrieval/ogc/chunking.py index 671deac9d..7aad91653 100644 --- a/dataretrieval/ogc/chunking.py +++ b/dataretrieval/ogc/chunking.py @@ -1,13 +1,21 @@ -"""Joint URL-byte chunking for the OGC getters. +"""URL-byte chunk planning and dispatch for the OGC getters. An OGC query has several chunkable axes: every multi-value list parameter (sites, parameter codes, …) plus the cql-text ``filter``, which splits along its top-level OR clauses. Any of them can fan the URL past the server's ~8 KB byte limit. ``ChunkPlan`` picks a fan-out for each axis that minimizes total sub-requests while keeping every -sub-request URL under the budget; ``ChunkedCall`` fetches the resulting -cartesian product of chunks. Requests that already fit get a trivial -single-step plan — ``ChunkedCall`` has one code path either way. +sub-request URL under the budget. Requests that already fit get a +trivial single-step plan — the executor has one code path either way. + +This module owns the OGC-specific half: the byte budget, the +``parallel_chunks`` dial, and the ``multi_value_chunked`` decorator that +ties a plan to a fetcher. Driving the resulting sub-requests to +completion — bounded concurrency, retry, failure precedence, resume — is +API-neutral and belongs to +:class:`dataretrieval.transport.fanout.FanOut`, which this module hands +its plan to. :class:`~dataretrieval.ogc.planning.ChunkPlan` satisfies +:class:`~dataretrieval.transport.fanout.FanOutPlan` structurally. Parallel chunks: the planner is conservative by default — it splits only as far as the byte limit forces. A caller who knows their result is large can opt @@ -15,50 +23,9 @@ the query out into ``n`` parallel sub-requests. ``n`` drives :meth:`ChunkPlan._refine`; see ``parallel_chunks`` for the why and the when. -This module owns the *execution* half — the event loop and bounded -concurrency that drive a plan to completion (``ChunkedCall``) plus the -public ``multi_value_chunked`` decorator. The neighboring concerns remain -separate: :mod:`~dataretrieval.ogc.planning` builds the -:class:`~dataretrieval.ogc.planning.ChunkPlan`; -:mod:`~dataretrieval.combining` assembles results; -:mod:`~dataretrieval.transport.retry` owns bounded retry policy; and -:mod:`~dataretrieval.ogc.interruptions` defines the resumable -:class:`~dataretrieval.ogc.interruptions.ChunkInterrupted` contract. - -Concurrency: ``multi_value_chunked`` fans every pending sub-request out -under one ``asyncio.gather`` sharing a single ``httpx.AsyncClient``. An -``asyncio.Semaphore`` — not the client's connection pool, which is -merely sized to match — caps the sub-requests in flight at ``N``. See -:meth:`ChunkedCall._run` for why the gate must be the semaphore rather -than the pool. ``API_USGS_CONCURRENT`` resolves ``N``: an integer N > 1 -allows N sub-requests in flight; ``1`` forces sequential dispatch (one -request at a time); the literal ``unbounded`` lifts the cap. ``N`` -bounds only how many of a chunked query's sub-requests are in flight at -once — a client-side trade-off between open connections and fan-out -latency. It does not affect the API rate limit: a chunked call issues -the same number of sub-requests regardless of ``N``, so ``N`` changes -their timing, not the total request volume. The USGS API rate-limits by -volume over time (HTTP 429), not by simultaneity; set ``API_USGS_PAT`` -to raise that quota. The default of 32 is a conservative cap that keeps -connection use modest. The fan-out runs in a short-lived worker thread -(an ``anyio`` blocking portal), so it works whether or not the caller is -already inside an event loop (Jupyter / IPython / async apps). - -Retries: each sub-request is retried on a transient failure (429, -5xx, connect/read timeout) with exponential backoff + full jitter, -honoring a server ``Retry-After`` when present. ``API_USGS_RETRIES`` -sets the cap (default 4; ``0`` disables). A ``Retry-After`` longer -than the per-call ceiling escalates to a resumable interruption. - -Interruption: any mid-stream transient failure — 429, 5xx, or a bare -transport error (connect/read timeout, oversize follow-up URL) — surfaces -as a ``ChunkInterrupted`` subclass: ``QuotaExhausted`` for 429, -``ServiceInterrupted`` for the rest. The exception carries ``.call``, a -``ChunkedCall`` handle that owns the already-completed sub-request -state (sparse-indexed, since gathered sub-requests complete out of -order). Call ``.call.resume()`` once the underlying condition clears; -only the still-pending sub-requests are re-issued. ``Retry-After`` (when -the server sets it) is surfaced on the exception as ``.retry_after``. +Concurrency, retries, and interruption semantics are documented on +:mod:`dataretrieval.transport.fanout`; ``API_USGS_CONCURRENT`` and +``API_USGS_RETRIES`` are read there. Dedup: list-axis chunks don't overlap; filter-axis chunks can, so ``_combine_chunk_frames`` dedupes by feature ``id``. ``properties``, @@ -69,33 +36,39 @@ from __future__ import annotations -import asyncio import functools -import os -from collections.abc import Awaitable, Callable, Iterator +from collections.abc import Callable, Iterator from contextlib import contextmanager -from contextvars import copy_context -from typing import Any, cast +from typing import Any import httpx import pandas as pd -from anyio.from_thread import start_blocking_portal -from dataretrieval import progress as _progress from dataretrieval._ambient import Ambient -from dataretrieval.combining import ( - _combine_chunk_frames, - _combine_chunk_responses, +from dataretrieval.transport.fanout import ( + FanOut, + _active_client, + _Fetch, + _Finalize, + _passthrough_result, + active_client, ) -from dataretrieval.exceptions import ConfigurationError -from dataretrieval.transport.http import open_async_client -from dataretrieval.transport.retry import _NO_RETRY, RetryPolicy -from dataretrieval.transport.retry import retry_async as _retry +from dataretrieval.transport.retry import RetryPolicy -from .interruptions import ChunkInterrupted from .planning import ChunkPlan from .policy import _require_positive_int -from .retry import _classify_chunk_error + +# Compatibility aliases. ``ChunkedCall`` was this module's executor before it +# moved down to transport as the API-neutral ``FanOut``; ``get_active_client`` +# and ``_chunked_client`` named its shared per-call client. Only the +# chunking/progress test modules still use these names, and the rename is not +# worth churning them over -- package code imports the canonical spellings from +# :mod:`dataretrieval.transport.fanout`. They are aliases, not copies: the +# ambient in particular must be the *same* object transport publishes, or a +# test reading it here would never see the running client. +ChunkedCall = FanOut +get_active_client = active_client +_chunked_client = _active_client # Empirically the API replies HTTP 414 above ~8200 bytes of full URL — # matches nginx's default ``large_client_header_buffers`` of 8 KB. 8000 @@ -105,73 +78,6 @@ _OGC_URL_BYTE_LIMIT = 8000 -# Fan-out concurrency cap, read at call time (not import) so test -# ``monkeypatch.setenv`` applies. Value grammar in :func:`_read_concurrency_env`; -# the concurrency model is in the module docstring. -_CONCURRENCY_ENV = "API_USGS_CONCURRENT" -_CONCURRENCY_DEFAULT = 32 -_CONCURRENCY_UNBOUNDED = "unbounded" - - -def _read_concurrency_env() -> int | None: - """ - Resolve the ``API_USGS_CONCURRENT`` env var to a parallelism cap. - - Returns - ------- - int or None - ``1`` for sequential dispatch (one sub-request at a time); an - integer >1 for bounded concurrency; ``None`` to disable the - per-call cap entirely (``unbounded`` keyword). Unset → default - of ``_CONCURRENCY_DEFAULT``. - """ - raw = os.environ.get(_CONCURRENCY_ENV) - if raw is None: - return _CONCURRENCY_DEFAULT - raw = raw.strip() - if raw == "": - return _CONCURRENCY_DEFAULT - if raw.lower() == _CONCURRENCY_UNBOUNDED: - return None - try: - value = int(raw) - except ValueError as exc: - raise ConfigurationError( - f"{_CONCURRENCY_ENV} must be a positive integer or " - f"'{_CONCURRENCY_UNBOUNDED}'; got {raw!r}." - ) from exc - if value < 1: - raise ConfigurationError( - f"{_CONCURRENCY_ENV} must be >= 1 (got {value}); use " - f"'{_CONCURRENCY_UNBOUNDED}' to disable the cap." - ) - return value - - -# Shared per-call ``httpx.AsyncClient``, scoped via ``with _chunked_client(c):`` -# during ``ChunkedCall._run`` so paginated-loop helpers (``_walk_pages``) reuse -# the same connection pool across every sub-request. ``None`` outside a chunked -# call — paginated helpers then open their own short-lived client. -_chunked_client: Ambient[httpx.AsyncClient | None] = Ambient("_chunked_client", None) - - -def get_active_client() -> httpx.AsyncClient | None: - """ - Return the chunker's currently-published client, or ``None``. - - Used by the paginated-loop helpers (e.g. - :func:`dataretrieval.transport.pagination._client_for`) to reuse the - per-call connection pool. - - Returns - ------- - httpx.AsyncClient or None - The client scoped via ``with _chunked_client(...)`` if currently inside - a :class:`ChunkedCall` run; ``None`` otherwise. - """ - return _chunked_client.get() - - # Parallel-chunks dial: opt-in to fan a query out *more finely* than the byte # limit alone requires. Scoped to a ``with parallel_chunks(...):`` block (a # ContextVar), deliberately NOT an env var (see :func:`parallel_chunks` for @@ -284,428 +190,11 @@ def parallel_chunks(n: int) -> Iterator[None]: yield -# --------------------------------------------------------------------------- -# Type aliases for the ChunkedCall contract. -# --------------------------------------------------------------------------- - -# The per-sub-request fetcher the decorator wraps and ``ChunkedCall`` drives: -# an ``async def fetch(args) -> (df, response)``. -_Fetch = Callable[[dict[str, Any]], Awaitable[tuple[pd.DataFrame, httpx.Response]]] - -# Caller-supplied transform applied to the combined chunk result, so a -# resumed call returns the same shape as an un-interrupted one rather than -# the chunker's raw ``(frame, httpx.Response)``. This keeps the chunker -# generic: the OGC getters inject their post-processing (type coercion, -# column arrangement, ``BaseMetadata``) through ``_finalize_ogc``. -# The default is identity, so direct ``ChunkedCall`` use is unaffected. -_Finalize = Callable[[pd.DataFrame, httpx.Response], tuple[pd.DataFrame, Any]] - - -def _passthrough_result( - frame: pd.DataFrame, response: httpx.Response -) -> tuple[pd.DataFrame, Any]: - """Default :data:`_Finalize`: return the raw combined pair unchanged.""" - return frame, response - - -class ChunkedCall: - """ - Stateful handle for a chunked call. - - Holds the in-flight state (per-sub-request frames and responses) - and the async fetcher. A single :meth:`resume` entry point drives - the call from wherever it is to completion — used both for the - first invocation (from :func:`multi_value_chunked`) and for subsequent - retries after a :class:`ChunkInterrupted`. - - :meth:`_run` gathers every pending sub-request over one shared - :class:`httpx.AsyncClient`, applies the failure-precedence rules, and - combines. :meth:`resume` drives :meth:`_run` through an ``anyio`` - blocking portal, so it works whether or not the caller is already - inside an event loop. Concurrency is bounded by a per-run - ``asyncio.Semaphore`` (see :meth:`_run`), so sequential dispatch - (``API_USGS_CONCURRENT=1``) is just a degenerate gather. - - A ``ChunkedCall`` is created internally when a :class:`ChunkPlan` - executes; callers reach it via :attr:`ChunkInterrupted.call` on - the exception raised by a mid-stream failure. - - :meth:`resume` is idempotent: :meth:`_run` iterates - :meth:`ChunkPlan.iter_sub_args` (deterministic order) and skips - any index whose result is already in ``self._chunks``. The - completion set is a sparse ``dict[int, (df, response)]`` so the - gather can record scattered completions (e.g. indices [0, 2, 5] - after siblings [1, 3, 4] failed) and a subsequent ``resume`` only - re-issues the missing indices. - - Parameters - ---------- - plan : ChunkPlan - The chunking plan to execute. - fetch : Callable - ``async def`` that issues a single sub-request, given the - substituted args dict, and returns ``(frame, response)``. - - Attributes - ---------- - plan : ChunkPlan - The plan being driven (read-only after construction). - fetch : Callable - The async per-sub-request fetch function. - finalize : Callable - Transform applied to the combined result (see :data:`_Finalize`) at - the terminal :meth:`_run` return, so a completed call yields the - caller's finished shape. The ``partial_*`` accessors deliberately - skip it and stay raw. - partial_frame : pandas.DataFrame - Raw combined frame of completed sub-requests (live; recomputed per - access). Not finalized — call :meth:`resume` for the finished shape. - partial_response : httpx.Response or None - Raw aggregate response (canonical URL restored), or ``None`` when - nothing has completed yet (live; recomputed per access). - """ - - def __init__( - self, - plan: ChunkPlan, - fetch: _Fetch, - retry_policy: RetryPolicy = _NO_RETRY, - finalize: _Finalize = _passthrough_result, - ) -> None: - self.plan = plan - self.fetch = fetch - self.retry_policy = retry_policy - self.finalize = finalize - # Snapshot the ambient context at construction time — i.e. inside the - # caller's ``with`` blocks (base URL, dialect, row cap, progress - # reporter). :meth:`resume` runs every drive inside this snapshot. So a - # *later* ``exc.call.resume()`` still rebuilds sub-requests against the - # original API's base URL/dialect rather than the process defaults, even - # though it fires after those ``with`` blocks have exited and reset - # their ContextVars. ``build_request`` reads those ContextVars when it - # reconstructs each sub-request, so the snapshot must outlive them. - self._ctx = copy_context() - # Completed (frame, response) pairs keyed by sub-args index; sparse - # (gathered sub-requests complete out of order — see class docstring). - # ``_run``'s ``track`` closure is the only writer, so ``dict`` insertion - # order is completion order (relied on by :meth:`_combine_raw`). - self._chunks: dict[int, tuple[pd.DataFrame, httpx.Response]] = {} - - def wrap_failure(self, exc: BaseException) -> ChunkInterrupted | None: - """ - Wrap ``exc`` as the matching :class:`ChunkInterrupted` carrying this call. - - Returns ``None`` when ``exc`` is not a recognized transient transport - failure, so the caller can re-raise it. Encapsulates the - ``classify → instantiate-with-call-state`` recipe so - :class:`ChunkedCall`'s private fields stay private. - - Parameters - ---------- - exc : BaseException - The exception raised by a sub-request. - - Returns - ------- - ChunkInterrupted or None - The matching :class:`ChunkInterrupted` subclass carrying this - call for a recognized transient failure; ``None`` otherwise. - """ - classification = _classify_chunk_error(exc) - if classification is None: - return None - interrupted_class, retry_after = classification - return interrupted_class( - completed_chunks=self.completed_chunks, - total_chunks=self.plan.total, - call=self, - retry_after=retry_after, - cause=exc, - ) - - @property - def completed_chunks(self) -> int: - """Number of sub-requests completed so far.""" - return len(self._chunks) - - def _combine_raw(self) -> tuple[pd.DataFrame, httpx.Response]: - """Assemble the raw ``(frame, response)`` from completed sub-requests. - - Runs before :attr:`finalize`. Frames concatenate in sub-args *index* - order (``sorted`` keys — deterministic, independent of parallel - completion order). The aggregated response takes its headers from the - response with the lowest reported ``x-ratelimit-remaining`` value. If no - response reports that header, the aggregate falls back to the last - completed response; ``self._chunks`` preserves completion order because - the ``track`` closure in :meth:`_run` is its only writer. - - Returns - ------- - tuple of (pandas.DataFrame, httpx.Response) - The concatenated frame and the aggregated response, before - :attr:`finalize` is applied. - """ - return self._combine_frames(), self._combine_responses() - - def _combine_frames(self) -> pd.DataFrame: - """Combine completed frames in deterministic sub-request order.""" - return _combine_chunk_frames([self._chunks[i][0] for i in sorted(self._chunks)]) - - def _combine_responses(self) -> httpx.Response: - """Aggregate completed responses under the canonical request URL.""" - responses = [response for _, response in self._chunks.values()] - return _combine_chunk_responses(responses, self.plan.canonical_url) - - @property - def partial_frame(self) -> pd.DataFrame: - """ - Raw combined frame of sub-requests that have completed so far. - - Live — recomputed on each access so it reflects current state - across resume attempts. Deliberately the *raw* combined frame - (``_combine_frames``), NOT the finalized result: this is a cheap, - side-effect-free snapshot for inspecting partial progress, so - reading it (or building a :class:`ChunkInterrupted` around it) - never triggers ``finalize`` work — which for OGC getters includes - a schema network fetch on an empty frame. Use ``call.resume()`` - for the finalized result. - - Returns - ------- - pandas.DataFrame - Combined frame of completed sub-requests, or an empty - ``DataFrame`` when nothing has completed. - """ - return self._combine_frames() if self._chunks else pd.DataFrame() - - @property - def partial_response(self) -> httpx.Response | None: - """ - Raw aggregate response with the canonical URL restored to the - user's full original query. - - Live — recomputed on each access. Like :attr:`partial_frame`, this - is the *raw* aggregate (an :class:`httpx.Response`), not the - finalized result, so inspecting it is side-effect-free. - - Returns - ------- - httpx.Response or None - Aggregated response when at least one sub-request has - completed, ``None`` otherwise. - """ - return self._combine_responses() if self._chunks else None - - def _pending(self) -> Iterator[tuple[int, dict[str, Any]]]: - """ - Yield ``(index, sub_args)`` for sub-requests not yet completed. - - Walks :meth:`ChunkPlan.iter_sub_args` in deterministic order - and skips any index already in ``self._chunks``. :meth:`_run` - uses this to pick up exactly the sub-requests it still owes — - first run and every resume alike. - - Yields - ------ - tuple of (int, dict) - The sub-args ``index`` and its ``sub_args`` dict for each - sub-request not yet completed. - """ - for index, sub_args in enumerate(self.plan.iter_sub_args()): - if index not in self._chunks: - yield index, sub_args - - def resume(self) -> tuple[pd.DataFrame, Any]: - """ - Drive the chunked call to completion and return the combined result. - - Runs :meth:`_run` through an ``anyio`` blocking portal (a - short-lived worker thread), so it works whether or not the caller - is already inside an event loop (Jupyter / IPython / async apps). - The portal copies the calling context, so the active progress - reporter still reaches the sub-requests. - - Idempotent: only sub-requests whose index isn't already in - ``self._chunks`` are re-issued. Sub-args order matches - :meth:`ChunkPlan.iter_sub_args` and is deterministic, so a - partial completion (sparse indices) resumes correctly. - - Returns - ------- - df : pandas.DataFrame - Combined data from every successful sub-request. - response - The finalized aggregate — a raw :class:`httpx.Response` - (canonical URL, headers from the response with the lowest reported - remaining quota, and summed response elapsed durations) by default, - or whatever :attr:`finalize` produces (e.g. ``BaseMetadata`` for - the OGC getters). - - Raises - ------ - ChunkInterrupted - On a mid-stream transient failure — 429, 5xx, or a bare - transport error: :class:`QuotaExhausted` for 429, - :class:`ServiceInterrupted` for the rest. The resumable - handle is on ``exc.call`` — wait for the underlying - condition to clear and call ``exc.call.resume()`` again. - """ - # Drive inside the snapshot taken at construction (see ``__init__``). - # ``start_blocking_portal`` copies the *calling* context into its - # worker thread, and running here means that calling context is the - # snapshot. So the base URL / dialect / row cap / progress reporter - # active when the call was created reach the rebuilt sub-requests, even - # when this is a resume fired long after the original ``with`` blocks - # exited. - return self._ctx.run(self._resume_in_context) - - def _resume_in_context(self) -> tuple[pd.DataFrame, Any]: - """Body of :meth:`resume`, run inside the captured context.""" - concurrency = _read_concurrency_env() - with start_blocking_portal() as portal: - # ``portal.call`` returns ``Any`` because ``functools.partial`` - # erases ``_run``'s return type; restore the declared tuple. - return cast( - "tuple[pd.DataFrame, Any]", - portal.call(functools.partial(self._run, concurrency)), - ) - - async def _run(self, max_concurrent: int | None) -> tuple[pd.DataFrame, Any]: - """ - Gather every pending sub-request; return the combined, finalized result. - - Pending sub-requests (:meth:`_pending`) fan out over one shared - :class:`httpx.AsyncClient` under ``asyncio.gather``, with - ``return_exceptions=True`` so completed sub-requests survive a - sibling's transient failure. On a recognized transient - (:class:`RateLimited`, :class:`ServiceUnavailable`, or a bare - ``httpx.HTTPError`` / ``httpx.InvalidURL``), a - :class:`ChunkInterrupted` subclass is raised carrying ``self`` on - ``.call``. ``exc.call.resume()`` then re-issues only the unfinished - indices through this same runner. - - The gather dispatches *every* pending sub-request at once, but an - ``asyncio.Semaphore`` caps the number of concurrent fetches at - ``N = max_concurrent`` — ``None`` lifts the cap, ``N=1`` runs them - one at a time. The connection pool is sized to the same ``N`` - (``httpx.Limits(max_connections=N, max_keepalive_connections=N)``) - so the in-flight fetches reuse keepalive connections. - - The semaphore, not the pool, is deliberately the throttle. If the - pool throttled instead, the excess sub-requests would queue - *inside* httpx waiting for a connection, and that wait counts - against the pool-acquire timeout (60 s, from ``HTTPX_ASYNC_DEFAULTS``). - A batch of slow pages that keeps every connection busy past that - window would then trip ``httpx.PoolTimeout`` on the queued tail — - a purely client-side failure that consumes the retry budget and - surfaces as a spurious resumable ``ServiceInterrupted``. Holding - sub-requests at the semaphore keeps them out of the pool until a - slot frees, so the pool timeout only fires for a genuinely stuck - connection. - - The shared client is published on :data:`_chunked_client` so - the paginated-loop helpers reuse its connection pool. - - Parameters - ---------- - max_concurrent : int or None - Maximum sub-requests in flight (the semaphore value, and the - connection-pool size). ``None`` lifts the cap entirely. - - Returns - ------- - df : pandas.DataFrame - Combined data from every sub-request. - response - The finalized aggregate — a raw :class:`httpx.Response` - (canonical URL, headers from the response with the lowest reported - remaining quota, and summed response elapsed durations) by default, - or whatever :attr:`finalize` produces (e.g. ``BaseMetadata`` for - OGC getters). - - Raises - ------ - ChunkInterrupted - On a transient sub-request failure. ``.call`` is ``self``, - holding the sparse completed sub-requests; ``.call.resume()`` - re-issues the unfinished ones. - """ - # The semaphore is the throttle; the pool is merely sized to match - # it. Left at httpx's default client limits (``max_connections=100``, - # keepalive 20) the pool would bottleneck a wider cap or churn - # connections by keeping too few alive. See the method docstring for - # why the gate can't be the pool itself. ``unbounded`` - # (``max_concurrent=None``) is a degenerate cap at the plan total — a - # semaphore that can never block — so gated is the only code path. - limits = httpx.Limits( - max_connections=max_concurrent, max_keepalive_connections=max_concurrent - ) - semaphore = asyncio.Semaphore( - self.plan.total if max_concurrent is None else max_concurrent - ) - - async with open_async_client(limits=limits) as client: - with _chunked_client(client): - reporter = _progress.current() - if reporter is not None: - reporter.set_chunks(self.plan.total) - - async def track( - index: int, args: dict[str, Any] - ) -> tuple[pd.DataFrame, httpx.Response]: - """One sub-request (with retry) + result-store + progress tick.""" - result = await _retry( - lambda: self.fetch(args), self.retry_policy, gate=semaphore - ) - self._chunks[index] = result - if reporter is not None: - # Chunks finish out of order under gather, so tick the - # completed *count* rather than a positional index. - reporter.start_chunk(self.completed_chunks) - return result - - # Dispatch every pending sub-request concurrently; the - # semaphore (held by ``_retry`` per attempt) is the only throttle. - # ``return_exceptions`` keeps completed pairs after a sibling - # fails, so partial state stays recoverable via :meth:`resume`. - # Failure precedence, in order: - # 1. Cancellation / interrupt signals (CancelledError, - # KeyboardInterrupt, SystemExit — non-Exception) propagate - # unmodified; wrapping them as a transient would swallow - # the user's stop signal. - # 2. A non-transient failure (a real bug — unrecognized by - # ``wrap_failure``) surfaces raw, so it isn't masked behind - # a resumable handle for a transient sibling that landed - # later. - # 3. Only when every failure is a recognized transient do we - # raise the first as a resumable ``ChunkInterrupted``. - results = await asyncio.gather( - *(track(index, args) for index, args in self._pending()), - return_exceptions=True, - ) - failures = [r for r in results if isinstance(r, BaseException)] - for exc in failures: - if not isinstance(exc, Exception): - raise exc - first_transient: tuple[ChunkInterrupted, BaseException] | None = None - for exc in failures: - interrupted = self.wrap_failure(exc) - if interrupted is None: - raise exc - if first_transient is None: - first_transient = (interrupted, exc) - if first_transient is not None: - interrupted, exc = first_transient - raise interrupted from exc - - return self.finalize(*self._combine_raw()) - - def multi_value_chunked( *, build_request: Callable[..., httpx.Request], url_limit: int | None = None, -) -> Callable[[_Fetch], Callable[..., tuple[pd.DataFrame, Any]]]: +) -> Callable[[_Fetch[dict[str, Any]]], Callable[..., tuple[pd.DataFrame, Any]]]: """ Decorate an async fetcher to transparently chunk over-budget requests. @@ -752,7 +241,9 @@ def multi_value_chunked( ChunkedCall : Per-sub-request execution and resume semantics. """ - def decorator(fetch: _Fetch) -> Callable[..., tuple[pd.DataFrame, Any]]: + def decorator( + fetch: _Fetch[dict[str, Any]], + ) -> Callable[..., tuple[pd.DataFrame, Any]]: @functools.wraps(fetch) def wrapper( args: dict[str, Any], @@ -772,7 +263,16 @@ def wrapper( # The concurrency cap is resolved inside ``resume()`` from # ``API_USGS_CONCURRENT``; ``1`` is a sequential gather, # ``total <= 1`` a one-element gather — no special branch. - return ChunkedCall(plan, fetch, retry_policy, finalize).resume() + return ChunkedCall( + plan, + fetch, + retry_policy, + finalize, + canonical_url=plan.canonical_url, + # The collection name, for the progress line the executor + # opens. ``get_ogc_data`` puts it in ``args``. + service=args.get("service"), + ).resume() return wrapper diff --git a/dataretrieval/ogc/context.py b/dataretrieval/ogc/context.py index ff2ea5a32..f416fd2a8 100644 --- a/dataretrieval/ogc/context.py +++ b/dataretrieval/ogc/context.py @@ -8,13 +8,19 @@ """ from dataretrieval._ambient import Ambient -from dataretrieval.ogc.policy import DEFAULT_DIALECT, OGC_API_URL, OgcDialect +from dataretrieval.ogc.policy import DEFAULT_DIALECT, OgcDialect # Optional cap on rows accumulated by one paginated request. _row_cap: Ambient[int | None] = Ambient("ogc_row_cap", None) -# OGC base URL targeted by request construction and schema lookup. -_ogc_base_url: Ambient[str] = Ambient("ogc_base_url", OGC_API_URL) +# OGC base URL targeted by request construction and schema lookup. Empty by +# default *on purpose*: this package is API-neutral, so the adapter naming the +# service is the one that sets it (``get_ogc_data(base_url=...)`` does, and a +# hand-built request path such as ``waterdata.get_cql`` enters this context +# itself). A default endpoint here would silently send a caller that forgot to +# set it -- e.g. an NGWMN path -- to whichever API happened to be the default; +# an unset value instead fails loudly on the malformed URL. +_ogc_base_url: Ambient[str] = Ambient("ogc_base_url", "") # Per-call request and response dialect. _dialect: Ambient[OgcDialect] = Ambient("ogc_dialect", DEFAULT_DIALECT) diff --git a/dataretrieval/ogc/engine.py b/dataretrieval/ogc/engine.py index 576af6d31..d9f5a0b0a 100644 --- a/dataretrieval/ogc/engine.py +++ b/dataretrieval/ogc/engine.py @@ -30,22 +30,16 @@ Awaitable, Callable, ) -from typing import Any, TypeVar, cast +from typing import TYPE_CHECKING, Any, TypeVar import httpx import pandas as pd import dataretrieval.ogc.chunking as chunking -import dataretrieval.progress as _progress -from dataretrieval._response_metadata import BaseMetadata -from dataretrieval.credentials import without_embedded_credentials -from dataretrieval.ogc.chunking import get_active_client -from dataretrieval.ogc.context import _row_cap +from dataretrieval.ogc.context import _dialect, _ogc_base_url, _row_cap from dataretrieval.ogc.errors import _raise_for_non_200 from dataretrieval.ogc.policy import ( - BASE_URL, # noqa: F401 — compatibility alias DEFAULT_DIALECT, - OGC_API_URL, OgcDialect, _require_positive_int, ) @@ -54,14 +48,17 @@ # the symbols its orchestration uses. from dataretrieval.ogc.requests import ( _construct_api_requests, - _dialect, - _ogc_base_url, _switch_arg_id, _switch_properties_id, ) from dataretrieval.ogc.shaping import GEOPANDAS, _finalize_ogc, _get_resp_data +from dataretrieval.transport.fanout import FanOut, active_client +from dataretrieval.transport.links import resolve_next_url from dataretrieval.transport.pagination import paginate -from dataretrieval.transport.sync import run_sync +from dataretrieval.transport.retry import RetryPolicy + +if TYPE_CHECKING: + from dataretrieval._response_metadata import BaseMetadata # Set up logger for this module logger = logging.getLogger(__name__) @@ -112,42 +109,14 @@ def _next_req_url( href = link.get("href") if not href: return None - # Refuse to follow a next-page link to a different host — - # the request's headers/auth were minted for the original - # host and shouldn't leak to whatever a poisoned response - # body might supply. Guarded against mock-shaped ``resp.url`` - # attributes (tests sometimes set strings or ``MagicMock``) - # by falling open when host extraction isn't reliable. - next_host: str | None - cur_host: str | None - next_url: httpx.URL | None - try: - next_url = httpx.URL(href) - next_host = next_url.host - resp_url = ( - resp.url - if isinstance(resp.url, httpx.URL) - else httpx.URL(str(resp.url)) - ) - cur_host = resp_url.host - except (httpx.InvalidURL, TypeError): - next_url = None - next_host = cur_host = None - if next_host and cur_host and next_host != cur_host: - raise RuntimeError( - f"Refusing to follow cross-host next-page URL: " - f"{next_host} != {cur_host}" - ) - # Matching hosts is not enough: a link may also carry ``user:pass@``, - # which httpx turns into an ``Authorization: Basic`` header on the - # follow-up request. The host check above passes in exactly that case, - # so strip it here rather than trusting the link we were handed. - if next_url is not None: - return str(without_embedded_credentials(next_url)) - # ``href`` comes from the JSON ``links`` array (typed ``Any``); the - # ``not href`` guard above already excluded empty/None, and it is a - # URL string (passed to ``httpx.URL`` above). - return cast("str", href) + # The link is response data: parsing it, resolving a relative + # reference, refusing a foreign host and stripping embedded + # credentials is one shared policy, so this walk cannot drift from + # the other two. ``RuntimeError`` rather than the taxonomy's + # ``DataRetrievalError`` because that is what this walk has always + # raised; retyping it is a released behavior change to make + # deliberately, not a side effect of sharing the check. + return resolve_next_url(href, resp, service="OGC", error=RuntimeError) return None @@ -163,12 +132,12 @@ async def _paginate( raise_for_status: Callable[[httpx.Response], None] = _raise_for_non_200, ) -> tuple[pd.DataFrame, httpx.Response]: """Compatibility wrapper around service-neutral cursor pagination.""" - active_client = client if client is not None else get_active_client() + session = client if client is not None else active_client() return await paginate( initial_req, parse_response=parse_response, follow_up=follow_up, - client=active_client, + client=session, raise_for_status=raise_for_status, row_cap=_row_cap.get(), ) @@ -251,7 +220,7 @@ def get_ogc_data( output_id: str, *, max_rows: int | None = None, - base_url: str = OGC_API_URL, + base_url: str | None = None, extra_id_cols: frozenset[str] | set[str] = frozenset(), dialect: OgcDialect | None = None, ) -> tuple[pd.DataFrame, BaseMetadata]: @@ -278,7 +247,11 @@ def get_ogc_data( fetches the full result. Intended for cheap previews of large, un-chunked tables (e.g. :func:`get_reference_table`). base_url : str, optional - OGC API base URL to target. Defaults to the main Water Data API. + OGC API base URL to target. Required in practice -- this package is + API-neutral and names no service of its own; each adapter passes its + own base (e.g. ``waterdata.utils.OGC_API_URL``, + ``ngwmn.NGWMN_OGC_API_URL``). Falls back to the base URL already in + scope for the current call. extra_id_cols : set or frozenset, optional Synthetic id columns to push to the end of a result frame (see :func:`_arrange_cols`). Defaults to an empty set. @@ -308,6 +281,8 @@ def get_ogc_data( if dialect is None: dialect = _DEFAULT_DIALECT + if base_url is None: + base_url = _ogc_base_url.get() args = args.copy() args["service"] = service @@ -339,12 +314,10 @@ def get_ogc_data( dialect=dialect, base_url=base_url, ) - 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) + # No progress block here: the executor that emits the events owns the line + # (see :meth:`~dataretrieval.transport.fanout.FanOut.resume`). + with _row_cap(max_rows), _ogc_base_url(base_url), _dialect(dialect): + return _fetch_once(args, finalize=finalize) @chunking.multi_value_chunked(build_request=_construct_api_requests) @@ -368,20 +341,6 @@ async def _fetch_once( return await _walk_pages(geopd=GEOPANDAS, req=req) -def _run_sync( - make_coro: Callable[[], Awaitable[tuple[pd.DataFrame, httpx.Response]]], - *, - service: str, - error_url: str | httpx.URL | None = None, -) -> tuple[pd.DataFrame, httpx.Response]: - """Compatibility wrapper around the service-neutral sync bridge.""" - return run_sync( - make_coro, - service=service, - error_url=error_url if error_url is not None else _ogc_base_url.get(), - ) - - def fetch_ogc_request( request: httpx.Request, *, @@ -392,8 +351,10 @@ def fetch_ogc_request( This is the facade-level entry point for generalized CQL requests: the caller builds its own :class:`httpx.Request` (e.g. via :func:`~dataretrieval.ogc.requests._construct_cql_request`) and hands it - here. Pagination, progress reporting, and error handling are identical to - the typed getters' path through :func:`_walk_pages`. + here. The request is driven as a one-item + :class:`~dataretrieval.transport.fanout.FanOut` -- the same executor the + typed getters use -- so pagination, retry, progress reporting, and error + handling are identical to their path through :func:`_walk_pages`. Parameters ---------- @@ -410,7 +371,13 @@ def fetch_ogc_request( Aggregated response metadata. """ - async def _coro() -> tuple[pd.DataFrame, httpx.Response]: - return await _walk_pages(geopd=GEOPANDAS, req=request) + async def _fetch(req: httpx.Request) -> tuple[pd.DataFrame, httpx.Response]: + return await _walk_pages(geopd=GEOPANDAS, req=req) - return _run_sync(_coro, service=service) + return FanOut( + [request], + _fetch, + RetryPolicy.from_env(), + canonical_url=str(request.url), + service=service, + ).resume() diff --git a/dataretrieval/ogc/errors.py b/dataretrieval/ogc/errors.py index a84ff7391..be8d8733f 100644 --- a/dataretrieval/ogc/errors.py +++ b/dataretrieval/ogc/errors.py @@ -11,7 +11,7 @@ import httpx from dataretrieval.exceptions import error_for_status -from dataretrieval.transport.retry import parse_retry_after as _parse_retry_after +from dataretrieval.exceptions import parse_retry_after as _parse_retry_after def _error_body(resp: httpx.Response) -> str: diff --git a/dataretrieval/ogc/interruptions.py b/dataretrieval/ogc/interruptions.py index a402fdc05..a9158fa78 100644 --- a/dataretrieval/ogc/interruptions.py +++ b/dataretrieval/ogc/interruptions.py @@ -1,177 +1,25 @@ -"""Resumable chunk-interruption exceptions — the public resume contract. - -When a transparently-chunked request fails mid-stream (a 429, a 5xx, or a -bare transport error), the work already completed is preserved and the call -is resumable. The raised exception carries a ``.call`` handle whose -``resume()`` re-issues only the still-pending sub-requests. These exception -types are that contract, re-exported at the top level -(``from dataretrieval import ChunkInterrupted``). The execution machinery -that raises and resumes them lives in :mod:`dataretrieval.ogc.chunking`. +"""Compatibility re-export: the interruption taxonomy moved to a top-level leaf. + +The resume contract is no longer OGC-specific — Water Use raises it too — so the +classes live in :mod:`dataretrieval.interruptions`, where the base class is +named :class:`~dataretrieval.interruptions.FanOutInterrupted`. This path is kept +because it is what existing code and tests import; new code should import from +the leaf, or the top level +(``from dataretrieval import FanOutInterrupted``). """ from __future__ import annotations -from typing import TYPE_CHECKING, Any, ClassVar - -import httpx -import pandas as pd - -from dataretrieval.exceptions import DataRetrievalError - -if TYPE_CHECKING: - from dataretrieval.ogc.chunking import ChunkedCall - - -class ChunkInterrupted(DataRetrievalError): - """ - Base class for mid-stream chunk failures whose completed work is resumable. - - A ``ChunkInterrupted`` subclass means: a sub-request failed, but - ``ChunkedCall`` still owns whatever completed successfully before - the failure. Call ``self.call.resume()`` to pick up where the - failure stopped you — only still-pending sub-requests are - re-issued. - - Subclasses describe *why* ``ChunkedCall`` stopped so callers can - pick a retry policy: :class:`QuotaExhausted` for 429 (wait for the - rate-limit window), :class:`ServiceInterrupted` for 5xx (wait for - the upstream to recover). The ``.call`` handle is the same object - across every interruption of a single chunked call — frames - accumulate across retries. - - Attributes - ---------- - call : ChunkedCall or None - Resumable handle into the ``ChunkedCall`` that raised this - exception. ``None`` only on hand-constructed exceptions (test - fixtures), where ``.call``-derived accessors degrade to - empty/``None``. - retry_after : float or None - Seconds the server suggested waiting (``Retry-After`` header). - ``None`` when the server gave no hint. - completed_chunks : int - Number of sub-requests successfully completed before the failure. - total_chunks : int - Total sub-requests in the plan. - partial_frame : pandas.DataFrame - Combined frame of work completed by the moment this exception - was raised. Snapshot at raise time — does NOT advance on a - later ``call.resume()`` (use ``exc.call.partial_frame`` for - the live view). - partial_response : httpx.Response or None - Raw aggregate response covering the completed sub-requests at - raise time; ``None`` if nothing had completed yet. Same snapshot - semantics as ``partial_frame``. (Raw, not finalized — use - ``exc.call.resume()`` for the finalized ``(df, metadata)`` result.) - - Examples - -------- - Retry on any transient interruption, honoring the server's - ``Retry-After`` hint when present and falling back to a fixed wait - otherwise. Each new interruption keeps the already-completed work - intact — only the still-pending sub-requests are re-issued. - - .. code-block:: python - - import time - from dataretrieval import ChunkInterrupted - - # ``getter`` is any chunked OGC getter — e.g. - # ``waterdata.get_daily`` or ``ngwmn.get_water_level``. - try: - df, md = getter(monitoring_location_id=long_list_of_sites) - except ChunkInterrupted as exc: - while True: - time.sleep(exc.retry_after or 5 * 60) - try: - df, md = exc.call.resume() - break - except ChunkInterrupted as next_exc: - exc = next_exc - """ - - # Subclasses override with a ``str.format`` template; the format - # call sees ``completed_chunks`` and ``total_chunks`` as kwargs. - _MESSAGE_TEMPLATE: ClassVar[str] = ( - "Chunked request interrupted after {completed_chunks}/" - "{total_chunks} sub-requests; call .call.resume() to continue." - ) - - def __init__( - self, - *, - completed_chunks: int, - total_chunks: int, - call: ChunkedCall | None = None, - retry_after: float | None = None, - cause: BaseException | None = None, - ) -> None: - message = self._MESSAGE_TEMPLATE.format( - completed_chunks=completed_chunks, total_chunks=total_chunks - ) - if cause is not None: - cause_msg = str(cause) or type(cause).__name__ - message = f"{message} Cause: {type(cause).__name__}: {cause_msg}" - super().__init__(message) - self.completed_chunks = completed_chunks - self.total_chunks = total_chunks - self.call = call - self.retry_after = retry_after - # Snapshot partial state at raise time so the exception stays a stable - # record of the failure moment: ``exc.partial_frame`` / - # ``.partial_response`` do NOT advance on a later ``call.resume()`` - # (that live view is on ``call.partial_frame`` / ``.partial_response``). - # This keeps each interruption in a resume loop a faithful record of - # what it saw, rather than every exception aliasing the shared call's - # advancing state. ``.copy()`` guards the single-chunk fast path, where - # the combined frame may be returned verbatim. - if call is None: - self.partial_frame: pd.DataFrame = pd.DataFrame() - self.partial_response: httpx.Response | None = None - else: - self.partial_frame = call.partial_frame.copy() - self.partial_response = call.partial_response - - def __getstate__(self) -> dict[str, Any]: - # Drop the live ChunkedCall before pickling: its ``.fetch`` is an - # undecorated module function pickle can't reference by name, so the - # interruption can't cross a process boundary with ``.call`` attached. - # The degraded ``call=None`` form keeps the counts, retry hint, and the - # snapshotted partial frame / response — plain instance attributes the - # base ``__getstate__`` already pickles. Only ``.resume()`` is lost, and - # cross-process resume was never possible anyway. - return {**super().__getstate__(), "call": None} - - -class QuotaExhausted(ChunkInterrupted): - """ - A sub-request returned HTTP 429 — the per-key rate-limit window is exhausted. - - Subclass of :class:`ChunkInterrupted`. The completed sub-requests are - preserved on ``.call``; once the rate-limit window resets, - ``.call.resume()`` re-issues only the still-pending work. - ``partial_frame`` holds what completed before the 429. - """ - - _MESSAGE_TEMPLATE = ( - "HTTP 429 after {completed_chunks}/{total_chunks} sub-requests; " - "catch QuotaExhausted (or ChunkInterrupted) to access " - ".partial_frame or .call.resume() once the rate-limit " - "window has rolled over." - ) - - -class ServiceInterrupted(ChunkInterrupted): - """ - A sub-request returned HTTP 5xx — the upstream service failed transiently. - - Subclass of :class:`ChunkInterrupted`. The completed sub-requests are - preserved on ``.call``; once the upstream recovers, ``.call.resume()`` - resumes only the still-pending work. - """ - - _MESSAGE_TEMPLATE = ( - "Service error after {completed_chunks}/{total_chunks} " - "sub-requests; catch ServiceInterrupted (or ChunkInterrupted) " - "and call .call.resume() once the upstream service recovers." - ) +from dataretrieval.interruptions import ( + ChunkInterrupted, + FanOutInterrupted, + QuotaExhausted, + ServiceInterrupted, +) + +__all__ = [ + "ChunkInterrupted", + "FanOutInterrupted", + "QuotaExhausted", + "ServiceInterrupted", +] diff --git a/dataretrieval/ogc/planning.py b/dataretrieval/ogc/planning.py index 03dc90736..d6e07ff7c 100644 --- a/dataretrieval/ogc/planning.py +++ b/dataretrieval/ogc/planning.py @@ -573,3 +573,14 @@ def iter_sub_args(self) -> Iterator[dict[str, Any]]: for axis, chunk in zip(self.axes, combo, strict=False): sub_args[axis.arg_key] = axis.render(chunk) yield sub_args + + # ``total`` and ``iter_sub_args`` are this class's domain vocabulary and + # stay as they are. The dunders are how a plan satisfies + # :class:`~dataretrieval.transport.fanout.FanOutPlan`, which asks for a + # sized iterable and nothing chunking-specific. They delegate rather than + # duplicate, so ``len(plan)`` cannot disagree with what iterating yields. + def __len__(self) -> int: + return self.total + + def __iter__(self) -> Iterator[dict[str, Any]]: + return self.iter_sub_args() diff --git a/dataretrieval/ogc/policy.py b/dataretrieval/ogc/policy.py index 04299a7e3..2f9f58c5a 100644 --- a/dataretrieval/ogc/policy.py +++ b/dataretrieval/ogc/policy.py @@ -1,10 +1,14 @@ -"""Low-level OGC policy: dialect, controls, and default endpoint constants. +"""Low-level OGC policy: the dialect type and control validation. This module is the single source of truth for the :class:`OgcDialect` type -(per-API quirks the generic request builder needs), OGC control validation, and -the default endpoint constants used by the Water Data OGC API. It depends only -on the stdlib and the credential-policy leaf, so any OGC submodule can import it -without creating cycles. +(per-API quirks the generic request builder needs) and OGC control validation. +It depends only on the stdlib, so any OGC submodule can import it without +creating cycles. + +It names no endpoint: which service an OGC call targets is the *adapter's* +policy, supplied per call as ``base_url`` (see +:data:`dataretrieval.ogc.context._ogc_base_url`). A default here would quietly +point every generic OGC caller at one API. It must NOT import engine, shaping, or any service adapter. """ @@ -14,8 +18,6 @@ import numbers from dataclasses import dataclass, field -from dataretrieval.credentials import WATERDATA_BASE_URL - def _require_positive_int( value: int, name: str, *, examples: str | None = None @@ -30,14 +32,6 @@ def _require_positive_int( raise ValueError(f"{name} must be a positive integer{eg} (got {value!r}).") -# --------------------------------------------------------------------------- -# Endpoint constants -# --------------------------------------------------------------------------- - -BASE_URL = WATERDATA_BASE_URL -OGC_API_VERSION = "v0" -OGC_API_URL = f"{BASE_URL}/ogcapi/{OGC_API_VERSION}" - # --------------------------------------------------------------------------- # Dialect type # --------------------------------------------------------------------------- diff --git a/dataretrieval/ogc/requests.py b/dataretrieval/ogc/requests.py index 944668031..882a8b5fa 100644 --- a/dataretrieval/ogc/requests.py +++ b/dataretrieval/ogc/requests.py @@ -120,6 +120,26 @@ def _partition_request_params( return get_params, {} +def _items_url(service: str) -> str: + """The OGC items endpoint for ``service`` under the active base URL.""" + return f"{_ogc_base_url.get()}/collections/{service}/items" + + +def _cql2_post_request( + service_url: str, *, content: str, params: dict[str, Any] +) -> httpx.Request: + """A POST/CQL2 request: the media type the API requires, in one place.""" + headers = _default_headers(service_url) + headers["Content-Type"] = "application/query-cql-json" + return httpx.Request( + method="POST", + url=service_url, + headers=headers, + content=content, + params=params, + ) + + def _construct_api_requests( service: str, properties: list[str] | None = None, @@ -129,7 +149,7 @@ def _construct_api_requests( **kwargs: Any, ) -> httpx.Request: """Construct an HTTP request object for the specified OGC API service.""" - service_url = f"{_ogc_base_url.get()}/collections/{service}/items" + service_url = _items_url(service) dialect = _dialect.get() for key in _DATE_RANGE_PARAMS: if key in kwargs: @@ -152,21 +172,14 @@ def _construct_api_requests( if "filter_lang" in params: params["filter-lang"] = params.pop("filter_lang") - headers = _default_headers(service_url) - if post_params: - headers["Content-Type"] = "application/query-cql-json" - return httpx.Request( - method="POST", - url=service_url, - headers=headers, - content=_cql2_param(post_params), - params=params, + return _cql2_post_request( + service_url, content=_cql2_param(post_params), params=params ) return httpx.Request( method="GET", url=service_url, - headers=headers, + headers=_default_headers(service_url), params=params, ) @@ -181,7 +194,7 @@ def _construct_cql_request( skip_geometry: bool | None = None, ) -> httpx.Request: """Build a POST/CQL2 request from a verbatim CQL2 body.""" - service_url = f"{_ogc_base_url.get()}/collections/{service}/items" + service_url = _items_url(service) params = _ogc_query_params( {}, properties=properties, @@ -189,25 +202,18 @@ def _construct_cql_request( limit=limit, skip_geometry=skip_geometry, ) - headers = _default_headers(service_url) - headers["Content-Type"] = "application/query-cql-json" - return httpx.Request( - method="POST", - url=service_url, - headers=headers, - content=cql_body, - params=params, - ) + return _cql2_post_request(service_url, content=cql_body, params=params) # --------------------------------------------------------------------------- # Argument normalization helpers # --------------------------------------------------------------------------- -# Default set of iterable-shaped params that ``_get_args`` must NOT push -# through ``_normalize_str_iterable`` (date-range params may carry -# ``pd.NaT``/None or interval strings; ``bbox`` is ``list[float]``). Callers -# with extra numeric params pass their own superset. +# Iterable-shaped params that ``_get_args`` must NOT push through +# ``_normalize_str_iterable`` (date-range params may carry ``pd.NaT``/None or +# interval strings; ``bbox`` is ``list[float]``). Every OGC caller gets these; +# an adapter with extra numeric params names only its extras via +# ``prepare_request_args(..., extra_no_normalize=...)``. _NO_NORMALIZE_PARAMS = _DATE_RANGE_PARAMS | {"bbox"} @@ -274,14 +280,20 @@ def prepare_request_args( local_vars: dict[str, Any], exclude: set[str] | None = None, *, - no_normalize: frozenset[str] | set[str] = _NO_NORMALIZE_PARAMS, + extra_no_normalize: frozenset[str] | set[str] = frozenset(), ) -> dict[str, Any]: """Build OGC request kwargs from a getter's ``locals()``. Internal bookkeeping keys, caller-supplied exclusions, and ``None`` values are omitted. Identifiers and properties are validated; other iterables are - normalized unless listed in ``no_normalize``. + normalized unless exempted. + + ``extra_no_normalize`` *adds* to the engine's own + :data:`_NO_NORMALIZE_PARAMS` rather than replacing it, so an adapter names + only the params it owns and cannot silently drop the date-range exemptions + by forgetting to union them back in. """ + no_normalize = _NO_NORMALIZE_PARAMS | frozenset(extra_no_normalize) to_exclude = {"service", "output_id"} if exclude: to_exclude.update(exclude) diff --git a/dataretrieval/ogc/retry.py b/dataretrieval/ogc/retry.py index 7eeafb44b..894395b5a 100644 --- a/dataretrieval/ogc/retry.py +++ b/dataretrieval/ogc/retry.py @@ -1,61 +1,20 @@ -"""OGC interruption classification over service-neutral transport retry policy. +"""Compatibility re-export: interruption classification moved to the taxonomy leaf. -Only the OGC-specific half of retry lives here: turning a transport failure into -the resumable :class:`~dataretrieval.ogc.interruptions.ChunkInterrupted` the -chunker reports. The policy itself -- backoff, bounds, classification of what is -transient -- belongs to :mod:`dataretrieval.transport.retry`, which callers -import directly; re-exporting its tunables here would hand out stale copies that -patching cannot reach. +Turning a transport failure into a resumable +:class:`~dataretrieval.interruptions.FanOutInterrupted` was never OGC-specific -- +it keys off the shared ``RateLimited``/``TransientError`` taxonomy and httpx -- +so it now lives beside the classes it produces, in +:mod:`dataretrieval.interruptions`. -"Should we retry this?" and "can the caller resume it?" are the same question -asked twice, so both answers come from one place in transport. Keeping a second -copy here is how they would end up disagreeing -- refusing to retry a failure -while still telling the caller it can be resumed. +The retry *policy* -- backoff, bounds, classification of what is transient -- +still belongs to :mod:`dataretrieval.transport.retry`, which callers import +directly; re-exporting its tunables here would hand out stale copies that +patching cannot reach. """ from __future__ import annotations -import httpx - -from dataretrieval.exceptions import RateLimited, TransientError -from dataretrieval.ogc.interruptions import ( - ChunkInterrupted, - QuotaExhausted, - ServiceInterrupted, -) -from dataretrieval.transport.retry import _deterministic_failure - - -def _classify_transient( - exc: BaseException, -) -> tuple[type[ChunkInterrupted], float | None] | None: - """Classify one failure as a resumable OGC interruption.""" - if isinstance(exc, RateLimited): - return QuotaExhausted, exc.retry_after - if isinstance(exc, TransientError): - return ServiceInterrupted, exc.retry_after - if isinstance(exc, (httpx.HTTPError, httpx.InvalidURL)): - # Some failures will fail the same way every time -- a bad scheme, a - # hostname that doesn't resolve. Offering to resume one would just - # hide the real error behind a retry that can never work. - if _deterministic_failure(exc): - return None - return ServiceInterrupted, None - return None - - -def _classify_chunk_error( - exc: BaseException, -) -> tuple[type[ChunkInterrupted], float | None] | None: - """Walk a wrapped pagination failure for a resumable transport cause.""" - current: BaseException | None = exc - while current is not None: - result = _classify_transient(current) - if result is not None: - return result - current = current.__cause__ - return None - +from dataretrieval.interruptions import _classify_chunk_error, _classify_transient __all__ = [ "_classify_chunk_error", diff --git a/dataretrieval/ogc/schema.py b/dataretrieval/ogc/schema.py index 6306dc4fc..c3b3fd232 100644 --- a/dataretrieval/ogc/schema.py +++ b/dataretrieval/ogc/schema.py @@ -10,21 +10,68 @@ from typing import Any, cast import httpx +import pandas as pd +from dataretrieval._response_metadata import BaseMetadata +from dataretrieval.ogc.context import _ogc_base_url from dataretrieval.ogc.errors import _raise_for_non_200 -from dataretrieval.ogc.policy import OGC_API_URL 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 def _check_ogc_requests( - endpoint: str, req_type: str = "queryables", *, base_url: str = OGC_API_URL + endpoint: str, req_type: str = "queryables", *, base_url: str | None = None ) -> tuple[dict[str, Any], httpx.Response]: - """Retrieve one collection's queryables or response schema.""" + """Retrieve one collection's queryables or response schema. + + ``base_url`` names the API to ask; it defaults to the one in scope for the + current call rather than to any particular service. + """ if req_type not in ("queryables", "schema"): raise ValueError(f"req_type must be 'queryables' or 'schema', got {req_type!r}") + if base_url is None: + base_url = _ogc_base_url.get() url = f"{base_url}/collections/{endpoint}/{req_type}" response = _get(url, headers=_default_headers(url), **HTTPX_DEFAULTS) _raise_for_non_200(response) return cast("dict[str, Any]", response.json()), response + + +def queryables_frame( + collection: str, *, base_url: str | None = None +) -> tuple[pd.DataFrame, BaseMetadata]: + """Tabulate one collection's queryable properties. + + Reading an OGC queryables document is protocol knowledge, not service + knowledge, so it lives here rather than in any one API's getters -- every + OGC adapter in the package can offer the same table. ``base_url`` names + the API to ask, defaulting to the one in scope for the current call. + + Returns + ------- + pd.DataFrame + One row per queryable, sorted by name, with columns ``queryable``, + ``type``, ``title``, and ``description``. + BaseMetadata + Metadata describing the request (URL, query time, response headers). + """ + # The OGC queryables document is a JSON Schema whose ``properties`` map each + # filterable property name to a ``{title, type, description}`` definition. + body, response = _check_ogc_requests( + endpoint=collection, req_type="queryables", base_url=base_url + ) + properties: dict[str, Any] = body.get("properties", {}) + df = pd.DataFrame( + [ + { + "queryable": name, + "type": prop.get("type"), + "title": prop.get("title"), + "description": (prop.get("description") or "").strip(), + } + for name, prop in sorted(properties.items()) + ], + columns=["queryable", "type", "title", "description"], + ) + return df, BaseMetadata(response) diff --git a/dataretrieval/ogc/shaping.py b/dataretrieval/ogc/shaping.py index 8e0e539c2..2cb1befc8 100644 --- a/dataretrieval/ogc/shaping.py +++ b/dataretrieval/ogc/shaping.py @@ -18,7 +18,8 @@ import pandas as pd from dataretrieval._response_metadata import BaseMetadata -from dataretrieval.ogc.policy import DEFAULT_DIALECT, OGC_API_URL, OgcDialect +from dataretrieval.ogc.context import _ogc_base_url +from dataretrieval.ogc.policy import DEFAULT_DIALECT, OgcDialect try: import geopandas as gpd @@ -65,6 +66,23 @@ def _attach_coordinates(df: pd.DataFrame, features: list[dict[str, Any]]) -> Non df["geometry"] = geoms +def _geo_feature_frame(features: list[dict[str, Any]]) -> pd.DataFrame: + """Build a ``GeoDataFrame`` from GeoJSON features, tolerating a missing + ``geometry`` key. + + ``GeoDataFrame.from_features`` indexes ``feature["geometry"]`` directly, so + collections that omit it (NGWMN observation collections, Water Data + statistics features) would raise ``KeyError``. Default the key to ``None`` + for only those features, so features that already carry geometry (the + common sites case) are passed through without a per-feature dict copy. + The single home for this upstream-schema workaround. + """ + return gpd.GeoDataFrame.from_features( + [f if "geometry" in f else {**f, "geometry": None} for f in features], + crs=_CRS, + ) + + def _get_resp_data( resp: httpx.Response, geopd: bool, @@ -132,16 +150,9 @@ def _get_resp_data( # Organize json into geodataframe and make sure id column comes along. # NGWMN observation collections (water levels, lithology, …) return - # features with no ``geometry`` key at all, which - # ``GeoDataFrame.from_features`` can't handle (it indexes - # ``feature["geometry"]`` directly). Default the key to ``None`` for only - # those features so the call is safe; the all-null check below then yields - # a plain DataFrame. Features that already carry geometry (the common - # sites case) are passed through without a per-feature dict copy. - df = gpd.GeoDataFrame.from_features( - [f if "geometry" in f else {**f, "geometry": None} for f in features], - crs=_CRS, - ) + # features with no ``geometry`` key at all; ``_geo_feature_frame`` absorbs + # that, and the all-null check below then yields a plain DataFrame. + df = _geo_feature_frame(features) # Mirror the non-geopandas branch's defensive ``f.get("id")`` so a feature # missing a top-level ``id`` yields None rather than a KeyError. df["id"] = [f.get("id") for f in features] @@ -160,7 +171,7 @@ def _deal_with_empty( properties: list[str] | None, service: str, *, - base_url: str = OGC_API_URL, + base_url: str | None = None, ) -> pd.DataFrame: """ Handles empty DataFrame results by returning a DataFrame with appropriate columns. @@ -179,7 +190,8 @@ def _deal_with_empty( service : str The service endpoint to query for schema properties if needed. base_url : str, optional - OGC API base URL to use for that schema query. + OGC API base URL to use for that schema query. Defaults to the base + URL in scope for the current call, not to any particular service. Returns ------- @@ -189,6 +201,8 @@ def _deal_with_empty( """ if return_list.empty: if not properties or all(pd.isna(properties)): + if base_url is None: + base_url = _ogc_base_url.get() # Schema lookup performs HTTP only for an empty result. from dataretrieval.ogc.schema import _check_ogc_requests @@ -353,7 +367,7 @@ def _finalize_ogc( max_rows: int | None = None, extra_id_cols: frozenset[str] | set[str] = frozenset(), dialect: OgcDialect | None = None, - base_url: str = OGC_API_URL, + base_url: str | None = None, ) -> tuple[pd.DataFrame, BaseMetadata]: """Shape a combined OGC result into the user-facing ``(df, md)``. @@ -378,6 +392,8 @@ def _finalize_ogc( """ if dialect is None: dialect = DEFAULT_DIALECT + if base_url is None: + base_url = _ogc_base_url.get() frame = _deal_with_empty(frame, properties, service, base_url=base_url) # Normalize to PEP-8 snake_case column names *first*, so the dialect's # ``time_cols``/``numerical_cols``/``sort_cols`` (all snake_case) match diff --git a/dataretrieval/progress.py b/dataretrieval/progress.py index b8ac14e28..9f51cf762 100644 --- a/dataretrieval/progress.py +++ b/dataretrieval/progress.py @@ -26,13 +26,13 @@ from __future__ import annotations -import contextvars import os import sys from collections.abc import Iterator from contextlib import contextmanager from typing import TYPE_CHECKING, TextIO +from dataretrieval._ambient import Ambient from dataretrieval.credentials import SIGNUP_URL, accepts_api_key, api_key if TYPE_CHECKING: @@ -53,9 +53,7 @@ def _group_int(value: str) -> str: # within one query, and an unrelated query in another context can't clobber its # state. (It does not give concurrent queries sharing one stderr separate # lines — they would still interleave.) -_active: contextvars.ContextVar[ProgressReporter | None] = contextvars.ContextVar( - "dataretrieval_progress", default=None -) +_active: Ambient[ProgressReporter | None] = Ambient("dataretrieval_progress", None) # Process-level latch so the "no API key" pointer is shown at most once. _api_key_hint_shown = False @@ -292,11 +290,10 @@ def progress_context( reporter = ProgressReporter( service=service, stream=stream, enabled=enabled, target_url=target_url ) - token = _active.set(reporter) try: - yield reporter + with _active(reporter): + yield reporter finally: - _active.reset(token) reporter.close() diff --git a/dataretrieval/streamstats.py b/dataretrieval/streamstats.py index af0763b09..22b3bf451 100644 --- a/dataretrieval/streamstats.py +++ b/dataretrieval/streamstats.py @@ -11,11 +11,13 @@ import httpx +from dataretrieval._querying import _get_with_retry from dataretrieval.transport.http import HTTPX_DEFAULTS -from dataretrieval.utils import _get_with_retry __all__ = ["download_workspace", "get_sample_watershed", "get_watershed", "Watershed"] +STREAMSTATS_URL = "https://streamstats.usgs.gov/streamstatsservices" + def download_workspace(workspaceID: str, format: str = "") -> httpx.Response: """Download a StreamStats workspace. @@ -37,7 +39,7 @@ def download_workspace(workspaceID: str, format: str = "") -> httpx.Response: """ payload = {"workspaceID": workspaceID, "format": format} - url = "https://streamstats.usgs.gov/streamstatsservices/download" + url = f"{STREAMSTATS_URL}/download" r = _get_with_retry(url, params=payload, **HTTPX_DEFAULTS) return r @@ -140,7 +142,7 @@ def get_watershed( "includefeatures": includefeatures, "simplify": simplify, } - url = "https://streamstats.usgs.gov/streamstatsservices/watershed.geojson" + url = f"{STREAMSTATS_URL}/watershed.geojson" r = _get_with_retry(url, params=payload, **HTTPX_DEFAULTS) @@ -189,7 +191,7 @@ def __init__(self, rcode: str, xlocation: float, ylocation: float) -> None: Parses the response onto this instance. """ response = cast( - httpx.Response, + "httpx.Response", get_watershed(rcode, xlocation, ylocation, format="geojson"), ) self._populate(json.loads(response.text)) diff --git a/dataretrieval/transport/env.py b/dataretrieval/transport/env.py new file mode 100644 index 000000000..2bfaf472b --- /dev/null +++ b/dataretrieval/transport/env.py @@ -0,0 +1,57 @@ +"""Environment parsing for the ``API_USGS_*`` numeric knobs. + +A dependency-free leaf: every transport setting read from the environment +shares one grammar and one error voice, and no policy module has to be +imported to get at the parser. +""" + +from __future__ import annotations + +import math +import os +from collections.abc import Callable +from typing import TypeVar + +from dataretrieval.exceptions import ConfigurationError + +_Number = TypeVar("_Number", int, float) + + +def _read_env_number( + name: str, + default: _Number, + cast: Callable[[str], _Number], + expected: str, + *, + minimum: float = 0, + hint: str = "", +) -> _Number: + """Read a bounded number from the environment, or ``default`` if unset. + + The single parser behind every ``API_USGS_*`` numeric knob, so they share + one grammar and one error voice rather than each adapter hand-rolling the + read-cast-validate sequence its own way. + + 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. ``hint`` appends + a sentence pointing at the fix when a setting has one (e.g. the keyword + that disables a cap). + """ + 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}).{hint}" + ) from exc + # ``nan`` passes every ordering test, so a bare ``< minimum`` 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}).{hint}") + if value < minimum: + raise ConfigurationError(f"{name} must be >= {minimum:g} (got {value}).{hint}") + return value diff --git a/dataretrieval/transport/fanout.py b/dataretrieval/transport/fanout.py new file mode 100644 index 000000000..80b17c729 --- /dev/null +++ b/dataretrieval/transport/fanout.py @@ -0,0 +1,710 @@ +"""Bounded, resumable fan-out execution over a plan of sub-requests. + +A fan-out is one logical query the service forces into several requests. Two +unrelated reasons produce one: + +- a Water Data / NGWMN query whose URL exceeds the server's byte limit, split + along its multi-value axes by :class:`dataretrieval.ogc.planning.ChunkPlan`; +- a Water Use query naming several locations, which the NWDC accepts only one + at a time. + +Chunking is how you divide the data structurally; fan-out is how you distribute +the work operationally. The two are orthogonal, and only the first is protocol +knowledge: dividing a query needs the byte budget, the CQL2 grammar, and which +parameters are list-valued, while distributing the pieces needs none of it. Only +the Water Data / NGWMN case above involves chunking at all — Water Use fans out +without dividing anything, because the caller's locations were never one body to +split. + +So this module owns distribution and nothing else: concurrency bounded by a +semaphore, per-attempt retry, deterministic failure precedence, sparse +completion tracking, and resume. It names no protocol concept — an adapter +supplies a :class:`FanOutPlan` (whatever structure it divided into, if any) and +an ``async def fetch(item) -> (df, response)``. + +Concurrency: :meth:`FanOut._run` dispatches every pending sub-request under one +``asyncio.gather`` sharing a single ``httpx.AsyncClient``. An +``asyncio.Semaphore`` -- not the client's connection pool, which is merely sized +to match -- caps the sub-requests in flight at ``N``; see :meth:`FanOut._run` +for why the gate must be the semaphore rather than the pool. +``API_USGS_CONCURRENT`` resolves ``N``: an integer N > 1 allows N sub-requests +in flight; ``1`` forces sequential dispatch; the literal ``unbounded`` lifts the +cap. ``N`` bounds only how many of a query's sub-requests are in flight at once +-- a client-side trade-off between open connections and fan-out latency. It does +not affect the API rate limit: a fanned-out call issues the same number of +sub-requests regardless of ``N``, so ``N`` changes their timing, not the total +request volume. The USGS API rate-limits by volume over time (HTTP 429), not by +simultaneity; set ``API_USGS_PAT`` to raise that quota. The default of 32 is a +conservative cap that keeps connection use modest. The fan-out runs in a +short-lived worker thread (an ``anyio`` blocking portal), so it works whether or +not the caller is already inside an event loop (Jupyter / IPython / async apps). + +Retries: each sub-request is retried on a transient failure (429, 5xx, +connect/read timeout) with exponential backoff + full jitter, honoring a server +``Retry-After`` when present. ``API_USGS_RETRIES`` sets the cap (default 4; +``0`` disables). A ``Retry-After`` longer than the per-call ceiling escalates to +a resumable interruption. + +Interruption: any mid-stream transient failure surfaces as a +:class:`~dataretrieval.interruptions.FanOutInterrupted` subclass carrying +``.call``, a :class:`FanOut` handle owning the already-completed sub-request +state. Call ``.call.resume()`` once the underlying condition clears; only the +still-pending sub-requests are re-issued. +""" + +from __future__ import annotations + +import asyncio +import functools +import os +from collections.abc import Awaitable, Callable, Iterator +from contextvars import copy_context +from typing import Any, Generic, Protocol, TypeVar, cast + +import httpx +import pandas as pd +from anyio.from_thread import start_blocking_portal + +from dataretrieval import progress as _progress +from dataretrieval._ambient import Ambient +from dataretrieval.combining import ( + _combine_chunk_frames, + _combine_chunk_responses, +) +from dataretrieval.interruptions import ( + FanOutInterrupted, + _classify_chunk_error, + _walk_causes, +) +from dataretrieval.transport.env import _read_env_number +from dataretrieval.transport.http import network_error, open_async_client +from dataretrieval.transport.retry import _NO_RETRY, RetryPolicy +from dataretrieval.transport.retry import retry_async as _retry + +#: One sub-request's description, as the adapter's ``fetch`` wants it. The +#: executor never inspects it — see :class:`FanOutPlan`. +_Item = TypeVar("_Item") +#: The same thing in :class:`FanOutPlan`, where it only ever comes *out* of the +#: plan. Covariant so a ``list[httpx.Request]`` satisfies a plan of any +#: supertype, the way ``Iterable`` is covariant for the same reason. +_ItemCo = TypeVar("_ItemCo", covariant=True) + +# Fan-out concurrency cap, read at call time (not import) so test +# ``monkeypatch.setenv`` applies. Value grammar in :func:`_read_concurrency_env`; +# the concurrency model is in the module docstring. +_CONCURRENCY_ENV = "API_USGS_CONCURRENT" +_CONCURRENCY_DEFAULT = 32 +_CONCURRENCY_UNBOUNDED = "unbounded" + + +def _resolve_concurrency(default: int = _CONCURRENCY_DEFAULT) -> int | None: + """ + Resolve the parallelism cap: the general setting, or a module's default. + + ``API_USGS_CONCURRENT`` is the general knob and applies to every fanned-out + call in the package. A module may pass a different ``default`` when its + service warrants one — Water Use ships a lower figure than the OGC getters, + because the NWDC is only stress-tested to that level. + + The ordering is deliberate: an explicitly set environment variable wins over + a module's default, never the reverse. A module that could override the + general setting would make ``API_USGS_CONCURRENT=1`` a lie — the user + dialing concurrency down to be polite to the service would find one adapter + quietly ignoring them, which is precisely the defect this consolidates away. + Module defaults express "absent instruction, this service prefers N"; they + do not express "this service knows better than you". + + Parameters + ---------- + default : int + Cap to use when ``API_USGS_CONCURRENT`` is unset or empty. + + Returns + ------- + int or None + ``1`` for sequential dispatch (one sub-request at a time); an + integer >1 for bounded concurrency; ``None`` to disable the + per-call cap entirely (the ``unbounded`` keyword). + """ + # Only the ``unbounded`` keyword is specific to this knob; the rest is the + # same read-cast-validate every ``API_USGS_*`` number gets, so it delegates + # rather than growing a third copy with its own error wording. + if os.environ.get(_CONCURRENCY_ENV, "").strip().lower() == _CONCURRENCY_UNBOUNDED: + return None + return _read_env_number( + _CONCURRENCY_ENV, + default, + int, + f"a positive integer or '{_CONCURRENCY_UNBOUNDED}'", + minimum=1, + hint=f" Use '{_CONCURRENCY_UNBOUNDED}' to disable the cap.", + ) + + +# --------------------------------------------------------------------------- +# The plan contract +# --------------------------------------------------------------------------- + + +class FanOutPlan(Protocol[_ItemCo]): + """ + A fan-out's shape: how many sub-requests, and what each one is. + + Deliberately the two standard protocols rather than bespoke members. A + plan is a sized, iterable collection of sub-request descriptions, which is + exactly ``__len__`` + ``__iter__`` — so a plain ``list`` of pre-built + requests satisfies this with no adapter class, and a real planner + satisfies it by delegating (see + :class:`~dataretrieval.ogc.planning.ChunkPlan`, whose domain vocabulary is + ``total`` / ``iter_sub_args``). Naming them ``total`` and + ``iter_sub_args`` here would mean two names for ``len`` that could report + different counts, and a shim class for every adapter whose sub-requests + are already a list. + + The item type is whatever an adapter's own ``fetch`` accepts: this executor + passes each item through untouched and never inspects it, so the OGC + getters yield kwargs dicts while Water Use yields ready + :class:`httpx.Request` objects. + + Iteration order is load-bearing: :meth:`FanOut.resume` keys completed work + by position, so a plan that yielded a different order on a second pass + would resume the wrong sub-requests. ``len`` must agree with the number of + items iteration yields — the usual contract for a sized collection. + + The identity of the query as a whole is *not* here: it is a value stamped + on the combined response, not a property of how the work divides, so it is + the ``canonical_url`` argument to :class:`FanOut`. + """ + + def __len__(self) -> int: ... + + def __iter__(self) -> Iterator[_ItemCo]: ... + + +# --------------------------------------------------------------------------- +# Shared per-call client +# --------------------------------------------------------------------------- + +# The per-call ``httpx.AsyncClient``, published for the duration of +# ``FanOut._run`` so paginated-loop helpers reuse the same connection pool +# across every sub-request. ``None`` outside a fan-out — paginated helpers then +# open their own short-lived client. Deliberately a plain ContextVar-backed +# ambient rather than a parameter: the fetch closure an adapter injects is often +# several frames below the client's owner. +_active_client: Ambient[httpx.AsyncClient | None] = Ambient("_fanout_client", None) + + +def active_client() -> httpx.AsyncClient | None: + """ + Return the fan-out's currently-published client, or ``None``. + + Used by paginated-loop helpers to reuse the per-call connection pool. + + Returns + ------- + httpx.AsyncClient or None + The client published for the duration of a :meth:`FanOut._run`; + ``None`` outside one. + """ + return _active_client.get() + + +# --------------------------------------------------------------------------- +# Type aliases for the FanOut contract +# --------------------------------------------------------------------------- + +# The per-sub-request fetcher an adapter injects and ``FanOut`` drives: an +# ``async def fetch(item) -> (df, response)``, where ``item`` is whatever the +# adapter's plan yields. +_Fetch = Callable[[_Item], Awaitable[tuple[pd.DataFrame, httpx.Response]]] + +# Caller-supplied transform applied to the combined result, so a resumed call +# returns the same shape as an un-interrupted one rather than the executor's raw +# ``(frame, httpx.Response)``. This keeps the executor generic: the OGC getters +# inject their post-processing (type coercion, column arrangement, +# ``BaseMetadata``) through ``_finalize_ogc``. The default is identity. +_Finalize = Callable[[pd.DataFrame, httpx.Response], tuple[pd.DataFrame, Any]] + + +def _passthrough_result( + frame: pd.DataFrame, response: httpx.Response +) -> tuple[pd.DataFrame, Any]: + """Default :data:`_Finalize`: return the raw combined pair unchanged.""" + return frame, response + + +class FanOut(Generic[_Item]): + """ + Stateful handle for a fanned-out call. + + Holds the in-flight state (per-sub-request frames and responses) + and the async fetcher. A single :meth:`resume` entry point drives + the call from wherever it is to completion — used both for the + first invocation and for subsequent retries after a + :class:`~dataretrieval.interruptions.FanOutInterrupted`. + + :meth:`_run` gathers every pending sub-request over one shared + :class:`httpx.AsyncClient`, applies the failure-precedence rules, and + combines; :meth:`resume` drives it through an ``anyio`` blocking + portal so it works whether or not the caller is already inside an + event loop. Concurrency is bounded by a per-run ``asyncio.Semaphore`` + (see :meth:`_run`), so sequential dispatch + (``API_USGS_CONCURRENT=1``) is just a degenerate gather. + + A ``FanOut`` is created internally when an adapter executes a plan; + callers reach it via ``FanOutInterrupted.call`` on the exception raised + by a mid-stream failure. + + :meth:`resume` is idempotent: :meth:`_run` iterates the plan + (deterministic order) and skips + any index whose result is already in ``self._chunks``. The + completion set is a sparse ``dict[int, (df, response)]`` so the + gather can record scattered completions (e.g. indices [0, 2, 5] + after siblings [1, 3, 4] failed) and a subsequent ``resume`` only + re-issues the missing indices. + + Parameters + ---------- + plan : FanOutPlan + The sub-requests to execute: anything sized and iterable, from a + :class:`~dataretrieval.ogc.planning.ChunkPlan` to a plain ``list`` of + pre-built requests. + fetch : Callable + ``async def`` that issues a single sub-request, given one item from + ``plan``, and returns ``(frame, response)``. + client_options : dict, optional + Extra ``httpx.AsyncClient`` options for the shared client this run + opens (e.g. ``{"verify": False}``). + default_concurrent : int, optional + This service's preferred in-flight cap when ``API_USGS_CONCURRENT`` + is unset. Defaults to 32. + canonical_url : str or None, optional + URL identifying the query as a whole, restored onto the combined + response so the caller sees the request they made rather than + whichever sub-request happened to land last. Also the destination + :meth:`resume` labels its progress line with. + service : str or None, optional + Human-facing name of what is being retrieved (e.g. ``"daily"``, + ``"wateruse"``), used to label the progress line :meth:`resume` + opens. ``None`` leaves the line unlabelled. + + Attributes + ---------- + plan : FanOutPlan + The plan being driven (read-only after construction). + fetch : Callable + The async per-sub-request fetch function. + finalize : Callable + Transform applied to the combined result (see :data:`_Finalize`) at + the terminal :meth:`_run` return, so a completed call yields the + caller's finished shape. The ``partial_*`` accessors deliberately + skip it and stay raw. + partial_frame : pandas.DataFrame + Raw combined frame of completed sub-requests (live; recomputed per + access). Not finalized — call :meth:`resume` for the finished shape. + partial_response : httpx.Response or None + Raw aggregate response (canonical URL restored), or ``None`` when + nothing has completed yet (live; recomputed per access). + """ + + def __init__( + self, + plan: FanOutPlan[_Item], + fetch: _Fetch[_Item], + retry_policy: RetryPolicy = _NO_RETRY, + finalize: _Finalize = _passthrough_result, + client_options: dict[str, Any] | None = None, + default_concurrent: int = _CONCURRENCY_DEFAULT, + *, + canonical_url: str | None = None, + service: str | None = None, + ) -> None: + self.plan = plan + self.fetch = fetch + self.retry_policy = retry_policy + self.finalize = finalize + self.canonical_url = canonical_url + # Label for the progress line :meth:`resume` opens. It lives here, next + # to ``canonical_url``, because this class is what emits the progress + # events — see :meth:`resume`. + self.service = service + # This service's preferred cap when the user has not set + # ``API_USGS_CONCURRENT``. Resolved at resume time, not here, so a + # test's ``monkeypatch.setenv`` still applies. See + # :func:`_resolve_concurrency` for why the env var outranks it. + self.default_concurrent = default_concurrent + # Extra ``httpx.AsyncClient`` options merged into the shared client this + # run opens (``verify`` for the Water Use ``ssl_check`` flag, say). The + # executor owns client lifecycle, so an adapter with a per-call client + # requirement has to hand it down rather than open its own — opening its + # own would defeat the shared connection pool. Empty for OGC, which + # exposes no such flag. + self.client_options = client_options or {} + # Snapshot the ambient context at construction time — i.e. inside the + # caller's ``with`` blocks (base URL, dialect, row cap, progress + # reporter). :meth:`resume` runs every drive inside this snapshot, so + # a *later* ``exc.call.resume()`` — which fires after those ``with`` + # blocks have exited and reset their ContextVars — still rebuilds + # sub-requests against the original API's base URL/dialect rather than + # the process defaults. The adapter's request builder reads those + # ContextVars when it reconstructs each sub-request, so the snapshot + # must outlive them. The mechanism is generic; which ambients matter is + # the adapter's business. + self._ctx = copy_context() + # Completed (frame, response) pairs keyed by sub-args index; sparse + # (gathered sub-requests complete out of order — see class docstring). + # ``_run``'s ``track`` closure is the only writer, so ``dict`` insertion + # order is completion order (relied on by :meth:`_combine_raw`). + self._chunks: dict[int, tuple[pd.DataFrame, httpx.Response]] = {} + + def wrap_failure(self, exc: BaseException) -> FanOutInterrupted | None: + """ + Build the matching :class:`FanOutInterrupted` carrying this + call when ``exc`` is a recognized transient transport failure; + return ``None`` for unrecognized failures so the caller can + re-raise. Encapsulates the + ``classify → instantiate-with-call-state`` recipe so + :class:`FanOut`'s private fields stay private. + + Parameters + ---------- + exc : BaseException + The exception raised by a sub-request. + + Returns + ------- + FanOutInterrupted or None + The matching :class:`FanOutInterrupted` subclass carrying this + call for a recognized transient failure; ``None`` otherwise. + """ + classification = _classify_chunk_error(exc) + if classification is None: + return None + interrupted_class, retry_after = classification + return interrupted_class( + completed_chunks=self.completed_chunks, + total_chunks=len(self.plan), + call=self, + retry_after=retry_after, + cause=exc, + ) + + def _normalize_failure(self, exc: BaseException) -> BaseException: + """Map an explicitly caused transport failure into the public taxonomy.""" + for current in _walk_causes(exc): + if isinstance(current, httpx.TransportError): + wrapped = network_error( + self.canonical_url or "unknown service", current + ) + wrapped.__cause__ = current + return wrapped + return exc + + @property + def completed_chunks(self) -> int: + """Number of sub-requests completed so far.""" + return len(self._chunks) + + def _combine_raw(self) -> tuple[pd.DataFrame, httpx.Response]: + """Assemble the raw ``(frame, response)`` from completed sub-requests, + before :attr:`finalize` runs. + + Frames concatenate in sub-args *index* order (``sorted`` keys — + deterministic, independent of parallel completion order). The + aggregated response takes its headers from the response with the + lowest reported ``x-ratelimit-remaining`` value. If no response + reports that header, it falls back to the last completed response; + ``self._chunks`` preserves completion order because the ``track`` + closure in :meth:`_run` is its only writer. + + Returns + ------- + tuple of (pandas.DataFrame, httpx.Response) + The concatenated frame and the aggregated response, before + :attr:`finalize` is applied. + """ + return self._combine_frames(), self._combine_responses() + + def _combine_frames(self) -> pd.DataFrame: + """Combine completed frames in deterministic sub-request order.""" + return _combine_chunk_frames([self._chunks[i][0] for i in sorted(self._chunks)]) + + def _combine_responses(self) -> httpx.Response: + """Aggregate completed responses under the canonical request URL.""" + responses = [response for _, response in self._chunks.values()] + return _combine_chunk_responses(responses, self.canonical_url) + + @property + def partial_frame(self) -> pd.DataFrame: + """ + Raw combined frame of sub-requests that have completed so far. + + Live — recomputed on each access so it reflects current state + across resume attempts. Deliberately the *raw* combined frame + (``_combine_frames``), NOT the finalized result: this is a cheap, + side-effect-free snapshot for inspecting partial progress, so + reading it (or building a :class:`FanOutInterrupted` around it) + never triggers ``finalize`` work — which for OGC getters includes + a schema network fetch on an empty frame. Use ``call.resume()`` + for the finalized result. + + Returns + ------- + pandas.DataFrame + Combined frame of completed sub-requests, or an empty + ``DataFrame`` when nothing has completed. + """ + return self._combine_frames() if self._chunks else pd.DataFrame() + + @property + def partial_response(self) -> httpx.Response | None: + """ + Raw aggregate response with the canonical URL restored to the + user's full original query. + + Live — recomputed on each access. Like :attr:`partial_frame`, this + is the *raw* aggregate (an :class:`httpx.Response`), not the + finalized result, so inspecting it is side-effect-free. + + Returns + ------- + httpx.Response or None + Aggregated response when at least one sub-request has + completed, ``None`` otherwise. + """ + return self._combine_responses() if self._chunks else None + + def _pending(self) -> Iterator[tuple[int, _Item]]: + """ + Yield ``(index, item)`` for sub-requests not yet completed. + + Iterates the plan in its deterministic order and skips any index + already in ``self._chunks``. :meth:`_run` uses this to pick up + exactly the sub-requests it still owes — the mechanism behind + idempotent resume. + """ + for index, item in enumerate(self.plan): + if index not in self._chunks: + yield index, item + + def resume(self) -> tuple[pd.DataFrame, Any]: + """ + Drive the call to completion and return the combined result. + + Opens the progress line for the drive and runs :meth:`_run` through + an ``anyio`` blocking portal (a short-lived worker thread), so it + works whether or not the caller is already inside an event loop + (Jupyter / IPython / async apps). The portal copies the calling + context, so the active progress reporter still reaches the + sub-requests. + + This executor is what emits progress events, so it is also what owns + the reporter's lifetime: an adapter that drives a ``FanOut`` gets the + line for free instead of having to remember a separate + ``with progress_context(...)`` block. A reporter already active + (a nested getter, or a caller's own context) is reused unchanged. + + Idempotent: only sub-requests whose index isn't already in + ``self._chunks`` are re-issued. Item order is the plan's own and + is deterministic, so a partial completion (sparse indices) + resumes correctly. + + Returns + ------- + df : pandas.DataFrame + Combined data from every successful sub-request. + response + The finalized aggregate — a raw :class:`httpx.Response` + (canonical URL, headers from the response with the lowest reported + remaining quota, and summed response elapsed durations) by default, + or whatever :attr:`finalize` produces (e.g. ``BaseMetadata`` for + the OGC getters). + + Raises + ------ + FanOutInterrupted + On a mid-stream transient failure — 429, 5xx, or a bare + transport error: :class:`~dataretrieval.interruptions.QuotaExhausted` + for 429, :class:`~dataretrieval.interruptions.ServiceInterrupted` + for the rest. The resumable handle is on ``exc.call`` — wait for + the underlying condition to clear and call ``exc.call.resume()`` + again. + """ + # Open the line out here, in the *calling* context, so an outer + # reporter (a nested getter, or the caller's own ``progress_context``) + # is the one found and reused; a drive that finds none gets a fresh + # line, which is what makes a resume long after the interruption + # report progress at all. + # + # Then drive inside the snapshot taken at construction (see + # ``__init__``). ``start_blocking_portal`` copies the *calling* context + # into its worker thread, and running here means that calling context + # is the snapshot — so the base URL / dialect / row cap active when the + # call was created reach the rebuilt sub-requests, even when this is a + # resume fired long after the original ``with`` blocks exited. The + # reporter is the one ambient that must NOT come from the snapshot: a + # reporter captured there belongs to a context that has since closed + # it, so this drive's reporter is republished over it. + with _progress.progress_context( + service=self.service, target_url=self.canonical_url + ): + return self._ctx.run(self._resume_in_context, _progress.current()) + + def _resume_in_context( + self, reporter: _progress.ProgressReporter | None + ) -> tuple[pd.DataFrame, Any]: + """Body of :meth:`resume`, run inside the captured context.""" + concurrency = _resolve_concurrency(self.default_concurrent) + with _progress._active(reporter), start_blocking_portal() as portal: + # ``portal.call`` returns ``Any`` because ``functools.partial`` + # erases ``_run``'s return type; restore the declared tuple. + return cast( + "tuple[pd.DataFrame, Any]", + portal.call(functools.partial(self._run, concurrency)), + ) + + async def _run(self, max_concurrent: int | None) -> tuple[pd.DataFrame, Any]: + """ + Gather every pending sub-request over one shared + :class:`httpx.AsyncClient` and return the combined, finalized result. + + Pending sub-requests (:meth:`_pending`) fan out under + ``asyncio.gather`` with ``return_exceptions=True`` so completed + sub-requests survive a sibling's transient failure. On a + recognized transient (:class:`~dataretrieval.exceptions.RateLimited`, + :class:`~dataretrieval.exceptions.ServiceUnavailable`, or a bare + ``httpx.HTTPError`` / ``httpx.InvalidURL``) a + :class:`FanOutInterrupted` subclass is raised carrying ``self`` on + ``.call``; ``exc.call.resume()`` then re-issues only the unfinished + indices through this same runner. + + The gather dispatches *every* pending sub-request at once, but an + ``asyncio.Semaphore`` caps the number of concurrent fetches at + ``N = max_concurrent`` — ``None`` lifts the cap, ``N=1`` runs them + one at a time. The connection pool is sized to the same ``N`` + (``httpx.Limits(max_connections=N, max_keepalive_connections=N)``) + so the in-flight fetches reuse keepalive connections. + + The semaphore, not the pool, is deliberately the throttle. If the + pool throttled instead, the excess sub-requests would queue + *inside* httpx waiting for a connection, and that wait counts + against the pool-acquire timeout (60 s, from ``HTTPX_ASYNC_DEFAULTS``). + A batch of slow pages that keeps every connection busy past that + window would then trip ``httpx.PoolTimeout`` on the queued tail — + a purely client-side failure that consumes the retry budget and + surfaces as a spurious resumable ``ServiceInterrupted``. Holding + sub-requests at the semaphore keeps them out of the pool until a + slot frees, so the pool timeout only fires for a genuinely stuck + connection. + + The shared client is published on :data:`_active_client` so + the paginated-loop helpers reuse its connection pool. + + Parameters + ---------- + max_concurrent : int or None + Maximum sub-requests in flight (the semaphore value, and the + connection-pool size). ``None`` lifts the cap entirely. + + Returns + ------- + df : pandas.DataFrame + Combined data from every sub-request. + response + The finalized aggregate — a raw :class:`httpx.Response` + (canonical URL, headers from the response with the lowest reported + remaining quota, and summed response elapsed durations) by default, + or whatever :attr:`finalize` produces. + + Raises + ------ + FanOutInterrupted + On a transient sub-request failure. ``.call`` is ``self``, + holding the sparse completed sub-requests; ``.call.resume()`` + re-issues the unfinished ones. + """ + # The semaphore is the throttle; the pool is merely sized to match + # it. Left at httpx's default client limits (``max_connections=100``, + # keepalive 20) the pool would bottleneck a wider cap or churn + # connections by keeping too few alive. See the method docstring for + # why the gate can't be the pool itself. ``unbounded`` + # (``max_concurrent=None``) is a degenerate cap at the plan total — a + # semaphore that can never block — so gated is the only code path. + limits = httpx.Limits( + max_connections=max_concurrent, max_keepalive_connections=max_concurrent + ) + semaphore = asyncio.Semaphore( + len(self.plan) if max_concurrent is None else max_concurrent + ) + + async with open_async_client(limits=limits, **self.client_options) as client: + with _active_client(client): + reporter = _progress.current() + if reporter is not None: + reporter.set_chunks(len(self.plan)) + + async def track( + index: int, item: _Item + ) -> tuple[pd.DataFrame, httpx.Response]: + """One sub-request (with retry) + result-store + progress tick.""" + result = await _retry( + lambda: self.fetch(item), self.retry_policy, gate=semaphore + ) + self._chunks[index] = result + if reporter is not None: + # Chunks finish out of order under gather, so tick the + # completed *count* rather than a positional index. + reporter.start_chunk(self.completed_chunks) + return result + + # Dispatch every pending sub-request concurrently; the + # semaphore (held by ``_retry`` per attempt) is the only throttle. + # ``return_exceptions`` keeps completed pairs after a sibling + # fails, so partial state stays recoverable via :meth:`resume`. + # Failure precedence, in order: + # 1. Cancellation / interrupt signals (CancelledError, + # KeyboardInterrupt, SystemExit — non-Exception) propagate + # unmodified; wrapping them as a transient would swallow + # the user's stop signal. + # 2. A non-transient failure (a real bug — unrecognized by + # ``wrap_failure``) surfaces raw, so it isn't masked behind + # a resumable handle for a transient sibling that landed + # later. + # 3. Only when every failure is a recognized transient do we + # raise the first as a resumable ``FanOutInterrupted``. + results = await asyncio.gather( + *(track(index, item) for index, item in self._pending()), + return_exceptions=True, + ) + failures = [r for r in results if isinstance(r, BaseException)] + for exc in failures: + if not isinstance(exc, Exception): + raise exc + # Classify first, build once. Every failure has to be + # examined -- a non-transient sibling must surface raw -- but + # only the first transient is ever raised. Asking + # ``wrap_failure`` per failure would snapshot the combined + # frame N times (a full concat over every completed + # sub-request) and discard all but one, which a batch of + # sub-requests failing together makes routine. + first_transient: BaseException | None = None + for exc in failures: + if _classify_chunk_error(exc) is None: + raise self._normalize_failure(exc) + if first_transient is None: + first_transient = exc + if first_transient is not None: + interrupted = self.wrap_failure(first_transient) + if interrupted is None: + # Unreachable: classified as transient just above. + raise self._normalize_failure(first_transient) + raise interrupted from first_transient + + return self.finalize(*self._combine_raw()) + + +__all__ = [ + "FanOut", + "FanOutPlan", + "active_client", +] diff --git a/dataretrieval/transport/links.py b/dataretrieval/transport/links.py new file mode 100644 index 000000000..c438f5492 --- /dev/null +++ b/dataretrieval/transport/links.py @@ -0,0 +1,118 @@ +"""One policy for the server-supplied next-page links every page walk follows. + +A ``next`` href is response *data*, not configuration: it arrives over the wire +from the service (or from whatever answered for it) and then becomes the URL of +our next request, carrying that request's headers and API key. Three page walks +-- the OGC engine's ``links`` array, the ratings STAC walk, and Water Use's +``Link:`` header -- each need the same things of it before it is trusted: parse +it, resolve a relative reference against the page it came from, refuse a host +the caller never asked for, and drop any embedded ``user:pass@`` (which +``httpx`` would otherwise turn into an ``Authorization: Basic`` header). + +They had three implementations of that policy, and the three disagreed: only two +resolved relative references, only two refused an unparseable link rather than +handing it back, and each worded its refusal differently. A security invariant +with three spellings is one that gets fixed in one place and stays broken in the +other two, so it lives here once and the walks pass in what genuinely differs -- +which hosts are acceptable, and whether an accepted host is rewritten. +""" + +from __future__ import annotations + +import httpx + +from dataretrieval.credentials import without_embedded_credentials +from dataretrieval.exceptions import DataRetrievalError + +__all__ = ["resolve_next_url"] + + +def _page_url(response: httpx.Response) -> httpx.URL: + """The URL *response* came from, as an ``httpx.URL``. + + Read lazily by :func:`resolve_next_url` (see there). The coercion is for + callers holding a response-shaped stand-in whose ``url`` is a plain string. + """ + url = response.url + return url if isinstance(url, httpx.URL) else httpx.URL(str(url)) + + +def resolve_next_url( + href: str, + response: httpx.Response, + *, + service: str, + allowed_hosts: frozenset[str] | None = None, + rewrite_host: str | None = None, + error: type[Exception] = DataRetrievalError, +) -> str: + """Return *href* as a URL safe to request, or raise if it is not. + + ``response.url`` is consulted only when it is actually needed -- to resolve a + relative reference, or as the default acceptable host. A walk that names its + own acceptable hosts and receives an absolute link never touches it, which + keeps this usable on a response whose request was never attached. + + Parameters + ---------- + href : str + The next-page link exactly as the service supplied it. + response : httpx.Response + The page the link arrived on. + service : str + Name of the service, used in the error messages (e.g. ``"ratings"``). + allowed_hosts : frozenset of str, optional + Hosts the link may name. Defaults to just the responding host; pass a + wider set only where the service is known to spell its own host several + ways. + rewrite_host : str, optional + Rewrite an accepted link to this host over ``https``, dropping any + explicit port. For a service whose links name a spelling of the host + that does not serve the API. + error : type of Exception, optional + Exception type to raise. Defaults to + :class:`~dataretrieval.exceptions.DataRetrievalError`; the OGC engine + passes ``RuntimeError`` to keep the type it has always raised, until + retyping it is a deliberate, released decision. + + Returns + ------- + str + An absolute URL on an acceptable host, carrying no embedded credentials. + """ + try: + target = httpx.URL(href) + except (httpx.InvalidURL, TypeError) as exc: + raise error( + f"The {service} service returned an unusable next-page link: " + f"{href!r}. The page walk cannot continue; report this if it " + f"persists." + ) from exc + if not target.is_absolute_url: + target = _page_url(response).join(target) + expected = ( + allowed_hosts + if allowed_hosts is not None + else frozenset({_page_url(response).host}) + ) + if target.host not in expected: + raise error( + f"Refusing to follow a cross-host next-page link: the {service} " + f"response points at {target.host} rather than " + f"{rewrite_host or ' or '.join(sorted(expected))}. Following it " + f"would send this request, and any credentials on it, to a host " + f"you did not ask for. Retrying will not help; report this if it " + f"persists." + ) + if rewrite_host is not None: + # The port goes with the scheme/host rewrite: one that went with the + # link's original scheme (``http://...:8080``) would otherwise survive + # into an https request and be dialed under TLS. ``userinfo`` goes for + # the reason below -- ``copy_with`` is doing both jobs at once here. + return str( + target.copy_with(scheme="https", host=rewrite_host, port=None, userinfo=b"") + ) + # A same-host link may still carry ``user:pass@``, which httpx turns into an + # ``Authorization: Basic`` header on the follow-up request. The host check + # passes in exactly that case, so strip it rather than trust the link. + return str(without_embedded_credentials(target)) diff --git a/dataretrieval/transport/retry.py b/dataretrieval/transport/retry.py index 577ed3347..9bf3a6edf 100644 --- a/dataretrieval/transport/retry.py +++ b/dataretrieval/transport/retry.py @@ -3,15 +3,10 @@ from __future__ import annotations import asyncio -import math -import os import random -import socket import time from collections.abc import Awaitable, Callable from dataclasses import dataclass -from datetime import datetime, timezone -from email.utils import parsedate_to_datetime from typing import NamedTuple, TypeVar import httpx @@ -22,6 +17,8 @@ NetworkError, TransientError, ) +from dataretrieval.interruptions import _deterministic_failure +from dataretrieval.transport.env import _read_env_number from dataretrieval.transport.liveness import ( credit_wait, elapsed_since_progress, @@ -50,98 +47,12 @@ # hint from waking together. Small on purpose: the server named the wait, so # jitter here decorrelates rather than extends it. _RETRY_AFTER_JITTER = 1.0 -# Resolver failures that will not resolve differently on a later attempt. The -# temporary ones (notably EAI_AGAIN -- "try again", raised while a resolver is -# still coming up, on VPN reconnect, or after a laptop wakes) are deliberately -# absent: those are worth another try. Looked up defensively because the EAI_* -# constants are platform-dependent; an unrecognized code stays retryable, since -# spending a few seconds on a retry is cheaper than dropping a recoverable call. -_PERMANENT_DNS_ERRORS = frozenset( - code - for code in ( - getattr(socket, name, None) for name in ("EAI_NONAME", "EAI_FAIL", "EAI_NODATA") - ) - if code is not None -) # Attempts the no-progress budget never withholds; see RetryPolicy.allows_wait. _STALL_EXEMPT_ATTEMPTS = 1 _STALL_TIMEOUT_ENV = "API_USGS_STALL_TIMEOUT" _STALL_TIMEOUT_DEFAULT = 60.0 _T = TypeVar("_T") -_Number = TypeVar("_Number", int, float) - - -def parse_retry_after(value: str | None) -> float | None: - """Parse a ``Retry-After`` header into seconds, or ``None`` for no usable hint. - - Both header forms mean the same thing and are treated the same way: the - seconds are returned as given, however large. A value past what a caller will - wait out inline stops the retry and surfaces a transient carrying the hint on - ``.retry_after``, so a long wait becomes the caller's decision (and, for a - chunked call, a resumable interruption) instead of being ignored. - - An over-long hint is honored rather than discarded. Dropping it would make - the client retry *harder* against a service that just asked for a long - pause, and would deny the caller the number it needs on ``.retry_after``. - Clock skew can inflate a date-form hint, but trusting one costs a - recoverable escalation while ignoring it costs hammering a service that is - already asking for room. - - A date that has *already* passed yields no hint at all rather than ``0.0``. - Read literally it says "retry now", but the likelier reading is that our - clock runs ahead of the server's -- and acting on it would re-send almost - immediately against a service that just asked for a pause. Falling back to - our own bounded backoff is right under either reading. (Delta-seconds is - clock-independent, so a literal ``Retry-After: 0`` is still honored as the - instruction it is, floored by :meth:`RetryPolicy.backoff`'s jitter.) - """ - if not value: - return None - raw = value.strip() - try: - seconds = float(raw) - except ValueError: - pass - else: - # ``inf``/``nan`` parse cleanly but poison every later comparison: an - # infinite hint would refuse retry forever and travel to the caller on - # ``.retry_after``. Treat them as no hint at all. - return max(0.0, seconds) if math.isfinite(seconds) else None - try: - retry_at = parsedate_to_datetime(raw) - except (TypeError, ValueError, OverflowError): - return None - if retry_at.tzinfo is None: - retry_at = retry_at.replace(tzinfo=timezone.utc) - delay = (retry_at - datetime.now(timezone.utc)).total_seconds() - return delay if delay > 0 else None - - -def _read_env_number( - name: str, default: _Number, cast: Callable[[str], _Number], expected: str -) -> _Number: - """Read a non-negative number from the environment, or ``default`` if unset. - - Raises :class:`~dataretrieval.exceptions.ConfigurationError` -- a - ``DataRetrievalError`` *and* a ``ValueError`` -- for an unusable value, so a - typo in the environment doesn't escape a request path as a bare - ``ValueError`` that ``except DataRetrievalError`` misses. - """ - raw = os.environ.get(name, "").strip() - if not raw: - return default - try: - value = cast(raw) - except ValueError as exc: - raise ConfigurationError(f"{name} must be {expected} (got {raw!r}).") from exc - # ``nan`` passes every ordering test, so a bare ``< 0`` guard lets it through - # and then silently makes each budget comparison false. - if not math.isfinite(value): - raise ConfigurationError(f"{name} must be {expected} (got {raw!r}).") - if value < 0: - raise ConfigurationError(f"{name} must be >= 0 (got {value}).") - return value @dataclass(frozen=True) @@ -294,44 +205,6 @@ def backoff(self, attempt: int, retry_after: float | None) -> float: _NO_RETRY = RetryPolicy(max_retries=0) -def _deterministic_failure(exc: BaseException) -> bool: - """Whether a transport failure would fail identically on every retry. - - An unsupported scheme or a request we built wrong is settled before a byte - goes out, and a hostname the resolver rejects outright won't be accepted on - the next attempt either -- so retrying only delays the error the caller - needs. A *temporary* resolver failure is not in that class and stays - retryable (see :data:`_PERMANENT_DNS_ERRORS`). - - The original failure is several layers down and not always an explicit - ``raise ... from``: a DNS failure reaches us as ``NetworkError`` -> - ``httpx.ConnectError`` -> ``httpcore.ConnectError`` -> ``socket.gaierror``, - linked by ``__context__`` (implicit chaining) rather than ``__cause__``. - - Both links of every frame are visited, not just the first one present. A - frame can carry an explicit ``__cause__`` *and* an unrelated ``__context__`` - (any ``raise X from Y`` inside an ``except`` block produces exactly that), so - following only the cause would walk off down the explicit branch and miss a - ``gaierror`` sitting on the implicit one -- spending the whole retry budget - on a hostname that will never resolve. The ``seen`` set keeps a chain that - rejoins itself, or points back at an ancestor, from looping. - """ - seen: set[int] = set() - pending: list[BaseException | None] = [exc] - while pending: - current = pending.pop() - if current is None or id(current) in seen: - continue - seen.add(id(current)) - if isinstance(current, (httpx.UnsupportedProtocol, httpx.LocalProtocolError)): - return True - if isinstance(current, socket.gaierror): - # Return, not continue: the first resolver code found settles the chain. - return current.errno in _PERMANENT_DNS_ERRORS - pending += [current.__cause__, current.__context__] - return False - - def _retryable( exc: BaseException, statuses: frozenset[int] = _RETRYABLE_STATUSES ) -> tuple[bool, float | None]: diff --git a/dataretrieval/transport/sync.py b/dataretrieval/transport/sync.py deleted file mode 100644 index f9799d2a4..000000000 --- a/dataretrieval/transport/sync.py +++ /dev/null @@ -1,29 +0,0 @@ -"""Synchronous dispatch over asynchronous retrieval internals.""" - -from __future__ import annotations - -from collections.abc import Awaitable, Callable -from typing import TypeVar, cast - -import httpx -from anyio.from_thread import start_blocking_portal - -from dataretrieval import progress as _progress -from dataretrieval.transport.http import network_error - -_T = TypeVar("_T") - - -def run_sync( - make_coro: Callable[[], Awaitable[_T]], - *, - service: str, - error_url: str | httpx.URL, -) -> _T: - """Run an async retrieval from synchronous code in a blocking portal.""" - with _progress.progress_context(service=service, target_url=error_url): - with start_blocking_portal() as portal: - try: - return cast("_T", portal.call(make_coro)) - except httpx.TransportError as exc: - raise network_error(error_url, exc) from exc diff --git a/dataretrieval/utils.py b/dataretrieval/utils.py index 59aeca15c..5e6bc09fd 100644 --- a/dataretrieval/utils.py +++ b/dataretrieval/utils.py @@ -1,82 +1,40 @@ -"""Useful utilities for data munging.""" +"""Data-shaping helpers, and the historical home of the legacy query path. + +What is *defined* here is frame munging that names no service: building a UTC +datetime column out of the separate date/time/zone columns a caller points at. +The one-shot HTTP query path that used to sit alongside it now lives in +:mod:`dataretrieval._querying`, and the WQX3 / legacy-WQP column conventions +live in :mod:`dataretrieval._wqx`; nothing here depends on either -- the names +below are re-exported so their documented ``dataretrieval.utils`` paths keep +resolving. + +By default, do not add new service-specific behavior here. +""" from __future__ import annotations import warnings -from collections.abc import Callable, Iterable -from typing import Any -import httpx import pandas as pd -import dataretrieval.credentials as _credentials +import dataretrieval._querying as _querying import dataretrieval.transport.http as _transport_http from dataretrieval._ambient import Ambient # noqa: F401 - compatibility re-export from dataretrieval._response_metadata import ( BaseMetadata, # noqa: F401 — compatibility re-export; defined there now ) from dataretrieval.codes import tz -from dataretrieval.exceptions import ( - NoSitesError, - URLTooLong, - error_for_status, -) -from dataretrieval.transport.retry import ( - _GATEWAY_STATUSES, - RetryPolicy, - parse_retry_after, - retry_sync, -) # Compatibility names retained at their historical utility paths. -_AUTHORIZED_API_KEY_HOST = _credentials._AUTHORIZED_API_KEY_HOST -HTTPX_ASYNC_DEFAULTS = _transport_http.HTTPX_ASYNC_DEFAULTS HTTPX_DEFAULTS = _transport_http.HTTPX_DEFAULTS USER_AGENT = _transport_http.USER_AGENT _default_headers = _transport_http.default_headers _get = _transport_http.get _network_error = _transport_http.network_error -_strip_api_key_from_untrusted_host = _transport_http.strip_api_key_from_untrusted_host -_strip_api_key_from_untrusted_host_async = ( - _transport_http.strip_api_key_from_untrusted_host_async -) - - -def to_str(listlike: object, delimiter: str = ",") -> str | None: - """Translate a list-like object into a delimited string. - - Parameters - ---------- - listlike: list-like object - A list, or a list-like object - (e.g. ``pandas.core.series.Series``). - delimiter: string, optional - String placed between entries of ``listlike`` when it is turned into a - string. Default value is a comma. - - Returns - ------- - listlike: string - The listlike object as a string separated by the delimiter. - - Examples - -------- - .. doctest:: - - >>> dataretrieval.utils.to_str([1, "a", 2]) - '1,a,2' - - >>> dataretrieval.utils.to_str([0, 10, 42], delimiter="+") - '0+10+42' - - """ - if isinstance(listlike, str): - return listlike - - if isinstance(listlike, Iterable): - return delimiter.join(map(str, listlike)) - - return None +# Public functions whose implementation moved to the private query module; this +# is the path they are documented at. +query = _querying.query +to_str = _querying.to_str def format_datetime( @@ -123,275 +81,3 @@ def format_datetime( ) return df - - -# (time-suffix, tz-suffix) pairs that follow a "Date" column. -_TIME_TZ_SUFFIXES = ( - # WQX3 / Samples, e.g. - # Activity_StartDate / Activity_StartTime / Activity_StartTimeZone - ("Time", "TimeZone"), - # Legacy WQP (slash-separated), e.g. - # ActivityStartDate / ActivityStartTime/Time / ActivityStartTime/TimeZoneCode - ("Time/Time", "Time/TimeZoneCode"), -) - - -def _build_utc_datetime( - date_series: pd.Series, time_series: pd.Series, tz_series: pd.Series -) -> pd.Series: - """Combine date + time + tz-abbreviation columns into a UTC pandas Series. - - Unknown timezone codes (and rows missing any of the three values) yield - ``NaT``. The input columns are not mutated. - """ - offsets = tz_series.map(tz) - combined = ( - date_series.astype("string") - + " " - + time_series.astype("string") - + " " - + offsets.astype("string") - ) - return pd.to_datetime( - combined, format="%Y-%m-%d %H:%M:%S %z", utc=True, errors="coerce" - ) - - -def _attach_datetime_columns(df: pd.DataFrame) -> pd.DataFrame: - """Append a UTC ``DateTime`` column per Date/Time/TimeZone triplet. - - Detects two naming patterns that appear in USGS Samples and Water Quality - Portal CSV responses: - - * **WQX3** — ``Date``, ``Time``, ``TimeZone`` - * **Legacy WQP** — ``Date``, ``Time/Time``, - ``Time/TimeZoneCode`` - - For every triplet present, a new ``DateTime`` column is appended - holding a UTC ``Timestamp`` (offsets resolved via - :data:`dataretrieval.codes.tz`). The original Date/Time/TimeZone columns - are left intact, and an existing ``DateTime`` column is never - overwritten. - - Rows are sorted (and the index reset) by the canonical activity-start - datetime when present — ``Activity_StartDateTime`` (WQX3) or - ``ActivityStartDateTime`` (legacy WQP) — falling back to the first - detected ``*Date`` column. Mirrors R ``dataRetrieval``'s - end-of-pipeline sort in ``importWQP.R``. - - Parameters - ---------- - df : ``pandas.DataFrame`` - DataFrame returned from a Samples or WQP CSV endpoint. - - Returns - ------- - df : ``pandas.DataFrame`` - A new DataFrame with derivable ``DateTime`` columns appended - and rows sorted by the activity-start datetime (if any date column - was detected). - """ - columns = set(df.columns) - new_columns = {} - first_date_col = None - for col in df.columns: - if not col.endswith("Date"): - continue - if first_date_col is None: - first_date_col = col - prefix = col.removesuffix("Date") - target = prefix + "DateTime" - if target in columns or target in new_columns: - continue - for time_suffix, tz_suffix in _TIME_TZ_SUFFIXES: - time_col = prefix + time_suffix - tz_col = prefix + tz_suffix - if time_col in columns and tz_col in columns: - new_columns[target] = _build_utc_datetime( - df[col], df[time_col], df[tz_col] - ) - break - if new_columns: - # Concat in one shot — per-column assignment on a wide CSV-derived - # frame triggers pandas' fragmentation PerformanceWarning. - df = pd.concat([df, pd.DataFrame(new_columns, index=df.index)], axis=1) - sort_key: str | None - if "Activity_StartDateTime" in df.columns: - sort_key = "Activity_StartDateTime" - elif "ActivityStartDateTime" in df.columns: - sort_key = "ActivityStartDateTime" - else: - sort_key = first_date_col - if sort_key is not None: - df = df.sort_values(by=sort_key, ignore_index=True) - return df - - -_URL_TOO_LONG_EXAMPLE = """ - # n is the number of chunks to divide the query into \n - split_list = np.array_split(site_list, n) - data_list = [] # list to store chunk results in \n - # loop through chunks and make requests \n - for site_list in split_list: \n - data = nwis.get_record(sites=site_list, service='dv', \n - start=start, end=end) \n - data_list.append(data) # append results to list""" - - -def _url_too_long_error(detail: str) -> URLTooLong: - return URLTooLong( - "Request URL too long. Modify your query to use fewer sites. " - f"{detail}. Pseudo-code example of how to split your query: " - f"\n {_URL_TOO_LONG_EXAMPLE}" - ) - - -def _raise_for_status( - response: httpx.Response, - *, - detail_from: Callable[[httpx.Response], str | None] | None = None, -) -> None: - """Raise the typed :class:`DataRetrievalError` for an HTTP error response. - - A success status returns ``None``. Shared by the legacy :func:`query` path - (and ``streamstats`` / ``wateruse``). Delegates the status-to-type mapping to - :func:`dataretrieval.exceptions.error_for_status`, except a too-long-URL - status (413 / 414): that gets the same actionable "split your query" - remediation as the client-side over-long-URL case below, rather than a bare - ``HTTP 414`` (both still raise :class:`~dataretrieval.exceptions.URLTooLong`). - - ``detail_from``, when given, is called *only on an error response* to pull an - API-specific detail string (e.g. a JSON error envelope's message) out of the - body; a truthy result is appended to the raised message. This lets callers - surface their API's error wording without re-implementing the status-to-type - mapping and message format. - """ - status = response.status_code - if status < 400: - return - if status in (413, 414): - raise _url_too_long_error(f"API response reason: {response.reason_phrase}") - message = f"HTTP {status} {response.reason_phrase}".rstrip() - detail = detail_from(response) if detail_from is not None else None - if detail: - message += f": {detail}" - message += f" (URL: {response.url})" - raise error_for_status( - status, - message, - retry_after=parse_retry_after(response.headers.get("Retry-After")), - ) - - -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_with_retry( - url: str, - payload: dict[str, Any], - delimiter: str = ",", - ssl_check: bool = True, - *, - retry_policy: RetryPolicy | None = None, -) -> httpx.Response: - """Send an active-service query with bounded transient retry by default.""" - - for key, value in payload.items(): - payload[key] = to_str(value, delimiter) - # httpx serializes None params as ``foo=``; USGS rejects with 400. - # 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": USER_AGENT} - - 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. - if response.text.startswith("No sites/data"): - raise NoSitesError(response.url) - - return response - - -def query( - url: str, - payload: dict[str, Any], - delimiter: str = ",", - ssl_check: bool = True, -) -> httpx.Response: - """Send a query. - - Wrapper for ``httpx.get`` that handles errors, converts listed query - parameters to comma-separated strings, and returns the response. - - Parameters - ---------- - url: string - URL to query. - payload: dict - Query parameters passed to ``httpx.get``. - delimiter: string - Delimiter to use with lists. - ssl_check: bool - Whether to check SSL certificates. Default is True. - - Returns - ------- - response: ``httpx.Response`` - The response from the API query ``httpx.get`` function call. - - Raises - ------ - DataRetrievalError - On an HTTP error response, the typed subclass for the status (see - :func:`dataretrieval.exceptions.error_for_status` for the mapping); or - :class:`~dataretrieval.exceptions.NoSitesError` when a 200 response - reports no data matched; or :class:`~dataretrieval.exceptions.NetworkError` - on a connection-level failure (timeout, DNS), with the underlying - ``httpx`` exception on ``__cause__``. - """ - return _query_with_retry( - url, - payload, - delimiter, - ssl_check, - retry_policy=RetryPolicy(max_retries=0), - ) diff --git a/dataretrieval/waterdata/cql.py b/dataretrieval/waterdata/cql.py index 8ba658fb1..f83a1f2b3 100644 --- a/dataretrieval/waterdata/cql.py +++ b/dataretrieval/waterdata/cql.py @@ -10,25 +10,31 @@ import json from collections.abc import Iterable -from typing import Any +from typing import TYPE_CHECKING, Any import pandas as pd -from dataretrieval._response_metadata import BaseMetadata from dataretrieval.ogc import fetch_ogc_request +from dataretrieval.ogc.context import _ogc_base_url from dataretrieval.ogc.requests import ( _as_str_list, _construct_cql_request, _switch_properties_id, ) -from dataretrieval.waterdata.types import ( - WATERDATA_SERVICES, -) +from dataretrieval.ogc.shaping import _finalize_ogc from dataretrieval.waterdata.utils import ( + _EXTRA_ID_COLS, _OUTPUT_ID_BY_SERVICE, - _finalize_ogc, + OGC_API_URL, + WATERDATA_DIALECT, ) +if TYPE_CHECKING: + from dataretrieval._response_metadata import BaseMetadata + from dataretrieval.waterdata.types import ( + WATERDATA_SERVICES, + ) + def get_cql( service: WATERDATA_SERVICES, @@ -148,25 +154,31 @@ def get_cql( # downstream), matching the typed getters. wire_properties = _switch_properties_id(properties_list, output_id, service) - req = _construct_cql_request( - service, - body, - properties=wire_properties, - bbox=bbox, - limit=limit, - skip_geometry=skip_geometry, - ) - - df, response = fetch_ogc_request(req, service=service) + # The OGC package names no service of its own, so this hand-built request + # path states the target itself -- request construction and the empty-result + # schema lookup in ``_finalize_ogc`` both read the base URL from here. + with _ogc_base_url(OGC_API_URL): + req = _construct_cql_request( + service, + body, + properties=wire_properties, + bbox=bbox, + limit=limit, + skip_geometry=skip_geometry, + ) - return _finalize_ogc( - df, - response, - properties=properties_list, - output_id=output_id, - convert_type=convert_type, - service=service, - ) + df, response = fetch_ogc_request(req, service=service) + + return _finalize_ogc( + df, + response, + properties=properties_list, + output_id=output_id, + convert_type=convert_type, + service=service, + extra_id_cols=_EXTRA_ID_COLS, + dialect=WATERDATA_DIALECT, + ) __all__ = ["get_cql"] diff --git a/dataretrieval/waterdata/endpoints.py b/dataretrieval/waterdata/endpoints.py new file mode 100644 index 000000000..7611d1204 --- /dev/null +++ b/dataretrieval/waterdata/endpoints.py @@ -0,0 +1,41 @@ +"""Every Water Data endpoint this package talks to, in one place. + +"Which services does Water Data reach, and at what URL" is a single question +with a single answer, so the answer lives in one file rather than being spelled +out again in each family module. The host is the authority of the credentials +leaf -- the host that serves these endpoints is the host that honors the API +key -- while the paths below stay here rather than importing OGC policy +internals. + +This module imports nothing but that leaf, so a family module can name its +endpoint without also taking on an OGC or transport edge. +""" + +from __future__ import annotations + +from dataretrieval.credentials import WATERDATA_BASE_URL + +#: Root of the modernized Water Data APIs. +BASE_URL = WATERDATA_BASE_URL + +#: OGC API - Features service backing the typed collection getters. +OGC_API_URL = f"{BASE_URL}/ogcapi/v0" + +#: Samples database (discrete water-quality results, WQX3 CSV). +SAMPLES_URL = f"{BASE_URL}/samples-data" + +#: Daily-statistics service (period-of-record and date-range normals). +STATISTICS_API_VERSION = "v0" +STATISTICS_API_URL = f"{BASE_URL}/statistics/{STATISTICS_API_VERSION}" + +#: STAC catalog serving NWIS rating-curve assets. +STAC_URL = f"{BASE_URL}/stac/v0" + +__all__ = [ + "BASE_URL", + "OGC_API_URL", + "SAMPLES_URL", + "STAC_URL", + "STATISTICS_API_URL", + "STATISTICS_API_VERSION", +] diff --git a/dataretrieval/waterdata/measurements.py b/dataretrieval/waterdata/measurements.py index 21d5a207f..e3aec5032 100644 --- a/dataretrieval/waterdata/measurements.py +++ b/dataretrieval/waterdata/measurements.py @@ -8,17 +8,19 @@ from __future__ import annotations from collections.abc import Iterable -from typing import Any +from typing import TYPE_CHECKING, Any import pandas as pd -from dataretrieval._response_metadata import BaseMetadata -from dataretrieval.ogc.filters import FILTER_LANG from dataretrieval.waterdata.utils import ( _get_args, get_ogc_data, ) +if TYPE_CHECKING: + from dataretrieval._response_metadata import BaseMetadata + from dataretrieval.ogc.filters import FILTER_LANG + def get_field_measurements( monitoring_location_id: str | Iterable[str] | None = None, diff --git a/dataretrieval/waterdata/metadata.py b/dataretrieval/waterdata/metadata.py index 0c99be544..c3e88ea02 100644 --- a/dataretrieval/waterdata/metadata.py +++ b/dataretrieval/waterdata/metadata.py @@ -9,18 +9,20 @@ from __future__ import annotations from collections.abc import Iterable -from typing import Any +from typing import TYPE_CHECKING, Any import pandas as pd -from dataretrieval._response_metadata import BaseMetadata -from dataretrieval.ogc.filters import FILTER_LANG from dataretrieval.waterdata.utils import ( _get_args, _with_state, get_ogc_data, ) +if TYPE_CHECKING: + from dataretrieval._response_metadata import BaseMetadata + from dataretrieval.ogc.filters import FILTER_LANG + def get_monitoring_locations( monitoring_location_id: str | Iterable[str] | None = None, diff --git a/dataretrieval/waterdata/nearest.py b/dataretrieval/waterdata/nearest.py index 13b4e4f21..b9c6b498d 100644 --- a/dataretrieval/waterdata/nearest.py +++ b/dataretrieval/waterdata/nearest.py @@ -6,13 +6,15 @@ from __future__ import annotations from collections.abc import Iterable -from typing import Any, Literal, get_args +from typing import TYPE_CHECKING, Any, Literal, get_args import pandas as pd -from dataretrieval._response_metadata import BaseMetadata from dataretrieval.waterdata.time_series import get_continuous +if TYPE_CHECKING: + from dataretrieval._response_metadata import BaseMetadata + __all__ = ["get_nearest_continuous"] diff --git a/dataretrieval/waterdata/ratings.py b/dataretrieval/waterdata/ratings.py index 9274067c3..ff852029e 100644 --- a/dataretrieval/waterdata/ratings.py +++ b/dataretrieval/waterdata/ratings.py @@ -17,7 +17,6 @@ import httpx import pandas as pd -from dataretrieval.credentials import WATERDATA_BASE_URL, without_embedded_credentials from dataretrieval.exceptions import DataRetrievalError from dataretrieval.ogc.dates import _DURATION_RE, _format_api_dates from dataretrieval.ogc.errors import _raise_for_non_200 @@ -33,14 +32,14 @@ from dataretrieval.transport.http import ( get as _get, ) +from dataretrieval.transport.links import resolve_next_url +from dataretrieval.waterdata.endpoints import STAC_URL __all__ = ["get_ratings"] logger = logging.getLogger(__name__) -STAC_URL = f"{WATERDATA_BASE_URL}/stac/v0" - RATING_FILE_TYPE = Literal["exsa", "base", "corr"] _VALID_FILE_TYPES = get_args(RATING_FILE_TYPE) @@ -245,28 +244,6 @@ def _search( one page isn't silently truncated. """ - def _checked_next_url(href: str, current: httpx.URL) -> str: - """Resolve a server-supplied ``next`` href into a URL safe to request.""" - try: - target = httpx.URL(href) - except (httpx.InvalidURL, TypeError) as exc: - raise DataRetrievalError( - f"The ratings service returned an unusable next-page link: " - f"{href!r}. The page walk cannot continue; report this if it " - f"persists." - ) from exc - if not target.is_absolute_url: - target = current.join(target) - if target.host != current.host: - raise DataRetrievalError( - f"Refusing to follow a ratings next-page link pointing at " - f"{target.host} rather than {current.host}. Following it would " - f"send this request, and any credentials on it, to a host you " - f"did not ask for. Retrying will not help; report this if it " - f"persists." - ) - return str(without_embedded_credentials(target)) - query_params: dict[str, Any] = {"limit": min(limit, 10000)} if filter_str is not None: query_params["filter"] = filter_str @@ -303,7 +280,11 @@ def _checked_next_url(href: str, current: httpx.URL) -> str: # carry this request's API key off the authorized host, and one carrying # ``user:pass@`` would mint an ``Authorization: Basic`` header the caller # never configured. - url = None if href is None else _checked_next_url(href, response.url) + url = ( + None + if href is None + else resolve_next_url(href, response, service="ratings") + ) params = None return features diff --git a/dataretrieval/waterdata/reference.py b/dataretrieval/waterdata/reference.py index adb5df4fd..7db5fb803 100644 --- a/dataretrieval/waterdata/reference.py +++ b/dataretrieval/waterdata/reference.py @@ -8,19 +8,22 @@ from __future__ import annotations -from typing import Any, get_args +from typing import TYPE_CHECKING, Any, get_args import pandas as pd -from dataretrieval._response_metadata import BaseMetadata -from dataretrieval.ogc.schema import _check_ogc_requests +from dataretrieval.ogc.schema import queryables_frame from dataretrieval.waterdata.types import ( METADATA_COLLECTIONS, ) from dataretrieval.waterdata.utils import ( + OGC_API_URL, get_ogc_data, ) +if TYPE_CHECKING: + from dataretrieval._response_metadata import BaseMetadata + def get_reference_table( collection: str, @@ -156,23 +159,9 @@ def get_queryables(collection: str) -> tuple[pd.DataFrame, BaseMetadata]: >>> df.set_index("queryable").loc["state_name", "type"] 'string' """ - # The OGC queryables document is a JSON Schema whose ``properties`` map each - # filterable property name to a ``{title, type, description}`` definition. - body, response = _check_ogc_requests(endpoint=collection, req_type="queryables") - properties: dict[str, Any] = body.get("properties", {}) - df = pd.DataFrame( - [ - { - "queryable": name, - "type": prop.get("type"), - "title": prop.get("title"), - "description": (prop.get("description") or "").strip(), - } - for name, prop in sorted(properties.items()) - ], - columns=["queryable", "type", "title", "description"], - ) - return df, BaseMetadata(response) + # Reading the queryables document is OGC protocol work; this getter only + # names the API to ask. + return queryables_frame(collection, base_url=OGC_API_URL) __all__ = ["get_reference_table", "get_queryables"] diff --git a/dataretrieval/waterdata/samples.py b/dataretrieval/waterdata/samples.py index 9738e1486..642fd05ac 100644 --- a/dataretrieval/waterdata/samples.py +++ b/dataretrieval/waterdata/samples.py @@ -18,7 +18,9 @@ import httpx import pandas as pd +from dataretrieval._querying import to_str from dataretrieval._response_metadata import BaseMetadata +from dataretrieval._wqx import _attach_datetime_columns from dataretrieval.ogc.errors import _raise_for_non_200 from dataretrieval.transport.http import ( HTTPX_DEFAULTS, @@ -29,16 +31,15 @@ from dataretrieval.transport.http import ( get as _get, ) -from dataretrieval.utils import _attach_datetime_columns, to_str from dataretrieval.waterdata.types import ( CODE_SERVICES, PROFILES, SERVICES, + _check_profiles, ) from dataretrieval.waterdata.utils import ( SAMPLES_URL, _accept_legacy_kwargs, - _check_profiles, _get_args, ) diff --git a/dataretrieval/waterdata/stats.py b/dataretrieval/waterdata/stats.py index 5fd6bb7b9..88ff751e0 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:`dataretrieval.transport.pagination.paginate` through the shared sync -bridge rather than going through ``multi_value_chunked``. The typed getters +:func:`dataretrieval.transport.pagination.paginate` as a one-item +:class:`~dataretrieval.transport.fanout.FanOut` 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. @@ -18,37 +19,22 @@ import pandas as pd from dataretrieval._response_metadata import BaseMetadata -from dataretrieval.credentials import WATERDATA_BASE_URL from dataretrieval.ogc.errors import _raise_for_non_200 from dataretrieval.ogc.shaping import ( - _CRS, GEOPANDAS, _attach_coordinates, _empty_feature_frame, + _geo_feature_frame, ) +from dataretrieval.transport.fanout import FanOut, active_client from dataretrieval.transport.http import default_headers from dataretrieval.transport.pagination import paginate -from dataretrieval.transport.sync import run_sync +from dataretrieval.transport.retry import RetryPolicy +from dataretrieval.waterdata.endpoints import STATISTICS_API_URL __all__ = ["get_data"] -# ``_handle_nesting``'s geopandas branch calls ``gpd.GeoDataFrame.from_features`` -# directly, so this module needs its own bound ``gpd`` name. Import it under the -# same guard the engine uses; when geopandas is absent ``gpd`` is left unbound -# (``GEOPANDAS`` is ``False``, so the stats path never touches it). The -# empty-page short-circuit instead delegates to ``shaping._empty_feature_frame``, -# which resolves ``shaping``'s ``gpd`` — so an empty-page test patches -# ``shaping.gpd`` while the populated geopandas branch uses ``stats.gpd``. -try: - import geopandas as gpd -except ImportError: # pragma: no cover - exercised only without geopandas - pass - -STATISTICS_API_VERSION = "v0" -STATISTICS_API_URL = f"{WATERDATA_BASE_URL}/statistics/{STATISTICS_API_VERSION}" - - def _handle_nesting( body: dict[str, Any], geopd: bool = False, @@ -111,14 +97,9 @@ def _handle_nesting( # consistent by NOT adding an id column. _attach_coordinates(df, features) else: - # Default a missing ``geometry`` key to ``None`` per feature so - # ``from_features`` (which indexes ``feature["geometry"]`` directly) - # can't ``KeyError`` on a stats feature that omits geometry — mirrors - # the guard in :func:`engine._get_resp_data`. - df = gpd.GeoDataFrame.from_features( - [f if "geometry" in f else {**f, "geometry": None} for f in features], - crs=_CRS, - ).drop(columns=["data"], errors="ignore") + # Stats features may omit ``geometry`` entirely; ``_geo_feature_frame`` + # is the shared home for that upstream-schema workaround. + df = _geo_feature_frame(features).drop(columns=["data"], errors="ignore") # Unnest json features, properties, data, and values while retaining necessary # metadata to merge with main dataframe. @@ -222,10 +203,11 @@ def get_data( parameters. The stats path doesn't go through ``multi_value_chunked`` (its query - 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. + shape has no chunkable list axes), so it drives transport pagination as a + one-item :class:`~dataretrieval.transport.fanout.FanOut`. The executor + runs the pagination loop in a short-lived worker thread, so this works + whether or not the caller is already inside an event loop, and the single + request gets the same retry and resume semantics as every other getter. Parameters ---------- @@ -240,12 +222,9 @@ def get_data( 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. 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. + Caller-borrowed async client. ``None`` (default) borrows the one this + call's :class:`~dataretrieval.transport.fanout.FanOut` opened, which + lives in the same event loop as the page walk. Primarily a test seam. Returns ------- @@ -260,9 +239,14 @@ def get_data( DataRetrievalError 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 + or :class:`~dataretrieval.exceptions.NetworkError` if the request + can't reach the service in a way retrying cannot fix (bad scheme, + a hostname that does not resolve), the ``httpx`` exception chained on ``__cause__``. + FanOutInterrupted + A transient failure (429 / 5xx / timeout) survived the built-in + retries. Resume with ``exc.call.resume()`` (see + :doc:`/userguide/errors`). """ url = f"{STATISTICS_API_URL}/{service}" @@ -289,16 +273,27 @@ async def follow_up(cursor: str, sess: httpx.AsyncClient) -> httpx.Response: method, url=url, params={**args, "next_token": cursor}, headers=headers ) - async def _run() -> tuple[pd.DataFrame, httpx.Response]: + async def _fetch(request: httpx.Request) -> tuple[pd.DataFrame, httpx.Response]: return await paginate( - req, + request, parse_response=parse_response, follow_up=follow_up, - client=client, + # Borrow the executor's shared client unless a caller injected one. + client=client if client is not None else active_client(), raise_for_status=_raise_for_non_200, ) - df, response = run_sync(_run, service=service, error_url=url) + # A one-item fan-out. Statistics has no chunkable axis, so the plan is a + # single request -- but running it through the same executor as every other + # getter is what gives it retry, the resumable interruption taxonomy, and + # the progress line, instead of a private sync bridge that had none of them. + df, response = FanOut( + [req], + _fetch, + RetryPolicy.from_env(), + canonical_url=str(req.url), + service=service, + ).resume() if expand_percentiles: df = _expand_percentiles(df) diff --git a/dataretrieval/waterdata/time_series.py b/dataretrieval/waterdata/time_series.py index 35f8d7d6d..755945497 100644 --- a/dataretrieval/waterdata/time_series.py +++ b/dataretrieval/waterdata/time_series.py @@ -12,12 +12,10 @@ from __future__ import annotations from collections.abc import Iterable -from typing import Any +from typing import TYPE_CHECKING, Any import pandas as pd -from dataretrieval._response_metadata import BaseMetadata -from dataretrieval.ogc.filters import FILTER_LANG from dataretrieval.waterdata import stats from dataretrieval.waterdata.utils import ( _get_args, @@ -25,6 +23,10 @@ get_ogc_data, ) +if TYPE_CHECKING: + from dataretrieval._response_metadata import BaseMetadata + from dataretrieval.ogc.filters import FILTER_LANG + def get_daily( monitoring_location_id: str | Iterable[str] | None = None, diff --git a/dataretrieval/waterdata/types.py b/dataretrieval/waterdata/types.py index 022627d03..ff13f42fe 100644 --- a/dataretrieval/waterdata/types.py +++ b/dataretrieval/waterdata/types.py @@ -1,4 +1,4 @@ -from typing import Literal +from typing import Literal, get_args __all__ = [ "CODE_SERVICES", @@ -102,3 +102,32 @@ "count", ], } + + +def _check_profiles( + service: SERVICES, + profile: PROFILES, +) -> None: + """Check whether a service profile is valid. + + Parameters + ---------- + service : string + One of the service names from the "services" list. + profile : string + One of the profile names from "results_profiles", + "locations_profiles", "activities_profiles", + "projects_profiles" or "organizations_profiles". + """ + valid_services = get_args(SERVICES) + if service not in valid_services: + raise ValueError( + f"Invalid service: '{service}'. Valid options are: {valid_services}." + ) + + valid_profiles = PROFILE_LOOKUP[service] + if profile not in valid_profiles: + raise ValueError( + f"Invalid profile: '{profile}' for service '{service}'. " + f"Valid options are: {valid_profiles}." + ) diff --git a/dataretrieval/waterdata/utils.py b/dataretrieval/waterdata/utils.py index db00d6245..dfacee5e4 100644 --- a/dataretrieval/waterdata/utils.py +++ b/dataretrieval/waterdata/utils.py @@ -1,7 +1,7 @@ """Water Data API layer over the generic OGC facade. This module is the Water-Data-specific adapter: it supplies the -service-to-id map, the CQL2/date-only dialect, profile validation, and a +service-to-id map, the CQL2/date-only dialect, and a thin ``get_ogc_data`` wrapper that injects the Water Data defaults. The statistics path lives in its own :mod:`dataretrieval.waterdata.stats` module. @@ -19,33 +19,20 @@ import functools import warnings from collections.abc import Callable, Mapping -from typing import Any, TypeVar, get_args +from typing import TYPE_CHECKING, Any, TypeVar -import httpx import pandas as pd -import dataretrieval.ogc.dates as _ogc_dates -import dataretrieval.ogc.shaping as _ogc_shaping -from dataretrieval._response_metadata import BaseMetadata from dataretrieval.codes.states import apply_state -from dataretrieval.credentials import WATERDATA_BASE_URL from dataretrieval.ogc import OgcDialect, prepare_request_args from dataretrieval.ogc import get_ogc_data as _facade_get_ogc_data -from dataretrieval.waterdata.types import ( - PROFILE_LOOKUP, - PROFILES, - SERVICES, -) -# --------------------------------------------------------------------------- -# Water Data endpoint constants. The authority comes from the credentials leaf -# -- the host that serves these endpoints is the host that honors the API key -- -# while the paths below stay local rather than importing OGC policy internals. -# --------------------------------------------------------------------------- +# Endpoint constants live in one place for the whole service; they are re-bound +# here because ``waterdata.utils.OGC_API_URL`` is a documented path. +from dataretrieval.waterdata.endpoints import BASE_URL, OGC_API_URL, SAMPLES_URL -BASE_URL = WATERDATA_BASE_URL -OGC_API_URL = f"{BASE_URL}/ogcapi/v0" -SAMPLES_URL = f"{BASE_URL}/samples-data" +if TYPE_CHECKING: + from dataretrieval._response_metadata import BaseMetadata # Maps each OGC waterdata service to its user-facing ``id`` column (the name the # typed getters rename the wire ``id`` to, e.g. ``daily`` -> ``daily_id``). @@ -108,23 +95,24 @@ sort_cols=("time", "monitoring_location_id"), ) -# Iterable-shaped params that ``_get_args`` must NOT push through -# ``_normalize_str_iterable`` (scalar non-string knobs are caught by runtime -# type, so only iterables with special handling need to be named here): -# - date-range params may contain ``pd.NaT``/None or interval strings -# - ``bbox``/``boundingBox`` are ``list[float]``, sometimes ``numpy.ndarray`` +# The Water-Data-specific *extras* on top of the engine's own no-normalize set +# (which already covers the date-range params and ``bbox``). Scalar non-string +# knobs are caught by runtime type, so only iterables with special handling +# need to be named here: +# - ``boundingBox`` is ``list[float]``, sometimes ``numpy.ndarray`` # - ``get_peaks``'s int-valued filters (``water_year`` etc.) are ``list[int]`` # - ``get_combined_metadata``'s ``thresholds`` is ``list[float]`` -_NO_NORMALIZE_PARAMS = _ogc_dates._DATE_RANGE_PARAMS | { - "bbox", - "boundingBox", - "water_year", - "year", - "month", - "day", - "peak_since", - "thresholds", -} +_NO_NORMALIZE_PARAMS = frozenset( + { + "boundingBox", + "water_year", + "year", + "month", + "day", + "peak_since", + "thresholds", + } +) def _flatten_queryables(local_vars: dict[str, Any]) -> dict[str, Any]: @@ -150,13 +138,15 @@ def _get_args( ) -> dict[str, Any]: """Water-Data wrapper over :func:`~dataretrieval.ogc.prepare_request_args`. - Supplies the Water Data API's extended ``no_normalize`` set (numeric - params such as ``water_year``, ``thresholds``, ``boundingBox``) so they - keep their element types. Also flattens any ``**queryables`` passthrough - (see :func:`_flatten_queryables`). + Adds the Water Data API's extra no-normalize params (numeric params such + as ``water_year``, ``thresholds``, ``boundingBox``) so they keep their + element types. Also flattens any ``**queryables`` passthrough (see + :func:`_flatten_queryables`). """ _flatten_queryables(local_vars) - return prepare_request_args(local_vars, exclude, no_normalize=_NO_NORMALIZE_PARAMS) + return prepare_request_args( + local_vars, exclude, extra_no_normalize=_NO_NORMALIZE_PARAMS + ) def _with_state(local_vars: dict[str, Any], *, to: str, into: str) -> dict[str, Any]: @@ -231,66 +221,6 @@ def get_ogc_data( ) -def _finalize_ogc( - frame: pd.DataFrame, - response: httpx.Response, - *, - properties: list[str] | None, - output_id: str, - convert_type: bool, - service: str, - max_rows: int | None = None, -) -> tuple[pd.DataFrame, BaseMetadata]: - """Water-Data wrapper over :func:`~dataretrieval.ogc.shaping._finalize_ogc`. - - Injects the Water Data ``extra_id_cols`` and ``dialect`` so a direct - call (e.g. from ``get_cql``) orders synthetic id columns and coerces/ - sorts result columns identically to the typed getters. See - :func:`~dataretrieval.ogc.shaping._finalize_ogc` for the full - result-shaping contract. - """ - return _ogc_shaping._finalize_ogc( - frame, - response, - properties=properties, - output_id=output_id, - convert_type=convert_type, - service=service, - max_rows=max_rows, - extra_id_cols=_EXTRA_ID_COLS, - dialect=WATERDATA_DIALECT, - ) - - -def _check_profiles( - service: SERVICES, - profile: PROFILES, -) -> None: - """Check whether a service profile is valid. - - Parameters - ---------- - service : string - One of the service names from the "services" list. - profile : string - One of the profile names from "results_profiles", - "locations_profiles", "activities_profiles", - "projects_profiles" or "organizations_profiles". - """ - valid_services = get_args(SERVICES) - if service not in valid_services: - raise ValueError( - f"Invalid service: '{service}'. Valid options are: {valid_services}." - ) - - valid_profiles = PROFILE_LOOKUP[service] - if profile not in valid_profiles: - raise ValueError( - f"Invalid profile: '{profile}' for service '{service}'. " - f"Valid options are: {valid_profiles}." - ) - - _R = TypeVar("_R") @@ -355,8 +285,6 @@ def wrapper(*args: Any, **kwargs: Any) -> _R: "_NO_NORMALIZE_PARAMS", "_OUTPUT_ID_BY_SERVICE", "_accept_legacy_kwargs", - "_check_profiles", - "_finalize_ogc", "_get_args", "_with_state", "get_ogc_data", diff --git a/dataretrieval/wateruse.py b/dataretrieval/wateruse.py index 2351ef07a..2ab24be21 100644 --- a/dataretrieval/wateruse.py +++ b/dataretrieval/wateruse.py @@ -41,7 +41,6 @@ from __future__ import annotations -import asyncio import io from collections.abc import Callable, Iterable from typing import Any @@ -49,28 +48,24 @@ import httpx import pandas as pd +from dataretrieval._querying import _raise_for_status, to_str from dataretrieval._response_metadata import BaseMetadata from dataretrieval.codes.states import to_state -from dataretrieval.combining import ( - _combine_chunk_frames, - _combine_chunk_responses, -) from dataretrieval.exceptions import DataRetrievalError -from dataretrieval.transport.http import default_headers, open_async_client +from dataretrieval.transport.fanout import FanOut, active_client +from dataretrieval.transport.http import default_headers, network_error +from dataretrieval.transport.links import resolve_next_url 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 _raise_for_status, to_str +from dataretrieval.transport.retry import RetryPolicy __all__ = [ "get_wateruse", "WATERUSE_URL", "MODELS", "TIME_RESOLUTIONS", - "MAX_CONCURRENT_REQUESTS", + "DEFAULT_CONCURRENT_REQUESTS", ] - 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 @@ -90,13 +85,15 @@ #: Temporal resolutions: monthly, annual calendar year, annual water year. TIME_RESOLUTIONS = ("monthly", "annualcy", "annualwy") -#: Maximum locations fetched concurrently when a list of state/county/huc -#: selectors is fanned out (one request per location). Kept conservative -#: because every location retries independently, so the burst a rate-limit -#: episode produces is this number times the retry count; the NWDC tolerates -#: this level of concurrency without rate-limit errors (verified by stress -#: test). Set ``wateruse.MAX_CONCURRENT_REQUESTS = 1`` for serial. -MAX_CONCURRENT_REQUESTS = 4 +#: This service's preferred in-flight cap when ``API_USGS_CONCURRENT`` is +#: unset. Lower than the package default of 32 because every location retries +#: independently, so a rate-limit episode bursts this number times the retry +#: count; the NWDC tolerates this level without rate-limit errors (verified by +#: stress test) and higher has not been tested. Setting ``API_USGS_CONCURRENT`` +#: overrides it -- see :func:`dataretrieval.transport.fanout._resolve_concurrency` +#: for why the general setting outranks a module's default rather than the +#: reverse. +DEFAULT_CONCURRENT_REQUESTS = 4 # Page responses carry the HUC12 identifier in this column; it must stay a # string so leading zeros (e.g. "010900020502") survive the round trip. @@ -128,8 +125,12 @@ def get_wateruse( Each selector also accepts a list of values. The NWDC queries one area per request, so a list is fanned out into one request per value — up to - :data:`MAX_CONCURRENT_REQUESTS` in parallel — and the results are - concatenated in the order given. + ``API_USGS_CONCURRENT`` in parallel, defaulting to + :data:`DEFAULT_CONCURRENT_REQUESTS` for this service — and the results are + concatenated in the order given. A fan-out interrupted by a rate limit or an + upstream fault raises a resumable + :class:`~dataretrieval.interruptions.FanOutInterrupted`, whose + ``.call.resume()`` re-issues only the locations that did not complete. Parameters ---------- @@ -199,9 +200,12 @@ def get_wateruse( is not five digits, or a HUC of invalid length). DataRetrievalError On an HTTP error response, the typed subclass for the status (see - :func:`dataretrieval.exceptions.error_for_status`); or - :class:`~dataretrieval.exceptions.NetworkError` on a connection-level - failure (timeout, DNS). + :func:`dataretrieval.exceptions.error_for_status`). A transient 429, + 5xx, or recoverable connection failure that exhausts inline retries is + raised as a resumable + :class:`~dataretrieval.interruptions.FanOutInterrupted`; a deterministic + connection failure (for example, a permanently unresolvable host) + remains a :class:`~dataretrieval.exceptions.NetworkError`. Examples -------- @@ -247,16 +251,7 @@ def get_wateruse( ) for location in _resolve_locations(state, county, huc) ] - # ``_run_sync`` drives the async fan-out via an anyio portal, so it is safe - # even inside an already-running event loop (e.g. a Jupyter notebook). - # ``error_url`` is the host reported in any connection-error message (this - # module builds its own requests, so it has no OGC request-builder base). - df, response = run_sync( - lambda: _fan_out(requests, headers, ssl_check), - service="wateruse", - error_url=WATERUSE_URL, - ) - return df, BaseMetadata(response) + return _fan_out(requests, headers, ssl_check) # Valid HUC code lengths (digits) → the hydrologic-unit level they query. @@ -341,18 +336,30 @@ def _validate_huc(value: object) -> str: return code -async def _fan_out( +def _fan_out( requests: list[httpx.Request], headers: dict[str, str], ssl_check: bool -) -> tuple[pd.DataFrame, httpx.Response]: - """Fetch every request (each paginated) concurrently over one shared client. +) -> tuple[pd.DataFrame, BaseMetadata]: + """Fetch every request (each paginated) over the shared fan-out executor. Each request is paginated by :func:`dataretrieval.transport.pagination.paginate` with NWDC strategies: parse a CSV page and read its ``Link`` header cursor (``parse``), follow that cursor (``follow``), and raise the typed error - carrying the NWDC ``detail`` (``raise_for_status``). Concurrency is bounded - by a semaphore at :data:`MAX_CONCURRENT_REQUESTS`, and ``asyncio.gather`` - preserves input order, so the concatenation is deterministic. The shared - :class:`httpx.AsyncClient` keeps connections alive across pages and requests. + carrying the NWDC ``detail`` (``raise_for_status``). + + Everything else -- bounded concurrency, per-attempt retry, failure + precedence, progress, and resumable interruption -- belongs to + :class:`~dataretrieval.transport.fanout.FanOut`, which Water Data and NGWMN + drive too. This function is now only the NWDC-specific half: what a + sub-request is, and how to read one. + + The plan is the request list itself. ``FanOut`` asks a plan only to be + sized and iterable, and the NWDC accepts one ``location=`` per request, so + the caller's locations arrive already separate -- there is nothing to + divide and so nothing for a plan class to hold. + + The broad retry status set is on purpose: NWDC reports a bad query as a 400 + with a ``{"detail": ...}`` envelope, so unlike WQP and StreamStats its 5xx + really is an upstream fault worth re-sending. """ def parse(response: httpx.Response) -> tuple[pd.DataFrame, str | None]: @@ -364,67 +371,43 @@ async def follow(cursor: str, sess: httpx.AsyncClient) -> httpx.Response: def raise_for_status(response: httpx.Response) -> None: _raise_for_status(response, detail_from=_nwdc_error_detail) - # The broad status set on purpose: NWDC reports a bad query as a 400 with a - # ``{"detail": ...}`` envelope, so unlike WQP and StreamStats its 5xx really - # is an upstream fault worth re-sending. Note the cost is multiplied by the - # fan-out -- see MAX_CONCURRENT_REQUESTS. - policy = RetryPolicy.from_env() - async with open_async_client(verify=ssl_check) as client: - semaphore = asyncio.Semaphore(max(1, MAX_CONCURRENT_REQUESTS)) - - async def _one(request: httpx.Request) -> tuple[pd.DataFrame, httpx.Response]: - async def attempt() -> tuple[pd.DataFrame, httpx.Response]: - return await paginate( - request, - parse_response=parse, - follow_up=follow, - client=client, - raise_for_status=raise_for_status, - ) - - # ``retry_async`` owns the gate: the slot is acquired per attempt, - # so a location backing off isn't holding one. A later-page failure - # is intentionally wrapped by ``paginate`` and propagates instead of - # restarting a partially completed walk. - return await retry_async(attempt, policy, gate=semaphore) - - # ``return_exceptions`` so every location is joined before the client - # block exits. Letting the first failure propagate out of the gather - # closed the shared client from under its still-running siblings: a - # location mid-page-walk (or asleep on a ``Retry-After`` backoff) then - # failed with "Cannot send a request, as the client has been closed" on - # a task nobody was awaiting any more -- a spurious error, and an - # unretrieved-exception warning, both attributable to our own teardown. - # The cost is that a fatal error waits for the slowest sibling; that is - # the price of not abandoning in-flight work mid-request. - results = await asyncio.gather( - *(_one(req) for req in requests), return_exceptions=True - ) - - # A cancellation or interrupt signal (``CancelledError``, - # ``KeyboardInterrupt`` -- non-``Exception``) wins over any request failure: - # gathering with ``return_exceptions`` captures it like any other result, and - # reporting a sibling's HTTP error instead would swallow the user's stop - # signal. Otherwise raise in input order, so which failure a caller sees - # stays deterministic rather than depending on which location lost the race. - # (Same precedence the chunked fan-out applies -- see ``ChunkedCall._run``.) - failures = [result for result in results if isinstance(result, BaseException)] - for failure in failures: - if not isinstance(failure, Exception): - raise failure - if failures: - raise failures[0] - pairs = [result for result in results if not isinstance(result, BaseException)] - - # Reuse the transport combine helpers: drop empty frames and concat, and fold - # the per-location responses into one (headers from the response with the - # lowest reported remaining quota plus summed response durations), keeping - # the first request's URL as the query identity. - frames = [frame for frame, _ in pairs] - responses = [resp for _, resp in pairs] - return _combine_chunk_frames(frames), _combine_chunk_responses( - responses, str(requests[0].url) - ) + async def fetch(request: httpx.Request) -> tuple[pd.DataFrame, httpx.Response]: + """One location's full page walk, over the executor's shared client. + + ``active_client()`` is the client :meth:`FanOut._run` published for this + run; borrowing it keeps every location's pages on one connection pool + instead of opening a client per location. + """ + try: + return await paginate( + request, + parse_response=parse, + follow_up=follow, + client=active_client(), + raise_for_status=raise_for_status, + ) + except httpx.TransportError as exc: + raise network_error(request.url, exc) from exc + + def finalize( + frame: pd.DataFrame, response: httpx.Response + ) -> tuple[pd.DataFrame, BaseMetadata]: + return frame, BaseMetadata(response) + + return FanOut( + requests, + fetch, + RetryPolicy.from_env(), + finalize=finalize, + client_options={"verify": ssl_check}, + default_concurrent=DEFAULT_CONCURRENT_REQUESTS, + # No single URL expresses "all of these locations" -- the service + # has no such request -- so the aggregate reports the first, + # matching what an un-fanned single-location call would show. + canonical_url=str(requests[0].url) if requests else None, + # Labels the progress line the executor opens for this drive. + service="wateruse", + ).resume() def _read_csv_page(response: httpx.Response) -> pd.DataFrame: @@ -456,31 +439,12 @@ def _next_page_url(response: httpx.Response) -> str | None: url = response.links.get("next", {}).get("url") if not url: return None - try: - target = httpx.URL(url) - except (httpx.InvalidURL, TypeError) as exc: - raise DataRetrievalError( - f"Water Use returned an unusable next-page link: {url!r}. The page " - f"walk cannot continue; report this if it persists." - ) from exc - if not target.is_absolute_url: - target = response.url.join(target) - if target.host not in _WATERUSE_HOST_ALIASES: - raise DataRetrievalError( - f"Refusing to follow a Water Use next-page link pointing at " - f"{target.host} rather than {_WATERUSE_HOST}. Following it would " - f"send this request, and any credentials on it, to a host you did " - f"not ask for. Retrying will not help; report this if it persists." - ) - # Drop any explicit port and any embedded userinfo along with the - # scheme/host rewrite. A port that went with the link's original scheme - # (``http://…:8080``) would otherwise survive into an https request and be - # dialed under TLS; userinfo (``http://user:pass@…``) would survive into an - # ``Authorization: Basic`` header that httpx derives from it and send a - # credential the caller never configured to the rewritten host -- the very - # thing the host check above exists to prevent. - return str( - target.copy_with(scheme="https", host=_WATERUSE_HOST, port=None, userinfo=b"") + return resolve_next_url( + url, + response, + service="Water Use", + allowed_hosts=_WATERUSE_HOST_ALIASES, + rewrite_host=_WATERUSE_HOST, ) diff --git a/dataretrieval/wqp.py b/dataretrieval/wqp.py index 1112cb73d..2a41fd968 100644 --- a/dataretrieval/wqp.py +++ b/dataretrieval/wqp.py @@ -18,7 +18,8 @@ from dataretrieval._response_metadata import BaseMetadata -from .utils import _attach_datetime_columns, _query_with_retry +from ._querying import _query_with_retry +from ._wqx import _attach_datetime_columns __all__ = [ "get_results", diff --git a/docs/source/architecture/decisions/0006-service-neutral-transport.rst b/docs/source/architecture/decisions/0006-service-neutral-transport.rst index 8eee96a62..67b3c2c6d 100644 --- a/docs/source/architecture/decisions/0006-service-neutral-transport.rst +++ b/docs/source/architecture/decisions/0006-service-neutral-transport.rst @@ -4,7 +4,9 @@ ADR 0006: Use a service-neutral transport layer Status ------ -Accepted +Accepted. The clause assigning resumable ``ChunkedCall`` state to OGC is +superseded by :doc:`0008-fan-out-execution`, which moves fan-out *execution* +into transport and leaves chunk *planning* in OGC. The rest stands. Context ------- diff --git a/docs/source/architecture/decisions/0008-fan-out-execution.rst b/docs/source/architecture/decisions/0008-fan-out-execution.rst new file mode 100644 index 000000000..c336d143b --- /dev/null +++ b/docs/source/architecture/decisions/0008-fan-out-execution.rst @@ -0,0 +1,144 @@ +ADR 0008: Separate fan-out execution from chunk planning +======================================================== + +Status +------ + +Accepted. Supersedes the clause of :doc:`0006-service-neutral-transport` +assigning "resumable ``ChunkedCall`` state" to OGC's protocol concerns; the rest +of ADR 0006 stands. + +Context +------- + +Two services turn one logical query into several requests, for unrelated +reasons. A Water Data or NGWMN query whose URL exceeds the server's byte limit +is split along its multi-value axes. A Water Use query naming several locations +is split because the NWDC accepts one ``location=`` per request -- its URLs run +around 63 bytes against an 8000-byte budget, so the byte limit has nothing to do +with it. + +Chunking is how you divide the data structurally; fan-out is how you distribute +the work operationally. The two are orthogonal, and only the first is protocol +knowledge: dividing a query needs the byte budget, the CQL2 grammar, and which +parameters are list-valued, while distributing the pieces needs none of it. + +The package had not drawn that line. ``ChunkPlan`` (division) and +``ChunkedCall`` (distribution) sat side by side in ``dataretrieval.ogc`` as +siblings, and ADR 0006 grouped them together deliberately. That grouping was +correct while a byte plan was the only thing anyone fanned out over. It stopped +being correct once Water Use fanned out too: unable to reach an OGC-internal +executor, ``wateruse._fan_out`` re-implemented the semaphore, the +``asyncio.gather``, and the cancellation-beats-HTTP-error failure precedence, +with a comment naming ``ChunkedCall._run`` as the original. One subtle rule, +two copies, synchronized by prose. + +The duplicate was not merely redundant. It lacked resume, so a rate limit +partway through discarded every location that had already succeeded -- against +an hourly quota, on fan-outs that reach into the hundreds. It reported no +progress. And it read its own module-global concurrency cap, so a user setting +``API_USGS_CONCURRENT`` to be polite to the service found one adapter ignoring +them. + +Decision +-------- + +``dataretrieval.transport.fanout`` owns fan-out execution for every service: +bounded concurrency, per-attempt retry, deterministic failure precedence, sparse +completion tracking, and resume. It names no protocol concept. An adapter +supplies a ``FanOutPlan`` and an ``async def fetch(item) -> (df, response)``. + +``FanOutPlan`` is a ``Protocol`` of ``__len__`` and ``__iter__``, generic in the +item type -- a sized, iterable collection of sub-request descriptions, and +nothing more. The executor passes each item to the adapter's own ``fetch`` +without inspecting it, so the item type is the adapter's business: the OGC +getters yield kwargs dicts, Water Use yields ready ``httpx.Request`` objects. + +The standard protocols, rather than bespoke members, are a deliberate choice. +A plan declaring ``total`` and ``iter_sub_args()`` would be stating ``len`` +twice under a private name: the two could then report different counts, and a +test would have to assert they agree. Every adapter whose sub-requests are +already a list would also need a wrapper class whose only job is renaming +``len``. + +With the standard names a plain ``list`` is a plan, which is exactly what Water +Use passes. ``ChunkPlan`` keeps ``total`` and ``iter_sub_args`` as its own +vocabulary and defines the dunders to delegate to them, so the two cannot +disagree. + +The protocol is structural rather than nominal for the original reason: +``ChunkPlan`` derives sub-requests from a byte budget over multi-value axes, +a list of requests derives nothing, and so there is no shared implementation an +abstract base could hold. + +The identity of the query as a whole is *not* part of the plan. ``canonical_url`` +is a value stamped on the combined response, not a property of how the work +divides, so it is an argument to ``FanOut``. ``ChunkPlan`` computes one while +planning and the OGC call site passes it through; Water Use passes its first +location's URL, since the service has no request expressing "all of these". + +``dataretrieval.ogc`` keeps chunk planning: the byte budget, the axis +partitioning, the CQL2 filter split, the ``parallel_chunks`` dial. Those are +division, and division is protocol-specific. + +The interruption taxonomy moves to ``dataretrieval.interruptions``, a top-level +leaf, for the reason ADR 0006 gives for ``combining``, ``progress``, and +``credentials``: adapters need it whether or not they went through transport, +and an exception taxonomy is not HTTP execution policy. Its base class is +renamed ``FanOutInterrupted``, since Water Use raises it without chunking +anything. ``ChunkInterrupted`` is retained as a permanent alias of the same +class object -- not a shim scheduled for deletion -- because it is the name +published in the user guide and caught in user code. The subclasses +(``QuotaExhausted``, ``ServiceInterrupted``) were already neutral and are +unchanged. + +Concurrency is one general setting with per-service defaults. +``API_USGS_CONCURRENT`` applies to every fanned-out call; a service may declare +a different default for when it is unset. The precedence is deliberate: an +explicitly set environment variable outranks a service default, never the +reverse. A service that could override the general setting would make +``API_USGS_CONCURRENT=1`` a lie. Service defaults say "absent instruction, this +service prefers N"; they do not say "this service knows better than you". + +Consequences +------------ + +- Water Use gains resume, progress reporting, and the shared concurrency + setting, and sheds roughly 75 lines of duplicated orchestration. +- One implementation of failure precedence, so cancellation-beats-error and + deterministic failure ordering cannot drift between services. +- **Breaking:** a Water Use fan-out interrupted by a 5xx, 429, or recoverable + connection failure now raises ``ServiceInterrupted`` / ``QuotaExhausted`` + rather than ``ServiceUnavailable`` / ``RateLimited`` / ``NetworkError``. All + remain ``DataRetrievalError``, so broad handlers are unaffected, but a narrow + handler around a Water Use call must widen. This is convergence, not novelty + -- it is what the OGC getters have always done -- and it is what makes the + failure resumable. Deterministic connection failures remain ``NetworkError``. +- **Breaking:** ``wateruse.MAX_CONCURRENT_REQUESTS`` is removed in favor of + ``API_USGS_CONCURRENT`` and ``wateruse.DEFAULT_CONCURRENT_REQUESTS``. +- Resume re-issues a failed location's entire page walk, so pages fetched before + the failure are fetched again. This already applied to OGC -- a partial walk + never enters the completion map -- and is a cost, not a correctness problem. +- Water Use frames carry ``huc12_id``, not ``id``, so ``_combine_chunk_frames`` + concatenates them without deduplicating. Correct, because locations partition + by construction, but the executor's dedup safety net does not apply there. +- ``transport`` is no longer purely leaf-shaped: ``fanout`` is a composite that + drives retry, pagination-borrowed clients, and combining. It remains HTTP + execution policy, which is the test the package applies. + +Compliance +---------- + +``tests/architecture_test.py`` asserts three things. That ``wateruse`` +contains no ``asyncio.gather``, ``Semaphore``, or ``TaskGroup``, so the +duplication cannot return. That both plan types are sized and *repeatably* +iterable -- resume keys completed work by position, so a generator mistaken for +a collection would re-issue the wrong sub-requests. And that an interruption +taxonomy does not reappear inside ``transport``. + +Adapter tests cover Water Use resume re-issuing only unfinished locations, +progress ticks, and the concurrency precedence rule. + +Conformance itself is left to the type checker rather than asserted at runtime: +with the protocol reduced to ``__len__`` and ``__iter__``, a missing member is a +``mypy`` error at the call site, not an ``AttributeError`` discovered mid-fan-out. diff --git a/docs/source/architecture/decisions/index.rst b/docs/source/architecture/decisions/index.rst index 450f29669..fb1315283 100644 --- a/docs/source/architecture/decisions/index.rst +++ b/docs/source/architecture/decisions/index.rst @@ -24,4 +24,5 @@ records sequentially. 0005-legacy-nwis 0006-service-neutral-transport 0007-adapter-facades + 0008-fan-out-execution template diff --git a/docs/source/architecture/index.rst b/docs/source/architecture/index.rst index 122e8189b..b58f4167d 100644 --- a/docs/source/architecture/index.rst +++ b/docs/source/architecture/index.rst @@ -105,23 +105,34 @@ Shared components state; ``requests`` owns argument normalization and HTTP request construction; ``schema`` executes queryables/schema requests; ``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 the package-wide fitness function in + ``planning`` determines chunk boundaries; ``chunking`` connects those plans + to the shared fan-out executor and retains compatibility aliases; + ``interruptions`` and ``retry`` re-export their moved compatibility + surfaces; and ``shaping``, ``dates``, ``filters``, and ``errors`` isolate + their named protocol concerns. The full runtime OGC graph, including the + facade, is acyclic -- enforced by the package-wide fitness function in ``tests/architecture_test.py``. ``dataretrieval.transport`` Internal service-neutral execution layer. Owns guarded client lifecycle and timeouts, host-scoped authentication, cursor pagination, bounded retry, - response aggregation, progress, and sync-over-async dispatch. Internally, - ``liveness`` is a stdlib-only leaf recording when data last arrived, so the - page loop that observes progress and the retry loop that acts on it both - depend on ``liveness`` rather than on each other. Transport imports no - service adapter or OGC protocol module, and it is not exposed as a public - framework API. + response aggregation, fan-out execution, progress integration, and + sync-over-async dispatch. ``fanout`` drives an injected plan and fetch + callback, owning bounded concurrency, deterministic failure precedence, + sparse completion state, resume, and the progress line. It is also the one + entry point from synchronous getter code into the async internals: a query + with nothing to divide runs as a one-item fan-out rather than crossing a + separate bridge. 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 depend on ``liveness`` rather than on each other. + Transport imports no service adapter or OGC protocol module, and is not + exposed as a public framework API. + +``dataretrieval.interruptions`` + Shared resumable fan-out failure contract. It owns + ``FanOutInterrupted`` and its subclasses; ``ChunkInterrupted`` remains a + permanent alias for compatibility. The module is outside transport because + adapters catch these errors independently of how they execute requests. ``dataretrieval.exceptions`` Stable error-policy leaf. It has no runtime third-party dependency, and @@ -141,11 +152,21 @@ Shared components import and resolves to this same class. ``dataretrieval.utils`` - Data-shaping helpers, legacy request composition, and compatibility imports - for names that historically lived here (including ``Ambient`` and - ``BaseMetadata``, so their original import paths keep working). OGC does not - depend on this mixed legacy module; by default, do not add new - service-specific behavior there. + Data-shaping helpers, plus compatibility imports for names that + historically lived here (including ``Ambient``, ``BaseMetadata``, ``query`` + and ``to_str``, so their original import paths keep working). OGC does not + depend on this legacy module; by default, do not add new service-specific + behavior there. + +``dataretrieval._querying`` + The one-shot HTTP query path the single-request adapters (``nwis``, + ``wqp``, ``nldi``, ``streamstats``, ``wateruse``) use: compose the URL, send + it, map the status, retry a transient. It left ``utils`` because the two + halves shared only a filename -- this one depends on ``exceptions`` and + ``transport``, the shaping half on ``codes`` and pandas, and no caller + wanted both. The implementation module is private; the established public + function paths remain ``dataretrieval.utils.query`` and + ``dataretrieval.utils.to_str``. ``dataretrieval.codes`` and ``dataretrieval.rdb`` State/time-zone code conversion and RDB parsing leaves. @@ -178,8 +199,10 @@ contracts; consistency alone is not sufficient reason for a breaking change. Failed requests derive from ``dataretrieval.DataRetrievalError``. Callers can inspect ``status_code``, ``retry_after``, and ``retryable`` without knowing the -concrete subtype. OGC calls may raise ``ChunkInterrupted`` subclasses carrying a -resumable call handle and completed partial state. +concrete subtype. A fanned-out call -- an over-large OGC request, or a Water Use +query naming several locations -- may raise ``FanOutInterrupted`` subclasses +(formerly, and still aliased as, ``ChunkInterrupted``) carrying a resumable call +handle and completed partial state. Package/module exports and documentation define the public surface. Underscore-prefixed symbols are implementation details even where existing @@ -230,15 +253,16 @@ A typical OGC-backed call follows this sequence:: -> shape columns and types -> return DataFrame and BaseMetadata -A transient failure after some chunks complete raises a resumable interruption. -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. +A transient failure after some subrequests complete raises a resumable +interruption. The retained ``FanOut`` (available as ``ChunkedCall`` on the OGC +compatibility path) reissues only missing work and applies the same finalization +path when resumed. Cancellation and non-transient programming errors take +precedence over retry/resume wrapping. 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. +matching semantics. Retry, cursor pagination, and fan-out remain explicit +adapter choices. Chunk planning remains OGC-specific; resumable fan-out is also +used by Water Use because the NWDC accepts only one location per request. Resource and configuration view ------------------------------- @@ -250,7 +274,8 @@ Resource and configuration view link to any other host, including external rating assets. ``API_USGS_CONCURRENT`` - OGC subrequest concurrency cap; defaults to 32, ``1`` is sequential, and + Fan-out concurrency cap; defaults to 32 for OGC and 4 for Water Use when + unset. An explicit value applies to every service, ``1`` is sequential, and ``unbounded`` removes the explicit cap. A semaphore, not pool waiting, is the execution throttle. @@ -291,9 +316,13 @@ Known architectural debt This view records categories and representative locations of debt. ``.importlinter`` is authoritative for exact current dependency allowlists. -- ``ogc/engine.py`` retains compatibility wrappers alongside OGC orchestration. -- ``utils.py`` combines metadata, shaping, ambient configuration, legacy - request composition, and transport compatibility imports. +- ``ogc/engine.py`` retains a compatibility pagination wrapper alongside OGC + orchestration. The sync-dispatch wrapper is gone: every retrieval path now + enters through ``transport.fanout.FanOut``. +- ``utils.py`` combines shaping with compatibility imports for metadata, + ambient configuration, transport, and the query path. +- ``waterdata/utils.py`` combines endpoint constants, argument normalization, + and the OGC engine wrappers. These are documented so guardrails distinguish accepted current dependencies from new erosion. They should be removed through small, test-protected changes, diff --git a/docs/source/reference/exceptions.rst b/docs/source/reference/exceptions.rst index 447da2263..7b5c29094 100644 --- a/docs/source/reference/exceptions.rst +++ b/docs/source/reference/exceptions.rst @@ -7,16 +7,23 @@ dataretrieval.exceptions :members: :show-inheritance: -Resumable chunk interruptions +Resumable fan-out interruptions ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ -These exceptions are raised when a transparently-chunked request is interrupted -mid-stream. The completed work is preserved, and ``exc.call.resume()`` continues -it. They are defined in ``dataretrieval.ogc.interruptions`` (they carry -pandas/httpx state), but you can import them from the top level, e.g. -``from dataretrieval import ChunkInterrupted``. +These are raised when a fanned-out request is interrupted mid-stream; the +completed work is preserved and ``exc.call.resume()`` continues it. They are +defined in ``dataretrieval.interruptions`` (they carry pandas/httpx state) but +are importable from the top level, e.g. +``from dataretrieval import FanOutInterrupted``. -.. autoclass:: dataretrieval.ChunkInterrupted +``ChunkInterrupted`` is a permanent alias of ``FanOutInterrupted`` -- the same +class object under the name it was first published as -- so ``except +ChunkInterrupted`` and ``except FanOutInterrupted`` are the same handler. The +base class is named for the fan-out rather than for chunking because a Water Use +call fans out without dividing anything: the NWDC simply accepts one location +per request. + +.. autoclass:: dataretrieval.FanOutInterrupted :members: :show-inheritance: diff --git a/docs/source/userguide/errors.rst b/docs/source/userguide/errors.rst index c8ed723c6..511d82174 100644 --- a/docs/source/userguide/errors.rst +++ b/docs/source/userguide/errors.rst @@ -19,11 +19,13 @@ clause handles any failure regardless of which service you called: except dataretrieval.DataRetrievalError: ... # any request failure: error status, connection loss, too-large, ... -Connection-level failures (timeouts, DNS, refused connections) are wrapped as -:class:`~dataretrieval.exceptions.NetworkError`, so the clause above covers them -too -- you never have to catch an ``httpx`` exception. A *no-data* result is **not** an -error: the modern getters return an empty ``DataFrame`` when nothing matches, so -check ``df.empty`` rather than catching anything. +Connection-level failures (timeouts, DNS, refused connections) remain inside +the package taxonomy, so the clause above covers them -- you never have to catch +an ``httpx`` exception. A deterministic connection failure is a +:class:`~dataretrieval.exceptions.NetworkError`; a recoverable one that exhausts +inline retries during fan-out is a resumable ``ServiceInterrupted``. A *no-data* +result is **not** an error: the modern getters return an empty ``DataFrame`` when +nothing matches, so check ``df.empty`` rather than catching anything. Branch without knowing the concrete type ========================================= @@ -71,31 +73,38 @@ and honors the server's ``Retry-After`` hint when present: raise time.sleep(e.retry_after or 2 ** attempt) -Resume a large Water Data request -================================= +Resume an interrupted request +============================= + +Some requests become several: the Water Data and NGWMN getters split an +over-large request into chunks, and a Water Use call with several locations +becomes one request per location. When a transient failure interrupts one +mid-stream, the work already completed is preserved: catch +``FanOutInterrupted`` and call ``exc.call.resume()`` once the condition clears +-- only the unfinished sub-requests are re-issued. -The Water Data getters transparently split an over-large request into chunks. -When a transient failure interrupts a chunk mid-stream, the work already -completed is preserved: catch ``ChunkInterrupted`` and call ``exc.call.resume()`` -once the condition clears -- only the unfinished sub-requests are re-issued. +(``ChunkInterrupted`` is the same class under its original name; either works.) .. code-block:: python import time - from dataretrieval import ChunkInterrupted + from dataretrieval import FanOutInterrupted from dataretrieval.waterdata import get_daily try: df, md = get_daily(monitoring_location_id=long_list_of_sites) - except ChunkInterrupted as exc: + except FanOutInterrupted as exc: while True: time.sleep(exc.retry_after or 5 * 60) try: df, md = exc.call.resume() break - except ChunkInterrupted as again: + except FanOutInterrupted as again: exc = again +The same loop works for ``wateruse.get_wateruse`` with a list of states, +counties, or HUCs. + Chunk a large request more finely ================================= diff --git a/pyproject.toml b/pyproject.toml index 124dea809..5a184141f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -118,6 +118,16 @@ select = [ "F", "E", "W", "I", "UP", "B", "Q", "SIM", "TID", "C90", # mccabe "E501", # line-length + # A first-party import used only in annotations is not a runtime + # dependency, and the project already says so twice: ``.importlinter`` sets + # ``exclude_type_checking_imports = True`` ("Contracts describe what runs") + # and ``tests/architecture_test.py`` skips ``TYPE_CHECKING`` blocks when it + # asserts the module seams. TC001 makes the code agree with the policy + # instead of leaving it to reviewer memory. TC002/TC003 (third-party and + # stdlib) are a separate import-cost concern with no policy behind them + # here, so they stay unselected. + "TC001", + "TC006", # quoted first argument to ``typing.cast`` ] ignore = [ "SIM105", # Use `contextlib.suppress(...)` instead of `try-except-pass` diff --git a/tests/architecture_test.py b/tests/architecture_test.py index acd9ccf46..595e2d305 100644 --- a/tests/architecture_test.py +++ b/tests/architecture_test.py @@ -223,11 +223,9 @@ def test_waterdata_utils_is_not_an_ogc_reexport_hub() -> None: if dependency == "dataretrieval.ogc" or dependency.startswith("dataretrieval.ogc.") } - assert ogc_deps == { - "dataretrieval.ogc", - "dataretrieval.ogc.dates", - "dataretrieval.ogc.shaping", - }, f"Water Data utils crossed its intended OGC seam: {ogc_deps}" + assert ogc_deps == {"dataretrieval.ogc"}, ( + f"Water Data utils crossed its intended OGC seam: {ogc_deps}" + ) exports = _literal_exports(path) @@ -298,6 +296,11 @@ def test_transport_is_execution_policy_only() -> None: misplaced = { "dataretrieval/transport/progress.py", "dataretrieval/transport/combining.py", + # An exception taxonomy is not HTTP execution policy either. ``fanout`` + # raises ``FanOutInterrupted`` and belongs here; defining it here would + # not, since adapters catch it whether or not they went through + # transport. + "dataretrieval/transport/interruptions.py", } present = { path @@ -508,3 +511,60 @@ def test_empty_result_shaping_consults_the_schema_endpoint() -> None: assert "dataretrieval.ogc.schema" in _runtime_imports( PACKAGE_ROOT / "ogc" / "shaping.py" ) + + +def test_wateruse_does_not_reimplement_fan_out_orchestration() -> None: + """Water Use must drive its locations through the shared fan-out executor. + + It previously ran its own ``asyncio.gather`` with a private semaphore and a + hand-copied failure-precedence rule, kept in sync with ``FanOut`` by a + comment. Two copies of that rule is how they drift, and the duplicate lost + resume, progress, and the shared concurrency setting. Assert the duplication + cannot quietly return. + """ + source = (PACKAGE_ROOT / "wateruse.py").read_text(encoding="utf-8") + tree = ast.parse(source) + offenders = { + f"{node.value.id}.{node.attr}" + for node in ast.walk(tree) + if isinstance(node, ast.Attribute) + and isinstance(node.value, ast.Name) + and node.value.id == "asyncio" + and node.attr in {"gather", "Semaphore", "wait", "TaskGroup"} + } + assert not offenders, ( + "Water Use re-implemented fan-out orchestration instead of using " + f"transport.fanout.FanOut: {sorted(offenders)}" + ) + + +def test_fan_out_plans_are_sized_and_repeatably_iterable() -> None: + """What ``FanOut`` needs of a plan, checked on both real plan types. + + ``FanOutPlan`` is the two standard protocols, so a ``list`` conforms with + no adapter class and ``isinstance`` against a ``runtime_checkable`` + protocol would prove only that the methods exist. What is actually + load-bearing and *not* guaranteed by the type is repeatability: resume + keys completed work by position, so a plan whose second pass differed -- + a generator mistaken for a collection, say -- would re-issue the wrong + sub-requests. + """ + import httpx + + from dataretrieval.ogc.planning import ChunkPlan + + def _build(**args: object) -> httpx.Request: + return httpx.Request("GET", "https://example.invalid/items", params=args) + + plans = [ + ChunkPlan({"sites": ["a", "b"]}, _build, url_limit=8000), + # Water Use hands its request list straight to ``FanOut``. + [httpx.Request("GET", "https://example.invalid/data")], + ] + for plan in plans: + name = type(plan).__name__ + first = list(plan) + assert len(first) == len(plan), ( + f"{name} yielded {len(first)} items but reports len {len(plan)}" + ) + assert list(plan) == first, f"{name} is not repeatably iterable" diff --git a/tests/headers_host_scoping_test.py b/tests/headers_host_scoping_test.py index 82dbbb22b..f5f31cd0f 100644 --- a/tests/headers_host_scoping_test.py +++ b/tests/headers_host_scoping_test.py @@ -8,8 +8,8 @@ import httpx import pytest +from dataretrieval.transport.http import HTTPX_ASYNC_DEFAULTS from dataretrieval.utils import ( - HTTPX_ASYNC_DEFAULTS, _default_headers, _get, ) diff --git a/tests/transport_test.py b/tests/transport_test.py index dd5e9278a..8008d9e7b 100644 --- a/tests/transport_test.py +++ b/tests/transport_test.py @@ -12,8 +12,10 @@ import pandas as pd import pytest +import dataretrieval.exceptions as exceptions import dataretrieval.transport.liveness as liveness import dataretrieval.transport.retry as retry +from dataretrieval._querying import _raise_for_status from dataretrieval.exceptions import ( ConfigurationError, DataRetrievalError, @@ -22,9 +24,8 @@ RateLimited, ServiceUnavailable, ) +from dataretrieval.transport.fanout import FanOut from dataretrieval.transport.pagination import paginate -from dataretrieval.transport.sync import run_sync -from dataretrieval.utils import _raise_for_status def _response( @@ -147,10 +148,23 @@ def test_shared_status_mapping_preserves_retry_after() -> None: def test_sync_bridge_runs_async_operation() -> None: - async def operation() -> str: - return "ok" + """A one-item fan-out is the package's only sync-to-async bridge. + + Every retrieval path now enters through ``FanOut``, so a single request + with nothing to chunk still reaches the network from synchronous caller + code through the executor's blocking portal. + """ + frame = pd.DataFrame({"value": ["ok"]}) + response = _response() - assert run_sync(operation, service="test", error_url="https://example.test") == "ok" + async def operation(item: str) -> tuple[pd.DataFrame, httpx.Response]: + assert item == "only" + return frame, response + + returned, aggregated = FanOut(["only"], operation).resume() + + assert returned["value"].tolist() == ["ok"] + assert aggregated is response def test_retry_tunables_have_a_single_home() -> None: @@ -174,13 +188,13 @@ def test_parse_retry_after_accepts_http_date() -> None: 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")) + parsed = exceptions.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 + assert exceptions.parse_retry_after("Wed, 21 Oct 2015 07:28:00 GMT") is None + assert exceptions.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 + assert exceptions.parse_retry_after("0") == 0.0 def test_both_retry_after_forms_are_honored_alike() -> None: @@ -195,9 +209,9 @@ def test_both_retry_after_forms_are_honored_alike() -> None: ) header = far_future.strftime("%a, %d %b %Y %H:%M:%S GMT") - parsed = retry.parse_retry_after(header) + parsed = exceptions.parse_retry_after(header) assert parsed is not None and 1750 < parsed <= 1800 - assert retry.parse_retry_after("1800") == 1800.0 + assert exceptions.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) @@ -560,3 +574,34 @@ def test_deterministic_failures_are_not_offered_as_resumable() -> None: temporary = _wrapped_dns_failure(socket.EAI_AGAIN) assert retry._retryable(temporary) == (True, None) assert _classify_chunk_error(temporary) is not None + + +def test_exception_chain_walk_terminates_on_a_self_referencing_chain() -> None: + """Every question asked of a failure chain shares one guarded traversal. + + ``raise ... from`` accepts an exception already in the chain, so a retry + loop that re-raises an earlier failure can close the cycle. Each of these + walks reaches an answer by inspecting links, so an unguarded one would spin + forever inside a request path rather than surface the failure. This fails by + hanging, not by asserting -- pytest's timeout is the real assertion. + """ + from dataretrieval.interruptions import ( + ServiceInterrupted, + _classify_chunk_error, + _deterministic_failure, + _walk_causes, + ) + + first = RuntimeError("first") + second = RuntimeError("second") + first.__cause__ = second + second.__cause__ = first + + assert {id(exc) for exc in _walk_causes(first)} == {id(first), id(second)} + assert _classify_chunk_error(first) is None + assert _deterministic_failure(first) is False + # The status hunt in the interruption constructor walks the same chain. + assert ( + ServiceInterrupted(completed_chunks=0, total_chunks=1, cause=first).status_code + is None + ) diff --git a/tests/utils_test.py b/tests/utils_test.py index 1649a3663..b461b019e 100644 --- a/tests/utils_test.py +++ b/tests/utils_test.py @@ -5,7 +5,7 @@ import pandas as pd import pytest -from dataretrieval import exceptions, nwis, utils +from dataretrieval import _querying, _wqx, exceptions, nwis, utils class Test_Ambient: @@ -89,7 +89,7 @@ def test_opt_in_retry_recovers_from_transient(self, httpx_mock, monkeypatch): monkeypatch.setattr(retry, "_RETRY_BASE_BACKOFF", 0.0) monkeypatch.setattr(retry, "_RETRY_MAX_BACKOFF", 0.0) - response = utils._query_with_retry(url, {"a": "1"}) + response = _querying._query_with_retry(url, {"a": "1"}) assert response.text == "ok" assert len(httpx_mock.get_requests()) == 2 @@ -370,7 +370,7 @@ def test_wqx3_triplet_resolves_to_utc(self): "Activity_StartTimeZone": ["PST", "EST"], } ) - df = utils._attach_datetime_columns(df) + df = _wqx._attach_datetime_columns(df) assert df["Activity_StartDateTime"][0] == pd.Timestamp( "2024-01-09 18:00:00", tz="UTC" ) @@ -387,7 +387,7 @@ def test_legacy_wqp_triplet_resolves_to_utc(self): "ActivityStartTime/TimeZoneCode": ["PST"], } ) - df = utils._attach_datetime_columns(df) + df = _wqx._attach_datetime_columns(df) assert df["ActivityStartDateTime"][0] == pd.Timestamp( "2024-01-09 18:00:00", tz="UTC" ) @@ -400,7 +400,7 @@ def test_unknown_timezone_is_NaT(self): "Activity_StartTimeZone": ["BOGUS"], } ) - df = utils._attach_datetime_columns(df) + df = _wqx._attach_datetime_columns(df) assert df["Activity_StartDateTime"].isna().all() def test_existing_datetime_column_not_overwritten(self): @@ -412,7 +412,7 @@ def test_existing_datetime_column_not_overwritten(self): "Activity_StartDateTime": ["preexisting"], } ) - df = utils._attach_datetime_columns(df) + df = _wqx._attach_datetime_columns(df) assert df["Activity_StartDateTime"].tolist() == ["preexisting"] @@ -471,10 +471,10 @@ def test_retrying_get_maps_invalid_url(monkeypatch): import httpx monkeypatch.setattr( - utils, + _querying, "_get", mock.Mock(side_effect=httpx.InvalidURL("invalid URL")), ) with pytest.raises(exceptions.URLTooLong): - utils._get_with_retry("https://example.invalid") + _querying._get_with_retry("https://example.invalid") diff --git a/tests/waterdata_chunking_test.py b/tests/waterdata_chunking_test.py index 06657821c..4b5213136 100644 --- a/tests/waterdata_chunking_test.py +++ b/tests/waterdata_chunking_test.py @@ -20,6 +20,7 @@ import contextvars import datetime import http.server +import socket import threading import time import warnings @@ -38,6 +39,7 @@ ) from dataretrieval.exceptions import ( DataRetrievalError, + NetworkError, RateLimited, ServiceUnavailable, TransientError, @@ -1913,7 +1915,7 @@ def test_retryable_skips_wrapped_midpagination_transient(): def test_retry_transient_then_recovers(monkeypatch): - monkeypatch.setattr(_chunking.asyncio, "sleep", _aiozero) + monkeypatch.setattr(_retry_mod.asyncio, "sleep", _aiozero) calls = {"n": 0} async def afn(): @@ -1928,7 +1930,7 @@ async def afn(): def test_retry_exhausted_reraises(monkeypatch): - monkeypatch.setattr(_chunking.asyncio, "sleep", _aiozero) + monkeypatch.setattr(_retry_mod.asyncio, "sleep", _aiozero) calls = {"n": 0} async def afn(): @@ -1943,7 +1945,7 @@ async def afn(): def test_retry_non_retryable_not_retried(monkeypatch): slept: list[float] = [] - monkeypatch.setattr(_chunking.asyncio, "sleep", _recording_sleep(slept)) + monkeypatch.setattr(_retry_mod.asyncio, "sleep", _recording_sleep(slept)) calls = {"n": 0} async def afn(): @@ -1958,7 +1960,7 @@ async def afn(): def test_retry_long_retry_after_escalates(monkeypatch): slept: list[float] = [] - monkeypatch.setattr(_chunking.asyncio, "sleep", _recording_sleep(slept)) + monkeypatch.setattr(_retry_mod.asyncio, "sleep", _recording_sleep(slept)) calls = {"n": 0} async def afn(): @@ -1974,7 +1976,7 @@ async def afn(): def test_retry_transient_then_success(monkeypatch): - monkeypatch.setattr(_chunking.asyncio, "sleep", _aiozero) + monkeypatch.setattr(_retry_mod.asyncio, "sleep", _aiozero) calls = {"n": 0} async def afn(): @@ -1994,7 +1996,7 @@ def test_chunker_retries_transient_then_completes(monkeypatch): """A transient on one sub-request is retried transparently; the decorated call completes with no ChunkInterrupted.""" monkeypatch.setenv("API_USGS_RETRIES", "3") - monkeypatch.setattr(_chunking.asyncio, "sleep", _aiozero) + monkeypatch.setattr(_retry_mod.asyncio, "sleep", _aiozero) state = {"failed": False} async def fetch(args): @@ -2033,7 +2035,7 @@ def test_chunker_exhausted_retries_still_resumable(monkeypatch): """When retries are exhausted the failure still surfaces as a resumable ChunkInterrupted — retries don't swallow the escape hatch.""" monkeypatch.setenv("API_USGS_RETRIES", "2") - monkeypatch.setattr(_chunking.asyncio, "sleep", _aiozero) + monkeypatch.setattr(_retry_mod.asyncio, "sleep", _aiozero) attempts = {"n": 0} async def fetch(args): @@ -2054,7 +2056,7 @@ def test_async_fan_out_retries_transient_then_completes(monkeypatch): """The parallel path retries a transient sub-request and completes.""" monkeypatch.setenv("API_USGS_RETRIES", "3") - monkeypatch.setattr(_chunking.asyncio, "sleep", _aiozero) + monkeypatch.setattr(_retry_mod.asyncio, "sleep", _aiozero) state = {"failed": False} async def fetch_async(args): @@ -2073,7 +2075,7 @@ def test_async_fan_out_surfaces_fatal_over_transient(monkeypatch): being masked behind a resumable interruption from a transient sibling.""" monkeypatch.setenv("API_USGS_RETRIES", "2") - monkeypatch.setattr(_chunking.asyncio, "sleep", _aiozero) + monkeypatch.setattr(_retry_mod.asyncio, "sleep", _aiozero) async def fetch_async(args): # One chunk carries a deterministic programmer error; the rest are @@ -2087,6 +2089,36 @@ async def fetch_async(args): fetch({"sites": ["S1" * 10, "S2" * 10, "S3" * 10, "S4" * 10]}) +def test_deterministic_transport_failure_is_normalized(): + """A permanent resolver failure must not leak a raw httpx exception.""" + + @multi_value_chunked(build_request=_fake_build, url_limit=8000) + async def fetch(_args): + resolution = socket.gaierror(socket.EAI_NONAME, "name not known") + failure = httpx.ConnectError("name not known") + failure.__context__ = resolution + raise failure + + with pytest.raises(NetworkError) as excinfo: + fetch({"sites": ["S1"]}) + + assert isinstance(excinfo.value.__cause__, httpx.ConnectError) + + +def test_transport_context_does_not_mask_fatal_error(): + """An unrelated implicit transport context cannot replace the fatal error.""" + + @multi_value_chunked(build_request=_fake_build, url_limit=8000) + async def fetch(_args): + try: + raise httpx.ConnectError("name not known") + except httpx.ConnectError: + raise ValueError("deterministic bug") # noqa: B904 - regression shape + + with pytest.raises(ValueError, match="deterministic bug"): + fetch({"sites": ["S1"]}) + + # --- finalize hook (resume finalizes; partials stay raw) ------------------- # # Regression for the bug where ``exc.call.resume()`` returned the chunker's diff --git a/tests/waterdata_progress_test.py b/tests/waterdata_progress_test.py index be81ffe84..de5618f37 100644 --- a/tests/waterdata_progress_test.py +++ b/tests/waterdata_progress_test.py @@ -371,6 +371,9 @@ def test_nested_context_reuses_outer_reporter(): def _resp(features, *, next_url=None, rate_remaining=None): resp = mock.MagicMock() + # A real response always carries an ``httpx.URL``; the next-page check + # resolves and host-checks the ``next`` link against it. + resp.url = httpx.URL("https://example.com/p1") links = [{"rel": "next", "href": next_url}] if next_url else [] resp.json.return_value = { "numberReturned": len(features), diff --git a/tests/waterdata_test.py b/tests/waterdata_test.py index 6f217ab19..5c54c9fee 100644 --- a/tests/waterdata_test.py +++ b/tests/waterdata_test.py @@ -11,6 +11,7 @@ import pytest from pandas import DataFrame +from dataretrieval.ogc.context import _ogc_base_url from dataretrieval.ogc.engine import _dialect from dataretrieval.ogc.requests import ( _check_monitoring_location_id, @@ -37,9 +38,10 @@ get_stats_por, get_time_series_metadata, ) +from dataretrieval.waterdata.types import _check_profiles from dataretrieval.waterdata.utils import ( + OGC_API_URL, WATERDATA_DIALECT, - _check_profiles, _get_args, ) @@ -58,14 +60,16 @@ @pytest.fixture(autouse=True) def _activate_waterdata_dialect(): - """Make the Water Data OGC dialect ambient for this module. - - The dialect (monitoring-locations -> POST/CQL2; daily -> date-only time - args) is normally set by ``get_ogc_data`` per call. The direct - ``_construct_api_requests`` unit tests here bypass it, so activate the - dialect module-wide so they exercise the real Water Data behavior. + """Make the Water Data OGC base URL and dialect ambient for this module. + + Both are normally set together by ``get_ogc_data`` per call: the base URL + (the OGC package names no service of its own) and the dialect + (monitoring-locations -> POST/CQL2; daily -> date-only time args). The + direct ``_construct_api_requests``/``_construct_cql_request`` unit tests + here bypass that entry point, so activate both module-wide so they + exercise the real Water Data behavior. """ - with _dialect(WATERDATA_DIALECT): + with _ogc_base_url(OGC_API_URL), _dialect(WATERDATA_DIALECT): yield diff --git a/tests/waterdata_utils_test.py b/tests/waterdata_utils_test.py index 30995debe..0d7dd6baa 100644 --- a/tests/waterdata_utils_test.py +++ b/tests/waterdata_utils_test.py @@ -1,5 +1,6 @@ import asyncio import datetime +import functools import json import logging from unittest import mock @@ -18,6 +19,7 @@ ServiceUnavailable, TransientError, ) +from dataretrieval.interruptions import ServiceInterrupted from dataretrieval.ogc.context import _row_cap from dataretrieval.ogc.dates import _format_api_dates from dataretrieval.ogc.engine import ( @@ -35,9 +37,21 @@ _get_resp_data, _to_snake_case, ) +from dataretrieval.ogc.shaping import _finalize_ogc as _ogc_finalize from dataretrieval.waterdata import get_stats_date_range, get_stats_por from dataretrieval.waterdata.stats import _handle_nesting, get_data -from dataretrieval.waterdata.utils import OGC_API_URL, _finalize_ogc, _get_args +from dataretrieval.waterdata.utils import ( + _EXTRA_ID_COLS, + OGC_API_URL, + WATERDATA_DIALECT, + _get_args, +) + +# The Water Data injection ``get_cql`` performs at its call site, so these tests +# exercise the same result shape the typed getters produce. +_finalize_ogc = functools.partial( + _ogc_finalize, extra_id_cols=_EXTRA_ID_COLS, dialect=WATERDATA_DIALECT +) _LOGGER_NAME = _utils_module.__name__ @@ -85,6 +99,9 @@ def test_get_args_empty(): def test_walk_pages_multiple_mocked(): # Setup mock responses resp1 = mock.MagicMock() + # A real response always carries an ``httpx.URL``; the next-page check + # resolves and host-checks the ``next`` link against it. + resp1.url = httpx.URL("https://example.com/page1") resp1.json.return_value = { "numberReturned": 1, "features": [{"id": "1", "properties": {"val": "a"}}], @@ -481,15 +498,26 @@ def test_get_data_raises_on_mid_pagination_failure(monkeypatch): same ``_paginate`` strategy helper, so error-routing behaviour is exercised by the ``_walk_pages`` triplet above. This single ``get_data`` mid-pagination case proves the stats-specific - follow-up callback is wired into ``_paginate`` correctly.""" - with pytest.raises(DataRetrievalError, match="Paginated request failed") as excinfo: + follow-up callback is wired into ``_paginate`` correctly. + + Statistics drives that page walk as a one-item ``FanOut``, the same + executor every other getter uses, so a transient mid-walk failure is + resumable here too rather than ending the call outright. + """ + with pytest.raises(ServiceInterrupted) as excinfo: _run_get_data_with_failure( httpx.ConnectError("stats-boom"), monkeypatch, ) - assert isinstance(excinfo.value.__cause__, httpx.ConnectError) - assert "stats-boom" in str(excinfo.value) + # The pagination wrapper is the direct cause and still names the failure. + paginated = excinfo.value.__cause__ + assert isinstance(paginated, DataRetrievalError) + assert "Paginated request failed" in str(paginated) + assert isinstance(paginated.__cause__, httpx.ConnectError) + assert "stats-boom" in str(paginated) + # Nothing completed, so there is nothing to hand back but the handle. + assert excinfo.value.call.completed_chunks == 0 def test_get_data_warning_includes_next_token(caplog, monkeypatch): @@ -986,7 +1014,7 @@ def test_check_ogc_requests_raises_typed_on_5xx(httpx_mock): json={"code": "ServiceUnavailable", "description": "maintenance window"}, ) with pytest.raises(ServiceUnavailable): - _check_ogc_requests(endpoint="daily", req_type="schema") + _check_ogc_requests(endpoint="daily", req_type="schema", base_url=OGC_API_URL) @pytest.mark.parametrize( diff --git a/tests/wateruse_test.py b/tests/wateruse_test.py index b6ee4e002..5e88f0060 100644 --- a/tests/wateruse_test.py +++ b/tests/wateruse_test.py @@ -4,6 +4,7 @@ """ import re +import socket from urllib.parse import parse_qs, urlsplit import httpx @@ -11,7 +12,9 @@ import pytest import dataretrieval +from dataretrieval import progress as _progress from dataretrieval import wateruse +from dataretrieval.transport import fanout as _fanout from dataretrieval.utils import BaseMetadata from dataretrieval.wateruse import _next_page_url, _resolve_locations, get_wateruse @@ -270,8 +273,8 @@ def test_multiple_states_fan_out_preserves_input_order(httpx_mock): def test_fan_out_is_serial_when_concurrency_is_one(httpx_mock, monkeypatch): - """``MAX_CONCURRENT_REQUESTS = 1`` still fans out correctly (serial path).""" - monkeypatch.setattr(wateruse, "MAX_CONCURRENT_REQUESTS", 1) + """``API_USGS_CONCURRENT=1`` still fans out correctly (serial path).""" + monkeypatch.setenv("API_USGS_CONCURRENT", "1") httpx_mock.add_response( method="GET", url=re.compile(r".*location=stateCd%3ARI.*"), text=_CSV_P1 ) @@ -311,7 +314,14 @@ def test_fan_out_surfaces_final_rate_limit_header(httpx_mock): def test_fan_out_failure_never_returns_partial_data(httpx_mock): - """A failed location aborts the call even when another location succeeded.""" + """A failed location aborts the call even when another location succeeded. + + The completed sibling is not returned as though the call had succeeded -- + it is carried on the raised interruption for ``resume()`` instead. Water Use + reports ``ServiceInterrupted`` rather than the bare ``ServiceUnavailable`` + it raised before sharing the fan-out executor: the same upstream 503, now + resumable. + """ httpx_mock.add_response( method="GET", url=re.compile(r".*location=stateCd%3ARI.*"), @@ -324,9 +334,18 @@ def test_fan_out_failure_never_returns_partial_data(httpx_mock): json={"detail": "temporarily unavailable"}, ) - with pytest.raises(dataretrieval.ServiceUnavailable): + with pytest.raises(dataretrieval.ServiceInterrupted) as excinfo: get_wateruse(model="wu-public-supply-wd", state=["RI", "WI"]) + # The 503 is still the reported cause, and the successful location survives + # on the exception rather than being passed off as the whole answer. + assert isinstance(excinfo.value.__cause__, dataretrieval.ServiceUnavailable) + assert excinfo.value.status_code == 503 + assert excinfo.value.retryable + assert excinfo.value.completed_chunks == 1 + assert excinfo.value.total_chunks == 2 + assert len(excinfo.value.partial_frame) == 2 + # --- _resolve_locations unit tests (no HTTP) ------------------------------- @@ -562,23 +581,15 @@ async def open_mock_client(**overrides): ) as client: yield client - monkeypatch.setattr(wateruse, "open_async_client", open_mock_client) + monkeypatch.setattr(_fanout, "open_async_client", open_mock_client) requests = [ httpx.Request("GET", wateruse.WATERUSE_URL, params={"location": location}) for location in ("stateCd:AA", "stateCd:BB") ] - async def drive() -> int: - with pytest.raises(dataretrieval.DataRetrievalError, match="Invalid model"): - await wateruse._fan_out(requests, {}, True) - # Nothing left running: the sibling was joined before teardown, so no - # task can later fail against the closed client. - return len( - [task for task in asyncio.all_tasks() if task is not asyncio.current_task()] - ) - - assert asyncio.run(drive()) == 0 + with pytest.raises(dataretrieval.DataRetrievalError, match="Invalid model"): + wateruse._fan_out(requests, {}, True) assert pages["n"] == 2, "the sibling finished its walk rather than being abandoned" @@ -591,3 +602,230 @@ def test_next_page_url_rejects_cross_host_link(): # other failure rather than seeing a bare RuntimeError. with pytest.raises(dataretrieval.DataRetrievalError, match="outside.example"): _next_page_url(response) + + +# --- capabilities Water Use gained by sharing the fan-out executor ---------- + + +def test_interrupted_fan_out_resumes_only_the_unfinished_locations(httpx_mock): + """A rate-limited location is resumable; completed siblings are not re-fetched. + + Before Water Use shared the executor, a 429 anywhere in the fan-out + discarded every location that had already succeeded. That is the whole + reason a multi-location pull needed re-running from scratch against an + hourly quota. + """ + httpx_mock.add_response( + method="GET", url=re.compile(r".*location=stateCd%3ARI.*"), text=_CSV_P1 + ) + # WI is rate-limited on the first pass, then succeeds once resumed. + httpx_mock.add_response( + method="GET", + url=re.compile(r".*location=stateCd%3AWI.*"), + status_code=429, + json={"detail": "rate limited"}, + is_reusable=False, + ) + httpx_mock.add_response( + method="GET", url=re.compile(r".*location=stateCd%3AWI.*"), text=_CSV_P2 + ) + + with pytest.raises(dataretrieval.QuotaExhausted) as excinfo: + get_wateruse(model="wu-public-supply-wd", state=["RI", "WI"]) + + interrupted = excinfo.value + assert interrupted.status_code == 429 + assert interrupted.retryable + assert interrupted.completed_chunks == 1 + assert interrupted.total_chunks == 2 + requests_before = len(httpx_mock.get_requests()) + + df, md = interrupted.call.resume() + + # Only WI was re-issued; RI's completed frame carried across the resume. + assert len(httpx_mock.get_requests()) == requests_before + 1 + assert len(df) == 3 + assert isinstance(md, BaseMetadata) + + +def test_fan_out_honors_the_general_concurrency_setting(monkeypatch): + """``API_USGS_CONCURRENT`` outranks this service's default. + + A user dialing concurrency down to be polite must not find Water Use + quietly ignoring them -- the defect that motivated consolidating the knob. + """ + monkeypatch.setenv("API_USGS_CONCURRENT", "7") + assert _fanout._resolve_concurrency(wateruse.DEFAULT_CONCURRENT_REQUESTS) == 7 + + monkeypatch.delenv("API_USGS_CONCURRENT", raising=False) + assert ( + _fanout._resolve_concurrency(wateruse.DEFAULT_CONCURRENT_REQUESTS) + == wateruse.DEFAULT_CONCURRENT_REQUESTS + ) + # The service default is deliberately below the package-wide 32. + assert wateruse.DEFAULT_CONCURRENT_REQUESTS < _fanout._CONCURRENCY_DEFAULT + + +def test_fan_out_reports_progress(httpx_mock, monkeypatch): + """The fan-out ticks the progress reporter, which it never did standalone.""" + seen = [] + + class _Recorder: + def set_chunks(self, total): + seen.append(("chunks", total)) + + def start_chunk(self, completed): + seen.append(("chunk", completed)) + + def set_rate_remaining(self, remaining, limit=None): + pass + + def add_page(self, rows): + seen.append(("page", rows)) + + monkeypatch.setattr(_fanout._progress, "current", lambda: _Recorder()) + httpx_mock.add_response( + method="GET", url=re.compile(r".*location=stateCd%3ARI.*"), text=_CSV_P1 + ) + httpx_mock.add_response( + method="GET", url=re.compile(r".*location=stateCd%3AWI.*"), text=_CSV_P2 + ) + + get_wateruse(model="wu-public-supply-wd", state=["RI", "WI"]) + + assert ("chunks", 2) in seen + assert ("chunk", 1) in seen and ("chunk", 2) in seen + + +def test_resume_uses_the_current_progress_reporter(httpx_mock, monkeypatch): + """Resume must not resurrect the reporter closed by the interrupted call.""" + created = [] + + class _Recorder: + def __init__(self, **_kwargs): + self.closed = False + self.events = [] + created.append(self) + + def _record(self, event): + assert not self.closed, "fan-out updated a closed progress reporter" + self.events.append(event) + + def set_chunks(self, total): + self._record(("chunks", total)) + + def start_chunk(self, completed): + self._record(("chunk", completed)) + + def set_rate_remaining(self, remaining, limit=None): + self._record(("remaining", remaining, limit)) + + def add_page(self, rows): + self._record(("page", rows)) + + def close(self): + self.closed = True + + monkeypatch.setattr(_progress, "ProgressReporter", _Recorder) + httpx_mock.add_response( + method="GET", url=re.compile(r".*location=stateCd%3ARI.*"), text=_CSV_P1 + ) + httpx_mock.add_response( + method="GET", + url=re.compile(r".*location=stateCd%3AWI.*"), + status_code=429, + is_reusable=False, + ) + httpx_mock.add_response( + method="GET", url=re.compile(r".*location=stateCd%3AWI.*"), text=_CSV_P2 + ) + + with pytest.raises(dataretrieval.QuotaExhausted) as excinfo: + get_wateruse(model="wu-public-supply-wd", state=["RI", "WI"]) + + assert created[0].closed + with _progress.progress_context(service="resume") as resumed_reporter: + _, md = excinfo.value.call.resume() + + assert isinstance(md, BaseMetadata) + assert resumed_reporter is created[1] + assert ("chunks", 2) in resumed_reporter.events + assert ("chunk", 2) in resumed_reporter.events + + +def test_permanent_transport_failure_remains_a_network_error(monkeypatch): + """A deterministic connection failure stays inside the public taxonomy.""" + + async def fail(*_args, **_kwargs): + resolution = socket.gaierror(socket.EAI_NONAME, "name not known") + failure = httpx.ConnectError("name not known") + failure.__context__ = resolution + raise failure + + monkeypatch.setattr(wateruse, "paginate", fail) + + with pytest.raises(dataretrieval.NetworkError) as excinfo: + get_wateruse(model="wu-public-supply-wd", state="RI") + + assert isinstance(excinfo.value.__cause__, httpx.ConnectError) + + +def test_permanent_later_page_failure_remains_a_network_error(httpx_mock): + """Normalization finds transport failures nested by pagination.""" + httpx_mock.add_response( + method="GET", + url=re.compile(r".*location=stateCd%3ARI(?!.*cursor).*"), + text=_CSV_P1, + headers={ + "Link": '; rel="next"' + }, + ) + resolution = socket.gaierror(socket.EAI_NONAME, "name not known") + failure = httpx.ConnectError("name not known") + failure.__context__ = resolution + httpx_mock.add_exception( + failure, + method="GET", + url=re.compile(r".*cursor=x.*"), + ) + + with pytest.raises(dataretrieval.NetworkError) as excinfo: + get_wateruse(model="wu-public-supply-wd", state="RI") + + assert isinstance(excinfo.value.__cause__, httpx.ConnectError) + + +def test_mid_page_walk_transient_is_still_resumable(httpx_mock): + """A 429 on page 2+ of a location must still be a resumable interruption. + + ``paginate`` re-wraps a later-page failure as a plain ``DataRetrievalError`` + (page 1's status check sits outside its ``try``), so the typed cause is only + reachable through ``__cause__``. ``_classify_chunk_error`` walks that chain + for exactly this reason; were it a single ``isinstance`` check, a mid-walk + rate limit would escape as a bare error and lose ``.call.resume()`` -- + inconsistently, since page 1 would still be resumable. + """ + httpx_mock.add_response( + method="GET", + url=re.compile(r".*location=stateCd%3ARI(?!.*cursor).*"), + text=_CSV_P1, + headers={ + "Link": '; rel="next"' + }, + ) + httpx_mock.add_response( + method="GET", + url=re.compile(r".*cursor=x.*"), + status_code=429, + json={"detail": "rate limited"}, + ) + httpx_mock.add_response( + method="GET", url=re.compile(r".*location=stateCd%3AWI.*"), text=_CSV_P2 + ) + + with pytest.raises(dataretrieval.QuotaExhausted) as excinfo: + get_wateruse(model="wu-public-supply-wd", state=["RI", "WI"]) + + assert excinfo.value.call is not None + assert excinfo.value.completed_chunks == 1 + assert excinfo.value.total_chunks == 2