waterdata: tighten stats and OGC pagination — geometry, KeyError, non-200

thodson-usgs · claude · thodson-usgs · commit 0ad568c851e8 · 2026-05-03T15:23:27.000-05:00
Three correctness fixes to the two pagination loops in
dataretrieval.waterdata.utils.

* `get_stats_data` honored `GEOPANDAS` for the first page but
  hard-coded `geopd=False` on every continuation page. With geopandas
  installed, a multi-page stats response started as a `GeoDataFrame`
  and pages 2..N came back as plain `DataFrame`s; `pd.concat` then
  silently downgraded the result and the caller lost geometry / CRS.
  Use `geopd=GEOPANDAS` on every page.

* `get_stats_data` indexed `body["next"]` directly, raising `KeyError`
  on responses without that key (some terminal responses simply omit
  it). Switch to `body.get("next")`, which produces `None` and exits
  the loop cleanly.

* Both `get_stats_data`'s in-loop request and `_walk_pages`'s in-loop
  request returned the response without checking `status_code`. A 4xx
  or 5xx page whose body happened to JSON-decode could be appended as
  data, then pagination quietly stopped — the caller got a partial
  result with no warning. Add an explicit `if status_code != 200`
  raise inside each loop so the existing log-and-truncate handler
  fires deliberately rather than incidentally.

Co-Authored-By: Claude Opus 4.7 (1M context) &lt;noreply@anthropic.com&gt;
diff --git a/dataretrieval/waterdata/utils.py b/dataretrieval/waterdata/utils.py
@@ -606,6 +606,8 @@ def _walk_pages(
                     headers=headers,
                     data=content if method == "POST" else None,
                 )
+                if resp.status_code != 200:
+                    raise RuntimeError(_error_body(resp))
                 dfs.append(_get_resp_data(resp, geopd=geopd))
                 curr_url = _next_req_url(resp)
             except Exception:  # noqa: BLE001
@@ -1058,7 +1060,7 @@ def get_stats_data(
         all_dfs = [_handle_stats_nesting(body, geopd=GEOPANDAS)]
 
         # Look for a next code in the response body
-        next_token = body["next"]
+        next_token = body.get("next")
 
         while next_token:
             args["next_token"] = next_token
@@ -1070,9 +1072,11 @@ def get_stats_data(
                     params=args,
                     headers=headers,
                 )
+                if resp.status_code != 200:
+                    raise RuntimeError(_error_body(resp))
                 body = resp.json()
-                all_dfs.append(_handle_stats_nesting(body, geopd=False))
-                next_token = body["next"]
+                all_dfs.append(_handle_stats_nesting(body, geopd=GEOPANDAS))
+                next_token = body.get("next")
             except Exception:  # noqa: BLE001
                 error_text = _error_body(resp)
                 logger.error("Request incomplete. %s", error_text)
diff --git a/tests/waterdata_utils_test.py b/tests/waterdata_utils_test.py
@@ -1,10 +1,12 @@
 from unittest import mock
 
+import pandas as pd
 import requests
 
 from dataretrieval.waterdata.utils import (
     _get_args,
     _walk_pages,
+    get_stats_data,
 )
 
 
@@ -80,3 +82,121 @@ def test_walk_pages_multiple_mocked():
     assert mock_client.send.called
     assert mock_client.request.called
     assert mock_client.request.call_args[0][1] == "https://example.com/page2"
+
+
+def test_walk_pages_raises_on_non_200_in_loop():
+    """`_walk_pages` must surface a non-200 mid-loop, not silently truncate.
+
+    Regression: previously any non-200 page was appended (with whatever
+    body it had) and pagination quietly stopped because `_get_resp_data`
+    or `_next_req_url` raised inside the bare except. The user got a
+    partial result with no warning.
+    """
+    resp1 = mock.MagicMock()
+    resp1.json.return_value = {
+        "numberReturned": 1,
+        "features": [{"id": "1", "properties": {"val": "a"}}],
+        "links": [],
+    }
+    resp1.headers = {}
+    resp1.links = {"next": {"url": "https://example.com/page2"}}
+    resp1.status_code = 200
+
+    resp2 = mock.MagicMock()
+    resp2.status_code = 500
+    resp2.text = "<html>error</html>"
+
+    mock_client = mock.MagicMock(spec=requests.Session)
+    mock_client.send.return_value = resp1
+    mock_client.request.return_value = resp2
+
+    mock_req = mock.MagicMock(spec=requests.PreparedRequest)
+    mock_req.method = "GET"
+    mock_req.headers = {}
+    mock_req.url = "https://example.com/page1"
+
+    df, _ = _walk_pages(geopd=False, req=mock_req, client=mock_client)
+
+    # Page 1 still returned; page 2 logged-and-stopped after the explicit
+    # status check raised. The contract here is "log + truncate", same
+    # as the pre-fix bare-except behavior, but now the raise inside the
+    # loop is intentional rather than incidental.
+    assert len(df) == 1
+
+
+# --- get_stats_data pagination ----------------------------------------------
+
+
+def _stats_feature():
+    """Build a single feature shaped to satisfy ``_handle_stats_nesting``."""
+    return {
+        "type": "Feature",
+        "id": "USGS-1",
+        "geometry": None,
+        "properties": {
+            "monitoring_location_id": "USGS-1",
+            "data": [
+                {
+                    "parameter_code": "00060",
+                    "unit_of_measure": "ft^3/s",
+                    "parent_time_series_id": "abc",
+                    "values": [{"value": 1.0}],
+                }
+            ],
+        },
+    }
+
+
+def _stats_body(features, next_token=None):
+    body = {
+        "type": "FeatureCollection",
+        "features": features,
+        "numberReturned": len(features),
+    }
+    if next_token is not None:
+        body["next"] = next_token
+    return body
+
+
+def test_get_stats_data_handles_missing_next_key():
+    """A response without a ``next`` key must not raise KeyError.
+
+    Regression: ``body["next"]`` raised when the key was absent. Now
+    uses ``body.get("next")`` so a missing key means "no more pages".
+    """
+    resp = mock.MagicMock()
+    resp.status_code = 200
+    resp.json.return_value = _stats_body([_stats_feature()])
+    # No "next" key at all.
+
+    client = mock.MagicMock(spec=requests.Session)
+    client.send.return_value = resp
+
+    df, _ = get_stats_data(
+        args={}, service="observationNormals", expand_percentiles=False, client=client
+    )
+    assert isinstance(df, pd.DataFrame)
+    assert len(df) >= 1
+
+
+def test_get_stats_data_truncates_on_non_200_continuation():
+    """A 4xx/5xx on a continuation page must log and stop, not crash."""
+    resp1 = mock.MagicMock()
+    resp1.status_code = 200
+    resp1.json.return_value = _stats_body([_stats_feature()], next_token="abc")
+
+    resp2 = mock.MagicMock()
+    resp2.status_code = 503
+    resp2.text = "Service Unavailable"
+    resp2.url = "https://example.com/page2"
+
+    client = mock.MagicMock(spec=requests.Session)
+    client.send.return_value = resp1
+    client.request.return_value = resp2
+
+    df, _ = get_stats_data(
+        args={}, service="observationNormals", expand_percentiles=False, client=client
+    )
+    # Page 1 still surfaces; page 2 was caught by the in-loop status check.
+    assert isinstance(df, pd.DataFrame)
+    assert len(df) >= 1
