diff --git a/dataretrieval/ogc/chunking.py b/dataretrieval/ogc/chunking.py index d0dae01c..78efcfea 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, @@ -104,6 +101,30 @@ _retry, ) +# --------------------------------------------------------------------------- +# 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 + + # Empirically the API replies HTTP 414 above ~8200 bytes of full URL — # matches nginx's default ``large_client_header_buffers`` of 8 KB. 8000 # leaves ~200 bytes for request-line framing and proxy variance. The decorator @@ -433,7 +454,7 @@ def _combine_raw(self) -> tuple[pd.DataFrame, httpx.Response]: frames = [self._chunks[i][0] for i in sorted(self._chunks)] responses = [response for _, response in self._chunks.values()] return ( - _combine_chunk_frames(frames), + _combine_chunk_frames(frames, dedup=self.plan.has_filter_axis), _combine_chunk_responses(responses, self.plan.canonical_url), ) diff --git a/dataretrieval/ogc/combining.py b/dataretrieval/ogc/combining.py new file mode 100644 index 00000000..ec87bc33 --- /dev/null +++ b/dataretrieval/ogc/combining.py @@ -0,0 +1,245 @@ +"""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: + 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 + ) + + +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_frames( + frames: list[pd.DataFrame], *, dedup: bool = True +) -> pd.DataFrame: + """ + Concatenate per-chunk frames, dropping empties and optionally deduping + by ``id``. + + Parameters + ---------- + frames : list[pandas.DataFrame] + One frame per completed sub-request. + dedup : bool, default True + Whether to drop duplicate rows keyed on the ``id`` column. + List-axis chunks (multi-value list parameters like + ``monitoring_location_id``) produce non-overlapping partitions, + so dedup is unnecessary and skipping it saves ~2 ms on a 50 k-row + result. Filter-axis chunks *can* overlap (a feature matching + multiple OR-clauses appears in each clause's chunk), so dedup is + required there. Callers that know their chunks don't overlap + (e.g. :meth:`ChunkedCall._combine_raw` when the plan has no + filter axis) pass ``dedup=False``. + + Returns + ------- + pandas.DataFrame + The concatenated (and optionally 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 dedup and "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 + + +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/engine.py b/dataretrieval/ogc/engine.py index 668fa7b7..3d22ac35 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, 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..45ae343e 100644 --- a/dataretrieval/ogc/planning.py +++ b/dataretrieval/ogc/planning.py @@ -1,31 +1,30 @@ -"""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:`_combine_chunk_frames`, +:func:`_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 +129,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: """ @@ -596,6 +549,28 @@ def total(self) -> int: """ return math.prod((len(self.chunks[ax.arg_key]) for ax in self.axes), start=1) + @property + def has_filter_axis(self) -> bool: + """Whether the plan splits along the cql-text ``filter`` axis. + + Filter-axis chunks can overlap (a feature matching multiple + OR-clauses appears in each clause's chunk), so deduplication is + required when combining their frames. List-axis chunks (the + common case — multi-value list parameters like + ``monitoring_location_id``) never overlap, so the combine step + can skip the ``drop_duplicates`` call. + + Returns + ------- + bool + ``True`` when the plan has a filter axis with >1 chunk; + ``False`` otherwise. + """ + return any( + ax.joiner != _LIST_SEP and len(self.chunks[ax.arg_key]) > 1 + for ax in self.axes + ) + def iter_sub_args(self) -> Iterator[dict[str, Any]]: """ Yield substituted args for each sub-request, in deterministic @@ -619,171 +594,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..8c3fab17 100644 --- a/dataretrieval/ogc/retry.py +++ b/dataretrieval/ogc/retry.py @@ -190,6 +190,39 @@ 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 a SINGLE exception as a known transient (resumable) failure. + + Does NOT walk the ``__cause__`` chain — inspects only the exception + passed in. Factored out of :func:`_classify_chunk_error` (which walks + the ``__cause__`` chain and calls this on each link) so the transient + taxonomy lives in one place. :func:`_retryable` deliberately does *not* + delegate here — it uses a narrower set (``httpx.InvalidURL`` and a bare + ``httpx.HTTPError`` are resumable but not worth an automatic retry). + + Parameters + ---------- + exc : BaseException + A single exception to classify. + + Returns + ------- + tuple[type[ChunkInterrupted], float or None] or None + ``(interrupted_class, retry_after)`` for a recognized transient + failure; ``None`` otherwise. + """ + if isinstance(exc, RateLimited): + return QuotaExhausted, exc.retry_after + if isinstance(exc, ServiceUnavailable): + 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,12 +264,9 @@ 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 @@ -262,6 +292,13 @@ def _retryable(exc: BaseException) -> tuple[bool, float | None]: ``(retryable, retry_after)`` — the server ``Retry-After`` hint (seconds) when the transient carried one, else ``None``. """ + # Only initial-page raw transients (TransientError, TransportError) are + # retryable — not a wrapped mid-pagination DataRetrievalError. This is a + # deliberately narrower taxonomy than _classify_transient's: RateLimited + # and ServiceUnavailable are caught via their TransientError base, but + # httpx.InvalidURL and bare httpx.HTTPError are intentionally excluded + # (a too-long cursor or non-transport HTTP error won't fix on retry), so + # this stays a separate check rather than delegating to _classify_transient. if isinstance(exc, TransientError): return True, exc.retry_after if isinstance(exc, httpx.TransportError): diff --git a/dataretrieval/ogc/shaping.py b/dataretrieval/ogc/shaping.py index 4df242f2..315be869 100644 --- a/dataretrieval/ogc/shaping.py +++ b/dataretrieval/ogc/shaping.py @@ -128,13 +128,17 @@ 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] + # Build the frame directly from the flat properties dicts rather than + # routing through ``pd.json_normalize``. The OGC API returns flat + # (non-nested) property objects, so ``json_normalize``'s recursive + # flattening is unnecessary overhead (~2.5× slower than plain + # ``pd.DataFrame`` for typical page sizes). The ``id`` key is merged + # into each row dict up-front so the column is always present (may be + # all-None) — ``_arrange_cols`` downstream relies on it for the rename + # to the service-specific output_id (daily_id, channel_measurements_id, + # …). + rows = [{**(f.get("properties") or {}), "id": f.get("id")} for f in features] + df = pd.DataFrame(rows) _attach_coordinates(df, features) return df @@ -369,7 +373,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..9d223f7c 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, diff --git a/tests/waterdata_chunking_test.py b/tests/waterdata_chunking_test.py index 600f05d8..12278f29 100644 --- a/tests/waterdata_chunking_test.py +++ b/tests/waterdata_chunking_test.py @@ -47,6 +47,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 +61,7 @@ _LIST_SEP, _NEVER_CHUNK, _OR_SEP, - _QUOTA_HEADER, ChunkPlan, - _combine_chunk_frames, - _combine_chunk_responses, _extract_axes, _request_bytes, _safe_request_bytes, 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.)