Support batch bearer token context

.github/workflows/testing.yml

@@ -14,7 +14,7 @@ jobs:
       - name: Set Up Python
         uses: actions/setup-python@v6
         with:
-          python-version: '3.x'
+          python-version: '3.12'
 
       # Unlike the code-check workflow, this job requires the dev dependencies to be
       # installed to make sure we have the necessary, tools, stub files, etc.

README.md

@@ -38,6 +38,23 @@ cwms.init_session(
 If both `token` and `api_key` are provided, `cwms-python` will use the token
 and log a warning.
 
+#### CWMS batch jobs
+
+- The runner can provide short-lived authentication through
+  environment variables.
+- Given an env `CDA_BEARER_TOKEN` is set, `cwms.init_session()`
+  uses it as a bearer token if no explicit `token=` is provided in the args. In addition, when
+  `BATCH_JOB_CONTEXT_TOKEN` is set, it is sent as `X-CWMS-Job-Context` so CDA can
+  apply the authorized batch run context.
+
+```python
+import cwms
+
+# Reads from env variables
+cwms.init_session(api_root="https://cwms-data.usace.army.mil/cwms-data/")
+profile = cwms.get_user_profile()
+```
+
 ## Getting Started
 
 ```python
@@ -93,7 +110,7 @@ print(json)
 ## TimeSeries Profile API Compatibility Warning
 
 Currently, the TimeSeries Profile API may not be fully supported
-until a new version of cwms-data-access is released with the updated 
+until a new version of cwms-data-access is released with the updated
 endpoint implementation.
 
 ## Contributing

compose_files/togglz/features.properties

@@ -0,0 +1 @@
+AUTH_RE_ENABLE_NON_HASH_KEY_SUPPORT=true

cwms/api.py

@@ -32,6 +32,7 @@
 import base64
 import json
 import logging
+import os
 from http import HTTPStatus
 from json import JSONDecodeError
 from typing import Any, Optional, cast
@@ -167,6 +168,7 @@ def init_session(
     api_root: Optional[str] = None,
     api_key: Optional[str] = None,
     token: Optional[str] = None,
+    job_context_token: Optional[str] = None,
     pool_connections: int = 100,
 ) -> BaseUrlSession:
     """Specify a root URL and authentication credentials for the CWMS Data API.
@@ -181,6 +183,8 @@ def init_session(
         api_key (optional): An authentication key.
         token (optional): A Keycloak access token. If both token and api_key are
             provided, token is used.
+        job_context_token (optional): A signed batch job context token. If not
+            provided, BATCH_JOB_CONTEXT_TOKEN is used when present.
 
     Returns:
         Returns the updated session object.
@@ -198,6 +202,10 @@ def init_session(
             max_retries=retry_strategy,
         )
         SESSION.mount("https://", adapter)
+    if token is None:
+        # Batch runner images provide CDA_BEARER_TOKEN for normal init_session
+        # calls, keeping office API keys out of scripts and job definitions.
+        token = os.getenv("CDA_BEARER_TOKEN")
     if token:
         if api_key:
             logging.warning(
@@ -211,6 +219,12 @@ def init_session(
         if api_key.startswith("apikey "):
             api_key = api_key.replace("apikey ", "")
         SESSION.headers.update({"Authorization": "apikey " + api_key})
+    if job_context_token is None:
+        # Keep batch run context separate from the OAuth token; CDA can validate
+        # it only when the job runner provides the signed fallback context.
+        job_context_token = os.getenv("BATCH_JOB_CONTEXT_TOKEN")
+    if job_context_token:
+        SESSION.headers.update({"X-CWMS-Job-Context": job_context_token})
 
     return SESSION
 

docker-compose.yml

@@ -56,8 +56,10 @@ services:
     restart: unless-stopped
     volumes:
       - ./compose_files/pki/certs:/conf/
+      - ./compose_files/togglz/features.properties:/conf/features.properties:ro
       - ./compose_files/tomcat/logging.properties:/usr/local/tomcat/conf/logging.properties:ro
     environment:
+      - JAVA_OPTS=-Dproperties.file=/conf/features.properties
       - CDA_JDBC_DRIVER=oracle.jdbc.driver.OracleDriver
       - CDA_JDBC_URL=jdbc:oracle:thin:@db/FREEPDB1
       - CDA_JDBC_USERNAME=s0webtest
@@ -140,4 +142,3 @@ services:
       - "traefik.enable=true"
       - "traefik.http.routers.traefik.rule=PathPrefix(`/traefik`)"
       - "traefik.http.routers.traefik.service=api@internal"
-

tests/cda/timeseries/timeseries_CDA_test.py

@@ -1,3 +1,4 @@
+import time
 from datetime import datetime, timedelta, timezone
 from unittest.mock import patch
 
@@ -7,6 +8,7 @@
 
 import cwms
 import cwms.timeseries.timeseries as ts
+from cwms.api import ApiError
 
 TEST_OFFICE = "MVP"
 TEST_LOCATION_ID = "pytest_ts"
@@ -137,6 +139,17 @@ def init_session():
     print("Initializing CWMS API session for timeseries tests...")
 
 
+def get_timeseries_with_retry(*args, **kwargs):
+    last_error = None
+    for _ in range(5):
+        try:
+            return ts.get_timeseries(*args, **kwargs)
+        except ApiError as exc:
+            last_error = exc
+            time.sleep(1)
+    raise last_error
+
+
 def test_store_timeseries():
     ts_json = {
         "name": TEST_TSID_STORE,
@@ -145,7 +158,9 @@ def test_store_timeseries():
         "values": [[EPOCH_MS, 99, 0]],
     }
     ts.store_timeseries(ts_json)
-    data = ts.get_timeseries(TEST_TSID_STORE, TEST_OFFICE, begin=BEGIN, end=END).json
+    data = get_timeseries_with_retry(
+        TEST_TSID_STORE, TEST_OFFICE, begin=BEGIN, end=END
+    ).json
     assert data["name"] == TEST_TSID_STORE
     assert data["office-id"] == TEST_OFFICE
     assert data["units"] == "ft"

tests/mock/api_test.py

@@ -43,13 +43,48 @@
     """Verify that the authentication key for the session can be set."""
 
     # Initialize a session with both an alternate root URL and an authentication key.
     session = init_session(api_root="https://example.com/", api_key="API_AUTH_KEY")
 
     # Both the URL and the auth key should be set on the session.
     assert session.base_url == "https://example.com/"
     assert session.headers["Authorization"] == "apikey API_AUTH_KEY"
 
 
+def test_session_init_token_from_env(monkeypatch):
+    monkeypatch.setenv("CDA_BEARER_TOKEN", "Bearer ENV_TOKEN")
+
+    session = init_session(api_root="https://example.com/")
+
+    assert session.headers["Authorization"] == "Bearer ENV_TOKEN"
+
+
+def test_session_init_explicit_token_precedes_env(monkeypatch):
+    monkeypatch.setenv("CDA_BEARER_TOKEN", "ENV_TOKEN")
+
+    session = init_session(api_root="https://example.com/", token="ARG_TOKEN")
+
+    assert session.headers["Authorization"] == "Bearer ARG_TOKEN"
+
+
+def test_session_init_batch_job_context_from_env(monkeypatch):
+    monkeypatch.setenv("BATCH_JOB_CONTEXT_TOKEN", "JOB_CONTEXT")
+
+    session = init_session(api_root="https://example.com/")
+
+    assert session.headers["X-CWMS-Job-Context"] == "JOB_CONTEXT"
+
+
+def test_session_init_explicit_batch_job_context_precedes_env(monkeypatch):
+    monkeypatch.setenv("BATCH_JOB_CONTEXT_TOKEN", "ENV_CONTEXT")
+
+    session = init_session(
+        api_root="https://example.com/",
+        job_context_token="ARG_CONTEXT",
+    )
+
+    assert session.headers["X-CWMS-Job-Context"] == "ARG_CONTEXT"
+
+
 def test_api_headers():
     """Verify that the API version headers are correct."""