river.io.usgs: move requests to the USGS Water Data APIs (WaterServices decommission)

mhkit/river/io/usgs.py

@@ -3,6 +3,10 @@
 Geological Survey (USGS) National Water Information System (NWIS). It enables access to
 river flow data and related measurements useful for hydrokinetic resource assessment.
 
+Requests go to the USGS Water Data APIs (https://api.waterdata.usgs.gov/ogcapi/v1/), which
+replace the legacy WaterServices endpoints (waterservices.usgs.gov) that USGS is
+decommissioning. Files saved from either service can still be read.
+
 """
 
 from typing import Dict, Union, Optional
@@ -14,6 +18,42 @@
 import xarray as xr
 from pandas import DataFrame
 from mhkit.utils.cache import handle_caching
+from mhkit.utils.type_handling import convert_to_dataset
+
+
+WATER_DATA_API = "https://api.waterdata.usgs.gov/ogcapi/v1"
+
+
+def _read_water_data_features(text: Dict) -> DataFrame:
+    """
+    Process a USGS Water Data API response (GeoJSON features) into a DataFrame
+    with the same layout as the legacy WaterServices reader: a UTC datetime
+    index and one numeric column per parameter, named by its description.
+
+    Parameters
+    ----------
+    text : dict
+        GeoJSON FeatureCollection from the Water Data API. The parameter
+        description used for the column name is read from
+        ``text["parameter_description"]`` when present, else the parameter code.
+
+    Returns
+    -------
+    data : pandas.DataFrame
+        Processed time series data
+    """
+    rows = [feature["properties"] for feature in text.get("features", [])]
+    if not rows:
+        return pd.DataFrame()
+    frame = pd.DataFrame(rows)
+    name = text.get("parameter_description") or str(frame["parameter_code"].iloc[0])
+    data = pd.DataFrame(
+        {name: pd.to_numeric(frame["value"], errors="coerce").values},
+        index=pd.to_datetime(frame["time"], utc=True),
+    )
+    data.index.name = None
+    data = data[~data.index.duplicated(keep="last")].sort_index()
+    return data
 
 
 def _read_usgs_json(text: Dict, to_pandas: bool = True) -> Union[DataFrame, xr.Dataset]:
@@ -23,7 +63,8 @@ def _read_usgs_json(text: Dict, to_pandas: bool = True) -> Union[DataFrame, xr.D
     Parameters
     ----------
     text : dict
-        JSON response from USGS API containing time series data
+        JSON response from USGS API containing time series data, either
+        legacy WaterServices JSON or Water Data API GeoJSON
     to_pandas : bool, optional
         Flag to output pandas instead of xarray. Default = True.
 
@@ -32,6 +73,10 @@ def _read_usgs_json(text: Dict, to_pandas: bool = True) -> Union[DataFrame, xr.D
     data : pandas.DataFrame or xarray.Dataset
         Processed time series data
     """
+    if "features" in text:
+        data = _read_water_data_features(text)
+        return data if to_pandas else convert_to_dataset(data)
+
     data = pd.DataFrame()
     for i in range(len(text["value"]["timeSeries"])):
         try:
@@ -53,7 +98,7 @@ def _read_usgs_json(text: Dict, to_pandas: bool = True) -> Union[DataFrame, xr.D
             continue
 
     if not to_pandas:
-        data = data.to_dataset()
+        data = convert_to_dataset(data)
 
     return data
 
@@ -62,7 +107,9 @@ def read_usgs_file(
     file_name: str, to_pandas: bool = True
 ) -> Union[DataFrame, xr.Dataset]:
     """
-    Reads a USGS JSON data file (from https://waterdata.usgs.gov/nwis)
+    Reads a USGS JSON data file: legacy WaterServices JSON (from
+    https://waterdata.usgs.gov/nwis), or Water Data API GeoJSON saved from
+    https://api.waterdata.usgs.gov/ogcapi/v1/
 
     Parameters
     ----------
@@ -97,10 +144,13 @@ def request_usgs_data(
     options: Optional[Dict] = None,
 ) -> Union[DataFrame, xr.Dataset]:
     """
-    Loads USGS data directly from https://waterdata.usgs.gov/nwis using a
-    GET request
+    Loads USGS data from the USGS Water Data APIs
+    (https://api.waterdata.usgs.gov/ogcapi/v1/) using GET requests: the
+    ``daily`` collection for 'Daily' and ``continuous`` for 'Instantaneous'.
 
-    The request URL prints to the screen.
+    The request URLs print to the screen. Instantaneous requests are split
+    into one query per year, as the ``continuous`` collection serves at most
+    three years per query.
 
     Parameters
     ----------
@@ -111,7 +161,9 @@ def request_usgs_data(
     start_date : str
         Start date in the format 'YYYY-MM-DD' (e.g. '2018-01-01')
     end_date : str
-        End date in the format 'YYYY-MM-DD' (e.g. '2018-12-31')
+        End date in the format 'YYYY-MM-DD' (e.g. '2018-12-31'). Both dates
+        are inclusive. For 'Instantaneous' data they bound the request in UTC;
+        the legacy WaterServices endpoint read them in the site's local time.
     options : dict, optional
         Dictionary containing optional parameters:
         - data_type: str
@@ -127,6 +179,10 @@ def request_usgs_data(
             Flag to output pandas instead of xarray. Default = True
         - timeout: int
             Timeout in seconds for the HTTP request. Default = 30
+        - api_key: str or None
+            USGS Water Data API key for higher rate limits
+            (https://api.waterdata.usgs.gov/signup/). Default = the
+            API_USGS_PAT environment variable, else None
 
     Returns
     -------
@@ -169,53 +225,18 @@ def request_usgs_data(
         return cached_data
 
     # If no cached data, proceed with the API request
-    if data_type == "Daily":
-        data_url = "https://waterservices.usgs.gov/nwis/dv"
-        api_query = (
-            "/?format=json&sites="
-            + station
-            + "&startDT="
-            + start_date
-            + "&endDT="
-            + end_date
-            + "&statCd=00003"
-            + "&parameterCd="
-            + parameter
-            + "&siteStatus=all"
-        )
-    else:
-        data_url = "https://waterservices.usgs.gov/nwis/iv"
-        api_query = (
-            "/?format=json&sites="
-            + station
-            + "&startDT="
-            + start_date
-            + "&endDT="
-            + end_date
-            + "&parameterCd="
-            + parameter
-            + "&siteStatus=all"
-        )
-
-    print("Data request URL: ", data_url + api_query)
-
-    max_retries = 3
-    retry_count = 0
-    while retry_count < max_retries:
-        try:
-            response = requests.get(
-                url=data_url + api_query, proxies=proxy, timeout=timeout, verify=True
-            )
-            text = json.loads(response.text)
-            break
-        except requests.exceptions.SSLError as e:
-            retry_count += 1
-            if retry_count == max_retries:
-                raise e
-            print(
-                f"SSL Error occurred, retrying... (Attempt {retry_count}/{max_retries})"
-            )
-            continue
+    text = _request_water_data(
+        station,
+        parameter,
+        start_date,
+        end_date,
+        {
+            "data_type": data_type,
+            "proxy": proxy,
+            "timeout": timeout,
+            "api_key": options.get("api_key", os.environ.get("API_USGS_PAT")),
+        },
+    )
 
     # handle_caching is only set-up for pandas, so force this data to output as pandas for now
     data = _read_usgs_json(text, True)
@@ -233,6 +254,111 @@ def request_usgs_data(
         shutil.copy(cache_filepath, write_json)
 
     if not to_pandas:
-        data = data.to_dataset()
+        data = convert_to_dataset(data)
 
     return data
+
+
+def _get_json(url: str, params: Optional[Dict], proxy, timeout) -> Dict:
+    """GET a URL and parse JSON, retrying SSL errors up to three times."""
+    max_retries = 3
+    for attempt in range(1, max_retries + 1):
+        try:
+            response = requests.get(
+                url=url, params=params, proxies=proxy, timeout=timeout, verify=True
+            )
+            response.raise_for_status()
+            return response.json()
+        except requests.exceptions.SSLError:
+            if attempt == max_retries:
+                raise
+            print(f"SSL Error occurred, retrying... (Attempt {attempt}/{max_retries})")
+    raise RuntimeError("unreachable")
+
+
+def _year_windows(start_date: str, end_date: str):
+    """Split [start_date, end_date] into consecutive windows of at most one year."""
+    start = pd.Timestamp(start_date, tz="UTC")
+    end = pd.Timestamp(end_date, tz="UTC") + pd.Timedelta(days=1)
+    while start < end:
+        stop = min(start + pd.DateOffset(years=1), end)
+        yield start, stop
+        start = stop
+
+
+def _request_water_data(
+    station: str, parameter: str, start_date: str, end_date: str, settings: Dict
+) -> Dict:
+    """
+    Request all features for one station and parameter from the Water Data
+    API, following "next" links, and return them as one FeatureCollection with
+    the parameter's description attached for column naming.
+    """
+    proxy, timeout, api_key = (
+        settings["proxy"],
+        settings["timeout"],
+        settings["api_key"],
+    )
+    base = {
+        "f": "json",
+        "monitoring_location_id": f"USGS-{station}",
+        "parameter_code": parameter,
+        "limit": 50000,
+    }
+    if api_key:
+        base["api_key"] = api_key
+
+    if settings["data_type"] == "Daily":
+        collection = "daily"
+        windows = [(start_date, end_date)]
+        base["statistic_id"] = "00003"
+    else:
+        collection = "continuous"
+        windows = [
+            (a.strftime("%Y-%m-%dT%H:%M:%SZ"), b.strftime("%Y-%m-%dT%H:%M:%SZ"))
+            for a, b in _year_windows(start_date, end_date)
+        ]
+    url = f"{WATER_DATA_API}/collections/{collection}/items"
+
+    features = []
+    for begin, finish in windows:
+        params = dict(base, time=f"{begin}/{finish}")
+        print(
+            "Data request URL: ",
+            requests.Request("GET", url, params=params).prepare().url,
+        )
+        page = _get_json(url, params, proxy, timeout)
+        features += page.get("features", [])
+        next_link = _next_link(page)
+        while next_link:
+            if api_key and "api_key=" not in next_link:
+                next_link += f"&api_key={api_key}"
+            page = _get_json(next_link, None, proxy, timeout)
+            features += page.get("features", [])
+            next_link = _next_link(page)
+
+    description = None
+    try:
+        code = _get_json(
+            f"{WATER_DATA_API}/collections/parameter-codes/items/{parameter}",
+            {"f": "json"},
+            proxy,
+            timeout,
+        )
+        description = code.get("properties", {}).get("parameter_description")
+    except (requests.exceptions.RequestException, ValueError):
+        description = None
+
+    return {
+        "type": "FeatureCollection",
+        "features": features,
+        "parameter_description": description,
+    }
+
+
+def _next_link(page: Dict) -> Optional[str]:
+    """Return the href of a page's "next" link, if any."""
+    for link in page.get("links", []):
+        if link.get("rel") == "next":
+            return link.get("href")
+    return None