Fix: Correctly calculate starting offset for retries of ranged reads

googleapis · Aug 7, 2024 · ae4bda1 · ae4bda1
1 parent 5d2091d
commit ae4bda1
Show file tree

Hide file tree

Showing 2 changed files with 74 additions and 2 deletions.
diff --git a/google/resumable_media/requests/download.py b/google/resumable_media/requests/download.py
@@ -221,7 +221,7 @@ def retriable_request():
             # received using a range request, and set object generation query param.
             if self._bytes_downloaded > 0:
                 _download.add_bytes_range(
-                    self._bytes_downloaded, self.end, self._headers
+                    (self.start or 0) + self._bytes_downloaded, self.end, self._headers
                 )
                 request_kwargs["headers"] = self._headers
 
@@ -426,7 +426,7 @@ def retriable_request():
             # received using a range request, and set object generation query param.
             if self._bytes_downloaded > 0:
                 _download.add_bytes_range(
-                    self._bytes_downloaded, self.end, self._headers
+                    (self.start or 0) + self._bytes_downloaded, self.end, self._headers
                 )
                 request_kwargs["headers"] = self._headers
 

diff --git a/tests/unit/requests/test_download.py b/tests/unit/requests/test_download.py
@@ -455,6 +455,42 @@ def test_consume_w_bytes_downloaded(self):
         range_bytes = "bytes={:d}-{:d}".format(offset, end)
         assert download._headers["range"] == range_bytes
 
+    def test_consume_w_bytes_downloaded_range_read(self):
+        stream = io.BytesIO()
+        chunks = (b"up down ", b"charlie ", b"brown")
+        start = 1024
+        end = 65536
+
+        download = download_mod.Download(
+            EXAMPLE_URL,
+            stream=stream,
+            start=start,
+            end=end,
+            headers=None,
+            checksum="md5",
+        )
+        transport = mock.Mock(spec=["request"])
+        transport.request.return_value = _mock_response(chunks=chunks, headers=None)
+
+        assert download._bytes_downloaded == 0
+
+        # Mock a retry operation with bytes already downloaded in the stream and checksum stored
+        offset = 256
+        download._bytes_downloaded = offset
+        download._expected_checksum = None
+        download._checksum_object = _helpers._DoNothingHash()
+        download.consume(transport)
+
+        called_kwargs = {
+            "data": None,
+            "headers": download._headers,
+            "timeout": EXPECTED_TIMEOUT,
+            "stream": True,
+        }
+        transport.request.assert_called_once_with("GET", EXAMPLE_URL, **called_kwargs)
+        range_bytes = "bytes={:d}-{:d}".format(offset + start, end)
+        assert download._headers["range"] == range_bytes
+
     def test_consume_gzip_reset_stream_w_bytes_downloaded(self):
         stream = io.BytesIO()
         chunks = (b"up down ", b"charlie ", b"brown")
@@ -904,6 +940,42 @@ def test_consume_w_bytes_downloaded(self):
         range_bytes = "bytes={:d}-{:d}".format(offset, end)
         assert download._headers["range"] == range_bytes
 
+    def test_consume_w_bytes_downloaded_range_read(self):
+        stream = io.BytesIO()
+        chunks = (b"up down ", b"charlie ", b"brown")
+        start = 1024
+        end = 65536
+
+        download = download_mod.RawDownload(
+            EXAMPLE_URL,
+            stream=stream,
+            start=start,
+            end=end,
+            headers=None,
+            checksum="md5",
+        )
+        transport = mock.Mock(spec=["request"])
+        transport.request.return_value = _mock_raw_response(chunks=chunks, headers=None)
+
+        assert download._bytes_downloaded == 0
+
+        # Mock a retry operation with bytes already downloaded in the stream and checksum stored
+        offset = 256
+        download._bytes_downloaded = offset
+        download._expected_checksum = None
+        download._checksum_object = _helpers._DoNothingHash()
+        download.consume(transport)
+
+        called_kwargs = {
+            "data": None,
+            "headers": download._headers,
+            "timeout": EXPECTED_TIMEOUT,
+            "stream": True,
+        }
+        transport.request.assert_called_once_with("GET", EXAMPLE_URL, **called_kwargs)
+        range_bytes = "bytes={:d}-{:d}".format(start + offset, end)
+        assert download._headers["range"] == range_bytes
+
     def test_consume_gzip_reset_stream_w_bytes_downloaded(self):
         stream = io.BytesIO()
         chunks = (b"up down ", b"charlie ", b"brown")