Merge branch 'main' into store-multi-ts-add-args

krowvin · web-flow · commit 5a757cb0833e · 2026-03-30T20:59:36.000-05:00
diff --git a/cwms/api.py b/cwms/api.py
@@ -37,6 +37,7 @@
 from typing import Any, Optional, cast
 
 from requests import Response, adapters
+from requests.exceptions import RetryError as RequestsRetryError
 from requests_toolbelt import sessions  # type: ignore
 from requests_toolbelt.sessions import BaseUrlSession  # type: ignore
 from urllib3.util.retry import Retry
@@ -55,12 +56,12 @@
     status_forcelist=[
         403,
         429,
-        500,
         502,
         503,
         504,
     ],  # Example: also retry on these HTTP status codes
     allowed_methods=["GET", "PUT", "POST", "PATCH", "DELETE"],  # Methods to retry
+    raise_on_status=False,
 )
 SESSION = sessions.BaseUrlSession(base_url=API_ROOT)
 adapter = adapters.HTTPAdapter(
@@ -140,6 +141,27 @@ class PermissionError(ApiError):
     """Raised when the CDA request is not authorized for the current caller."""
 
 
+def _unwrap_retry_error(error: RequestsRetryError) -> Exception:
+    """Return the original retry cause when requests wraps it in RetryError."""
+
+    current: Exception = error
+    cause = error.__cause__
+    while isinstance(cause, Exception):
+        current = cause
+        cause = cause.__cause__
+
+    if current is error and error.args:
+        first_arg = error.args[0]
+        if isinstance(first_arg, Exception):
+            current = first_arg
+            reason = getattr(current, "reason", None)
+            while isinstance(reason, Exception):
+                current = reason
+                reason = getattr(current, "reason", None)
+
+    return current
+
+
 def init_session(
     *,
     api_root: Optional[str] = None,
@@ -308,11 +330,14 @@ def get(
     """
 
     headers = {"Accept": api_version_text(api_version)}
-    with SESSION.get(endpoint, params=params, headers=headers) as response:
-        if not response.ok:
-            logging.error(f"CDA Error: response={response}")
-            raise ApiError(response)
-        return _process_response(response)
+    try:
+        with SESSION.get(endpoint, params=params, headers=headers) as response:
+            if not response.ok:
+                logging.error(f"CDA Error: response={response}")
+                raise ApiError(response)
+            return _process_response(response)
+    except RequestsRetryError as error:
+        raise _unwrap_retry_error(error) from None
 
 
 def get_with_paging(
@@ -367,11 +392,16 @@ def _post_function(
     headers = {"accept": "*/*", "Content-Type": api_version_text(api_version)}
     if isinstance(data, dict) or isinstance(data, list):
         data = json.dumps(data)
-    with SESSION.post(endpoint, params=params, headers=headers, data=data) as response:
-        if not response.ok:
-            logging.error(f"CDA Error: response={response}")
-            raise ApiError(response)
-        return response
+    try:
+        with SESSION.post(
+            endpoint, params=params, headers=headers, data=data
+        ) as response:
+            if not response.ok:
+                logging.error(f"CDA Error: response={response}")
+                raise ApiError(response)
+            return response
+    except RequestsRetryError as error:
+        raise _unwrap_retry_error(error) from None
 
 
 def post(
@@ -461,10 +491,15 @@ def patch(
 
     if data and isinstance(data, dict) or isinstance(data, list):
         data = json.dumps(data)
-    with SESSION.patch(endpoint, params=params, headers=headers, data=data) as response:
-        if not response.ok:
-            logging.error(f"CDA Error: response={response}")
-            raise ApiError(response)
+    try:
+        with SESSION.patch(
+            endpoint, params=params, headers=headers, data=data
+        ) as response:
+            if not response.ok:
+                logging.error(f"CDA Error: response={response}")
+                raise ApiError(response)
+    except RequestsRetryError as error:
+        raise _unwrap_retry_error(error) from None
 
 
 def delete(
@@ -488,7 +523,10 @@ def delete(
     """
 
     headers = {"Accept": api_version_text(api_version)}
-    with SESSION.delete(endpoint, params=params, headers=headers) as response:
-        if not response.ok:
-            logging.error(f"CDA Error: response={response}")
-            raise ApiError(response)
+    try:
+        with SESSION.delete(endpoint, params=params, headers=headers) as response:
+            if not response.ok:
+                logging.error(f"CDA Error: response={response}")
+                raise ApiError(response)
+    except RequestsRetryError as error:
+        raise _unwrap_retry_error(error) from None
diff --git a/cwms/catalog/catalog.py b/cwms/catalog/catalog.py
@@ -67,7 +67,9 @@ def get_locations_catalog(
         "location-kind-like": location_kind_like,
     }
 
-    response = api.get(endpoint=endpoint, params=params, api_version=2)
+    response = api.get_with_paging(
+        endpoint=endpoint, selector="entries", params=params, api_version=2
+    )
     return Data(response, selector="entries")
 
 
@@ -131,7 +133,9 @@ def get_timeseries_catalog(
         "include-extents": include_extents,
     }
 
-    response = api.get(endpoint=endpoint, params=params, api_version=2)
+    response = api.get_with_paging(
+        endpoint=endpoint, selector="entries", params=params, api_version=2
+    )
     return Data(response, selector="entries")
 
 
diff --git a/cwms/timeseries/timeseries.py b/cwms/timeseries/timeseries.py
@@ -233,6 +233,14 @@ def combine_timeseries_results(results: List[Data]) -> Data:
     )
     combined_df["date-time"] = combined_df["date-time"].astype("Int64")
     combined_df = combined_df.reindex(columns=["date-time", "value", "quality-code"])
+
+    # Replace NaN in value column with None so they serialize as JSON null
+    # rather than the invalid JSON literal NaN.
+    combined_df["value"] = (
+        combined_df["value"]
+        .astype(object)
+        .where(combined_df["value"].notna(), other=None)
+    )
     # Update the "values" key in the JSON to include the combined data
     combined_json["values"] = combined_df.values.tolist()
 
@@ -438,8 +446,11 @@ def timeseries_df_to_json(
         pd.Timestamp.isoformat
     )
     df = df.reindex(columns=["date-time", "value", "quality-code"])
-    if df.isnull().values.any():
-        raise ValueError("Null/NaN data must be removed from the dataframe")
+
+    # Replace NaN/NA/NaT in value column with None so they serialize as JSON
+    # null rather than the invalid JSON literal NaN.
+    df["value"] = df["value"].astype(object).where(df["value"].notna(), other=None)
+
     if version_date:
         version_date_iso = version_date.isoformat()
     else:
diff --git a/cwms/users/users.py b/cwms/users/users.py
@@ -35,18 +35,66 @@ def get_user_profile() -> dict[str, Any]:
     return dict(response)
 
 
+def filter_users_by_office(data: dict[str, Any], office: str) -> dict[str, Any]:
+    """
+    Filter users JSON to only include users that have roles for the specified office.
+    Each user's roles dict will only contain the entry for that office.
+
+    Args:
+        data:   The full users JSON as a Python dict.
+        office: The office key to filter by (e.g., 'MVP', 'LRL').
+
+    Returns:
+        A new dict with the same structure, filtered to the specified office.
+    """
+    filtered_users = []
+
+    for user in data.get("users", []):
+        roles = user.get("roles", {})
+
+        if office in roles:
+            # Build a copy of the user with only the target office's roles
+            filtered_user = {k: v for k, v in user.items() if k != "roles"}
+            filtered_user["roles"] = {office: roles[office]}
+            filtered_users.append(filtered_user)
+
+    return {
+        "page": data.get("page"),
+        "page-size": data.get("page-size"),
+        "total": len(filtered_users),
+        "users": filtered_users,
+    }
+
+
 def get_users(
     office_id: Optional[str] = None,
+    username_like: Optional[str] = None,
+    include_roles: Optional[bool] = None,
     page: Optional[str] = None,
-    page_size: Optional[int] = None,
+    page_size: Optional[int] = 5000,
 ) -> Data:
     """Retrieve users with optional office and paging filters."""
 
-    params = {"office": office_id, "page": page, "page-size": page_size}
+    endpoint = "users"
+    params = {
+        "office": office_id,
+        "username-like": username_like,
+        "include-roles": include_roles,
+        "page": page,
+        "page-size": page_size,
+    }
     try:
-        response = api.get("users", params=params, api_version=1)
+        response = api.get_with_paging(
+            endpoint=endpoint, selector="users", params=params, api_version=1
+        )
     except api.ApiError as error:
         _raise_user_management_error(error, "User list lookup")
+
+    # filter by office if office_id is provided since the API does not
+    # currently support filtering by office on the backend. This is a
+    # temporary workaround until the API supports office filtering.
+    if office_id:
+        response = filter_users_by_office(response, office_id)
     return Data(response, selector="users")
 
 
diff --git a/examples/08_cwms_ts_shift_examples.ipynb b/examples/08_cwms_ts_shift_examples.ipynb
@@ -0,0 +1,77 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d9b66ebe",
+   "metadata": {
+    "vscode": {
+     "languageId": "plaintext"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "import pandas as pd\n",
+    "import cwms\n",
+    "import os\n",
+    "import logging\n",
+    "\n",
+    "\n",
+    "TS_NAME = \"KEYS.Elev.Inst.1Hour.0.Ccp-Rev\"\n",
+    "OFFICE = os.getenv(\"OFFICE\", \"SWT\")\n",
+    "API_KEY_DEV = os.getenv(\"CDA_API_KEY_DEV\")\n",
+    "API_ROOT_DEV = os.getenv(\"CDA_API_ROOT_DEV\")\n",
+    "\n",
+    "logging.basicConfig(\n",
+    "    level=logging.DEBUG, format=\"%(asctime)s - %(levelname)s - %(message)s\"\n",
+    ")\n",
+    "\n",
+    "\n",
+    "if not API_ROOT_DEV:\n",
+    "    logging.error(\n",
+    "        \"CDA_API_KEY environment variable is not set. Please set it to your API key.\"\n",
+    "    )\n",
+    "\n",
+    "if not API_KEY_DEV:\n",
+    "    logging.error(\n",
+    "        \"CDA_API_ROOT environment variable is not set. Please set it to your API root URL.\"\n",
+    "    )\n",
+    "\n",
+    "logging.debug(\"Starting Script!\")\n",
+    "\n",
+    "\n",
+    "logging.info(f\"Getting timeseries data for {TS_NAME} from office {OFFICE}\")\n",
+    "\n",
+    "cwms.init_session()\n",
+    "ts = cwms.get_timeseries(ts_id=TS_NAME, office_id=OFFICE)\n",
+    "shifted_df = ts.df.copy()\n",
+    "\n",
+    "if isinstance(shifted_df.index, pd.DatetimeIndex):\n",
+    "    shifted_df.index = shifted_df.index + pd.Timedelta(hours=10)\n",
+    "else:\n",
+    "    shifted_df[\"date-time\"] = pd.to_datetime(shifted_df[\"date-time\"]) + pd.Timedelta(\n",
+    "        hours=10\n",
+    "    )\n",
+    "\n",
+    "ts_json = cwms.timeseries_df_to_json(\n",
+    "    data=shifted_df,\n",
+    "    office_id=OFFICE,\n",
+    "    ts_id=TS_NAME,\n",
+    "    units=ts.json.get(\"units\"),\n",
+    ")\n",
+    "cwms.init_session(api_root=API_ROOT_DEV, api_key=API_KEY_DEV)\n",
+    "cwms.store_timeseries(data=ts_json)\n",
+    "\n",
+    "\n",
+    "logging.debug(\"Script Done\")\n"
+   ]
+  }
+ ],
+ "metadata": {
+  "language_info": {
+   "name": "python"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}
diff --git a/pyproject.toml b/pyproject.toml
@@ -2,8 +2,7 @@
 name = "cwms-python"
 repository = "https://github.com/HydrologicEngineeringCenter/cwms-python"
 
-version = "1.0.4"
-
+version = "1.0.7"
 
 packages = [
     { include = "cwms" },
diff --git a/tests/cda/timeseries/timeseries_CDA_test.py b/tests/cda/timeseries/timeseries_CDA_test.py
diff --git a/tests/cda/users/users_CDA_test.py b/tests/cda/users/users_CDA_test.py
diff --git a/tests/mock/api_test.py b/tests/mock/api_test.py
