Introduce job result downloading with item walking and ref rewriting

.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

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
 

openeo/_version.py

@@ -1 +1 @@
-__version__ = "0.52.0a3"
+__version__ = "0.52.0a4"

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,17 +11,50 @@
 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.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,
+    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__)
 
 # Default timeouts for requests
 # 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 +296,60 @@ 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()
+            ensure_parent_dir_for(target)
+            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:
+        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)
+                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

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,114 @@ def build_capabilities(
         "links": [],
     }
     return capabilities
+
+
+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):
+        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,
+        linked_docs: Iterable[dict] = (),
+    ) -> BatchJob:
+        collection_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, item_data=item_data, assets=assets)
+            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=collection_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_error(self, href, error: dict):
+        self.requests_mock.get(
+            href,
+            status_code=error.get("status", 500),
+            text=error.get("message", "Unspecified error"),
+        )
+
+    def setup_item(self, *, job_id: str, item_id: str, item_data: dict, assets: dict) -> dict:
+        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:
+            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:
+        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)
+        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"),
+        )
+
+    def setup_linked_document(self, *, job_id: str, doc: dict):
+        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:
+            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}