From f8e4efd58c109f51c9706e6a8a8754a54de2ad0c Mon Sep 17 00:00:00 2001 From: thodson-usgs Date: Sun, 2 Aug 2026 07:59:43 -0500 Subject: [PATCH] refactor(ogc): clarify chunking module boundaries Extract response and frame recombination into combining.py and colocate the private fetch/finalize contracts with ChunkedCall. Centralize transient classification while preserving retry, resumability, exception snapshot, deduplication, nested GeoJSON, and response metadata contracts. Add regression coverage for those behaviors. Align Ruff 0.16.1 across CI, pre-commit, and test metadata, and apply its Markdown code-fence formatting. --- .github/workflows/python-package.yml | 3 +- .pre-commit-config.yaml | 2 +- CONTRIBUTING.md | 21 +-- README.md | 65 +++---- dataretrieval/ogc/chunking.py | 60 ++++--- dataretrieval/ogc/combining.py | 206 ++++++++++++++++++++++ dataretrieval/ogc/engine.py | 8 +- dataretrieval/ogc/interruptions.py | 41 ++--- dataretrieval/ogc/planning.py | 245 ++------------------------- dataretrieval/ogc/retry.py | 67 ++++---- dataretrieval/ogc/shaping.py | 32 ++-- dataretrieval/wateruse.py | 9 +- pyproject.toml | 2 +- tests/waterdata_chunking_test.py | 48 +++++- tests/waterdata_utils_test.py | 19 +++ tests/wateruse_test.py | 2 +- 16 files changed, 445 insertions(+), 385 deletions(-) create mode 100644 dataretrieval/ogc/combining.py diff --git a/.github/workflows/python-package.yml b/.github/workflows/python-package.yml index f4a27f07..26131f96 100644 --- a/.github/workflows/python-package.yml +++ b/.github/workflows/python-package.yml @@ -20,7 +20,8 @@ jobs: python-version: "3.14" cache: "pip" - name: Install ruff - run: pip install ruff + # Keep this version aligned with the ruff-pre-commit revision. + run: pip install ruff==0.16.1 - name: Lint with ruff run: | ruff check . --output-format=github diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 0813f1fe..8dd44167 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -21,7 +21,7 @@ repos: - id: debug-statements - repo: https://github.com/astral-sh/ruff-pre-commit - rev: v0.15.15 + rev: v0.16.1 hooks: - id: ruff-check args: [--fix] diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 2fbd526b..10b1ef4c 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -149,19 +149,20 @@ via any automated processes or pipelines. * Example: ``` python + LIGHT_MESSAGES = { + "English": "There are %(number_of_lights)s lights.", + "Pirate": "Arr! Thar be %(number_of_lights)s lights.", + } - LIGHT_MESSAGES = { - 'English': "There are %(number_of_lights)s lights.", - 'Pirate': "Arr! Thar be %(number_of_lights)s lights." - } - def lights_message(language, number_of_lights): - """Return a language-appropriate string reporting the light count.""" - return LIGHT_MESSAGES[language] % locals() + def lights_message(language, number_of_lights): + """Return a language-appropriate string reporting the light count.""" + return LIGHT_MESSAGES[language] % locals() - def is_pirate(message): - """Return True if the given message sounds piratical.""" - return re.search(r"(?i)(arr|avast|yohoho)!", message) is not None + + def is_pirate(message): + """Return True if the given message sounds piratical.""" + return re.search(r"(?i)(arr|avast|yohoho)!", message) is not None ``` --- diff --git a/README.md b/README.md index a3f494fe..88b076ab 100644 --- a/README.md +++ b/README.md @@ -47,6 +47,7 @@ and set it as an environment variable: ```python import os + os.environ["API_USGS_PAT"] = "your_api_key_here" ``` @@ -59,9 +60,9 @@ from dataretrieval import waterdata # Get daily streamflow data (returns DataFrame and metadata) df, metadata = waterdata.get_daily( - monitoring_location_id='USGS-01646500', - parameter_code='00060', # Discharge - time='2024-10-01/2025-09-30' + monitoring_location_id="USGS-01646500", + parameter_code="00060", # Discharge + time="2024-10-01/2025-09-30", ) print(f"Retrieved {len(df)} records") @@ -72,9 +73,9 @@ Retrieve streamflow at multiple locations from October 1, 2024 to the present: ```python df, metadata = waterdata.get_daily( - monitoring_location_id=["USGS-13018750","USGS-13013650"], - parameter_code='00060', - time='2024-10-01/..' + monitoring_location_id=["USGS-13018750", "USGS-13013650"], + parameter_code="00060", + time="2024-10-01/..", ) print(f"Retrieved {len(df)} records") @@ -85,8 +86,8 @@ stream sites in Maryland: ```python # Get monitoring location information df, metadata = waterdata.get_monitoring_locations( - state='Maryland', # full name, postal code ('MD'), or FIPS ('24') - site_type_code='ST' # Stream sites + state="Maryland", # full name, postal code ('MD'), or FIPS ('24') + site_type_code="ST", # Stream sites ) print(f"Found {len(df)} stream monitoring locations in Maryland") @@ -98,9 +99,9 @@ windows to avoid timeouts and other issues: ```python # Get continuous data for a single monitoring location and water year df, metadata = waterdata.get_continuous( - monitoring_location_id='USGS-01646500', - parameter_code='00065', # Gage height - time='2024-10-01/2025-09-30' + monitoring_location_id="USGS-01646500", + parameter_code="00065", # Gage height + time="2024-10-01/2025-09-30", ) print(f"Retrieved {len(df)} continuous gage height measurements") ``` @@ -125,10 +126,10 @@ from dataretrieval import waterdata # enough to span many pages, so it profits from a finer split. sites, _ = waterdata.get_monitoring_locations(state="Ohio", site_type_code="ST") -with waterdata.parallel_chunks(32): # fan out into 32 sub-requests +with waterdata.parallel_chunks(32): # fan out into 32 sub-requests df, md = waterdata.get_daily( monitoring_location_id=sites["monitoring_location_id"], - parameter_code="00060", # discharge + parameter_code="00060", # discharge time="2004-01-01/2023-12-31", ) ``` @@ -167,6 +168,7 @@ API — enable debug-level ```python import logging + logging.basicConfig(level=logging.DEBUG) ``` @@ -181,14 +183,14 @@ from dataretrieval import ngwmn # Find the groundwater monitoring sites in a state # (state accepts a full name, a postal code like 'WI', or a FIPS code like '55') -sites, metadata = ngwmn.get_sites(state='Wisconsin') +sites, metadata = ngwmn.get_sites(state="Wisconsin") print(f"Found {len(sites)} NGWMN sites in Wisconsin") # Pull water levels from the first twenty sites over a time window. water_levels, metadata = ngwmn.get_water_level( - monitoring_location_id=sites['monitoring_location_id'][:20], - datetime=['2022-01-01', '2024-01-01'] + monitoring_location_id=sites["monitoring_location_id"][:20], + datetime=["2022-01-01", "2024-01-01"], ) print(f"Retrieved {len(water_levels)} water-level observations") @@ -203,16 +205,15 @@ from dataretrieval import wqp # Find water quality monitoring sites (returns a DataFrame and metadata) sites, metadata = wqp.what_sites( - statecode='US:55', # Wisconsin - siteType='Stream' + statecode="US:55", # Wisconsin + siteType="Stream", ) print(f"Found {len(sites)} stream monitoring sites in Wisconsin") # Get water quality results results, metadata = wqp.get_results( - siteid='USGS-05427718', - characteristicName='Temperature, water' + siteid="USGS-05427718", characteristicName="Temperature, water" ) print(f"Retrieved {len(results)} temperature measurements") @@ -227,18 +228,18 @@ from dataretrieval import nldi # Get watershed basin for a stream reach basin = nldi.get_basin( - feature_source='comid', - feature_id='13293474' # NHD reach identifier + feature_source="comid", + feature_id="13293474", # NHD reach identifier ) print(f"Basin contains {len(basin)} feature(s)") # Find upstream flowlines flowlines = nldi.get_flowlines( - feature_source='comid', - feature_id='13293474', - navigation_mode='UT', # Upstream tributaries - distance=50 # km + feature_source="comid", + feature_id="13293474", + navigation_mode="UT", # Upstream tributaries + distance=50, # km ) print(f"Found {len(flowlines)} upstream tributaries within 50km") @@ -255,17 +256,17 @@ from dataretrieval import wateruse # Monthly public-supply withdrawals for Rhode Island, split into # groundwater and surface-water sources (returns a DataFrame and metadata). df, metadata = wateruse.get_wateruse( - model='wu-public-supply-wd', - variable=['pswdtot', 'pswdgw', 'pswdsw'], - state='RI', # name/postal/FIPS; pass a list to fan out over several areas - start_date='2020-01', - time_resolution='monthly', + model="wu-public-supply-wd", + variable=["pswdtot", "pswdgw", "pswdsw"], + state="RI", # name/postal/FIPS; pass a list to fan out over several areas + start_date="2020-01", + time_resolution="monthly", ) print(f"Retrieved {len(df)} records across {df['huc12_id'].nunique()} watersheds") # Aggregate the HUC12 grid to a statewide monthly total (million gallons/day) -statewide = df.groupby('year_month')['pswdtot_mgd'].sum() +statewide = df.groupby("year_month")["pswdtot_mgd"].sum() print(statewide.head()) ``` diff --git a/dataretrieval/ogc/chunking.py b/dataretrieval/ogc/chunking.py index d0dae01c..b46c17df 100644 --- a/dataretrieval/ogc/chunking.py +++ b/dataretrieval/ogc/chunking.py @@ -74,7 +74,7 @@ import asyncio import functools import os -from collections.abc import Callable, Iterator +from collections.abc import Awaitable, Callable, Iterator from contextlib import contextmanager from contextvars import copy_context from typing import Any, cast @@ -86,17 +86,14 @@ from dataretrieval.utils import HTTPX_DEFAULTS, Ambient, _require_positive_int from . import progress as _progress -from .interruptions import ( - ChunkInterrupted, - _Fetch, - _Finalize, - _passthrough_result, -) -from .planning import ( - ChunkPlan, +from .combining import ( _combine_chunk_frames, _combine_chunk_responses, ) +from .interruptions import ( + ChunkInterrupted, +) +from .planning import ChunkPlan from .retry import ( _NO_RETRY, RetryPolicy, @@ -291,6 +288,30 @@ 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. @@ -417,12 +438,11 @@ def _combine_raw(self) -> tuple[pd.DataFrame, httpx.Response]: Frames concatenate in sub-args *index* order (``sorted`` keys — deterministic, independent of parallel completion order). The - aggregated response takes its headers from the most-recently- - *completed* sub-request: the ``track`` closure in :meth:`_run` - is the only writer of ``self._chunks`` and ``dict`` preserves - insertion order, so the chunks' natural order is completion - order and the last one carries the freshest - ``x-ratelimit-remaining``. + 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 ------- @@ -521,8 +541,9 @@ def resume(self) -> tuple[pd.DataFrame, Any]: Combined data from every successful sub-request. response The finalized aggregate — a raw :class:`httpx.Response` - (canonical URL, most-recently-completed sub-request's headers, - cumulative elapsed time) by default, or whatever + (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). @@ -603,8 +624,9 @@ async def _run(self, max_concurrent: int | None) -> tuple[pd.DataFrame, Any]: Combined data from every sub-request. response The finalized aggregate — a raw :class:`httpx.Response` - (canonical URL, most-recently-completed sub-request's headers, - cumulative elapsed time) by default, or whatever + (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 diff --git a/dataretrieval/ogc/combining.py b/dataretrieval/ogc/combining.py new file mode 100644 index 00000000..be4366c1 --- /dev/null +++ b/dataretrieval/ogc/combining.py @@ -0,0 +1,206 @@ +"""Result recombination: merge per-chunk frames and responses (no I/O). + +These utilities assemble the output of a chunked/fan-out call from its +individual per-sub-request results. They have no event-loop, retry, or +network state — they're pure data transforms imported by both the +chunked-call execution (:mod:`dataretrieval.ogc.chunking`) and the +per-page pagination (:mod:`dataretrieval.ogc.engine`). + +Separated from :mod:`dataretrieval.ogc.planning` so that module stays +focused on *what* to split, while this module owns *how* to reassemble. +""" + +from __future__ import annotations + +import copy +from datetime import timedelta + +import httpx +import pandas as pd + +# Response header USGS uses to advertise remaining hourly quota. Lives in this +# module so every layer (the combine helpers below, the engine's per-page +# progress reporter) reads it from one place rather than hard-coding the string. +_QUOTA_HEADER = "x-ratelimit-remaining" + + +def _safe_elapsed(response: httpx.Response) -> timedelta: + """ + Read ``response.elapsed``, falling back to ``timedelta(0)`` when + the attribute hasn't been populated. + + httpx only writes ``.elapsed`` when a response is closed through + its normal transport path. ``MockTransport`` (used by + ``pytest-httpx``) and hand-constructed ``httpx.Response`` objects + leave the attribute unset, so accessing it raises ``RuntimeError``. + Combining responses across chunks needs a defined duration, so we + treat the missing attribute as zero elapsed. + """ + try: + elapsed: object = response.elapsed + except RuntimeError: + return timedelta(0) + return elapsed if isinstance(elapsed, timedelta) else timedelta(0) + + +def _set_response_url(response: httpx.Response, url: str | httpx.URL) -> None: + """ + Overwrite the URL surfaced by a response without back-propagating + the change into any aliased original. + + Lightweight test doubles expose ``.url`` as a writable attribute. Real + :class:`httpx.Response` objects resolve it through a bound request, so swap + in a fresh request carrying the new URL; mutating the existing request would + leak through any shallow copy that shares it. + """ + if not isinstance(response, httpx.Response): + # Lightweight test doubles expose ``url`` as a writable attribute. + response.url = url + return + + target = httpx.URL(str(url)) + try: + old = response.request + except RuntimeError: + # No request bound (some hand-built httpx.Response fixtures); + # synthesize a minimal one to hold the URL. + response.request = httpx.Request("GET", target) + return + response.request = httpx.Request(method=old.method, url=target, headers=old.headers) + + +def _lowest_remaining(responses: list[httpx.Response]) -> httpx.Response: + """The response reporting the lowest ``x-ratelimit-remaining``. + + Within a rate-limit window, the counter decreases monotonically, so the + smallest value observed is the most conservative value to surface. Under + concurrent fan-out, the last response *by index* need not be the one the + server processed last. Fall back to the last response when none reports + the header. + """ + best: httpx.Response | None = None + best_remaining: int | None = None + for response in responses: + try: + remaining = int(response.headers[_QUOTA_HEADER]) + except (KeyError, ValueError): + continue + if best_remaining is None or remaining < best_remaining: + best, best_remaining = response, remaining + return best if best is not None else responses[-1] + + +def _merge_response( + base: httpx.Response, + *, + headers_from: httpx.Response, + elapsed: timedelta, + url: str | httpx.URL | None = None, +) -> httpx.Response: + """Fold several responses into one: a shallow copy of ``base`` whose + ``.headers`` are rebuilt as a fresh ``httpx.Headers`` from ``headers_from``, + ``.elapsed`` set to ``elapsed``, and ``.url`` overridden when ``url`` is + given. ``base`` and ``headers_from`` are never mutated, and the fresh + ``httpx.Headers`` means downstream mutations don't back-propagate into any + underlying response — so callers may re-fold idempotently. This is the one + low-level merge behind both pagination + (:func:`~dataretrieval.ogc.engine._paginate`) and the chunked / fan-out + aggregation (:func:`_combine_chunk_responses`).""" + merged = copy.copy(base) + merged.headers = httpx.Headers(headers_from.headers) + merged.elapsed = elapsed + if url is not None: + _set_response_url(merged, url) + return merged + + +def _combine_chunk_frames(frames: list[pd.DataFrame]) -> pd.DataFrame: + """Concatenate per-chunk frames and deduplicate IDs across chunks. + + Empty frames are ignored before concatenation so an empty plain + :class:`pandas.DataFrame` cannot downgrade a real ``GeoDataFrame`` and + strip its geometry or CRS. When every frame is empty, the first frame is + returned to preserve its concrete type. + + When multiple non-empty frames are combined, non-null feature IDs are + deduplicated regardless of the plan axis. Filter clauses can match the same + feature, and list inputs can contain repeated values or otherwise select + overlapping records. Rows without an ``id`` are preserved verbatim: pandas + treats null values as duplicates, so deduplicating them would silently lose + data. + """ + non_empty = [frame for frame in frames if not frame.empty] + if not non_empty: + return frames[0] if frames else pd.DataFrame() + if len(non_empty) == 1: + return non_empty[0].copy() + + combined = pd.concat(non_empty, ignore_index=True) + if "id" not in combined.columns: + return combined + + has_id = combined["id"].notna() + if has_id.all(): + return combined.drop_duplicates(subset="id", ignore_index=True) + if has_id.any(): + id_rows = combined[has_id].drop_duplicates(subset="id") + no_id_rows = combined[~has_id] + return pd.concat([id_rows, no_id_rows], ignore_index=True) + return combined + + +def _combine_chunk_responses( + responses: list[httpx.Response], canonical_url: str | None +) -> httpx.Response: + """ + Fold per-sub-request responses into a single aggregated response. + + For a multi-response input, returns a shallow copy of + ``responses[0]`` with ``.headers`` set to those of the response reporting + the lowest ``x-ratelimit-remaining`` value (the most conservative quota + observation; see :func:`_lowest_remaining`), ``.elapsed`` set to the sum of + the per-response elapsed durations, and ``.url`` set to the + canonical original-query URL (when supplied) so ``BaseMetadata`` + reflects the user's full request rather than the first chunk. + + For a single-response input with no canonical-URL override, + ``responses[0]`` is returned unchanged to skip the copy on the + passthrough hot path. + + Parameters + ---------- + responses : list[httpx.Response] + One response per completed sub-request, in caller-provided order. + canonical_url : str or None + URL of the unchunked original request. ``None`` skips the URL + override — used by the passthrough path (the fetcher's + response already carries the original-query URL) and by the + worst-case overflow path (no buildable canonical URL exists). + + Returns + ------- + httpx.Response + A shallow copy of the first response with aggregated + ``headers``, ``elapsed``, and ``url``. The function is + idempotent (the input responses' ``headers`` / ``elapsed`` / + ``url`` are never mutated), so it's safe to call repeatedly + via :attr:`ChunkedCall.partial_response` during error + inspection or resume retries. ``headers`` on the returned + object is a fresh ``httpx.Headers``, so mutations there don't + back-propagate into any chunk's underlying response. + """ + if len(responses) == 1 and canonical_url is None: + return responses[0] + + # Headers come from the response with the lowest reported remaining quota; + # ``_lowest_remaining`` returns the lone response as-is + # for a single-element list). ``_merge_response`` re-sums elapsed onto a + # fresh copy, so repeated calls (e.g. via ``ChunkedCall.partial_response`` + # during resume) stay idempotent. + elapsed = sum((_safe_elapsed(r) for r in responses), start=timedelta()) + return _merge_response( + responses[0], + headers_from=_lowest_remaining(responses), + elapsed=elapsed, + url=canonical_url, + ) diff --git a/dataretrieval/ogc/engine.py b/dataretrieval/ogc/engine.py index 668fa7b7..5fda67e5 100644 --- a/dataretrieval/ogc/engine.py +++ b/dataretrieval/ogc/engine.py @@ -47,9 +47,9 @@ from dataretrieval.ogc import chunking from dataretrieval.ogc import progress as _progress from dataretrieval.ogc.chunking import get_active_client +from dataretrieval.ogc.combining import _QUOTA_HEADER, _merge_response, _safe_elapsed from dataretrieval.ogc.dates import _DATE_RANGE_PARAMS, _format_api_dates from dataretrieval.ogc.errors import _paginated_failure_message, _raise_for_non_200 -from dataretrieval.ogc.planning import _QUOTA_HEADER, _merge_response, _safe_elapsed from dataretrieval.ogc.shaping import GEOPANDAS, _finalize_ogc, _get_resp_data from dataretrieval.utils import ( HTTPX_DEFAULTS, @@ -658,9 +658,9 @@ async def _paginate( response : httpx.Response A shallow copy of the first-page response, with ``.headers`` rebuilt as a fresh ``httpx.Headers`` reflecting the last page and - ``.elapsed`` set to cumulative wall-clock. The canonical URL is - preserved from the first page. The original first-page response - is not mutated. + ``.elapsed`` set to the sum of the per-page response durations. The + canonical URL is preserved from the first page. The original first-page + response is not mutated. Raises ------ diff --git a/dataretrieval/ogc/interruptions.py b/dataretrieval/ogc/interruptions.py index 18398654..8cb5723c 100644 --- a/dataretrieval/ogc/interruptions.py +++ b/dataretrieval/ogc/interruptions.py @@ -11,7 +11,6 @@ from __future__ import annotations -from collections.abc import Awaitable, Callable from typing import TYPE_CHECKING, Any, ClassVar import httpx @@ -23,27 +22,6 @@ from dataretrieval.ogc.chunking import ChunkedCall -# ``_Fetch`` is 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 ``utils._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 ChunkInterrupted(DataRetrievalError): """ Base class for mid-stream chunk failures whose completed work is @@ -140,10 +118,14 @@ def __init__( self.total_chunks = total_chunks self.call = call self.retry_after = retry_after - # Snapshot partial state at raise time so the exception's view stays - # stable across later ``call.resume()`` advances (the live view is on - # ``call.partial_frame`` / ``.partial_response``). ``.copy()`` guards - # the single-chunk fast path, where the frame may be returned verbatim. + # 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 @@ -155,9 +137,10 @@ 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 - # partial frame / response; only ``.resume()`` is lost (cross-process - # resume was never possible anyway). + # 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} diff --git a/dataretrieval/ogc/planning.py b/dataretrieval/ogc/planning.py index e80d0643..15b397a0 100644 --- a/dataretrieval/ogc/planning.py +++ b/dataretrieval/ogc/planning.py @@ -1,31 +1,32 @@ -"""Pure URL-byte chunk planning and result recombination (no I/O). - -This module holds the side-effect-free half of the chunker: deciding how -to split one over-budget OGC request into URL-fitting sub-requests -(:class:`ChunkPlan` and the axis/byte-accounting helpers) and reassembling -their per-chunk frames and responses (:func:`_combine_chunk_frames`, -:func:`_combine_chunk_responses`). It has no event loop, retry policy, or -network state — those live in :mod:`dataretrieval.ogc.chunking` (execution) -and :mod:`dataretrieval.ogc.retry` (retry policy), which import the plan and -drive it. Keeping the planning/combination logic here -makes it unit-testable without an HTTP client and gives the two concerns -separate reasons to change. +"""Pure URL-byte chunk planning (no I/O). + +This module holds the side-effect-free planning half of the chunker: +deciding how to split one over-budget OGC request into URL-fitting +sub-requests (:class:`ChunkPlan` and the axis/byte-accounting helpers). +It has no event loop, retry policy, or network state — those live in +:mod:`dataretrieval.ogc.chunking` (execution) and +:mod:`dataretrieval.ogc.retry` (retry policy), which import the plan and +drive it. + +Result recombination — reassembling the per-chunk frames and responses +back into one result +(:func:`~dataretrieval.ogc.combining._combine_chunk_frames`, +:func:`~dataretrieval.ogc.combining._combine_chunk_responses`, etc.) — lives in +the sibling :mod:`dataretrieval.ogc.combining` module, which callers import +directly. """ from __future__ import annotations -import copy import itertools import math from collections.abc import Callable, Iterator from contextlib import suppress from dataclasses import dataclass -from datetime import timedelta from typing import Any from urllib.parse import quote_plus import httpx -import pandas as pd from dataretrieval.exceptions import Unchunkable from dataretrieval.ogc.filters import ( @@ -130,52 +131,6 @@ def _safe_request_bytes( return _request_bytes(req) -def _safe_elapsed(response: httpx.Response) -> timedelta: - """ - Read ``response.elapsed``, falling back to ``timedelta(0)`` when - the attribute hasn't been populated. - - httpx only writes ``.elapsed`` when a response is closed through - its normal transport path. ``MockTransport`` (used by - ``pytest-httpx``) and hand-constructed ``httpx.Response`` objects - leave the attribute unset, so accessing it raises ``RuntimeError``. - Combining responses across chunks needs a defined duration, so we - treat the missing attribute as zero elapsed. - """ - try: - return response.elapsed - except RuntimeError: - return timedelta(0) - - -def _set_response_url(response: httpx.Response, url: str | httpx.URL) -> None: - """ - Overwrite the URL surfaced by a response without back-propagating - the change into any aliased original. - - Try the direct assignment first: on lightweight test mocks ``.url`` - is a plain writable attribute. On real ``httpx.Response`` it's - read-only (it resolves through the bound request), so swap in a - fresh :class:`httpx.Request` carrying the new URL — mutating the - existing one would leak through any shallow copy that shares the - same ``.request``. - """ - try: - response.url = url # type: ignore[misc, assignment] - except AttributeError: - target = httpx.URL(str(url)) - try: - old = response.request - except RuntimeError: - # No request bound (some hand-built httpx.Response fixtures); - # synthesize a minimal one to hold the URL. - response.request = httpx.Request("GET", target) - return - response.request = httpx.Request( - method=old.method, url=target, headers=old.headers - ) - - @dataclass(frozen=True) class _Axis: """ @@ -619,171 +574,3 @@ 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 - - -def _combine_chunk_frames(frames: list[pd.DataFrame]) -> pd.DataFrame: - """ - Concatenate per-chunk frames, dropping empties and deduping by ``id``. - - Parameters - ---------- - frames : list[pandas.DataFrame] - One frame per completed sub-request. - - Returns - ------- - pandas.DataFrame - The concatenated, deduplicated result. Empty when every input - frame is empty. - - Notes - ----- - An empty chunk can be a plain ``pd.DataFrame()`` (no geopandas); - concatenating it with real ``GeoDataFrame``s downgrades the result - to plain ``DataFrame`` and strips geometry/CRS, so empties are - dropped first. Dedup on the pre-rename feature ``id`` keeps - overlapping user OR-clauses from producing duplicate rows across - chunks. - - Dedup is restricted to rows whose ``id`` is non-null. ``pandas`` - treats NaN==NaN as a duplicate for ``drop_duplicates``, so a - blanket call would collapse every id-less row into a single one — - silent data loss if any chunk emits features without an - ``id`` field. - """ - non_empty = [f for f in frames if not f.empty] - if not non_empty: - # Preserve the frame type (GeoDataFrame vs DataFrame) of the - # input even when every chunk is empty — ``_get_resp_data`` - # returns ``gpd.GeoDataFrame()`` on empty geopd responses, and - # returning a plain ``pd.DataFrame()`` here would downgrade - # the type in a downstream ``pd.concat([result, geo_page])`` to - # a plain DataFrame and strip geometry/CRS. - return frames[0] if frames else pd.DataFrame() - if len(non_empty) == 1: - # Single-completed-chunk fast path. Return a copy so callers - # who treat ``ChunkedCall.partial_frame`` as a fresh result - # (the property docstring says "live; recomputed per access") - # don't accidentally mutate ``_chunks[0][0]`` in place. - return non_empty[0].copy() - combined = pd.concat(non_empty, ignore_index=True) - if "id" in combined.columns: - has_id = combined["id"].notna() - if has_id.all(): - combined = combined.drop_duplicates(subset="id", ignore_index=True) - elif has_id.any(): - # Mixed: dedupe only the id-bearing rows; preserve id-less - # rows verbatim (their order relative to id-bearing rows - # may shift, which is acceptable — dedup can't be id-keyed - # for rows without an id). - id_rows = combined[has_id].drop_duplicates(subset="id") - no_id_rows = combined[~has_id] - combined = pd.concat([id_rows, no_id_rows], ignore_index=True) - return combined - - -# Response header USGS uses to advertise remaining hourly quota. Lives in this -# base module so every layer (planning's ``_lowest_remaining``, the engine's -# per-page progress) reads it from one place rather than hard-coding the string. -_QUOTA_HEADER = "x-ratelimit-remaining" - - -def _lowest_remaining(responses: list[httpx.Response]) -> httpx.Response: - """The response reporting the lowest ``x-ratelimit-remaining``. - - The rate-limit counter decreases monotonically within a window, so the - smallest value any sub-request saw is the most-current "quota left after - this call" — the right thing to surface. Under concurrent fan-out the - last response *by index* need not be the one the server processed last, so - pick the minimum (falling back to the last response if none report it). - """ - best: httpx.Response | None = None - best_remaining: int | None = None - for response in responses: - try: - remaining = int(response.headers[_QUOTA_HEADER]) - except (KeyError, ValueError): - continue - if best_remaining is None or remaining < best_remaining: - best, best_remaining = response, remaining - return best if best is not None else responses[-1] - - -def _merge_response( - base: httpx.Response, - *, - headers_from: httpx.Response, - elapsed: timedelta, - url: str | httpx.URL | None = None, -) -> httpx.Response: - """Fold several responses into one: a shallow copy of ``base`` whose - ``.headers`` are rebuilt as a fresh ``httpx.Headers`` from ``headers_from``, - ``.elapsed`` set to ``elapsed``, and ``.url`` overridden when ``url`` is - given. ``base`` and ``headers_from`` are never mutated, and the fresh - ``httpx.Headers`` means downstream mutations don't back-propagate into any - underlying response — so callers may re-fold idempotently. This is the one - low-level merge behind both pagination (:func:`_paginate`) and the chunked / - fan-out aggregation (:func:`_combine_chunk_responses`).""" - merged = copy.copy(base) - merged.headers = httpx.Headers(headers_from.headers) - merged.elapsed = elapsed - if url is not None: - _set_response_url(merged, url) - return merged - - -def _combine_chunk_responses( - responses: list[httpx.Response], canonical_url: str | None -) -> httpx.Response: - """ - Fold per-sub-request responses into a single aggregated response. - - For a multi-response input, returns a shallow copy of - ``responses[0]`` with ``.headers`` set to those of the most-depleted - response (lowest ``x-ratelimit-remaining`` — the quota actually left - after the fan-out; see :func:`_lowest_remaining`), ``.elapsed`` set - to total wall-clock across every response, and ``.url`` set to the - canonical original-query URL (when supplied) so ``BaseMetadata`` - reflects the user's full request rather than the first chunk. - - For a single-response input with no canonical-URL override, - ``responses[0]`` is returned unchanged to skip the copy on the - passthrough hot path. - - Parameters - ---------- - responses : list[httpx.Response] - One response per completed sub-request, in execution order. - canonical_url : str or None - URL of the unchunked original request. ``None`` skips the URL - override — used by the passthrough path (the fetcher's - response already carries the original-query URL) and by the - worst-case overflow path (no buildable canonical URL exists). - - Returns - ------- - httpx.Response - A shallow copy of the first response with aggregated - ``headers``, ``elapsed``, and ``url``. The function is - idempotent (the input responses' ``headers`` / ``elapsed`` / - ``url`` are never mutated), so it's safe to call repeatedly - via :attr:`ChunkedCall.partial_response` during error - inspection or resume retries. ``headers`` on the returned - object is a fresh ``httpx.Headers``, so mutations there don't - back-propagate into any chunk's underlying response. - """ - if len(responses) == 1 and canonical_url is None: - return responses[0] - - # Headers come from the most-depleted response (lowest quota left after a - # concurrent fan-out; ``_lowest_remaining`` returns the lone response as-is - # for a single-element list). ``_merge_response`` re-sums elapsed onto a - # fresh copy, so repeated calls (e.g. via ``ChunkedCall.partial_response`` - # during resume) stay idempotent. - elapsed = sum((_safe_elapsed(r) for r in responses), start=timedelta()) - return _merge_response( - responses[0], - headers_from=_lowest_remaining(responses), - elapsed=elapsed, - url=canonical_url, - ) diff --git a/dataretrieval/ogc/retry.py b/dataretrieval/ogc/retry.py index b81ce60e..37ea040a 100644 --- a/dataretrieval/ogc/retry.py +++ b/dataretrieval/ogc/retry.py @@ -19,7 +19,7 @@ import httpx import pandas as pd -from dataretrieval.exceptions import RateLimited, ServiceUnavailable, TransientError +from dataretrieval.exceptions import RateLimited, TransientError from dataretrieval.ogc import progress as _progress from dataretrieval.ogc.interruptions import ( ChunkInterrupted, @@ -190,6 +190,28 @@ def backoff(self, attempt: int, retry_after: float | None) -> float: _NO_RETRY = RetryPolicy(max_retries=0) +def _classify_transient( + exc: BaseException, +) -> tuple[type[ChunkInterrupted], float | None] | None: + """Classify one exception as a transient, resumable failure. + + This function owns the shared exception taxonomy; it deliberately does not + walk ``__cause__``. :func:`_classify_chunk_error` walks wrapped pagination + failures, while :func:`_retryable` applies the narrower automatic-retry + policy to this classification. + """ + if isinstance(exc, RateLimited): + return QuotaExhausted, exc.retry_after + if isinstance(exc, TransientError): + # Every typed transient other than a rate-limit error is a service + # interruption. This fallback keeps future TransientError subclasses + # resumable after their inline retries are exhausted. + return ServiceInterrupted, exc.retry_after + if isinstance(exc, (httpx.HTTPError, httpx.InvalidURL)): + return ServiceInterrupted, None + return None + + def _classify_chunk_error( exc: BaseException, ) -> tuple[type[ChunkInterrupted], float | None] | None: @@ -231,41 +253,28 @@ def _classify_chunk_error( """ cur: BaseException | None = exc while cur is not None: - if isinstance(cur, RateLimited): - return QuotaExhausted, cur.retry_after - if isinstance(cur, ServiceUnavailable): - return ServiceInterrupted, cur.retry_after - if isinstance(cur, (httpx.HTTPError, httpx.InvalidURL)): - return ServiceInterrupted, None + result = _classify_transient(cur) + if result is not None: + return result cur = cur.__cause__ return None def _retryable(exc: BaseException) -> tuple[bool, float | None]: - """ - Decide whether ``exc`` is a transient worth an automatic retry. - - Only the *top-level* exception is inspected — unlike - :func:`_classify_chunk_error`, which walks the ``__cause__`` chain. - The distinction matters because ``_paginate`` raises an - initial-request transient (429 / 5xx / :class:`httpx.TransportError`) - *raw*, but wraps a mid-pagination failure as a base ``DataRetrievalError``. - So a raw transient means a sub-request that made no progress and is cheap to - re-issue, whereas a mid-pagination failure is left to escalate to a - resumable :class:`ChunkInterrupted` rather than re-walked from page 1 - (which would re-spend the quota just exhausted). ``httpx.InvalidURL`` - is never retried — a too-long cursor won't fix on a retry. + """Decide whether a top-level transient is worth an automatic retry. - Returns - ------- - tuple[bool, float or None] - ``(retryable, retry_after)`` — the server ``Retry-After`` hint - (seconds) when the transient carried one, else ``None``. + Wrapped mid-pagination failures are not retried from page one; they instead + escalate to a resumable :class:`ChunkInterrupted`. ``httpx.InvalidURL`` and + non-transport ``httpx.HTTPError`` instances are classified as resumable but + excluded from automatic retry by policy. """ - if isinstance(exc, TransientError): - return True, exc.retry_after - if isinstance(exc, httpx.TransportError): - return True, None + classification = _classify_transient(exc) + if classification is None: + return False, None + + _, retry_after = classification + if isinstance(exc, (TransientError, httpx.TransportError)): + return True, retry_after return False, None diff --git a/dataretrieval/ogc/shaping.py b/dataretrieval/ogc/shaping.py index 4df242f2..5870c63b 100644 --- a/dataretrieval/ogc/shaping.py +++ b/dataretrieval/ogc/shaping.py @@ -101,17 +101,13 @@ def _get_resp_data( Notes ----- - The non-geopandas branch builds the frame directly from each - feature's ``properties`` dict, plus the top-level ``id`` and - ``geometry.coordinates`` columns — the ``id`` column is always - added (so the downstream rename to the service-specific output id - works even on an all-None id), while the ``geometry`` column is - added only when at least one feature carries geometry. This skips - the GeoJSON envelope entirely, so - newly-added Feature-level fields (e.g. ``geometry.type`` after - USGS migrated to full GeoJSON geometry objects) can't leak into - the result frame; no reactive drop-list needs maintenance every - time the upstream schema grows. + The non-geopandas branch normalizes each feature's ``properties`` object, + flattening nested dictionaries with an underscore separator, then adds the + top-level ``id`` and a ``geometry`` column containing the coordinates. The + ``id`` column is always added so the downstream service-specific rename + works even when all IDs are missing; ``geometry`` is added only when + coordinates are present. Feature-level envelope fields are deliberately + excluded. """ if body is None: body = resp.json() @@ -128,13 +124,11 @@ def _get_resp_data( return _empty_feature_frame(geopd) if not geopd: - df = pd.json_normalize([f.get("properties") or {} for f in features], sep="_") - # Always materialize the ``id`` column (may be all-None) so - # ``_arrange_cols``'s ``df.rename(columns={"id": output_id})`` - # produces the documented service-specific output_id column - # (daily_id, channel_measurements_id, …) even if the upstream - # response carried no feature-level id. - df["id"] = [f.get("id") for f in features] + properties = [feature.get("properties") or {} for feature in features] + df = pd.json_normalize(properties, sep="_") + # Always materialize the feature-level ID (possibly all-None) so + # ``_arrange_cols`` can perform the documented service-specific rename. + df["id"] = [feature.get("id") for feature in features] _attach_coordinates(df, features) return df @@ -369,7 +363,7 @@ def _finalize_ogc( as :class:`~dataretrieval.utils.BaseMetadata`. Injected into the chunker as its ``finalize`` hook (see - :data:`~dataretrieval.ogc.interruptions._Finalize`) so the + :data:`~dataretrieval.ogc.chunking._Finalize`) so the un-interrupted return *and* a resumed ``ChunkInterrupted.call.resume()`` produce the same post-processed ``(DataFrame, BaseMetadata)`` shape, not the chunker's raw frame and bare ``httpx.Response``. diff --git a/dataretrieval/wateruse.py b/dataretrieval/wateruse.py index 635a97fb..5dd5cb9c 100644 --- a/dataretrieval/wateruse.py +++ b/dataretrieval/wateruse.py @@ -53,8 +53,8 @@ from dataretrieval.codes.states import to_state from dataretrieval.exceptions import DataRetrievalError +from dataretrieval.ogc.combining import _combine_chunk_frames, _combine_chunk_responses from dataretrieval.ogc.engine import _paginate, _run_sync -from dataretrieval.ogc.planning import _combine_chunk_frames, _combine_chunk_responses from dataretrieval.utils import ( HTTPX_DEFAULTS, BaseMetadata, @@ -365,8 +365,9 @@ async def _one(request: httpx.Request) -> tuple[pd.DataFrame, httpx.Response]: results = await asyncio.gather(*(_one(req) for req in requests)) # Reuse the engine's combine helpers: drop empty frames and concat, and fold - # the per-location responses into one (lowest-remaining rate-limit headers + - # cumulative elapsed), keeping the first request's URL as the query identity. + # the per-location responses into one (headers from the response with the + # lowest reported remaining quota plus summed response durations), keeping + # the first request's URL as the query identity. frames = [frame for frame, _ in results] responses = [resp for _, resp in results] return _combine_chunk_frames(frames), _combine_chunk_responses( @@ -398,7 +399,7 @@ def _next_page_url(response: httpx.Response) -> str | None: url = response.links.get("next", {}).get("url") if not url: return None - return url.replace("https://water.usgs.gov", "https://api.water.usgs.gov", 1) + return str(url).replace("https://water.usgs.gov", "https://api.water.usgs.gov", 1) def _nwdc_error_detail(response: httpx.Response) -> str | None: diff --git a/pyproject.toml b/pyproject.toml index e261707b..3e6bfbe7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -54,7 +54,7 @@ test = [ "pytest-rerunfailures", "coverage", "pytest-httpx", - "ruff", + "ruff==0.16.1", "dataretrieval[type-check]", # mypy, pinned once in the type-check extra ] doc = [ diff --git a/tests/waterdata_chunking_test.py b/tests/waterdata_chunking_test.py index 600f05d8..4a86bf17 100644 --- a/tests/waterdata_chunking_test.py +++ b/tests/waterdata_chunking_test.py @@ -35,6 +35,7 @@ DataRetrievalError, RateLimited, ServiceUnavailable, + TransientError, Unchunkable, ) from dataretrieval.ogc import chunking as _chunking @@ -47,6 +48,11 @@ multi_value_chunked, parallel_chunks, ) +from dataretrieval.ogc.combining import ( + _QUOTA_HEADER, + _combine_chunk_frames, + _combine_chunk_responses, +) from dataretrieval.ogc.interruptions import ( ChunkInterrupted, QuotaExhausted, @@ -56,10 +62,7 @@ _LIST_SEP, _NEVER_CHUNK, _OR_SEP, - _QUOTA_HEADER, ChunkPlan, - _combine_chunk_frames, - _combine_chunk_responses, _extract_axes, _request_bytes, _safe_request_bytes, @@ -1124,15 +1127,29 @@ def test_combine_chunk_frames_does_not_collapse_none_ids(): def test_combine_chunk_frames_still_dedupes_overlapping_ids(): - """The original dedup contract — overlapping OR-clause partitions - that produce duplicate-id rows across chunks must still collapse - to one row — has to keep working when ids ARE present.""" + """Duplicate feature IDs collapse regardless of why chunks overlap.""" df_a = pd.DataFrame({"id": ["x", "y"], "val": [1, 2]}) df_b = pd.DataFrame({"id": ["y", "z"], "val": [2, 3]}) combined = _combine_chunk_frames([df_a, df_b]) assert sorted(combined["id"].tolist()) == ["x", "y", "z"] +def test_list_axis_chunks_dedupe_repeated_feature_ids(): + """Repeated list values can select the same feature in separate chunks.""" + + @multi_value_chunked(build_request=_fake_build, url_limit=8000) + async def fetch(args): + return ( + pd.DataFrame({"id": ["feature-1"], "site": [args["sites"][0]]}), + _quota_response(500), + ) + + with parallel_chunks(2): + frame, _ = fetch({"sites": ["A", "A"]}) + + assert frame.to_dict(orient="records") == [{"id": "feature-1", "site": "A"}] + + def test_retry_after_surfaces_on_quota_exhausted(): """If the 429 response includes a ``Retry-After`` header, that delay must travel from the typed transport exception @@ -1997,6 +2014,25 @@ async def fetch(args): assert sorted(df["sites"]) == sorted(sites) # all recovered despite the 429 +def test_future_transient_subclass_remains_resumable(monkeypatch): + """The shared taxonomy handles new typed transients by default.""" + + class MaintenanceWindow(TransientError): + _DEFAULT_STATUS = 503 + + monkeypatch.setenv("API_USGS_RETRIES", "0") + + async def fetch(args): + raise MaintenanceWindow("planned maintenance", retry_after=30.0) + + decorated = multi_value_chunked(build_request=_fake_build, url_limit=240)(fetch) + with pytest.raises(ServiceInterrupted) as excinfo: + decorated({"sites": ["S1" * 10, "S2" * 10]}) + + assert excinfo.value.retry_after == 30.0 + assert isinstance(excinfo.value.__cause__, MaintenanceWindow) + + 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.""" diff --git a/tests/waterdata_utils_test.py b/tests/waterdata_utils_test.py index 95aabe30..431cf7f4 100644 --- a/tests/waterdata_utils_test.py +++ b/tests/waterdata_utils_test.py @@ -676,6 +676,25 @@ def test_get_resp_data_always_materializes_id_column(): assert df["id"].isna().all() +def test_get_resp_data_flattens_nested_properties(): + """Nested GeoJSON properties keep underscore-separated column names.""" + resp = mock.MagicMock() + resp.json.return_value = { + "features": [ + { + "id": "feature-1", + "properties": {"station": {"code": "A"}, "value": 1}, + } + ] + } + + df = _get_resp_data(resp, geopd=False) + + assert df.to_dict(orient="records") == [ + {"value": 1, "station_code": "A", "id": "feature-1"} + ] + + # --- _arrange_cols ---------------------------------------------------------- diff --git a/tests/wateruse_test.py b/tests/wateruse_test.py index 031aad74..00a843c8 100644 --- a/tests/wateruse_test.py +++ b/tests/wateruse_test.py @@ -306,7 +306,7 @@ def test_fan_out_surfaces_final_rate_limit_header(httpx_mock): assert md.header["x-ratelimit-remaining"] == "850" -# (response aggregation now reuses ogc.planning._combine_chunk_responses; the +# (response aggregation now reuses ogc.combining._combine_chunk_responses; the # integration test above pins the rate-limit-header behavior end-to-end.)