diff --git a/setup.py b/setup.py index 869afb4..980af90 100644 --- a/setup.py +++ b/setup.py @@ -1,6 +1,6 @@ from setuptools import setup, find_packages -VERSION = "1.5.23" +VERSION = "1.5.24" with open("requirements.txt") as f: requirements = f.read().splitlines() diff --git a/surge/async_jobs.py b/surge/async_jobs.py new file mode 100644 index 0000000..23dd7ca --- /dev/null +++ b/surge/async_jobs.py @@ -0,0 +1,67 @@ +"""Polling helper for async-job status endpoints.""" + +import time + +DEFAULT_IN_PROGRESS_STATUSES = frozenset({"IN_PROGRESS", "CREATING"}) +DEFAULT_COMPLETED_STATUSES = frozenset({"COMPLETED", "READY"}) +DEFAULT_ERROR_STATUSES = frozenset({"ERROR", "FAILED"}) + + +class AsyncJobError(Exception): + """Raised when an async job fails or times out. + + The ``status`` dict carries the underlying error; timeouts use the + synthetic ``status="TIMEOUT"`` sentinel so callers can distinguish + them from server-side failures. + """ + + def __init__(self, status): + self.status = status + error = status.get("error") or status.get("type") or "async job failed" + super().__init__(error) + + +def poll_async_job( + check_status, + *, + poll_time, + poll_interval, + initial_status=None, + in_progress_statuses=DEFAULT_IN_PROGRESS_STATUSES, + completed_statuses=DEFAULT_COMPLETED_STATUSES, + error_statuses=DEFAULT_ERROR_STATUSES, +): + """Poll ``check_status()`` until it returns a terminal status. + + Arguments: + check_status: zero-arg callable returning a dict with a ``status`` key. + poll_time: seconds budgeted for waiting between polls; when exhausted, + AsyncJobError is raised with the synthetic ``status="TIMEOUT"`` + sentinel. + poll_interval: seconds to sleep between polls. + initial_status: optional status dict consumed in place of the first + ``check_status()`` call — useful when the caller already has the + response from a kick-off endpoint. + in_progress_statuses: status strings that mean "keep polling". + completed_statuses: status strings that mean "return the status dict". + error_statuses: status strings that mean "raise AsyncJobError". + + Returns the terminal status dict. + """ + deadline = time.monotonic() + poll_time + status = initial_status if initial_status is not None else check_status() + while True: + value = status.get("status") + if value in completed_statuses: + return status + if value in error_statuses: + raise AsyncJobError(status) + if value not in in_progress_statuses: + message = f"unexpected status {value!r}" + raise AsyncJobError({**status, "error": message}) + remaining = deadline - time.monotonic() + if remaining <= 0: + message = f"async job did not complete within {poll_time} seconds" + raise AsyncJobError({"status": "TIMEOUT", "error": message}) + time.sleep(min(poll_interval, remaining)) + status = check_status() diff --git a/surge/reports.py b/surge/reports.py index 06579d0..ae3dd68 100644 --- a/surge/reports.py +++ b/surge/reports.py @@ -1,5 +1,4 @@ import gzip -from time import sleep, monotonic import urllib import tempfile import shutil @@ -8,6 +7,7 @@ import warnings from surge.api_resource import REPORTS_ENDPOINT, APIResource +from surge.async_jobs import AsyncJobError, poll_async_job from surge.errors import SurgeRequestError @@ -18,7 +18,7 @@ def __init__(self, **kwargs): self.__dict__.update(kwargs) def __str__(self): - return f"" + return "" def __repr__(self): return f"" @@ -35,6 +35,8 @@ def save_report( poll_time=5 * 60, api_key: str = None, poll_interval: float = 2, + *, + extra_params: dict = None, ): """ Request creation of a report, poll until the report is generated, and save the data to a file all in one call. @@ -57,35 +59,40 @@ def save_report( filepath (string or IO or None): Location to save the results file. If not specified, will save to "project_{project_id}_results.{csv/json} poll_time (int): Maximum number of seconds to wait for the report to be generated poll_interval (int or float): Seconds to wait between status checks + extra_params (dict, optional): Additional params merged into the + initial report request body. See ``request``. """ - response = cls.request(project_id=project_id, - type=type, - api_key=api_key) - # Capture the job_id from the initial CREATING response and - # reuse it across polls. check_status's IN_PROGRESS response - # does not include job_id; only RETRYING does (and means the - # server kicked off a new underlying job). - job_id = getattr(response, "job_id", None) - deadline = monotonic() + poll_time - while response.status in ("CREATING", "IN_PROGRESS", "RETRYING"): - if monotonic() >= deadline: - raise Exception( - "Report failed to generate within {poll_time} seconds". - format(poll_time=poll_time)) - sleep(poll_interval) - response = cls.check_status(project_id, job_id, api_key=api_key) - if response.status == "RETRYING": - job_id = response.job_id - - if response.status not in ("READY", "COMPLETED"): - raise ValueError("Report failed to generate with status {}".format( - response.status)) + initial = cls.request(project_id=project_id, + type=type, + api_key=api_key, + extra_params=extra_params) + # job_id swaps mid-poll on RETRYING; IN_PROGRESS responses + # omit it, so only update when present. + state = {"job_id": getattr(initial, "job_id", None)} + + def _check(): + r = cls.check_status(project_id, state["job_id"], api_key=api_key) + if getattr(r, "job_id", None): + state["job_id"] = r.job_id + return vars(r) + + try: + terminal = poll_async_job( + _check, + poll_time=poll_time, + poll_interval=poll_interval, + initial_status=vars(initial), + in_progress_statuses=("IN_PROGRESS", "CREATING", "RETRYING"), + ) + except AsyncJobError as e: + raise ValueError(f"Report failed to generate: {e}") from e + url = terminal["url"] file_ext = "csv" if "csv" in type else "json" default_file_name = "project_{project_id}_results.{file_ext}".format( project_id=project_id, file_ext=file_ext) target = filepath or default_file_name - with urllib.request.urlopen(response.url) as remote: + with urllib.request.urlopen(url) as remote: with tempfile.NamedTemporaryFile() as tmp_file: shutil.copyfileobj(remote, tmp_file) tmp_file.flush() @@ -129,13 +136,28 @@ def download_json(cls, return json.load(bytesio) @classmethod - def request(cls, project_id: str, type: str, api_key: str = None): + def request( + cls, + project_id: str, + type: str, + api_key: str = None, + *, + extra_params: dict = None, + ): """ - Request creation of a report for the given type. Note that reports are generated - asychronously so the response may include a `job_id` which needs to be used with - the `status` method to get the job status. In the event that the report has is - already generated and current, the report URL will be returned. Note that the URL - is a presigned URL which is active for only a limited duration. + Request creation of a report for the given type. + + Most callers want `save_report` instead — it kicks off the + request, polls the specific job, and saves the result to a + file in a single call. Use this raw method for fire-and-forget + batches or custom poll loops. + + Reports are generated asychronously so the response may + include a `job_id` which needs to be used with `check_status` + to get the job status. In the event that the report has is + already generated and current, the report URL will be returned. + Note that the URL is a presigned URL which is active for only + a limited duration. Type may be one of these types: * `export_json` @@ -164,12 +186,21 @@ def request(cls, project_id: str, type: str, api_key: str = None): Arguments: project_id (str): ID of project. report_type (str): report type + extra_params (dict, optional): Additional params merged into the + report request body. Use for backend-supported params that + aren't first-class kwargs on this method; refer to the API + reference for valid keys. Returns: status: Report status object which includes report id """ + extra = extra_params or {} + if "report_type" in extra: + raise ValueError( + "`report_type` is set via the `type` argument and cannot " + "be passed in `extra_params`") endpoint = f"{REPORTS_ENDPOINT}/{project_id}/report" - params = {"report_type": type} + params = {"report_type": type, **extra} response_json = cls.post(endpoint, params, api_key=api_key) if "error" in response_json: raise SurgeRequestError(response_json["error"]) @@ -221,7 +252,12 @@ def status(cls, project_id: str, job_id: str, api_key: str = None): @classmethod def check_status(cls, project_id: str, job_id: str, api_key: str = None): """ - Checks the status of a given report job. The response will be of one of these shapes: + Checks the status of a given report job. + + Most callers want `save_report` instead — this method is for + custom poll loops over a `job_id` returned by the raw `request`. + + The response will be of one of these shapes: 1) Report is still being generated (HTTP 202): diff --git a/tests/test_async_jobs.py b/tests/test_async_jobs.py new file mode 100644 index 0000000..1434b9b --- /dev/null +++ b/tests/test_async_jobs.py @@ -0,0 +1,168 @@ +from unittest.mock import patch + +import pytest + +from surge.async_jobs import AsyncJobError, poll_async_job + + +def test_returns_on_first_completed_status(): + statuses = iter([{"status": "COMPLETED", "url": "u"}]) + result = poll_async_job(lambda: next(statuses), + poll_time=10, + poll_interval=0) + assert result == {"status": "COMPLETED", "url": "u"} + + +def test_polls_through_in_progress_then_returns_completed(): + statuses = iter([ + { + "status": "IN_PROGRESS" + }, + { + "status": "IN_PROGRESS" + }, + { + "status": "COMPLETED", + "url": "u" + }, + ]) + with patch("time.sleep"): + result = poll_async_job(lambda: next(statuses), + poll_time=10, + poll_interval=0) + assert result == {"status": "COMPLETED", "url": "u"} + + +def test_accepts_creating_and_ready_status_vocabulary(): + statuses = iter([ + { + "status": "CREATING" + }, + { + "status": "READY", + "url": "u" + }, + ]) + with patch("time.sleep"): + result = poll_async_job(lambda: next(statuses), + poll_time=10, + poll_interval=0) + assert result == {"status": "READY", "url": "u"} + + +def test_initial_status_is_consumed_in_place_of_first_check_call(): + """When given, initial_status acts as iter 1; check_status is invoked + starting from iter 2.""" + check_calls = [] + + def check(): + check_calls.append(True) + return {"status": "COMPLETED", "url": "u"} + + result = poll_async_job( + check, + poll_time=10, + poll_interval=0, + initial_status={ + "status": "READY", + "url": "from-initial" + }, + ) + assert result == {"status": "READY", "url": "from-initial"} + assert check_calls == [] + + +def test_initial_status_falls_through_to_polling_when_in_progress(): + """In-progress initial_status hands off to check_status from iter 2.""" + statuses = iter([{"status": "COMPLETED", "url": "u"}]) + with patch("time.sleep"): + result = poll_async_job( + lambda: next(statuses), + poll_time=10, + poll_interval=0, + initial_status={"status": "CREATING"}, + ) + assert result == {"status": "COMPLETED", "url": "u"} + + +def test_raises_async_job_error_on_error_status(): + with pytest.raises(AsyncJobError, match="kaboom"): + poll_async_job( + lambda: { + "status": "ERROR", + "error": "kaboom" + }, + poll_time=10, + poll_interval=0, + ) + + +def test_raises_async_job_error_on_unknown_status(): + with pytest.raises(AsyncJobError, match="weird"): + poll_async_job(lambda: {"status": "weird"}, + poll_time=10, + poll_interval=0) + + +def test_raises_timeout_when_poll_time_elapses(): + with patch("time.sleep"): + with pytest.raises(AsyncJobError) as excinfo: + poll_async_job( + lambda: {"status": "IN_PROGRESS"}, + poll_time=0, + poll_interval=0, + ) + assert excinfo.value.status.get("status") == "TIMEOUT" + + +def test_caps_sleep_to_remaining_poll_time(): + """Sleep should never exceed the remaining budget — even with a + poll_interval much larger than poll_time.""" + + statuses = iter([ + { + "status": "IN_PROGRESS" + }, + { + "status": "COMPLETED", + "url": "u" + }, + ]) + + with patch("time.sleep") as mock_sleep: + poll_async_job( + lambda: next(statuses), + poll_time=5, + poll_interval=60, + ) + + assert mock_sleep.call_count == 1 + (slept_seconds, ) = mock_sleep.call_args.args + assert slept_seconds <= 5 + + +def test_raises_timeout_instead_of_returning_completed_after_deadline(): + """When poll_interval > poll_time the loop sleeps past the deadline. + The second status must not be honored — even if it would be COMPLETED.""" + + statuses = iter([ + { + "status": "IN_PROGRESS" + }, + { + "status": "COMPLETED", + "url": "u" + }, + ]) + + # Simulate the deadline elapsing immediately after the first poll: + # call 1 sets the deadline, calls 2+ all report well past it. + with patch("time.sleep"), patch("time.monotonic", + side_effect=[0.0, 1e9, 1e9, 1e9]): + with pytest.raises(AsyncJobError) as excinfo: + poll_async_job( + lambda: next(statuses), + poll_time=1, + poll_interval=2, + ) + assert excinfo.value.status.get("status") == "TIMEOUT" diff --git a/tests/test_reports.py b/tests/test_reports.py index 2b0341a..4b407f5 100644 --- a/tests/test_reports.py +++ b/tests/test_reports.py @@ -26,10 +26,12 @@ def test_save_report_downloads_when_request_returns_ready(): payload = b'[{"a": 1}]' ready = Report(status="READY", url="https://signed.example/report.gz") sink = io.BytesIO() - with mock.patch.object(Report, "request", - return_value=ready) as mock_request, \ - mock.patch.object(Report, "check_status") as mock_check, \ - mock.patch("urllib.request.urlopen") as mock_urlopen: + with ( + mock.patch.object(Report, "request", return_value=ready) as + mock_request, + mock.patch.object(Report, "check_status") as mock_check, + mock.patch("urllib.request.urlopen") as mock_urlopen, + ): mock_urlopen.return_value.__enter__.return_value = io.BytesIO( _gzipped(payload)) Report.save_report("proj-123", "export_json", filepath=sink) @@ -51,12 +53,16 @@ def test_save_report_polls_check_status_for_returned_job_id(): completed = Report(status="COMPLETED", url="https://signed.example/report.gz") sink = io.BytesIO() - with mock.patch.object(Report, "request", - return_value=creating) as mock_request, \ - mock.patch.object(Report, "check_status", - side_effect=[in_progress, completed]) as mock_check, \ - mock.patch("urllib.request.urlopen") as mock_urlopen, \ - mock.patch("surge.reports.sleep") as mock_sleep: + with ( + mock.patch.object(Report, "request", + return_value=creating) as mock_request, + mock.patch.object(Report, + "check_status", + side_effect=[in_progress, completed]) as + mock_check, + mock.patch("urllib.request.urlopen") as mock_urlopen, + mock.patch("time.sleep") as mock_sleep, + ): mock_urlopen.return_value.__enter__.return_value = io.BytesIO( _gzipped(payload)) Report.save_report("proj-123", "export_csv", filepath=sink) @@ -70,18 +76,22 @@ def test_save_report_polls_check_status_for_returned_job_id(): def test_save_report_switches_job_id_on_retrying(): """RETRYING responses include a new job_id; subsequent polls use it.""" - payload = b'[]' + payload = b"[]" creating = Report(status="CREATING", job_id="job-abc") retrying = Report(status="RETRYING", job_id="job-xyz") in_progress = Report(status="IN_PROGRESS") completed = Report(status="COMPLETED", url="https://signed.example/report.gz") sink = io.BytesIO() - with mock.patch.object(Report, "request", return_value=creating), \ - mock.patch.object(Report, "check_status", - side_effect=[retrying, in_progress, completed]) as mock_check, \ - mock.patch("urllib.request.urlopen") as mock_urlopen, \ - mock.patch("surge.reports.sleep"): + with ( + mock.patch.object(Report, "request", return_value=creating), + mock.patch.object(Report, + "check_status", + side_effect=[retrying, in_progress, completed]) + as mock_check, + mock.patch("urllib.request.urlopen") as mock_urlopen, + mock.patch("time.sleep"), + ): mock_urlopen.return_value.__enter__.return_value = io.BytesIO( _gzipped(payload)) Report.save_report("proj-123", "export_json", filepath=sink) @@ -92,10 +102,12 @@ def test_save_report_switches_job_id_on_retrying(): def test_save_report_raises_on_unexpected_status(): creating = Report(status="CREATING", job_id="job-abc") error = Report(status="ERROR", type="Report generation error") - with mock.patch.object(Report, "request", return_value=creating), \ - mock.patch.object(Report, "check_status", return_value=error), \ - mock.patch("surge.reports.sleep"): - with pytest.raises(ValueError, match="ERROR"): + with ( + mock.patch.object(Report, "request", return_value=creating), + mock.patch.object(Report, "check_status", return_value=error), + mock.patch("time.sleep"), + ): + with pytest.raises(ValueError, match="Report generation error"): Report.save_report("proj-123", "export_json", filepath=io.BytesIO()) @@ -104,13 +116,93 @@ def test_save_report_raises_on_unexpected_status(): def test_save_report_times_out_when_job_never_completes(): creating = Report(status="CREATING", job_id="job-abc") in_progress = Report(status="IN_PROGRESS", job_id="job-abc") - with mock.patch.object(Report, "request", return_value=creating), \ + with ( + mock.patch.object(Report, "request", return_value=creating), mock.patch.object(Report, "check_status", - return_value=in_progress), \ - mock.patch("surge.reports.sleep"), \ - mock.patch("surge.reports.monotonic", - side_effect=[0.0, 1000.0, 2000.0]): + return_value=in_progress), + mock.patch("time.sleep"), + mock.patch("time.monotonic", side_effect=[0.0, 1000.0, 2000.0]), + ): with pytest.raises(Exception, match="within 300 seconds"): Report.save_report("proj-123", "export_json", filepath=io.BytesIO()) + + +def test_request_merges_extra_params_into_post_body(): + with mock.patch.object(Report, "post") as mock_post: + mock_post.return_value = {"status": "CREATING", "job_id": "j"} + Report.request("proj-1", "export_json", extra_params={"item_id": "i"}) + endpoint, params = mock_post.call_args.args[:2] + assert "report" in endpoint + assert params == {"report_type": "export_json", "item_id": "i"} + + +def test_request_without_extra_params_sends_only_report_type(): + with mock.patch.object(Report, "post") as mock_post: + mock_post.return_value = {"status": "CREATING", "job_id": "j"} + Report.request("proj-1", "export_json") + _, params = mock_post.call_args.args[:2] + assert params == {"report_type": "export_json"} + + +def test_save_report_raises_when_initial_request_returns_error(): + """If request() returns ERROR directly, raise without polling check_status.""" + error = Report(status="ERROR", type="Report generation error") + with ( + mock.patch.object(Report, "request", return_value=error), + mock.patch.object(Report, "check_status") as mock_check, + ): + with pytest.raises(ValueError, match="Report generation error"): + Report.save_report("proj-e", "export_json", filepath=io.BytesIO()) + mock_check.assert_not_called() + + +def test_save_report_treats_completed_initial_response_as_terminal(): + """If request() returns COMPLETED (not just READY), skip polling.""" + payload = b"[]" + completed = Report(status="COMPLETED", + url="https://signed.example/report.gz") + sink = io.BytesIO() + with ( + mock.patch.object(Report, "request", return_value=completed) as + mock_request, + mock.patch.object(Report, "check_status") as mock_check, + mock.patch("urllib.request.urlopen") as mock_urlopen, + ): + mock_urlopen.return_value.__enter__.return_value = io.BytesIO( + _gzipped(payload)) + Report.save_report("proj-c", "export_json", filepath=sink) + mock_request.assert_called_once() + mock_check.assert_not_called() + assert sink.getvalue() == payload + + +def test_request_rejects_report_type_in_extra_params(): + """report_type is set via the `type` argument; allowing it in + extra_params would silently override the explicit `type`.""" + with mock.patch.object(Report, "post") as mock_post: + with pytest.raises(ValueError, match="report_type"): + Report.request("proj-1", + "export_csv", + extra_params={"report_type": "export_json"}) + mock_post.assert_not_called() + + +def test_save_report_forwards_extra_params_to_request(): + payload = b"[]" + ready = Report(status="READY", url="https://signed.example/report.gz") + with ( + mock.patch.object(Report, "request", return_value=ready) as + mock_request, + mock.patch("urllib.request.urlopen") as mock_urlopen, + ): + mock_urlopen.return_value.__enter__.return_value = io.BytesIO( + _gzipped(payload)) + Report.save_report( + "proj-1", + "export_json", + filepath=io.BytesIO(), + extra_params={"item_id": "i"}, + ) + assert mock_request.call_args.kwargs["extra_params"] == {"item_id": "i"}