From 63509a9e19c8648b4774ab545a9731d80c2e5c26 Mon Sep 17 00:00:00 2001 From: Stefaan Lippens Date: Tue, 11 Aug 2026 22:16:13 +0200 Subject: [PATCH 1/3] Introduce job result downloading with item walking and ref rewriting ref #931 --- openeo/rest/_connection.py | 87 ++++++++++++- openeo/rest/_testing.py | 61 +++++++++- openeo/rest/job.py | 244 ++++++++++++++++++++++++++----------- openeo/testing/stac.py | 3 +- tests/rest/test_job.py | 197 +++++++++++++++++++++++++++++- tests/rest/test_testing.py | 49 +++++++- 6 files changed, 562 insertions(+), 79 deletions(-) diff --git a/openeo/rest/_connection.py b/openeo/rest/_connection.py index f7a0d7fa9..647cccbda 100644 --- a/openeo/rest/_connection.py +++ b/openeo/rest/_connection.py @@ -2,6 +2,7 @@ import logging import sys +from pathlib import Path from typing import Iterable, Optional, Union import requests @@ -10,10 +11,25 @@ from requests.auth import AuthBase import openeo -from openeo.rest import OpenEoApiError, OpenEoApiPlainError, OpenEoRestError +from openeo.rest import ( + DEFAULT_DOWNLOAD_CHUNK_SIZE, + DEFAULT_DOWNLOAD_RANGE_SIZE, + OpenEoApiError, + OpenEoApiPlainError, + OpenEoRestError, +) from openeo.rest.auth.auth import NullAuth from openeo.util import ContextTimer, ensure_list, str_truncate, url_join -from openeo.utils.http import HTTP_502_BAD_GATEWAY, session_with_retries +from openeo.utils.http import ( + HTTP_408_REQUEST_TIMEOUT, + HTTP_429_TOO_MANY_REQUESTS, + HTTP_500_INTERNAL_SERVER_ERROR, + HTTP_501_NOT_IMPLEMENTED, + HTTP_502_BAD_GATEWAY, + HTTP_503_SERVICE_UNAVAILABLE, + HTTP_504_GATEWAY_TIMEOUT, + session_with_retries, +) _log = logging.getLogger(__name__) @@ -21,6 +37,18 @@ # TODO: get default_timeout from config? DEFAULT_TIMEOUT = 20 * 60 +MAX_DOWNLOAD_RETRIES_PER_RANGE = 3 + +RETRIABLE_DOWNLOAD_STATUSCODES = [ + HTTP_408_REQUEST_TIMEOUT, + HTTP_429_TOO_MANY_REQUESTS, + HTTP_500_INTERNAL_SERVER_ERROR, + HTTP_501_NOT_IMPLEMENTED, + HTTP_502_BAD_GATEWAY, + HTTP_503_SERVICE_UNAVAILABLE, + HTTP_504_GATEWAY_TIMEOUT, +] + class RestApiConnection: """Base connection class implementing generic REST API request functionality""" @@ -262,3 +290,58 @@ def put(self, path: str, headers: Optional[dict] = None, data: Optional[dict] = def __repr__(self): return "<{c} to {r!r} with {a}>".format(c=type(self).__name__, r=self._root_url, a=type(self.auth).__name__) + + def download_url( + self, + url: str, + target: Path, + *, + chunk_size: int = DEFAULT_DOWNLOAD_CHUNK_SIZE, + range_size: int = DEFAULT_DOWNLOAD_RANGE_SIZE, + ) -> None: + head = self.head(url, stream=True) + if head.ok and head.headers.get("Accept-Ranges") == "bytes" and "Content-Length" in head.headers: + file_size = int(head.headers["Content-Length"]) + self._download_ranged( + url=url, target=target, file_size=file_size, chunk_size=chunk_size, range_size=range_size + ) + else: + self._download_all_at_once(url=url, target=target, chunk_size=chunk_size) + + def _download_all_at_once(self, url: str, target: Path, *, chunk_size: int = DEFAULT_DOWNLOAD_CHUNK_SIZE) -> None: + with self.get(path=url, stream=True) as r: + r.raise_for_status() + with target.open("wb") as f: + for block in r.iter_content(chunk_size=chunk_size): + f.write(block) + + def _download_ranged( + self, + url: str, + target: Path, + file_size: int, + *, + chunk_size: int = DEFAULT_DOWNLOAD_CHUNK_SIZE, + range_size: int = DEFAULT_DOWNLOAD_RANGE_SIZE, + ) -> None: + with target.open("wb") as f: + for from_byte_index in range(0, file_size, range_size): + to_byte_index = min(from_byte_index + range_size - 1, file_size - 1) + tries_left = MAX_DOWNLOAD_RETRIES_PER_RANGE + while tries_left > 0: + try: + range_headers = {"Range": f"bytes={from_byte_index}-{to_byte_index}"} + with self.get(path=url, headers=range_headers, stream=True) as r: + r.raise_for_status() + for block in r.iter_content(chunk_size=chunk_size): + f.write(block) + break + except OpenEoApiPlainError as error: + tries_left -= 1 + if tries_left > 0 and error.http_status_code in RETRIABLE_DOWNLOAD_STATUSCODES: + _log.warning( + f"Failed to retrieve chunk {from_byte_index}-{to_byte_index} from {url} (status {error.http_status_code}) - retrying" + ) + continue + else: + raise error diff --git a/openeo/rest/_testing.py b/openeo/rest/_testing.py index 6e8423c26..4e8cda9b5 100644 --- a/openeo/rest/_testing.py +++ b/openeo/rest/_testing.py @@ -14,8 +14,9 @@ Union, ) -from openeo import Connection, DataCube +from openeo import BatchJob, Connection, DataCube from openeo.rest.vectorcube import VectorCube +from openeo.testing.stac import StacDummyBuilder from openeo.utils.http import HTTP_201_CREATED, HTTP_202_ACCEPTED, HTTP_204_NO_CONTENT OPENEO_BACKEND = "https://openeo.test/" @@ -488,3 +489,61 @@ def build_capabilities( "links": [], } return capabilities + + +class JobResultCollectionMocker: + """ + Helper to mock job result metadata (openEO 1.1 Collection style) + with items and assets. + """ + + def __init__(self, *, requests_mock, connection: Connection): + self.requests_mock = requests_mock + self.connection = connection + + def setup_job_results( + self, *, job_id: str = "job-123", items: dict, add_collection_assets: bool = True + ) -> BatchJob: + links = [] + collection_assets = {} + for item_id, item_data in items.items(): + assets = {} + for asset_key, asset_data in item_data.get("assets", {}).items(): + asset = self.setup_asset(job_id=job_id, asset_data=asset_data) + assets[asset_key] = asset + collection_assets[f"{item_id}-{asset_key}"] = asset + + item_href = self.setup_item(job_id=job_id, item_id=item_id, assets=assets) + links.append({"rel": "item", "href": item_href}) + + collection_href = self.connection.build_url(f"/jobs/{job_id}/results") + collection_doc = StacDummyBuilder.collection( + id=f"{job_id}-results", + stac_version="1.1.0", + links=links, + assets=collection_assets if add_collection_assets else {}, + ) + self.requests_mock.get(collection_href, json=collection_doc) + + job = BatchJob(job_id, connection=self.connection) + return job + + def setup_item(self, *, job_id: str, item_id: str, assets: dict) -> dict: + href = self.connection.build_url(f"/j/{job_id}/r/i/{item_id}.json") + doc = StacDummyBuilder.item( + id=item_id, + stac_version="1.1.0", + assets=assets, + ) + self.requests_mock.get(href, json=doc) + return href + + def setup_asset(self, *, job_id: str, asset_data: dict) -> dict: + href = self.connection.build_url(f"/j/{job_id}/r/a/{asset_data.get('path', 'asset.tiff')}") + content = asset_data.get("content", b"TIFF-DUMMY-DATA") + self.requests_mock.head(href, headers={"Content-Length": f"{len(content)}"}) + self.requests_mock.get(href, content=content) + return StacDummyBuilder.asset( + href=href, + type=asset_data.get("type", "image/tiff; application=geotiff"), + ) diff --git a/openeo/rest/job.py b/openeo/rest/job.py index 39d9e2093..3106a553e 100644 --- a/openeo/rest/job.py +++ b/openeo/rest/job.py @@ -6,9 +6,9 @@ import re import time import typing +import urllib.parse from pathlib import Path -from typing import Dict, List, Optional, Union -from urllib.parse import unquote, urlparse +from typing import Dict, Iterable, List, Optional, Union import requests @@ -48,16 +48,6 @@ DEFAULT_JOB_RESULTS_FILENAME = "job-results.json" -MAX_RETRIES_PER_RANGE = 3 -RETRIABLE_STATUSCODES = [ - HTTP_408_REQUEST_TIMEOUT, - HTTP_429_TOO_MANY_REQUESTS, - HTTP_500_INTERNAL_SERVER_ERROR, - HTTP_501_NOT_IMPLEMENTED, - HTTP_502_BAD_GATEWAY, - HTTP_503_SERVICE_UNAVAILABLE, - HTTP_504_GATEWAY_TIMEOUT, -] class BatchJob: @@ -409,6 +399,21 @@ def _sanitize_filename(s: str, replacement: str = "") -> str: return FILENAME_UNSAFE_REGEX.sub(replacement, s) +def _filename_from_url(url: str, *, full: bool = False) -> str: + """ + Try to extract a filename from a URL (based on the path), + with sanitization of risky characters, + and option to only get the final part (basename) or the full path. + """ + parsed = urllib.parse.urlparse(url) + path = urllib.parse.unquote(parsed.path) + parts = path.strip("/").split("/") + if not full: + parts = parts[-1:] + parts = [_sanitize_filename(p) for p in parts if p] + return "/".join(parts) + + _MEDIA_TYPE_EXTENSION_MAP = { "image/tiff": ".tiff", "image/tiff; application=geotiff": ".tiff", @@ -475,8 +480,7 @@ def _make_filename(self) -> str: # Build filename from key, href's path (if any) # and guess extension from media type if necessary sanitized_key = _sanitize_filename(self.key) - href_path = unquote(urlparse(str(self.href)).path) - href_basename = _sanitize_filename(Path(href_path).name) + href_basename = _filename_from_url(self.href, full=False) filename = f"{sanitized_key}-{href_basename}" if not re.fullmatch(r".*\.[a-zA-Z0-9]{1,10}$", filename): @@ -507,7 +511,7 @@ def download( target = target / self._make_filename() ensure_dir(target.parent) logger.info(f"Downloading job result asset {self.key!r} from {self.href!s} to {target!s}") - self._download_to_file(url=self.href, target=target, chunk_size=chunk_size, range_size=range_size) + self.job.connection.download_url(url=self.href, target=target, chunk_size=chunk_size, range_size=range_size) return target def _get_response(self, stream=True) -> requests.Response: @@ -525,61 +529,6 @@ def load_bytes(self) -> bytes: # TODO: more `load` methods e.g.: load GTiff asset directly as numpy array - def _download_to_file( - self, - url: str, - target: Path, - *, - chunk_size: int = DEFAULT_DOWNLOAD_CHUNK_SIZE, - range_size: int = DEFAULT_DOWNLOAD_RANGE_SIZE, - ): - head = self.job.connection.head(url, stream=True) - if head.ok and head.headers.get("Accept-Ranges") == "bytes" and "Content-Length" in head.headers: - file_size = int(head.headers["Content-Length"]) - self._download_ranged( - url=url, target=target, file_size=file_size, chunk_size=chunk_size, range_size=range_size - ) - else: - self._download_all_at_once(url=url, target=target, chunk_size=chunk_size) - - def _download_ranged( - self, - url: str, - target: Path, - file_size: int, - *, - chunk_size: int = DEFAULT_DOWNLOAD_CHUNK_SIZE, - range_size: int = DEFAULT_DOWNLOAD_RANGE_SIZE, - ): - with target.open("wb") as f: - for from_byte_index in range(0, file_size, range_size): - to_byte_index = min(from_byte_index + range_size - 1, file_size - 1) - tries_left = MAX_RETRIES_PER_RANGE - while tries_left > 0: - try: - range_headers = {"Range": f"bytes={from_byte_index}-{to_byte_index}"} - with self.job.connection.get(path=url, headers=range_headers, stream=True) as r: - r.raise_for_status() - for block in r.iter_content(chunk_size=chunk_size): - f.write(block) - break - except OpenEoApiPlainError as error: - tries_left -= 1 - if tries_left > 0 and error.http_status_code in RETRIABLE_STATUSCODES: - logger.warning( - f"Failed to retrieve chunk {from_byte_index}-{to_byte_index} from {url} (status {error.http_status_code}) - retrying" - ) - continue - else: - raise error - - def _download_all_at_once(self, url: str, target: Path, *, chunk_size: int = DEFAULT_DOWNLOAD_CHUNK_SIZE): - with self.job.connection.get(path=url, stream=True) as r: - r.raise_for_status() - with target.open("wb") as f: - for block in r.iter_content(chunk_size=chunk_size): - f.write(block) - class MultipleAssetException(OpenEoClientException): pass @@ -722,6 +671,161 @@ def download_files( return downloaded + def download_as_collection( + self, + target: Union[Path, str, None] = None, + *, + rewrite_references: bool = True, + download_derived_from: bool = False, + download_collection_assets: bool = False, + json_dumping: Optional[dict] = None, + ) -> List[Path]: + """ + Download the job results as a self-contained STAC collection: + + - job result metadata (the root STAC collection) + - linked items containing the result assets + - additionally linked metadata + + + :param target: folder path to download to + :param rewrite_references: whether to rewrite (item/asset/...) references + in the downloaded STAC collection to point to the local files + instead of the original URLs + :param download_derived_from: whether to download + additional "derived_from" documents linked from the STAC collection. + :param download_collection_assets: whether to download + the STAC Collection level assets in addition to assets from linked STAC Items + :param json_dumping: kwargs to finetune json.dump when writing STAC metadata files + """ + downloader = _JobResultDownloader( + job=self._job, + target=target, + rewrite_references=rewrite_references, + json_dumping=json_dumping, + ) + return downloader.download_collection( + download_derived_from=download_derived_from, + download_collection_assets=download_collection_assets, + ) + + +class _JobResultDownloader: + """ + Helper class to download batch job results as a STAC collection (openEO API 1.1 style): + recursively walking through items, assets and additional linked metadata. + """ + + # TODO: make this a public API that users can implement for custom download behavior (e.g. download to S3, ...) + # TODO: strategy to handle download failures: retry, warn, ignore, error, ... + # TODO: API to warn about or skip existing/previously downloaded files? + # TODO: dedicated request session (with appropriate retry strategy) for downloading? + + def __init__( + self, + *, + job: BatchJob, + target: Union[Path, str, None] = None, + rewrite_references: bool = True, + json_dumping: Optional[dict] = None, + ): + self._job = job + self._connection = job.connection + self._root_dir = Path(target or Path.cwd() / job.job_id) + if self._root_dir.exists() and not self._root_dir.is_dir(): + raise OpenEoClientException(f"Download target {self._root_dir} exists but isn't a folder.") + self._rewrite_references = rewrite_references + # TODO: also support passing a `json.dump`-style callable to customize json dumping + self._json_dumping = {"ensure_ascii": False, **(json_dumping or {})} + self._downloaded: List[Path] = [] + + def _write_json_file(self, data: dict, path: Union[str, Path]) -> Path: + path = Path(path) + ensure_dir(path.parent) + with open(path, mode="w", encoding="utf-8") as f: + json.dump(obj=data, fp=f, **self._json_dumping) + return path + + def download_collection( + self, + *, + download_derived_from: bool = False, + download_collection_assets: bool = False, + ) -> List[Path]: + """ + Download the job results as a self-contained STAC collection. + """ + result_metadata = self._connection.get(self._job.get_results_metadata_url(), expected_status=200).json() + if result_metadata.get("type") != "Collection": + raise OpenEoClientException( + f"Result metadata is not a STAC Collection (openEO API 1.1 style), but {result_metadata.get('type')}" + ) + + result_metadata_path = self._root_dir / DEFAULT_JOB_RESULTS_FILENAME + # Initial write of metadata, will possibly be updated later if rewrite_references is True + self._write_json_file(data=result_metadata, path=result_metadata_path) + + extra_rels = ["derived_from"] if download_derived_from else [] + for link in result_metadata["links"]: + if link["rel"] == "item": + path = self._download_item(href=link["href"]) + if self._rewrite_references: + link["href"] = path.relative_to(result_metadata_path.parent).as_posix() + elif link["rel"] in extra_rels: + path = self._root_dir / (_filename_from_url(link["href"], full=False) or link["rel"]) + self._connection.download_url(url=link["href"], target=path) + if self._rewrite_references: + link["href"] = path.relative_to(result_metadata_path.parent).as_posix() + + if download_collection_assets: + for asset_key, asset_metadata in result_metadata.get("assets", {}).items(): + path = self._download_asset( + asset_key=asset_key, asset_href=asset_metadata["href"], asset_metadata=asset_metadata, item_id=None + ) + if self._rewrite_references: + asset_metadata["href"] = path.relative_to(result_metadata_path.parent).as_posix() + + if self._rewrite_references: + # Rewrite the root collection metadata with updated references + self._write_json_file(data=result_metadata, path=result_metadata_path) + + self._downloaded.append(result_metadata_path) + + return self._downloaded + + def _download_item(self, href: str) -> Path: + item_metadata = self._connection.get(href, expected_status=200).json() + # TODO: sanitize item id to be safe as filename? + # TODO: different strategy to build structure: tree vs flat + metadata_path = self._root_dir / item_metadata["id"] / (item_metadata["id"] + ".json") + self._write_json_file(data=item_metadata, path=metadata_path) + + for asset_key, asset in item_metadata.get("assets", {}).items(): + asset_path = self._download_asset( + asset_key=asset_key, asset_href=asset["href"], asset_metadata=asset, item_id=item_metadata["id"] + ) + if self._rewrite_references: + asset["href"] = asset_path.relative_to(metadata_path.parent).as_posix() + + if self._rewrite_references: + self._write_json_file(data=item_metadata, path=metadata_path) + + self._downloaded.append(metadata_path) + + return metadata_path + + def _download_asset( + self, *, asset_key: str, asset_href: str, asset_metadata: dict, item_id: Optional[str] = None + ) -> Path: + asset = ResultAsset(job=self._job, key=asset_key, href=asset_href, metadata=asset_metadata) + path = self._root_dir + if item_id: + path = path / item_id + path = path / (_filename_from_url(asset_href, full=False) or asset_key) + asset.download(target=path) + self._downloaded.append(path) + return path + @deprecated(reason="Use :py:class:`JobResults` instead", version="0.4.10") class _Result: diff --git a/openeo/testing/stac.py b/openeo/testing/stac.py index e64121667..cace25e95 100644 --- a/openeo/testing/stac.py +++ b/openeo/testing/stac.py @@ -21,6 +21,7 @@ def item( properties: Optional[dict] = None, cube_dimensions: Optional[dict] = None, stac_extensions: Optional[List[str]] = None, + assets: Union[dict, None] = None, **kwargs, ) -> dict: """Create a STAC Item represented as dictionary.""" @@ -38,7 +39,7 @@ def item( "geometry": None, "properties": properties, "links": [], - "assets": {}, + "assets": assets or {}, **kwargs, } diff --git a/tests/rest/test_job.py b/tests/rest/test_job.py index cc6ff834b..9ebbbcbb2 100644 --- a/tests/rest/test_job.py +++ b/tests/rest/test_job.py @@ -4,9 +4,10 @@ import logging import re from pathlib import Path -from typing import Callable, Optional +from typing import Any, Callable, Dict, List, Optional, Union from unittest import mock +import dirty_equals import httpretty import pytest import requests @@ -14,10 +15,18 @@ import openeo import openeo.rest.job from openeo.rest import JobFailedException, OpenEoApiPlainError, OpenEoClientException -from openeo.rest.job import BatchJob, ResultAsset +from openeo.rest._testing import JobResultCollectionMocker +from openeo.rest.job import ( + BatchJob, + ResultAsset, + _filename_from_url, + _JobResultDownloader, + _sanitize_filename, +) from openeo.rest.models.general import Link from openeo.rest.models.logs import LogEntry -from openeo.util import dict_no_none +from openeo.testing.stac import StacDummyBuilder +from openeo.util import dict_no_none, load_json from openeo.utils.events import EVENTS from openeo.utils.http import ( HTTP_402_PAYMENT_REQUIRED, @@ -682,6 +691,7 @@ def test_get_results_metadata_url_full(con100): def job_with_results_mocker(con100, requests_mock) -> Callable: """ Helper to set up a job with downloadable assets + (STAC Item style) """ def setup(*, job_id="jj1", assets: dict, media_type: str = "image/tiff; application=geotiff"): @@ -707,6 +717,7 @@ def setup(*, job_id="jj1", assets: dict, media_type: str = "image/tiff; applicat return setup + @pytest.fixture def job_with_1_asset(job_with_results_mocker) -> BatchJob: return job_with_results_mocker(job_id="jj1", assets={"1.tiff": "/dl/jjr1.tiff"}) @@ -1078,7 +1089,6 @@ def download_tiff(request, context): class TestResultAsset: - @pytest.fixture def job(self, con100): return BatchJob("jj", connection=con100) @@ -1200,3 +1210,182 @@ def get_jobs(request, context): assert jobs.links == [Link(rel="next", href="https://oeo.test/jobs?limit=2&offset=2")] assert jobs.ext_federation_missing() == ["oeob"] assert "Partial job listing: missing federation components: ['oeob']." in caplog.text + + +def test_sanitize_filename(): + assert _sanitize_filename("foo/bar.txt") == "foobar.txt" + assert _sanitize_filename("foo/bar.txt", replacement="_") == "foo_bar.txt" + assert _sanitize_filename(r"foo\bar.txt", replacement="_") == "foo_bar.txt" + assert _sanitize_filename("foo\nbar.txt", replacement="_") == "foo_bar.txt" + assert _sanitize_filename("foo$bar.txt", replacement="_") == "foo_bar.txt" + assert _sanitize_filename("foo%bar.txt", replacement="_") == "foo_bar.txt" + + +def test_filename_from_url(): + assert _filename_from_url("https://example.com/foo/bar.txt") == "bar.txt" + assert _filename_from_url("https://example.com/foo/bar.txt?q=1&r=2#frag") == "bar.txt" + assert _filename_from_url("https://example.com/foo/ba%CF%83.txt") == "baσ.txt" + assert _filename_from_url("https://example.com/foo/bar") == "bar" + assert _filename_from_url("https://example.com/foo/bar/") == "bar" + assert _filename_from_url("https://example.com/") == "" + + # Full mode + assert _filename_from_url("https://example.com/foo/bar.txt", full=True) == "foo/bar.txt" + assert _filename_from_url("https://example.com/foo/bar.txt?q=1&r=2#frag", full=True) == "foo/bar.txt" + assert _filename_from_url("https://example.com/fo%CF%83/ba%CF%83", full=True) == "foσ/baσ" + assert _filename_from_url("https://example.com/foo/bar", full=True) == "foo/bar" + assert _filename_from_url("https://example.com/foo/bar/", full=True) == "foo/bar" + assert _filename_from_url("https://example.com/", full=True) == "" + + # Relative href + assert _filename_from_url("foo/bar.txt") == "bar.txt" + assert _filename_from_url("/foo/bar.txt") == "bar.txt" + assert _filename_from_url("foo/bar.txt", full=True) == "foo/bar.txt" + assert _filename_from_url("/foo/bar.txt", full=True) == "foo/bar.txt" + + +class TestJobResultDownloader: + + @pytest.fixture + def result_mocker(self, con100, requests_mock) -> JobResultCollectionMocker: + return JobResultCollectionMocker(requests_mock=requests_mock, connection=con100) + + def check_expected_downloads(self, downloaded: List[Path], expected: Dict[str, Any], tmp_path: Path): + expected_paths = set(tmp_path / k for k in expected.keys()) + assert set(downloaded) == expected_paths + assert set(p for p in tmp_path.glob("**/*") if p.is_file()) == expected_paths + + for path, expected_value in expected.items(): + actual = tmp_path / path + if actual.suffix == ".json": + assert load_json(actual) == expected_value + elif actual.suffix in {".tif", ".tiff"}: + assert actual.read_bytes() == expected_value + else: + raise ValueError(f"Unsupported {path=} {expected_value=}") + + def test_basic(self, result_mocker, tmp_path): + job = result_mocker.setup_job_results( + items={ + "item1": {"assets": {"asset1": {"path": "asset1.tiff"}}}, + } + ) + downloader = _JobResultDownloader(job=job, target=tmp_path) + downloaded = downloader.download_collection() + expected = { + "job-results.json": dirty_equals.IsPartialDict( + { + "id": "job-123-results", + "type": "Collection", + "stac_version": "1.1.0", + "links": [{"rel": "item", "href": "item1/item1.json"}], + } + ), + "item1/item1.json": dirty_equals.IsPartialDict( + { + "id": "item1", + "type": "Feature", + "stac_version": "1.1.0", + "assets": { + "asset1": dirty_equals.IsPartialDict(href="asset1.tiff"), + }, + } + ), + "item1/asset1.tiff": b"TIFF-DUMMY-DATA", + } + self.check_expected_downloads(downloaded=downloaded, expected=expected, tmp_path=tmp_path) + + def test_one_item_multiple_assets(self, result_mocker, tmp_path): + job = result_mocker.setup_job_results( + items={ + "item1": { + "assets": { + "asset1": {"path": "asset1.tiff"}, + "asset2": {"path": "asset2.tiff"}, + "asset3": {"path": "asset3.tiff"}, + } + } + } + ) + downloader = _JobResultDownloader(job=job, target=tmp_path) + downloaded = downloader.download_collection() + expected = { + "job-results.json": dirty_equals.IsPartialDict( + { + "id": "job-123-results", + "type": "Collection", + "stac_version": "1.1.0", + "links": [{"rel": "item", "href": "item1/item1.json"}], + } + ), + "item1/item1.json": dirty_equals.IsPartialDict( + { + "id": "item1", + "type": "Feature", + "stac_version": "1.1.0", + "assets": { + "asset1": dirty_equals.IsPartialDict(href="asset1.tiff"), + "asset2": dirty_equals.IsPartialDict(href="asset2.tiff"), + "asset3": dirty_equals.IsPartialDict(href="asset3.tiff"), + }, + } + ), + "item1/asset1.tiff": b"TIFF-DUMMY-DATA", + "item1/asset2.tiff": b"TIFF-DUMMY-DATA", + "item1/asset3.tiff": b"TIFF-DUMMY-DATA", + } + self.check_expected_downloads(downloaded=downloaded, expected=expected, tmp_path=tmp_path) + + def test_multiple_items(self, result_mocker, tmp_path): + job = result_mocker.setup_job_results( + items={ + "item1": { + "assets": { + "asset1": {"path": "asset1.tiff"}, + }, + }, + "item2": { + "assets": { + "asset2": {"path": "asset2.tiff"}, + "asset3": {"path": "asset3.tiff"}, + } + }, + } + ) + downloader = _JobResultDownloader(job=job, target=tmp_path) + downloaded = downloader.download_collection() + expected = { + "job-results.json": dirty_equals.IsPartialDict( + { + "id": "job-123-results", + "type": "Collection", + "links": [ + {"rel": "item", "href": "item1/item1.json"}, + {"rel": "item", "href": "item2/item2.json"}, + ], + } + ), + "item1/item1.json": dirty_equals.IsPartialDict( + { + "id": "item1", + "type": "Feature", + "stac_version": "1.1.0", + "assets": {"asset1": dirty_equals.IsPartialDict(href="asset1.tiff")}, + } + ), + "item2/item2.json": dirty_equals.IsPartialDict( + { + "id": "item2", + "type": "Feature", + "stac_version": "1.1.0", + "assets": { + "asset2": dirty_equals.IsPartialDict(href="asset2.tiff"), + "asset3": dirty_equals.IsPartialDict(href="asset3.tiff"), + }, + } + ), + "item1/asset1.tiff": b"TIFF-DUMMY-DATA", + "item2/asset2.tiff": b"TIFF-DUMMY-DATA", + "item2/asset3.tiff": b"TIFF-DUMMY-DATA", + } + self.check_expected_downloads(downloaded=downloaded, expected=expected, tmp_path=tmp_path) diff --git a/tests/rest/test_testing.py b/tests/rest/test_testing.py index 589dda3dc..79f5d898e 100644 --- a/tests/rest/test_testing.py +++ b/tests/rest/test_testing.py @@ -1,9 +1,10 @@ import re +import dirty_equals import pytest from openeo.rest import OpenEoApiError -from openeo.rest._testing import DummyBackend +from openeo.rest._testing import DummyBackend, JobResultCollectionMocker @pytest.fixture @@ -104,3 +105,49 @@ def test_setup_job_start_failure(self, dummy_backend): with pytest.raises(OpenEoApiError, match=re.escape("[500] Internal: No job starting for you, buddy")): job.start() assert job.status() == "error" + + +class TestJobResultCollectionMocker: + def test_basic(self, requests_mock, con120): + result_mocker = JobResultCollectionMocker(requests_mock=requests_mock, connection=con120) + result_mocker.setup_job_results( + job_id="job-456", + items={"item-567": {"assets": {"asset-678": {"path": "assets/asset-678.tif"}}}}, + ) + + job = con120.job("job-456") + assert job.get_results().get_metadata() == dirty_equals.IsPartialDict( + { + "type": "Collection", + "stac_version": "1.1.0", + "id": "job-456-results", + "links": [ + { + "rel": "item", + "href": "https://oeo.test/j/job-456/r/i/item-567.json", + } + ], + "assets": { + "item-567-asset-678": { + "href": "https://oeo.test/j/job-456/r/a/assets/asset-678.tif", + "roles": ["data"], + "type": "image/tiff; application=geotiff", + } + }, + } + ) + assert con120.get("https://oeo.test/j/job-456/r/i/item-567.json").json() == dirty_equals.IsPartialDict( + { + "type": "Feature", + "stac_version": "1.1.0", + "id": "item-567", + "assets": { + "asset-678": { + "href": "https://oeo.test/j/job-456/r/a/assets/asset-678.tif", + "roles": ["data"], + "type": "image/tiff; application=geotiff", + } + }, + } + ) + assert con120.get("https://oeo.test/j/job-456/r/a/assets/asset-678.tif").content == b"TIFF-DUMMY-DATA" From 8856f7fbf0816a8757ec864af43302f6cb62f801 Mon Sep 17 00:00:00 2001 From: Stefaan Lippens Date: Wed, 12 Aug 2026 08:44:25 +0200 Subject: [PATCH 2/3] Run pytest in more verbose mode --- .github/workflows/unittests.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/unittests.yml b/.github/workflows/unittests.yml index 5df7499f6..210c3431b 100644 --- a/.github/workflows/unittests.yml +++ b/.github/workflows/unittests.yml @@ -47,4 +47,4 @@ jobs: python -m pip install --upgrade pip python -m pip install .[tests] - name: Test with pytest - run: pytest + run: pytest -v From d7c8d370cefd661760f31ca35aa914377cfc68d2 Mon Sep 17 00:00:00 2001 From: Stefaan Lippens Date: Wed, 12 Aug 2026 15:23:18 +0200 Subject: [PATCH 3/3] _JobResultDownloader: add `on_download_failure` ref #931 --- openeo/rest/_testing.py | 36 +++++++++++++------ openeo/rest/job.py | 74 ++++++++++++++++++++++++++------------ tests/rest/test_job.py | 52 +++++++++++++++++++++++++++ tests/rest/test_testing.py | 28 ++++++++++++--- 4 files changed, 151 insertions(+), 39 deletions(-) diff --git a/openeo/rest/_testing.py b/openeo/rest/_testing.py index 4e8cda9b5..edc41f99f 100644 --- a/openeo/rest/_testing.py +++ b/openeo/rest/_testing.py @@ -513,7 +513,7 @@ def setup_job_results( assets[asset_key] = asset collection_assets[f"{item_id}-{asset_key}"] = asset - item_href = self.setup_item(job_id=job_id, item_id=item_id, assets=assets) + item_href = self.setup_item(job_id=job_id, item_id=item_id, item_data=item_data, assets=assets) links.append({"rel": "item", "href": item_href}) collection_href = self.connection.build_url(f"/jobs/{job_id}/results") @@ -528,21 +528,35 @@ def setup_job_results( job = BatchJob(job_id, connection=self.connection) return job - def setup_item(self, *, job_id: str, item_id: str, assets: dict) -> dict: - href = self.connection.build_url(f"/j/{job_id}/r/i/{item_id}.json") - doc = StacDummyBuilder.item( - id=item_id, - stac_version="1.1.0", - assets=assets, + def setup_error(self, href, error: dict): + self.requests_mock.get( + href, + status_code=error.get("status", 500), + text=error.get("message", "Unspecified error"), ) - self.requests_mock.get(href, json=doc) + + def setup_item(self, *, job_id: str, item_id: str, item_data: dict, assets: dict) -> dict: + href = self.connection.build_url(f"/j/{job_id}/r/i/{item_id}.json") + if error := item_data.get("error"): + self.setup_error(href, error=error) + else: + doc = StacDummyBuilder.item( + id=item_id, + stac_version="1.1.0", + assets=assets, + ) + self.requests_mock.get(href, json=doc) return href def setup_asset(self, *, job_id: str, asset_data: dict) -> dict: href = self.connection.build_url(f"/j/{job_id}/r/a/{asset_data.get('path', 'asset.tiff')}") - content = asset_data.get("content", b"TIFF-DUMMY-DATA") - self.requests_mock.head(href, headers={"Content-Length": f"{len(content)}"}) - self.requests_mock.get(href, content=content) + if error := asset_data.get("error"): + self.requests_mock.head(href, headers={}) + self.setup_error(href, error=error) + else: + content = asset_data.get("content", b"TIFF-DUMMY-DATA") + self.requests_mock.head(href, headers={"Content-Length": f"{len(content)}"}) + self.requests_mock.get(href, content=content) return StacDummyBuilder.asset( href=href, type=asset_data.get("type", "image/tiff; application=geotiff"), diff --git a/openeo/rest/job.py b/openeo/rest/job.py index 3106a553e..4b6118ca1 100644 --- a/openeo/rest/job.py +++ b/openeo/rest/job.py @@ -1,5 +1,7 @@ from __future__ import annotations +import contextlib +import copy import datetime import json import logging @@ -710,10 +712,16 @@ def download_as_collection( ) +class JobResultDownloadException(OpenEoClientException): + pass + + class _JobResultDownloader: """ Helper class to download batch job results as a STAC collection (openEO API 1.1 style): recursively walking through items, assets and additional linked metadata. + + Experimental API, subject to change. """ # TODO: make this a public API that users can implement for custom download behavior (e.g. download to S3, ...) @@ -728,6 +736,7 @@ def __init__( target: Union[Path, str, None] = None, rewrite_references: bool = True, json_dumping: Optional[dict] = None, + on_download_failure: str = "warn", ): self._job = job self._connection = job.connection @@ -738,6 +747,7 @@ def __init__( # TODO: also support passing a `json.dump`-style callable to customize json dumping self._json_dumping = {"ensure_ascii": False, **(json_dumping or {})} self._downloaded: List[Path] = [] + self._on_download_failure = on_download_failure def _write_json_file(self, data: dict, path: Union[str, Path]) -> Path: path = Path(path) @@ -746,6 +756,17 @@ def _write_json_file(self, data: dict, path: Union[str, Path]) -> Path: json.dump(obj=data, fp=f, **self._json_dumping) return path + @contextlib.contextmanager + def _download_attempt_context(self, name: str): + try: + yield + except Exception as e: + message = f"Failed to download {name} ({e=})" + if self._on_download_failure == "warn": + logger.warning(message, exc_info=True) + else: + raise JobResultDownloadException(message) from e + def download_collection( self, *, @@ -760,6 +781,8 @@ def download_collection( raise OpenEoClientException( f"Result metadata is not a STAC Collection (openEO API 1.1 style), but {result_metadata.get('type')}" ) + # Make a copy of the metadata, as we will rewrite references + result_metadata = copy.deepcopy(result_metadata) result_metadata_path = self._root_dir / DEFAULT_JOB_RESULTS_FILENAME # Initial write of metadata, will possibly be updated later if rewrite_references is True @@ -768,22 +791,26 @@ def download_collection( extra_rels = ["derived_from"] if download_derived_from else [] for link in result_metadata["links"]: if link["rel"] == "item": - path = self._download_item(href=link["href"]) - if self._rewrite_references: - link["href"] = path.relative_to(result_metadata_path.parent).as_posix() + with self._download_attempt_context(name=f"item {link=}"): + path = self._download_item(href=link["href"]) + if self._rewrite_references: + link["href"] = path.relative_to(result_metadata_path.parent).as_posix() + elif link["rel"] in extra_rels: - path = self._root_dir / (_filename_from_url(link["href"], full=False) or link["rel"]) - self._connection.download_url(url=link["href"], target=path) - if self._rewrite_references: - link["href"] = path.relative_to(result_metadata_path.parent).as_posix() + with self._download_attempt_context(name=f"link {link=}"): + path = self._root_dir / (_filename_from_url(link["href"], full=False) or link["rel"]) + self._connection.download_url(url=link["href"], target=path) + if self._rewrite_references: + link["href"] = path.relative_to(result_metadata_path.parent).as_posix() if download_collection_assets: - for asset_key, asset_metadata in result_metadata.get("assets", {}).items(): - path = self._download_asset( - asset_key=asset_key, asset_href=asset_metadata["href"], asset_metadata=asset_metadata, item_id=None - ) - if self._rewrite_references: - asset_metadata["href"] = path.relative_to(result_metadata_path.parent).as_posix() + for asset_key, asset in result_metadata.get("assets", {}).items(): + with self._download_attempt_context(name=f"collection asset {asset_key=} {asset=}"): + path = self._download_asset( + asset_key=asset_key, asset_href=asset["href"], asset_metadata=asset, item_id=None + ) + if self._rewrite_references: + asset["href"] = path.relative_to(result_metadata_path.parent).as_posix() if self._rewrite_references: # Rewrite the root collection metadata with updated references @@ -794,21 +821,22 @@ def download_collection( return self._downloaded def _download_item(self, href: str) -> Path: - item_metadata = self._connection.get(href, expected_status=200).json() + item: dict = self._connection.get(href, expected_status=200).json() # TODO: sanitize item id to be safe as filename? # TODO: different strategy to build structure: tree vs flat - metadata_path = self._root_dir / item_metadata["id"] / (item_metadata["id"] + ".json") - self._write_json_file(data=item_metadata, path=metadata_path) + metadata_path = self._root_dir / item["id"] / (item["id"] + ".json") + self._write_json_file(data=item, path=metadata_path) - for asset_key, asset in item_metadata.get("assets", {}).items(): - asset_path = self._download_asset( - asset_key=asset_key, asset_href=asset["href"], asset_metadata=asset, item_id=item_metadata["id"] - ) - if self._rewrite_references: - asset["href"] = asset_path.relative_to(metadata_path.parent).as_posix() + for asset_key, asset in item.get("assets", {}).items(): + with self._download_attempt_context(name=f"item asset {asset_key=} {asset=}"): + asset_path = self._download_asset( + asset_key=asset_key, asset_href=asset["href"], asset_metadata=asset, item_id=item["id"] + ) + if self._rewrite_references: + asset["href"] = asset_path.relative_to(metadata_path.parent).as_posix() if self._rewrite_references: - self._write_json_file(data=item_metadata, path=metadata_path) + self._write_json_file(data=item, path=metadata_path) self._downloaded.append(metadata_path) diff --git a/tests/rest/test_job.py b/tests/rest/test_job.py index 9ebbbcbb2..c22cab43f 100644 --- a/tests/rest/test_job.py +++ b/tests/rest/test_job.py @@ -18,6 +18,7 @@ from openeo.rest._testing import JobResultCollectionMocker from openeo.rest.job import ( BatchJob, + JobResultDownloadException, ResultAsset, _filename_from_url, _JobResultDownloader, @@ -1389,3 +1390,54 @@ def test_multiple_items(self, result_mocker, tmp_path): "item2/asset3.tiff": b"TIFF-DUMMY-DATA", } self.check_expected_downloads(downloaded=downloaded, expected=expected, tmp_path=tmp_path) + + @pytest.mark.parametrize( + ["on_download_failure", "items_setup", "expected"], + [ + ( + "warn", + { + "item1": { + "assets": {"asset1": {"path": "asset1.tiff", "error": {"message": "Nope no asset1 for you"}}}, + } + }, + "Failed to download item asset..*Nope no asset1 for you", + ), + ( + "error", + { + "item1": { + "assets": {"asset1": {"path": "asset1.tiff", "error": {"message": "Nope no asset1 for you"}}}, + } + }, + "Failed to download item asset.*Nope no asset1 for you", + ), + ( + "warn", + {"item1": {"error": {"message": "Nope no item1 for you"}}}, + "Failed to download item.*Nope no item1 for you", + ), + ( + "error", + {"item1": {"error": {"message": "Nope no item1 for you"}}}, + "Failed to download item.*Nope no item1 for you", + ), + ], + ) + def test_warn_or_error_on_download_fail( + self, result_mocker, tmp_path, caplog, on_download_failure, items_setup, expected + ): + job = result_mocker.setup_job_results(items=items_setup) + + expected = re.compile(expected) + if on_download_failure == "error": + context = pytest.raises(JobResultDownloadException, match=expected) + else: + context = contextlib.nullcontext() + + downloader = _JobResultDownloader(job=job, target=tmp_path, on_download_failure=on_download_failure) + with context: + downloader.download_collection() + + if on_download_failure == "warn": + assert expected.search(caplog.text) diff --git a/tests/rest/test_testing.py b/tests/rest/test_testing.py index 79f5d898e..c505b2e4e 100644 --- a/tests/rest/test_testing.py +++ b/tests/rest/test_testing.py @@ -3,7 +3,7 @@ import dirty_equals import pytest -from openeo.rest import OpenEoApiError +from openeo.rest import OpenEoApiError, OpenEoRestError from openeo.rest._testing import DummyBackend, JobResultCollectionMocker @@ -112,7 +112,7 @@ def test_basic(self, requests_mock, con120): result_mocker = JobResultCollectionMocker(requests_mock=requests_mock, connection=con120) result_mocker.setup_job_results( job_id="job-456", - items={"item-567": {"assets": {"asset-678": {"path": "assets/asset-678.tif"}}}}, + items={"item-567": {"assets": {"asset-678": {"path": "asset-678.tif"}}}}, ) job = con120.job("job-456") @@ -129,7 +129,7 @@ def test_basic(self, requests_mock, con120): ], "assets": { "item-567-asset-678": { - "href": "https://oeo.test/j/job-456/r/a/assets/asset-678.tif", + "href": "https://oeo.test/j/job-456/r/a/asset-678.tif", "roles": ["data"], "type": "image/tiff; application=geotiff", } @@ -143,11 +143,29 @@ def test_basic(self, requests_mock, con120): "id": "item-567", "assets": { "asset-678": { - "href": "https://oeo.test/j/job-456/r/a/assets/asset-678.tif", + "href": "https://oeo.test/j/job-456/r/a/asset-678.tif", "roles": ["data"], "type": "image/tiff; application=geotiff", } }, } ) - assert con120.get("https://oeo.test/j/job-456/r/a/assets/asset-678.tif").content == b"TIFF-DUMMY-DATA" + assert con120.get("https://oeo.test/j/job-456/r/a/asset-678.tif").content == b"TIFF-DUMMY-DATA" + + def test_item_error(self, requests_mock, con120): + result_mocker = JobResultCollectionMocker(requests_mock=requests_mock, connection=con120) + result_mocker.setup_job_results( + job_id="job-456", + items={"item-567": {"error": {"message": "Nope!"}}}, + ) + with pytest.raises(OpenEoRestError, match=re.escape("[500] Nope!")): + con120.get("https://oeo.test/j/job-456/r/i/item-567.json") + + def test_asset_error(self, requests_mock, con120): + result_mocker = JobResultCollectionMocker(requests_mock=requests_mock, connection=con120) + result_mocker.setup_job_results( + job_id="job-456", + items={"item-567": {"assets": {"asset-678": {"path": "asset-678.tif", "error": {"message": "Nope!"}}}}}, + ) + with pytest.raises(OpenEoRestError, match=re.escape("[500] Nope!")): + con120.get("https://oeo.test/j/job-456/r/a/asset-678.tif")