From 11a6a6c10251a5228824538f23a9c77d01ce0e0d Mon Sep 17 00:00:00 2001 From: Sujen Shah Date: Fri, 16 Oct 2026 15:26:42 -0700 Subject: [PATCH 1/2] fix: add timeout, retry and clear errors to get_maap_pgt_token - 5s connect / 30s read timeout; up to 6 attempts with exponential backoff and jitter on connect errors, timeouts and 5xx - 4xx and malformed responses fail immediately with PgtTokenError - on failure, message goes to stderr (stdout is captured into MAAP_PGT), _alt_error.txt is written for the HySDS job `error` field, and the script exits 78; dps_wrapper.sh propagates the exit code - resolve MAAP_API_HOST at call time instead of import time Closes #8 Claude-Session: https://claude.ai/code/session_0171HvL7yCdP2wCuxh7UngYP --- dps_wrapper.sh | 2 +- get_maap_pgt_token.py | 78 +++++++++++++++++++-- tests/test_get_maap_pgt_token.py | 117 +++++++++++++++++++++++++++++++ 3 files changed, 189 insertions(+), 8 deletions(-) create mode 100644 tests/test_get_maap_pgt_token.py diff --git a/dps_wrapper.sh b/dps_wrapper.sh index 72edadb..63bda48 100644 --- a/dps_wrapper.sh +++ b/dps_wrapper.sh @@ -4,7 +4,7 @@ USERNAME=$(python /app/get_username.py) source /maap-dps.env export MAAP_API_HOST export DPS_MACHINE_TOKEN -MAAP_PGT=$(python /app/get_maap_pgt_token.py) +MAAP_PGT=$(python /app/get_maap_pgt_token.py) || exit $? export MAAP_PGT unset DPS_MACHINE_TOKEN diff --git a/get_maap_pgt_token.py b/get_maap_pgt_token.py index 28d3546..6a89492 100644 --- a/get_maap_pgt_token.py +++ b/get_maap_pgt_token.py @@ -1,14 +1,77 @@ import json import os +import random +import sys +import time + import requests +TIMEOUT = (5, 30) # (connect, read) seconds +MAX_ATTEMPTS = 6 +BACKOFF_BASE = 1.0 # seconds; doubles each attempt +BACKOFF_CAP = 30.0 +EXIT_CODE_PGT_FAILURE = 78 # distinct job exit status for PGT token failures + + +class PgtTokenError(Exception): + """Raised when the MAAP PGT token cannot be obtained.""" + -def get_maap_pgt_token(username, maap_api_host=os.environ.get("MAAP_API_HOST")): +def _sleep_time(attempt): + """Exponential backoff with full jitter for the given 0-based attempt.""" + return random.uniform(0, min(BACKOFF_CAP, BACKOFF_BASE * 2 ** attempt)) + + +def get_maap_pgt_token(username, maap_api_host=None): + maap_api_host = maap_api_host or os.environ.get("MAAP_API_HOST") url = f"https://{maap_api_host}/api/members/{username}" headers = {"dps-token": os.environ.get("DPS_MACHINE_TOKEN")} - response = requests.get(url, headers=headers) - maap_pgt_token = response.json()["session_key"] - return maap_pgt_token + + last_error = None + for attempt in range(MAX_ATTEMPTS): + if attempt: + time.sleep(_sleep_time(attempt - 1)) + try: + response = requests.get(url, headers=headers, timeout=TIMEOUT) + except (requests.ConnectionError, requests.Timeout) as e: + last_error = f"{type(e).__name__}: {e}" + continue + + if response.status_code >= 500: + last_error = f"HTTP {response.status_code}" + continue + + try: + response.raise_for_status() + except requests.HTTPError as e: + # 4xx is permanent (e.g. bad token, unknown user): do not retry + raise PgtTokenError( + f"Request for user '{username}' to {url} failed: HTTP {response.status_code}" + ) from e + + try: + return response.json()["session_key"] + except (ValueError, KeyError, TypeError) as e: + raise PgtTokenError( + f"Response from {url} for user '{username}' has no usable 'session_key'" + ) from e + + raise PgtTokenError( + f"Could not get MAAP_PGT for user '{username}' from {url} " + f"after {MAX_ATTEMPTS} attempts; last error: {last_error}" + ) + + +def report_failure(error, job_dir="."): + """Surface the failure to HySDS. + + stderr lands in _stderr.txt; _alt_error.txt overrides the job's `error` + field (job_worker reads it from the job dir). Keep it a short single line. + """ + message = " ".join(str(error).split()) + print(f"Unable to get user MAAP_PGT token: {message}", file=sys.stderr) + with open(os.path.join(job_dir, "_alt_error.txt"), "w") as f: + f.write(f"MAAP_PGT token failure: {message}\n") if __name__ == '__main__': @@ -19,6 +82,7 @@ def get_maap_pgt_token(username, maap_api_host=os.environ.get("MAAP_API_HOST")): try: maap_pgt_token = get_maap_pgt_token(username) print(maap_pgt_token) - except Exception as e: - print(f"Unable to get user MAAP_PGT token: {e}") - raise e + except PgtTokenError as e: + # stdout is captured into MAAP_PGT by dps_wrapper.sh, so report on stderr + report_failure(e) + sys.exit(EXIT_CODE_PGT_FAILURE) diff --git a/tests/test_get_maap_pgt_token.py b/tests/test_get_maap_pgt_token.py new file mode 100644 index 0000000..ad4a3d4 --- /dev/null +++ b/tests/test_get_maap_pgt_token.py @@ -0,0 +1,117 @@ +import sys +from pathlib import Path +from unittest.mock import Mock, patch + +import pytest +import requests + +sys.path.insert(0, str(Path(__file__).parent.parent)) + +import get_maap_pgt_token as mod +from get_maap_pgt_token import get_maap_pgt_token, PgtTokenError + + +def _resp(status=200, body=None): + r = Mock() + r.status_code = status + r.json.return_value = body if body is not None else {} + if status >= 400: + r.raise_for_status.side_effect = requests.HTTPError(f"{status}") + return r + + +@pytest.fixture(autouse=True) +def no_sleep(): + with patch.object(mod.time, "sleep") as s: + yield s + + +@pytest.fixture(autouse=True) +def env(monkeypatch): + monkeypatch.setenv("MAAP_API_HOST", "api.example.org") + monkeypatch.setenv("DPS_MACHINE_TOKEN", "tok") + + +def test_success_passes_timeout_and_header(): + with patch.object(mod.requests, "get", return_value=_resp(200, {"session_key": "abc"})) as g: + assert get_maap_pgt_token("alice") == "abc" + g.assert_called_once_with( + "https://api.example.org/api/members/alice", + headers={"dps-token": "tok"}, + timeout=mod.TIMEOUT, + ) + + +def test_retries_flaky_endpoint_then_succeeds(no_sleep): + side = [ + requests.ConnectionError("boom"), + _resp(503), + _resp(200, {"session_key": "abc"}), + ] + with patch.object(mod.requests, "get", side_effect=side) as g: + assert get_maap_pgt_token("alice") == "abc" + assert g.call_count == 3 + assert no_sleep.call_count == 2 + + +def test_retries_on_timeout(): + side = [requests.Timeout("slow"), _resp(200, {"session_key": "abc"})] + with patch.object(mod.requests, "get", side_effect=side): + assert get_maap_pgt_token("alice") == "abc" + + +def test_permanent_5xx_gives_clear_error_after_max_attempts(): + with patch.object(mod.requests, "get", return_value=_resp(500)) as g: + with pytest.raises(PgtTokenError, match=r"after 6 attempts.*HTTP 500"): + get_maap_pgt_token("alice") + assert g.call_count == mod.MAX_ATTEMPTS + + +def test_4xx_is_not_retried(): + with patch.object(mod.requests, "get", return_value=_resp(401)) as g: + with pytest.raises(PgtTokenError, match="HTTP 401"): + get_maap_pgt_token("alice") + assert g.call_count == 1 + + +def test_missing_session_key_gives_clear_error(): + with patch.object(mod.requests, "get", return_value=_resp(200, {"other": 1})): + with pytest.raises(PgtTokenError, match="session_key"): + get_maap_pgt_token("alice") + + +def test_non_json_body_gives_clear_error(): + r = _resp(200) + r.json.side_effect = ValueError("not json") + with patch.object(mod.requests, "get", return_value=r): + with pytest.raises(PgtTokenError, match="session_key"): + get_maap_pgt_token("alice") + + +def test_backoff_is_bounded_and_grows(): + for attempt in range(10): + assert 0 <= mod._sleep_time(attempt) <= min(mod.BACKOFF_CAP, mod.BACKOFF_BASE * 2 ** attempt) + + +def test_report_failure_writes_stderr_and_alt_error(tmp_path, capsys): + mod.report_failure(PgtTokenError("HTTP 503\nafter 6 attempts"), job_dir=str(tmp_path)) + assert "HTTP 503 after 6 attempts" in capsys.readouterr().err + assert (tmp_path / "_alt_error.txt").read_text() == "MAAP_PGT token failure: HTTP 503 after 6 attempts\n" + + +def test_main_exits_with_custom_code_and_keeps_stdout_clean(tmp_path, monkeypatch): + import subprocess, json + (tmp_path / "_job.json").write_text(json.dumps({"username": "alice"})) + # port 9 on localhost refuses connections: retries are exhausted quickly with backoff patched out + code = ( + "import sys, runpy, time; time.sleep = lambda s: None; " + f"runpy.run_path({mod.__file__!r}, run_name='__main__')" + ) + r = subprocess.run( + [sys.executable, "-c", code], cwd=tmp_path, capture_output=True, text=True, + env={"MAAP_API_HOST": "localhost:9", "DPS_MACHINE_TOKEN": "t", "PATH": ""}, + ) + assert r.returncode == mod.EXIT_CODE_PGT_FAILURE + assert r.stdout == "" + assert "Unable to get user MAAP_PGT token" in r.stderr + assert (tmp_path / "_alt_error.txt").exists() From f03241269fbb5964f0734a18beb7dcfee2cbb929 Mon Sep 17 00:00:00 2001 From: Sujen Shah Date: Fri, 16 Oct 2026 15:40:05 -0700 Subject: [PATCH 2/2] refactor: use urllib3 Retry for get_maap_pgt_token retries Replace the hand-rolled retry loop with a requests session mounting an HTTPAdapter(Retry(...)): exponential backoff with jitter (backoff_jitter needs urllib3>=2), retries on connect errors, timeouts and 5xx. Pin urllib3>=2 in the Dockerfile and CI. Retry behaviour is tested against a local flaky HTTP server. Claude-Session: https://claude.ai/code/session_0171HvL7yCdP2wCuxh7UngYP --- .github/workflows/docker-build.yml | 2 +- docker/Dockerfile | 3 +- get_maap_pgt_token.py | 92 ++++++++++--------- tests/test_get_maap_pgt_token.py | 140 ++++++++++++++++++++--------- 4 files changed, 151 insertions(+), 86 deletions(-) diff --git a/.github/workflows/docker-build.yml b/.github/workflows/docker-build.yml index cc1fa98..f3a67a2 100644 --- a/.github/workflows/docker-build.yml +++ b/.github/workflows/docker-build.yml @@ -23,7 +23,7 @@ jobs: with: python-version: '3.12' - name: Install dependencies - run: pip install pytest cwltool==3.1.20250110105449 cwl_utils==0.35 PyYAML pystac==1.15.2 maap-py requests + run: pip install pytest cwltool==3.1.20250110105449 cwl_utils==0.35 PyYAML pystac==1.15.2 maap-py requests "urllib3>=2" - name: Run tests run: pytest tests diff --git a/docker/Dockerfile b/docker/Dockerfile index b106fb2..9c64f14 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -23,7 +23,8 @@ RUN pip install --no-cache-dir \ awscli \ pystac==1.15.2 \ maap-py \ - requests + requests \ + "urllib3>=2" # Copy only what you need COPY docker_stats/docker-stats-on-exit-shim /docker-stats-on-exit-shim diff --git a/get_maap_pgt_token.py b/get_maap_pgt_token.py index 6a89492..cd06f94 100644 --- a/get_maap_pgt_token.py +++ b/get_maap_pgt_token.py @@ -1,15 +1,17 @@ import json import os -import random import sys -import time import requests +from requests.adapters import HTTPAdapter +from urllib3.util.retry import Retry # backoff_jitter requires urllib3>=2 TIMEOUT = (5, 30) # (connect, read) seconds -MAX_ATTEMPTS = 6 -BACKOFF_BASE = 1.0 # seconds; doubles each attempt -BACKOFF_CAP = 30.0 +MAX_RETRIES = 5 # i.e. up to 6 attempts +BACKOFF_FACTOR = 1.0 # sleeps ~ factor * 2**(n-1) seconds between retries, plus jitter +BACKOFF_JITTER = 1.0 +BACKOFF_MAX = 30.0 +RETRY_STATUSES = tuple(range(500, 600)) EXIT_CODE_PGT_FAILURE = 78 # distinct job exit status for PGT token failures @@ -17,9 +19,21 @@ class PgtTokenError(Exception): """Raised when the MAAP PGT token cannot be obtained.""" -def _sleep_time(attempt): - """Exponential backoff with full jitter for the given 0-based attempt.""" - return random.uniform(0, min(BACKOFF_CAP, BACKOFF_BASE * 2 ** attempt)) +def _session(): + """Session retrying connect errors, timeouts and 5xx with exponential backoff and jitter.""" + retry = Retry( + total=MAX_RETRIES, + backoff_factor=BACKOFF_FACTOR, + backoff_jitter=BACKOFF_JITTER, + backoff_max=BACKOFF_MAX, + status_forcelist=RETRY_STATUSES, + raise_on_status=False, # hand back the last response so we can raise a clear error + ) + session = requests.Session() + adapter = HTTPAdapter(max_retries=retry) + session.mount("https://", adapter) + session.mount("http://", adapter) + return session def get_maap_pgt_token(username, maap_api_host=None): @@ -27,39 +41,35 @@ def get_maap_pgt_token(username, maap_api_host=None): url = f"https://{maap_api_host}/api/members/{username}" headers = {"dps-token": os.environ.get("DPS_MACHINE_TOKEN")} - last_error = None - for attempt in range(MAX_ATTEMPTS): - if attempt: - time.sleep(_sleep_time(attempt - 1)) - try: - response = requests.get(url, headers=headers, timeout=TIMEOUT) - except (requests.ConnectionError, requests.Timeout) as e: - last_error = f"{type(e).__name__}: {e}" - continue - - if response.status_code >= 500: - last_error = f"HTTP {response.status_code}" - continue - - try: - response.raise_for_status() - except requests.HTTPError as e: - # 4xx is permanent (e.g. bad token, unknown user): do not retry - raise PgtTokenError( - f"Request for user '{username}' to {url} failed: HTTP {response.status_code}" - ) from e - - try: - return response.json()["session_key"] - except (ValueError, KeyError, TypeError) as e: - raise PgtTokenError( - f"Response from {url} for user '{username}' has no usable 'session_key'" - ) from e - - raise PgtTokenError( - f"Could not get MAAP_PGT for user '{username}' from {url} " - f"after {MAX_ATTEMPTS} attempts; last error: {last_error}" - ) + try: + with _session() as session: + response = session.get(url, headers=headers, timeout=TIMEOUT) + except requests.RequestException as e: + raise PgtTokenError( + f"Could not get MAAP_PGT for user '{username}' from {url} " + f"after {MAX_RETRIES + 1} attempts; last error: {type(e).__name__}: {e}" + ) from e + + if response.status_code >= 500: + raise PgtTokenError( + f"Could not get MAAP_PGT for user '{username}' from {url} " + f"after {MAX_RETRIES + 1} attempts; last error: HTTP {response.status_code}" + ) + + try: + response.raise_for_status() + except requests.HTTPError as e: + # 4xx is permanent (e.g. bad token, unknown user) and was not retried + raise PgtTokenError( + f"Request for user '{username}' to {url} failed: HTTP {response.status_code}" + ) from e + + try: + return response.json()["session_key"] + except (ValueError, KeyError, TypeError) as e: + raise PgtTokenError( + f"Response from {url} for user '{username}' has no usable 'session_key'" + ) from e def report_failure(error, job_dir="."): diff --git a/tests/test_get_maap_pgt_token.py b/tests/test_get_maap_pgt_token.py index ad4a3d4..a20ef1b 100644 --- a/tests/test_get_maap_pgt_token.py +++ b/tests/test_get_maap_pgt_token.py @@ -1,6 +1,10 @@ +import json +import subprocess import sys +import threading +from http.server import BaseHTTPRequestHandler, HTTPServer from pathlib import Path -from unittest.mock import Mock, patch +from unittest.mock import MagicMock, patch import pytest import requests @@ -12,7 +16,7 @@ def _resp(status=200, body=None): - r = Mock() + r = MagicMock() r.status_code = status r.json.return_value = body if body is not None else {} if status >= 400: @@ -20,10 +24,12 @@ def _resp(status=200, body=None): return r -@pytest.fixture(autouse=True) -def no_sleep(): - with patch.object(mod.time, "sleep") as s: - yield s +def _mock_session(get): + """Patch mod._session to return a context-manager session whose .get is `get`.""" + session = MagicMock() + session.__enter__.return_value = session + session.get = get + return patch.object(mod, "_session", return_value=session) @pytest.fixture(autouse=True) @@ -32,50 +38,101 @@ def env(monkeypatch): monkeypatch.setenv("DPS_MACHINE_TOKEN", "tok") +@pytest.fixture +def fast_backoff(monkeypatch): + monkeypatch.setattr(mod, "BACKOFF_FACTOR", 0.01) + monkeypatch.setattr(mod, "BACKOFF_JITTER", 0.01) + + +@pytest.fixture +def flaky_server(): + """Local HTTP server; set .statuses to the sequence of codes to return, then 200.""" + class Handler(BaseHTTPRequestHandler): + def do_GET(self): + self.server.hits += 1 + code = self.server.statuses.pop(0) if self.server.statuses else 200 + body = json.dumps({"session_key": "abc"} if code == 200 else {}).encode() + self.send_response(code) + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def log_message(self, *a): + pass + + server = HTTPServer(("127.0.0.1", 0), Handler) + server.hits, server.statuses = 0, [] + threading.Thread(target=server.serve_forever, daemon=True).start() + server.url = f"http://127.0.0.1:{server.server_port}/" + yield server + server.shutdown() + + +# --- retry behaviour, against a real server through the real session --- + +def test_session_retries_5xx_then_succeeds(flaky_server, fast_backoff): + flaky_server.statuses = [503, 502] + with mod._session() as s: + r = s.get(flaky_server.url, timeout=mod.TIMEOUT) + assert r.status_code == 200 + assert flaky_server.hits == 3 + + +def test_session_gives_up_after_max_retries_returning_last_response(flaky_server, fast_backoff): + flaky_server.statuses = [500] * 20 + with mod._session() as s: + r = s.get(flaky_server.url, timeout=mod.TIMEOUT) + assert r.status_code == 500 + assert flaky_server.hits == mod.MAX_RETRIES + 1 + + +def test_session_does_not_retry_4xx(flaky_server, fast_backoff): + flaky_server.statuses = [401, 200] + with mod._session() as s: + r = s.get(flaky_server.url, timeout=mod.TIMEOUT) + assert r.status_code == 401 + assert flaky_server.hits == 1 + + +def test_session_retries_connection_errors(fast_backoff): + with mod._session() as s: + with pytest.raises(requests.ConnectionError): + s.get("http://127.0.0.1:9/", timeout=mod.TIMEOUT) + + +# --- get_maap_pgt_token logic --- + def test_success_passes_timeout_and_header(): - with patch.object(mod.requests, "get", return_value=_resp(200, {"session_key": "abc"})) as g: + get = MagicMock(return_value=_resp(200, {"session_key": "abc"})) + with _mock_session(get): assert get_maap_pgt_token("alice") == "abc" - g.assert_called_once_with( + get.assert_called_once_with( "https://api.example.org/api/members/alice", headers={"dps-token": "tok"}, timeout=mod.TIMEOUT, ) -def test_retries_flaky_endpoint_then_succeeds(no_sleep): - side = [ - requests.ConnectionError("boom"), - _resp(503), - _resp(200, {"session_key": "abc"}), - ] - with patch.object(mod.requests, "get", side_effect=side) as g: - assert get_maap_pgt_token("alice") == "abc" - assert g.call_count == 3 - assert no_sleep.call_count == 2 - - -def test_retries_on_timeout(): - side = [requests.Timeout("slow"), _resp(200, {"session_key": "abc"})] - with patch.object(mod.requests, "get", side_effect=side): - assert get_maap_pgt_token("alice") == "abc" +def test_exhausted_5xx_gives_clear_error(): + with _mock_session(MagicMock(return_value=_resp(503))): + with pytest.raises(PgtTokenError, match=r"after 6 attempts.*HTTP 503"): + get_maap_pgt_token("alice") -def test_permanent_5xx_gives_clear_error_after_max_attempts(): - with patch.object(mod.requests, "get", return_value=_resp(500)) as g: - with pytest.raises(PgtTokenError, match=r"after 6 attempts.*HTTP 500"): +def test_exhausted_connection_errors_give_clear_error(): + with _mock_session(MagicMock(side_effect=requests.ConnectionError("boom"))): + with pytest.raises(PgtTokenError, match=r"after 6 attempts.*ConnectionError: boom"): get_maap_pgt_token("alice") - assert g.call_count == mod.MAX_ATTEMPTS -def test_4xx_is_not_retried(): - with patch.object(mod.requests, "get", return_value=_resp(401)) as g: +def test_4xx_gives_clear_error(): + with _mock_session(MagicMock(return_value=_resp(401))): with pytest.raises(PgtTokenError, match="HTTP 401"): get_maap_pgt_token("alice") - assert g.call_count == 1 def test_missing_session_key_gives_clear_error(): - with patch.object(mod.requests, "get", return_value=_resp(200, {"other": 1})): + with _mock_session(MagicMock(return_value=_resp(200, {"other": 1}))): with pytest.raises(PgtTokenError, match="session_key"): get_maap_pgt_token("alice") @@ -83,15 +140,12 @@ def test_missing_session_key_gives_clear_error(): def test_non_json_body_gives_clear_error(): r = _resp(200) r.json.side_effect = ValueError("not json") - with patch.object(mod.requests, "get", return_value=r): + with _mock_session(MagicMock(return_value=r)): with pytest.raises(PgtTokenError, match="session_key"): get_maap_pgt_token("alice") -def test_backoff_is_bounded_and_grows(): - for attempt in range(10): - assert 0 <= mod._sleep_time(attempt) <= min(mod.BACKOFF_CAP, mod.BACKOFF_BASE * 2 ** attempt) - +# --- failure reporting to HySDS --- def test_report_failure_writes_stderr_and_alt_error(tmp_path, capsys): mod.report_failure(PgtTokenError("HTTP 503\nafter 6 attempts"), job_dir=str(tmp_path)) @@ -99,19 +153,19 @@ def test_report_failure_writes_stderr_and_alt_error(tmp_path, capsys): assert (tmp_path / "_alt_error.txt").read_text() == "MAAP_PGT token failure: HTTP 503 after 6 attempts\n" -def test_main_exits_with_custom_code_and_keeps_stdout_clean(tmp_path, monkeypatch): - import subprocess, json +def test_main_exits_with_custom_code_and_keeps_stdout_clean(tmp_path): (tmp_path / "_job.json").write_text(json.dumps({"username": "alice"})) - # port 9 on localhost refuses connections: retries are exhausted quickly with backoff patched out + # nothing listens on localhost:9; sleep is stubbed so retries are exhausted quickly code = ( - "import sys, runpy, time; time.sleep = lambda s: None; " + "import runpy, time; time.sleep = lambda s: None; " f"runpy.run_path({mod.__file__!r}, run_name='__main__')" ) r = subprocess.run( [sys.executable, "-c", code], cwd=tmp_path, capture_output=True, text=True, - env={"MAAP_API_HOST": "localhost:9", "DPS_MACHINE_TOKEN": "t", "PATH": ""}, + env={"MAAP_API_HOST": "localhost:9", "DPS_MACHINE_TOKEN": "t", "PATH": "", + "PYTHONPATH": str(Path(mod.__file__).parent)}, ) - assert r.returncode == mod.EXIT_CODE_PGT_FAILURE + assert r.returncode == mod.EXIT_CODE_PGT_FAILURE, r.stderr assert r.stdout == "" assert "Unable to get user MAAP_PGT token" in r.stderr assert (tmp_path / "_alt_error.txt").exists()