From 63509a9e19c8648b4774ab545a9731d80c2e5c26 Mon Sep 17 00:00:00 2001 From: Stefaan Lippens Date: Tue, 11 Aug 2026 22:16:13 +0200 Subject: [PATCH 1/8] 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/8] 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 e4e37cec55f1a70f3afc048ddfe0ccbd387afc10 Mon Sep 17 00:00:00 2001 From: Stefaan Lippens Date: Wed, 12 Aug 2026 15:23:18 +0200 Subject: [PATCH 3/8] _JobResultDownloader: add `on_download_failure` ref #931 --- openeo/rest/_testing.py | 36 ++++++++++++------ openeo/rest/job.py | 75 ++++++++++++++++++++++++++------------ tests/rest/test_job.py | 52 ++++++++++++++++++++++++++ tests/rest/test_testing.py | 28 +++++++++++--- 4 files changed, 151 insertions(+), 40 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..5ed585505 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,14 +712,19 @@ 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, ...) - # 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? @@ -728,6 +735,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 +746,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 +755,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 +780,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 +790,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 +820,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") From 2152e82cf0d9db4943aa286198f27c9a4027fe68 Mon Sep 17 00:00:00 2001 From: Stefaan Lippens Date: Wed, 12 Aug 2026 19:05:45 +0200 Subject: [PATCH 4/8] _JobResultDownloader: support download path templating --- openeo/rest/_connection.py | 10 ++- openeo/rest/_testing.py | 27 ++++++-- openeo/rest/job.py | 88 +++++++++++++++++++++----- openeo/util.py | 4 ++ tests/rest/test_job.py | 123 +++++++++++++++++++++++++++++++++++++ 5 files changed, 232 insertions(+), 20 deletions(-) diff --git a/openeo/rest/_connection.py b/openeo/rest/_connection.py index 647cccbda..6c17c3693 100644 --- a/openeo/rest/_connection.py +++ b/openeo/rest/_connection.py @@ -19,7 +19,13 @@ OpenEoRestError, ) from openeo.rest.auth.auth import NullAuth -from openeo.util import ContextTimer, ensure_list, str_truncate, url_join +from openeo.util import ( + ContextTimer, + ensure_list, + ensure_parent_dir_for, + str_truncate, + url_join, +) from openeo.utils.http import ( HTTP_408_REQUEST_TIMEOUT, HTTP_429_TOO_MANY_REQUESTS, @@ -311,6 +317,7 @@ def download_url( 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() + ensure_parent_dir_for(target) with target.open("wb") as f: for block in r.iter_content(chunk_size=chunk_size): f.write(block) @@ -324,6 +331,7 @@ def _download_ranged( chunk_size: int = DEFAULT_DOWNLOAD_CHUNK_SIZE, range_size: int = DEFAULT_DOWNLOAD_RANGE_SIZE, ) -> None: + ensure_parent_dir_for(target) 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) diff --git a/openeo/rest/_testing.py b/openeo/rest/_testing.py index edc41f99f..d5f387a34 100644 --- a/openeo/rest/_testing.py +++ b/openeo/rest/_testing.py @@ -7,6 +7,7 @@ Callable, Dict, Iterable, + List, Mapping, Optional, Sequence, @@ -502,9 +503,14 @@ def __init__(self, *, requests_mock, connection: Connection): self.connection = connection def setup_job_results( - self, *, job_id: str = "job-123", items: dict, add_collection_assets: bool = True + self, + *, + job_id: str = "job-123", + items: dict, + add_collection_assets: bool = True, + linked_docs: Iterable[dict] = (), ) -> BatchJob: - links = [] + collection_links = [] collection_assets = {} for item_id, item_data in items.items(): assets = {} @@ -514,13 +520,16 @@ def setup_job_results( collection_assets[f"{item_id}-{asset_key}"] = asset 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_links.append({"rel": "item", "href": item_href}) + + for doc in linked_docs: + collection_links.append(self.setup_linked_document(job_id=job_id, doc=doc)) 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, + links=collection_links, assets=collection_assets if add_collection_assets else {}, ) self.requests_mock.get(collection_href, json=collection_doc) @@ -561,3 +570,13 @@ def setup_asset(self, *, job_id: str, asset_data: dict) -> dict: href=href, type=asset_data.get("type", "image/tiff; application=geotiff"), ) + + def setup_linked_document(self, *, job_id: str, doc: dict): + href = self.connection.build_url(f"/j/{job_id}/r/d/{doc.get('path', 'doc.txt')}") + if "json" in doc: + text = json.dumps(doc["json"]) + else: + text = doc.get("text", "hello world") + self.requests_mock.head(href, headers={"Content-Length": f"{len(text)}"}) + self.requests_mock.get(href, text=text) + return {"rel": doc.get("rel", "doc"), "href": href} diff --git a/openeo/rest/job.py b/openeo/rest/job.py index 5ed585505..576818ebe 100644 --- a/openeo/rest/job.py +++ b/openeo/rest/job.py @@ -5,6 +5,7 @@ import datetime import json import logging +import os.path import re import time import typing @@ -30,7 +31,7 @@ ) from openeo.rest.models.general import LogsResponse from openeo.rest.models.logs import log_level_name -from openeo.util import ensure_dir +from openeo.util import ensure_dir, ensure_parent_dir_for from openeo.utils.events import EVENTS from openeo.utils.http import ( HTTP_408_REQUEST_TIMEOUT, @@ -681,6 +682,7 @@ def download_as_collection( download_derived_from: bool = False, download_collection_assets: bool = False, json_dumping: Optional[dict] = None, + path_templates: Optional[dict] = None, ) -> List[Path]: """ Download the job results as a self-contained STAC collection: @@ -699,12 +701,14 @@ def download_as_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 + :param path_templates: optional tempalte overrides for download paths. """ downloader = _JobResultDownloader( job=self._job, target=target, rewrite_references=rewrite_references, json_dumping=json_dumping, + path_templates=path_templates, ) return downloader.download_collection( download_derived_from=download_derived_from, @@ -736,6 +740,7 @@ def __init__( rewrite_references: bool = True, json_dumping: Optional[dict] = None, on_download_failure: str = "warn", + path_templates: Optional[dict] = None, ): self._job = job self._connection = job.connection @@ -747,10 +752,18 @@ def __init__( self._json_dumping = {"ensure_ascii": False, **(json_dumping or {})} self._downloaded: List[Path] = [] self._on_download_failure = on_download_failure + self._path_templates = { + "collection": "job-results.json", + "item": "{item_id}/{item_id}.json", + "asset": "{item_id}/{asset_filename}", + "collection-asset": "{asset_filename}", + "generic-link": "{filename}", + **(path_templates or {}), + } def _write_json_file(self, data: dict, path: Union[str, Path]) -> Path: path = Path(path) - ensure_dir(path.parent) + ensure_parent_dir_for(path) with open(path, mode="w", encoding="utf-8") as f: json.dump(obj=data, fp=f, **self._json_dumping) return path @@ -766,6 +779,51 @@ def _download_attempt_context(self, name: str): else: raise JobResultDownloadException(message) from e + def _check_download_path(self, path: Path): + if path in self._downloaded: + raise JobResultDownloadException("Download collision: {path} already downloaded") + + def build_path_collection(self, *, collection_id: str) -> Path: + """Build path for the root STAC collection metadata file (job results metadata)""" + vars = { + "job_id": _sanitize_filename(self._job.job_id), + "collection_id": _sanitize_filename(collection_id), + } + return self._root_dir / self._path_templates["collection"].format(**vars) + + def build_path_item(self, *, item_id: str) -> Path: + """Build path for a STAC item metadata file (job result item)""" + vars = { + "job_id": _sanitize_filename(self._job.job_id), + "item_id": _sanitize_filename(item_id), + } + return self._root_dir / self._path_templates["item"].format(**vars) + + def build_path_asset(self, *, asset_key: str, asset_href: str, item_id: Optional[str] = None) -> Path: + """Build path for a STAC asset file (job result asset)""" + vars = { + "job_id": _sanitize_filename(self._job.job_id), + "asset_key": _sanitize_filename(asset_key), + "asset_filename": _filename_from_url(asset_href, full=False), + } + if item_id: + vars["item_id"] = _sanitize_filename(item_id) + return self._root_dir / self._path_templates["asset"].format(**vars) + else: + return self._root_dir / self._path_templates["collection-asset"].format(**vars) + + def build_path_generic_link(self, *, rel: str, href: str) -> Path: + vars = { + "job_id": _sanitize_filename(self._job.job_id), + "rel": _sanitize_filename(rel), + "filename": _filename_from_url(href, full=False), + } + return self._root_dir / self._path_templates["generic-link"].format(**vars) + + def _relative_to(self, target: Path, doc: Path) -> str: + """Get relative reference to target to be used from given document""" + return Path(os.path.relpath(target, start=doc.parent)).as_posix() + def download_collection( self, *, @@ -783,7 +841,8 @@ def download_collection( # 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 + result_metadata_path = self.build_path_collection(collection_id=result_metadata.get("id")) + self._check_download_path(result_metadata_path) # 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) @@ -793,14 +852,16 @@ def download_collection( 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() + link["href"] = self._relative_to(target=path, doc=result_metadata_path) elif link["rel"] in extra_rels: with self._download_attempt_context(name=f"link {link=}"): - path = self._root_dir / (_filename_from_url(link["href"], full=False) or link["rel"]) + path = self.build_path_generic_link(rel=link["rel"], href=link["href"]) + self._check_download_path(path) self._connection.download_url(url=link["href"], target=path) + self._downloaded.append(path) if self._rewrite_references: - link["href"] = path.relative_to(result_metadata_path.parent).as_posix() + link["href"] = self._relative_to(target=path, doc=result_metadata_path) if download_collection_assets: for asset_key, asset in result_metadata.get("assets", {}).items(): @@ -809,7 +870,7 @@ def download_collection( 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() + asset["href"] = self._relative_to(target=path, doc=result_metadata_path) if self._rewrite_references: # Rewrite the root collection metadata with updated references @@ -821,9 +882,8 @@ def download_collection( def _download_item(self, href: str) -> Path: 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["id"] / (item["id"] + ".json") + metadata_path = self.build_path_item(item_id=item["id"]) + self._check_download_path(metadata_path) self._write_json_file(data=item, path=metadata_path) for asset_key, asset in item.get("assets", {}).items(): @@ -832,7 +892,7 @@ def _download_item(self, href: str) -> Path: 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() + asset["href"] = self._relative_to(target=asset_path, doc=metadata_path) if self._rewrite_references: self._write_json_file(data=item, path=metadata_path) @@ -845,10 +905,8 @@ 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) + path = self.build_path_asset(asset_key=asset_key, asset_href=asset_href, item_id=item_id) + self._check_download_path(path) asset.download(target=path) self._downloaded.append(path) return path diff --git a/openeo/util.py b/openeo/util.py index fd2517652..79cf184bf 100644 --- a/openeo/util.py +++ b/openeo/util.py @@ -247,6 +247,10 @@ def ensure_dir(path: Union[str, Path]) -> Path: return path +def ensure_parent_dir_for(path: Union[str, Path]) -> Path: + return ensure_dir(Path(path).parent) + + def ensure_list(x): """Convert given data structure to a list.""" try: diff --git a/tests/rest/test_job.py b/tests/rest/test_job.py index c22cab43f..f0329e830 100644 --- a/tests/rest/test_job.py +++ b/tests/rest/test_job.py @@ -1441,3 +1441,126 @@ def test_warn_or_error_on_download_fail( if on_download_failure == "warn": assert expected.search(caplog.text) + + @pytest.mark.parametrize( + [ + "download_derived_from", + "expected_links", + "expected_downloads_extra", + ], + [ + ( + False, + [ + {"rel": "item", "href": "item1/item1.json"}, + {"rel": "derived_from", "href": "https://oeo.test/j/job-123/r/d/derived_from.json"}, + ], + {}, + ), + ( + True, + [ + {"rel": "item", "href": "item1/item1.json"}, + {"rel": "derived_from", "href": "derived_from.json"}, + ], + { + "derived_from.json": {"hello": "world"}, + }, + ), + ], + ) + def test_download_derived_from_link( + self, result_mocker, tmp_path, download_derived_from, expected_links, expected_downloads_extra + ): + job = result_mocker.setup_job_results( + items={"item1": {}}, + linked_docs=[ + {"rel": "derived_from", "path": "derived_from.json", "json": {"hello": "world"}}, + ], + ) + downloader = _JobResultDownloader(job=job, target=tmp_path) + downloaded = downloader.download_collection(download_derived_from=download_derived_from) + expected = { + "job-results.json": dirty_equals.IsPartialDict( + { + "type": "Collection", + "links": expected_links, + } + ), + "item1/item1.json": dirty_equals.IsPartialDict({"id": "item1", "type": "Feature"}), + **expected_downloads_extra, + } + self.check_expected_downloads(downloaded=downloaded, expected=expected, tmp_path=tmp_path) + + @pytest.mark.parametrize( + ["path_templates", "expected"], + [ + ( + # Flat structure + { + "collection": "CO_{job_id}.json", + "item": "IT_{item_id}.json", + "asset": "AS_{item_id}-{asset_key}-{asset_filename}", + "collection-asset": "CA_{asset_key}", + "generic-link": "GL_{filename}", + }, + { + "CO_job-123.json": dirty_equals.IsPartialDict( + { + "links": [ + {"rel": "item", "href": "IT_item1.json"}, + {"rel": "derived_from", "href": "GL_derived_from.json"}, + ] + } + ), + "IT_item1.json": dirty_equals.IsPartialDict( + {"assets": {"a1": dirty_equals.IsPartialDict(href="AS_item1-a1-asset1.tiff")}} + ), + "AS_item1-a1-asset1.tiff": b"TIFF-DUMMY-DATA", + "GL_derived_from.json": {"hello": "world"}, + }, + ), + ( + # folder organisation per type + { + "collection": "collections/{job_id}.json", + "item": "items/{job_id}-{item_id}/item.json", + "asset": "assets/{job_id}-{item_id}-{asset_key}/{asset_filename}", + "collection-asset": "assets/{job_id}-{asset_key}/{asset_filename}", + "generic-link": "docs/{job_id}/{filename}", + }, + { + "collections/job-123.json": dirty_equals.IsPartialDict( + { + "links": [ + {"rel": "item", "href": "../items/job-123-item1/item.json"}, + {"rel": "derived_from", "href": "../docs/job-123/derived_from.json"}, + ] + } + ), + "items/job-123-item1/item.json": dirty_equals.IsPartialDict( + {"assets": {"a1": dirty_equals.IsPartialDict(href="../../assets/job-123-item1-a1/asset1.tiff")}} + ), + "assets/job-123-item1-a1/asset1.tiff": b"TIFF-DUMMY-DATA", + "docs/job-123/derived_from.json": {"hello": "world"}, + }, + ), + ], + ) + def test_custom_file_tree_structure(self, result_mocker, tmp_path, path_templates, expected): + job = result_mocker.setup_job_results( + items={ + "item1": {"assets": {"a1": {"path": "asset1.tiff"}}}, + }, + linked_docs=[ + {"rel": "derived_from", "path": "derived_from.json", "json": {"hello": "world"}}, + ], + ) + downloader = _JobResultDownloader( + job=job, + target=tmp_path, + path_templates=path_templates, + ) + downloaded = downloader.download_collection(download_derived_from=True) + + self.check_expected_downloads(downloaded=downloaded, expected=expected, tmp_path=tmp_path) From f8db2ef94f9a79e2dc9cc9958d3fdb9b03b358dc Mon Sep 17 00:00:00 2001 From: Stefaan Lippens Date: Thu, 13 Aug 2026 09:58:58 +0200 Subject: [PATCH 5/8] Add CHANGELOG entry about JobResult.download_as_collection ref #931 --- CHANGELOG.md | 1 + openeo/_version.py | 2 +- openeo/rest/job.py | 10 +++++++++- 3 files changed, 11 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index f92231a73..f97ca9012 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,6 +12,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - Added `AGENTS.md` with guidance for AI coding agents contributing to this project, including a request to disclose AI assistance in PRs ([#923](https://github.com/Open-EO/openeo-python-client/issues/923)) - Add a `py.typed` to indicate to type checkers that the package contains type annotations. - Support document based "derived_from" links in `openeo.testing.results` ([#928](https://github.com/Open-EO/openeo-python-client/issues/928)) +- Add `JobResults.download_as_collection()` (experimental) to download job results as a self-contained STAC collection with rewritten hrefs ([#931](https://github.com/Open-EO/openeo-python-client/issues/931)) ### Changed diff --git a/openeo/_version.py b/openeo/_version.py index 6204e6fa2..746f84b44 100644 --- a/openeo/_version.py +++ b/openeo/_version.py @@ -1 +1 @@ -__version__ = "0.52.0a3" +__version__ = "0.52.0a4" diff --git a/openeo/rest/job.py b/openeo/rest/job.py index 576818ebe..5a7f5783d 100644 --- a/openeo/rest/job.py +++ b/openeo/rest/job.py @@ -691,6 +691,7 @@ def download_as_collection( - linked items containing the result assets - additionally linked metadata + .. warning:: this is an experimental API, subject to change. :param target: folder path to download to :param rewrite_references: whether to rewrite (item/asset/...) references @@ -702,6 +703,8 @@ def download_as_collection( 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 :param path_templates: optional tempalte overrides for download paths. + + .. versionadded:: 0.52.0 """ downloader = _JobResultDownloader( job=self._job, @@ -725,12 +728,17 @@ 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. + .. warning:: this is an experimental API, subject to change. + + .. versionadded:: 0.52.0 """ # TODO: make this a public API that users can implement for custom download behavior (e.g. download to S3, ...) # TODO: API to warn about or skip existing/previously downloaded files? # TODO: dedicated request session (with appropriate retry strategy) for downloading? + # TODO: expose chunk_size/range_size from ResultAsset.download + # TODO: verbose mode to log/print each downloaded file to allow showing progress on large result sets + # TODO: download STAC documents (root collection, items) to file before parsing, instead of parsing in memory an re-json-encode them to file def __init__( self, From c22d189b8ff09579db77a0798527deabe0b0dce3 Mon Sep 17 00:00:00 2001 From: Stefaan Lippens Date: Thu, 13 Aug 2026 10:01:12 +0200 Subject: [PATCH 6/8] address Claude code review #931 --- openeo/rest/_testing.py | 1 - openeo/rest/job.py | 49 +++++++++++++++++++++++------------------ tests/rest/test_job.py | 42 +++++++++++++++++++++++++++++++++++ 3 files changed, 70 insertions(+), 22 deletions(-) diff --git a/openeo/rest/_testing.py b/openeo/rest/_testing.py index d5f387a34..3ef100794 100644 --- a/openeo/rest/_testing.py +++ b/openeo/rest/_testing.py @@ -7,7 +7,6 @@ Callable, Dict, Iterable, - List, Mapping, Optional, Sequence, diff --git a/openeo/rest/job.py b/openeo/rest/job.py index 5a7f5783d..dbdfca3c9 100644 --- a/openeo/rest/job.py +++ b/openeo/rest/job.py @@ -11,7 +11,7 @@ import typing import urllib.parse from pathlib import Path -from typing import Dict, Iterable, List, Optional, Union +from typing import Container, Dict, List, Literal, Optional, Union import requests @@ -34,13 +34,8 @@ from openeo.util import ensure_dir, ensure_parent_dir_for from openeo.utils.events import EVENTS 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, ) if typing.TYPE_CHECKING: @@ -393,13 +388,18 @@ class RESTJob(BatchJob): FILENAME_UNSAFE_REGEX = re.compile(r"[^\w_.-]+") -def _sanitize_filename(s: str, replacement: str = "") -> str: +def _sanitize_filename( + name: str, *, replacement: str = "", invalid: Container[str] = frozenset(("", ".", "..")) +) -> str: """ - Sanitize a filename (strip/replace risky characters) - so that it can be safely used as a filename. + Sanitize a string (strip/replace risky characters) + so that it can be safely used as file or folder name. """ - s = str(s).strip() - return FILENAME_UNSAFE_REGEX.sub(replacement, s) + sanitized = str(name).strip() + sanitized = FILENAME_UNSAFE_REGEX.sub(replacement, sanitized) + if sanitized in invalid: + raise ValueError(f"Invalid file/folder name {sanitized!r} (sanitized from {name!r})") + return sanitized def _filename_from_url(url: str, *, full: bool = False) -> str: @@ -682,27 +682,29 @@ def download_as_collection( download_derived_from: bool = False, download_collection_assets: bool = False, json_dumping: Optional[dict] = None, + on_download_failure: Literal["warn", "raise"] = "warn", path_templates: 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 + - linked STAC items: metadata and the result assets + - additionally linked metadata (e.g. "derived_from" documents) .. warning:: this is an experimental API, subject to change. :param target: folder path to download to - :param rewrite_references: whether to rewrite (item/asset/...) references + :param rewrite_references: whether to rewrite (item/asset/...) HREFs in the STAC documents. in the downloaded STAC collection to point to the local files - instead of the original URLs + 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 - :param path_templates: optional tempalte overrides for download paths. + 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. + :param on_download_failure: how to handle download failures, one of "warn" or "raise". + :param path_templates: optional template overrides for download paths. .. versionadded:: 0.52.0 """ @@ -711,6 +713,7 @@ def download_as_collection( target=target, rewrite_references=rewrite_references, json_dumping=json_dumping, + on_download_failure=on_download_failure, path_templates=path_templates, ) return downloader.download_collection( @@ -747,7 +750,7 @@ def __init__( target: Union[Path, str, None] = None, rewrite_references: bool = True, json_dumping: Optional[dict] = None, - on_download_failure: str = "warn", + on_download_failure: Literal["warn", "raise"] = "warn", path_templates: Optional[dict] = None, ): self._job = job @@ -782,14 +785,18 @@ def _download_attempt_context(self, name: str): yield except Exception as e: message = f"Failed to download {name} ({e=})" - if self._on_download_failure == "warn": + if self._on_download_failure in {"warn"}: logger.warning(message, exc_info=True) else: + # TODO: other handling strategies? + # e.g. collect all failures and raise a single exception at the end, + if self._on_download_failure != "raise": + logger.warning(f"Unknown on_download_failure strategy {self._on_download_failure!r}") raise JobResultDownloadException(message) from e def _check_download_path(self, path: Path): if path in self._downloaded: - raise JobResultDownloadException("Download collision: {path} already downloaded") + raise JobResultDownloadException(f"Download collision: {path} already downloaded") def build_path_collection(self, *, collection_id: str) -> Path: """Build path for the root STAC collection metadata file (job results metadata)""" diff --git a/tests/rest/test_job.py b/tests/rest/test_job.py index f0329e830..86aa5a705 100644 --- a/tests/rest/test_job.py +++ b/tests/rest/test_job.py @@ -1222,6 +1222,17 @@ def test_sanitize_filename(): assert _sanitize_filename("foo%bar.txt", replacement="_") == "foo_bar.txt" +def test_sanitize_filename_invalid(): + for filename in ["", " ", ".", " . ", "..", " .. "]: + with pytest.raises(ValueError): + _sanitize_filename(filename) + + # This is still fine however + assert _sanitize_filename(".config") == ".config" + # Weird, but you do you + assert _sanitize_filename("..config") == "..config" + + 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" @@ -1245,6 +1256,37 @@ def test_filename_from_url(): assert _filename_from_url("/foo/bar.txt", full=True) == "foo/bar.txt" +@pytest.mark.parametrize( + ["url", "full", "expected"], + [ + # Base cases + ("https://example.com/foo/bar.txt", False, "bar.txt"), + ("https://example.com/foo/bar.txt", True, "foo/bar.txt"), + # Period usage + ("https://example.com/foo/./bar.txt", False, "bar.txt"), + ("https://example.com/foo/./bar.txt", True, ValueError(r"Invalid file/folder name '\.'")), + ("https://example.com/foo/%2E/bar.txt", True, ValueError(r"Invalid file/folder name '\.'")), + ("https://example.com/foo/.", False, ValueError(r"Invalid file/folder name '\.'")), + # Double period usage + ("https://example.com/foo/../bar.txt", False, "bar.txt"), + ("https://example.com/foo/../bar.txt", True, ValueError(r"Invalid file/folder name '\.\.'")), + ("https://example.com/foo/%2E%2E/bar.txt", True, ValueError(r"Invalid file/folder name '\.\.'")), + ("https://example.com/foo/..", False, ValueError(r"Invalid file/folder name '\.\.'")), + # Empty path parts + ("https://example.com/foo/bar/", False, "bar"), + ("https://example.com/foo/bar/", True, "foo/bar"), + ("https://example.com/foo//bar.txt", False, "bar.txt"), + ("https://example.com/foo//bar.txt", True, "foo/bar.txt"), + ], +) +def test_test_filename_from_url_invalid_parts(url, full, expected): + if isinstance(expected, Exception): + with pytest.raises(type(expected), match=str(expected)): + _filename_from_url(url, full=full) + else: + assert _filename_from_url(url, full=full) == expected + + class TestJobResultDownloader: @pytest.fixture From b2c4400ca075ca9484b1c25b031faadfda80d3f1 Mon Sep 17 00:00:00 2001 From: Stefaan Lippens Date: Thu, 13 Aug 2026 12:50:55 +0200 Subject: [PATCH 7/8] _JobResultDownloader: improve download path tracking ref #931 --- openeo/rest/_testing.py | 9 ++++-- openeo/rest/job.py | 67 ++++++++++++++++++++++++++--------------- tests/rest/test_job.py | 27 +++++++++++++++++ 3 files changed, 76 insertions(+), 27 deletions(-) diff --git a/openeo/rest/_testing.py b/openeo/rest/_testing.py index 3ef100794..a4509f664 100644 --- a/openeo/rest/_testing.py +++ b/openeo/rest/_testing.py @@ -544,7 +544,8 @@ def setup_error(self, href, error: dict): ) 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") + path = item_data.get("full_path") or f"/j/{job_id}/r/i/{item_id}.json" + href = self.connection.build_url(path) if error := item_data.get("error"): self.setup_error(href, error=error) else: @@ -557,7 +558,8 @@ def setup_item(self, *, job_id: str, item_id: str, item_data: dict, assets: dict 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')}") + path = asset_data.get("full_path") or f"/j/{job_id}/r/a/{asset_data.get('path', 'asset.tiff')}" + href = self.connection.build_url(path) if error := asset_data.get("error"): self.requests_mock.head(href, headers={}) self.setup_error(href, error=error) @@ -571,7 +573,8 @@ def setup_asset(self, *, job_id: str, asset_data: dict) -> dict: ) def setup_linked_document(self, *, job_id: str, doc: dict): - href = self.connection.build_url(f"/j/{job_id}/r/d/{doc.get('path', 'doc.txt')}") + path = doc.get("full_path") or f"/j/{job_id}/r/d/{doc.get('path', 'doc.txt')}" + href = self.connection.build_url(path) if "json" in doc: text = json.dumps(doc["json"]) else: diff --git a/openeo/rest/job.py b/openeo/rest/job.py index dbdfca3c9..d99ab745f 100644 --- a/openeo/rest/job.py +++ b/openeo/rest/job.py @@ -695,8 +695,8 @@ def download_as_collection( .. warning:: this is an experimental API, subject to change. :param target: folder path to download to - :param rewrite_references: whether to rewrite (item/asset/...) HREFs in the STAC documents. - in the downloaded STAC collection to point to the local files + :param rewrite_references: whether to rewrite (item/asset/...) HREFs + in the downloaded STAC documents 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. @@ -726,6 +726,27 @@ class JobResultDownloadException(OpenEoClientException): pass +class _DownloadTracker: + """Simple tracker of download paths to avoid unintended collisions or double downloads.""" + + # TODO: also track the origin of a download for better error reporting? + + __slots__ = ("paths",) + + def __init__(self): + self.paths: List[Path] = [] + + def assert_new(self, path: Path): + """Check that download path is not known already, to avoid download collisions.""" + if path in self.paths: + raise JobResultDownloadException(f"Download collision, already downloaded {path}") + + def register(self, path: Path): + """Register a path as downloaded.""" + # TODO: add verification if path actually exists? + self.paths.append(path) + + class _JobResultDownloader: """ Helper class to download batch job results as a STAC collection (openEO API 1.1 style): @@ -742,6 +763,15 @@ class _JobResultDownloader: # TODO: expose chunk_size/range_size from ResultAsset.download # TODO: verbose mode to log/print each downloaded file to allow showing progress on large result sets # TODO: download STAC documents (root collection, items) to file before parsing, instead of parsing in memory an re-json-encode them to file + # TODO: give STAC collection/item docs an ".inprogress" suffix on initial write, before rewriting is done + + DEFAULT_PATH_TEMPLATES = { + "collection": "job-results.json", + "item": "{item_id}/{item_id}.json", + "asset": "{item_id}/{asset_filename}", + "collection-asset": "{asset_filename}", + "generic-link": "{filename}", + } def __init__( self, @@ -761,16 +791,9 @@ def __init__( 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] = [] + self._download_tracker = _DownloadTracker() self._on_download_failure = on_download_failure - self._path_templates = { - "collection": "job-results.json", - "item": "{item_id}/{item_id}.json", - "asset": "{item_id}/{asset_filename}", - "collection-asset": "{asset_filename}", - "generic-link": "{filename}", - **(path_templates or {}), - } + self._path_templates = {**self.DEFAULT_PATH_TEMPLATES, **(path_templates or {})} def _write_json_file(self, data: dict, path: Union[str, Path]) -> Path: path = Path(path) @@ -794,10 +817,6 @@ def _download_attempt_context(self, name: str): logger.warning(f"Unknown on_download_failure strategy {self._on_download_failure!r}") raise JobResultDownloadException(message) from e - def _check_download_path(self, path: Path): - if path in self._downloaded: - raise JobResultDownloadException(f"Download collision: {path} already downloaded") - def build_path_collection(self, *, collection_id: str) -> Path: """Build path for the root STAC collection metadata file (job results metadata)""" vars = { @@ -857,7 +876,7 @@ def download_collection( result_metadata = copy.deepcopy(result_metadata) result_metadata_path = self.build_path_collection(collection_id=result_metadata.get("id")) - self._check_download_path(result_metadata_path) + self._download_tracker.assert_new(result_metadata_path) # 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) @@ -872,9 +891,9 @@ def download_collection( elif link["rel"] in extra_rels: with self._download_attempt_context(name=f"link {link=}"): path = self.build_path_generic_link(rel=link["rel"], href=link["href"]) - self._check_download_path(path) + self._download_tracker.assert_new(path) self._connection.download_url(url=link["href"], target=path) - self._downloaded.append(path) + self._download_tracker.register(path) if self._rewrite_references: link["href"] = self._relative_to(target=path, doc=result_metadata_path) @@ -891,14 +910,14 @@ def download_collection( # 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) + self._download_tracker.register(result_metadata_path) - return self._downloaded + return self._download_tracker.paths def _download_item(self, href: str) -> Path: item: dict = self._connection.get(href, expected_status=200).json() metadata_path = self.build_path_item(item_id=item["id"]) - self._check_download_path(metadata_path) + self._download_tracker.assert_new(metadata_path) self._write_json_file(data=item, path=metadata_path) for asset_key, asset in item.get("assets", {}).items(): @@ -912,7 +931,7 @@ def _download_item(self, href: str) -> Path: if self._rewrite_references: self._write_json_file(data=item, path=metadata_path) - self._downloaded.append(metadata_path) + self._download_tracker.register(metadata_path) return metadata_path @@ -921,9 +940,9 @@ def _download_asset( ) -> Path: asset = ResultAsset(job=self._job, key=asset_key, href=asset_href, metadata=asset_metadata) path = self.build_path_asset(asset_key=asset_key, asset_href=asset_href, item_id=item_id) - self._check_download_path(path) + self._download_tracker.assert_new(path) asset.download(target=path) - self._downloaded.append(path) + self._download_tracker.register(path) return path diff --git a/tests/rest/test_job.py b/tests/rest/test_job.py index 86aa5a705..558a3ae73 100644 --- a/tests/rest/test_job.py +++ b/tests/rest/test_job.py @@ -1606,3 +1606,30 @@ def test_custom_file_tree_structure(self, result_mocker, tmp_path, path_template downloaded = downloader.download_collection(download_derived_from=True) self.check_expected_downloads(downloaded=downloaded, expected=expected, tmp_path=tmp_path) + + def test_download_collision_default(self, result_mocker, tmp_path, caplog): + """Download collisions are logged as warning by default (on_download_failure="warn")""" + job = result_mocker.setup_job_results( + items={ + "item1": {"assets": {"a": {"full_path": "data/item1/asset.tiff"}}}, + "item2": {"assets": {"a": {"full_path": "data/item2/asset.tiff"}}}, + }, + ) + downloader = _JobResultDownloader(job=job, target=tmp_path, path_templates={"asset": "assets/{asset_filename}"}) + downloader.download_collection() + assert caplog.text == dirty_equals.IsStr( + regex=r".*Download collision, already downloaded.*asset\.tiff.*", regex_flags=re.DOTALL + ) + + def test_download_collision_with_raise(self, result_mocker, tmp_path): + job = result_mocker.setup_job_results( + items={ + "item1": {"assets": {"a": {"full_path": "data/item1/asset.tiff"}}}, + "item2": {"assets": {"a": {"full_path": "data/item2/asset.tiff"}}}, + }, + ) + downloader = _JobResultDownloader( + job=job, target=tmp_path, path_templates={"asset": "assets/{asset_filename}"}, on_download_failure="raise" + ) + with pytest.raises(JobResultDownloadException, match=r"Download collision, already downloaded.*asset\.tiff"): + downloader.download_collection() From ffd378ee59a3f5bb9c1fb7557631c6b87d4663e1 Mon Sep 17 00:00:00 2001 From: Stefaan Lippens Date: Thu, 13 Aug 2026 14:03:07 +0200 Subject: [PATCH 8/8] Add basic test for JobResults.download_as_collection ref #931 --- openeo/rest/_testing.py | 18 ++++++++++++++++++ tests/rest/test_job.py | 40 +++++++++++++++++++++++++++++++++++++++- 2 files changed, 57 insertions(+), 1 deletion(-) diff --git a/openeo/rest/_testing.py b/openeo/rest/_testing.py index a4509f664..84056a904 100644 --- a/openeo/rest/_testing.py +++ b/openeo/rest/_testing.py @@ -495,6 +495,24 @@ class JobResultCollectionMocker: """ Helper to mock job result metadata (openEO 1.1 Collection style) with items and assets. + + Usage: + + - Define a fixture to create an instance with injected `requests_mock` + and `connection`. E.g.: + + @pytest.fixture + def result_collection_mocker(requests_mock, con) -> JobResultCollectionMocker: + return JobResultCollectionMocker(requests_mock=requests_mock, connection=con) + + - Call `setup_job_results` to mock the job results collection, + items, assets, ... E.g.: + + job = result_collection_mocker.setup_job_results( + items={ + "item1": {"assets": {"asset1": {"path": "asset1.tiff"}}}, + } + ) """ def __init__(self, *, requests_mock, connection: Connection): diff --git a/tests/rest/test_job.py b/tests/rest/test_job.py index 558a3ae73..61e28c667 100644 --- a/tests/rest/test_job.py +++ b/tests/rest/test_job.py @@ -1089,6 +1089,43 @@ def download_tiff(request, context): assert f.read() == TIFF_CONTENT +class TestJobResults: + # TODO: move all "job.get_results()" based tests + # inside this class for cleaner test structure + + @pytest.fixture + def result_collection_mocker(self, con100, requests_mock) -> JobResultCollectionMocker: + """helper to mock collection-style job results""" + return JobResultCollectionMocker(requests_mock=requests_mock, connection=con100) + + def test_download_as_collection_basic(self, result_collection_mocker, tmp_path): + job = result_collection_mocker.setup_job_results( + items={"item1": {"assets": {"asset1": {"path": "asset1.tiff"}}}} + ) + downloaded = job.get_results().download_as_collection(target=tmp_path) + + expected = { + "job-results.json": dirty_equals.IsPartialDict( + { + "type": "Collection", + "links": [{"rel": "item", "href": "item1/item1.json"}], + } + ), + "item1/item1.json": dirty_equals.IsPartialDict( + { + "id": "item1", + "type": "Feature", + "assets": { + "asset1": dirty_equals.IsPartialDict(href="asset1.tiff"), + }, + } + ), + "item1/asset1.tiff": b"TIFF-DUMMY-DATA", + } + + TestJobResultDownloader.check_expected_downloads(downloaded=downloaded, expected=expected, tmp_path=tmp_path) + + class TestResultAsset: @pytest.fixture def job(self, con100): @@ -1293,7 +1330,8 @@ class TestJobResultDownloader: 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): + @staticmethod + def check_expected_downloads(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