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 -vv

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,44 @@
 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__)
 
 # 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 +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

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"),
+        )