Issue #747 fixed chunking

Author: dsamaey

openeo/rest/__init__.py

@@ -4,7 +4,7 @@
 
 # TODO: get from config file
 DEFAULT_DOWNLOAD_CHUNK_SIZE = 10_000_000  # 10MB
-
+DEFAULT_DOWNLOAD_RANGE_SIZE = 500_000_000 # 500MB
 
 DEFAULT_JOB_STATUS_POLL_INTERVAL_MAX = 60
 DEFAULT_JOB_STATUS_POLL_CONNECTION_RETRY_INTERVAL = 30

openeo/rest/job.py

@@ -16,6 +16,7 @@
 from openeo.internal.warnings import deprecated, legacy_alias
 from openeo.rest import (
     DEFAULT_DOWNLOAD_CHUNK_SIZE,
+    DEFAULT_DOWNLOAD_RANGE_SIZE,
     DEFAULT_JOB_STATUS_POLL_CONNECTION_RETRY_INTERVAL,
     DEFAULT_JOB_STATUS_POLL_INTERVAL_MAX,
     DEFAULT_JOB_STATUS_POLL_SOFT_ERROR_MAX,
@@ -36,7 +37,7 @@
 
 
 DEFAULT_JOB_RESULTS_FILENAME = "job-results.json"
-MAX_RETRIES_PER_CHUNK = 3
+MAX_RETRIES_PER_RANGE = 3
 RETRIABLE_STATUSCODES = [408, 429, 500, 501, 502, 503, 504]
 
 class BatchJob:
@@ -384,7 +385,7 @@ def __repr__(self):
         )
 
     def download(
-        self, target: Optional[Union[Path, str]] = None, *, chunk_size: int = DEFAULT_DOWNLOAD_CHUNK_SIZE
+        self, target: Optional[Union[Path, str]] = None, *, chunk_size: int = DEFAULT_DOWNLOAD_CHUNK_SIZE, range_size: int=DEFAULT_DOWNLOAD_RANGE_SIZE
     ) -> Path:
         """
         Download asset to given location
@@ -399,7 +400,7 @@ def download(
             target = target / self.name
         ensure_dir(target.parent)
         logger.info("Downloading Job result asset {n!r} from {h!s} to {t!s}".format(n=self.name, h=self.href, t=target))
-        self._download_to_file(url=self.href, target=target, chunk_size=chunk_size)
+        self._download_to_file(url=self.href, target=target, chunk_size=chunk_size, range_size=range_size)
         return target
 
     def _get_response(self, stream=True) -> requests.Response:
@@ -418,26 +419,26 @@ 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):
+    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":
+        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_chunked(url=url, target=target, file_size=file_size, chunk_size=chunk_size)
+            self._download_ranged(url=url, target=target, file_size=file_size, chunk_size=chunk_size, range_size=range_size)
         else:
-            self._download_unchunked(url=url, target=target)
+            self._download_all_at_once(url=url, target=target, chunk_size=chunk_size)
 
 
-    def _download_chunked(self, url: str, target: Path, file_size: int, chunk_size: int):
+    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, chunk_size):
-                to_byte_index = min(from_byte_index + chunk_size - 1, file_size - 1)
-                tries_left = MAX_RETRIES_PER_CHUNK
+            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()
-                            shutil.copyfileobj(r.raw, f)
+                            shutil.copyfileobj(fsrc=r.raw, fdst=f, length=chunk_size)
                         break
                     except OpenEoApiPlainError as error:
                         tries_left -= 1
@@ -448,11 +449,11 @@ def _download_chunked(self, url: str, target: Path, file_size: int, chunk_size:
                             raise error
 
 
-    def _download_unchunked(self, url: str, target: Path):
+    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:
-                shutil.copyfileobj(r.raw, f)
+                shutil.copyfileobj(fsrc=r.raw, fdst=f, length=chunk_size)
 
 
 class MultipleAssetException(OpenEoClientException):
@@ -532,7 +533,7 @@ def get_asset(self, name: str = None) -> ResultAsset:
                     "No asset {n!r} in: {a}".format(n=name, a=[a.name for a in assets])
                 )
 
-    def download_file(self, target: Union[Path, str] = None, name: str = None, chunk_size=DEFAULT_DOWNLOAD_CHUNK_SIZE) -> Path:
+    def download_file(self, target: Union[Path, str] = None, name: str = None, *, chunk_size=DEFAULT_DOWNLOAD_CHUNK_SIZE, range_size: int=DEFAULT_DOWNLOAD_RANGE_SIZE) -> Path:
         """
         Download single asset. Can be used when there is only one asset in the
         :py:class:`JobResults`, or when the desired asset name is given explicitly.
@@ -544,7 +545,7 @@ def download_file(self, target: Union[Path, str] = None, name: str = None, chunk
         :return: path of downloaded asset
         """
         try:
-            return self.get_asset(name=name).download(target=target, chunk_size=chunk_size)
+            return self.get_asset(name=name).download(target=target, chunk_size=chunk_size, range_size=range_size)
         except MultipleAssetException:
             raise OpenEoClientException(
                 "Can not use `download_file` with multiple assets. Use `download_files` instead.")

tests/rest/test_job.py

@@ -600,10 +600,10 @@ def test_get_results_download_file(job_with_1_asset: BatchJob, tmp_path):
     with target.open("rb") as f:
         assert f.read() == TIFF_CONTENT
 
-def test_get_results_download_chunked_file_using_head(job_with_chunked_asset_using_head: BatchJob, tmp_path):
+def test_get_results_download_file_ranged(job_with_chunked_asset_using_head: BatchJob, tmp_path):
     job = job_with_chunked_asset_using_head
     target = tmp_path / "result.tiff"
-    res = job.get_results().download_file(target, chunk_size=1000)
+    res = job.get_results().download_file(target, range_size=1000)
     assert res == target
     with target.open("rb") as f:
         assert f.read() == TIFF_CONTENT