Merge pull request #194 from orange-kao/orange-incomplete-read-error

Re-reaise IncompleteReadError as MaybeRecoverableError
Aiven-Open · Oct 29, 2024 · 7ed9186 · 7ed9186
2 parents 47d3ed7 + 95c329e
commit 7ed9186
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 4 deletions.
diff --git a/rohmu/object_storage/azure.py b/rohmu/object_storage/azure.py
@@ -5,12 +5,13 @@
 
 from __future__ import annotations
 
-from azure.core.exceptions import HttpResponseError, ResourceExistsError
+from azure.core.exceptions import HttpResponseError, IncompleteReadError, ResourceExistsError
 from azure.storage.blob import BlobServiceClient, ContentSettings
 from rohmu.common.statsd import StatsdConfig
 from rohmu.errors import (
     FileNotFoundFromStorageError,
     InvalidConfigurationError,
+    MaybeRecoverableError,
     StorageError,
     TransferObjectStoreInitializationError,
     TransferObjectStoreMissingError,
@@ -379,6 +380,8 @@ def get_contents_to_fileobj(
             self._stream_blob(path, fileobj_to_store_to, byte_range, progress_callback)
         except azure.core.exceptions.ResourceNotFoundError as ex:
             raise FileNotFoundFromStorageError(path) from ex
+        except IncompleteReadError as ex:
+            raise MaybeRecoverableError("IncompleteReadError") from ex
 
         if progress_callback:
             progress_callback(1, 1)

diff --git a/rohmu/object_storage/google.py b/rohmu/object_storage/google.py
@@ -889,8 +889,7 @@ def next_chunk(self) -> tuple[MediaDownloadProgress, bool]:
         headers = self._headers.copy()
         chunk_start = self._num_bytes_downloaded + self._start_position
         chunk_end = chunk_start + self._chunksize - 1
-        if self._end_position < chunk_end:
-            chunk_end = self._end_position
+        chunk_end = min(self._end_position, chunk_end)
         headers["range"] = f"bytes={chunk_start}-{chunk_end}"
         resp, content = self._http.request(self._uri, "GET", headers=headers)
 

diff --git a/rohmu/object_storage/s3.py b/rohmu/object_storage/s3.py
@@ -15,6 +15,7 @@
     ConcurrentUploadError,
     FileNotFoundFromStorageError,
     InvalidConfigurationError,
+    MaybeRecoverableError,
     StorageError,
     TransferObjectStoreInitializationError,
     TransferObjectStoreMissingError,
@@ -424,7 +425,11 @@ def _read_object_to_fileobj(
         while data_read < body_length:
             read_amount = body_length - data_read
             read_amount = min(read_amount, READ_BLOCK_SIZE)
-            data = streaming_body.read(amt=read_amount)
+            try:
+                data = streaming_body.read(amt=read_amount)
+            except botocore.exceptions.IncompleteReadError as ex:
+                raise MaybeRecoverableError("botocore.exceptions.IncompleteReadError") from ex
+
             fileobj.write(data)
             data_read += len(data)
             if cb: