fix: add a check for partial response data (#435)

* chore: update checksum mismatch error message with additional context

* trigger retry if lib detects incomplete read

* update condition and message

* move check to before checksum mismatch errors

Author: cojenco

packages/google-resumable-media/google/resumable_media/requests/download.py

@@ -44,6 +44,14 @@
 Please restart the download.
 """
 
+_RESPONSE_HEADERS_INFO = """\
+
+The X-Goog-Stored-Content-Length is {}. The X-Goog-Stored-Content-Encoding is {}.
+
+The download request read {} bytes of data.
+If the download was incomplete, please check the network connection and restart the download.
+"""
+
 
 class Download(_request_helpers.RequestsMixin, _download.Download):
     """Helper to manage downloading a resource from a Google API.
@@ -131,14 +139,32 @@ def _write_to_stream(self, response):
             and response.status_code != http.client.PARTIAL_CONTENT
         ):
             actual_checksum = _helpers.prepare_checksum_digest(checksum_object.digest())
+
             if actual_checksum != expected_checksum:
-                msg = _CHECKSUM_MISMATCH.format(
-                    self.media_url,
-                    expected_checksum,
-                    actual_checksum,
-                    checksum_type=self.checksum.upper(),
+                headers = self._get_headers(response)
+                x_goog_encoding = headers.get("x-goog-stored-content-encoding")
+                x_goog_length = headers.get("x-goog-stored-content-length")
+                content_length_msg = _RESPONSE_HEADERS_INFO.format(
+                    x_goog_length, x_goog_encoding, self._bytes_downloaded
                 )
-                raise common.DataCorruption(response, msg)
+                if (
+                    x_goog_length
+                    and self._bytes_downloaded < int(x_goog_length)
+                    and x_goog_encoding != "gzip"
+                ):
+                    # The library will attempt to trigger a retry by raising a ConnectionError, if
+                    # (a) bytes_downloaded is less than response header x-goog-stored-content-length, and
+                    # (b) the object is not gzip-compressed when stored in Cloud Storage.
+                    raise ConnectionError(content_length_msg)
+                else:
+                    msg = _CHECKSUM_MISMATCH.format(
+                        self.media_url,
+                        expected_checksum,
+                        actual_checksum,
+                        checksum_type=self.checksum.upper(),
+                    )
+                    msg += content_length_msg
+                    raise common.DataCorruption(response, msg)
 
     def consume(
         self,
@@ -321,13 +347,30 @@ def _write_to_stream(self, response):
             actual_checksum = _helpers.prepare_checksum_digest(checksum_object.digest())
 
             if actual_checksum != expected_checksum:
-                msg = _CHECKSUM_MISMATCH.format(
-                    self.media_url,
-                    expected_checksum,
-                    actual_checksum,
-                    checksum_type=self.checksum.upper(),
+                headers = self._get_headers(response)
+                x_goog_encoding = headers.get("x-goog-stored-content-encoding")
+                x_goog_length = headers.get("x-goog-stored-content-length")
+                content_length_msg = _RESPONSE_HEADERS_INFO.format(
+                    x_goog_length, x_goog_encoding, self._bytes_downloaded
                 )
-                raise common.DataCorruption(response, msg)
+                if (
+                    x_goog_length
+                    and self._bytes_downloaded < int(x_goog_length)
+                    and x_goog_encoding != "gzip"
+                ):
+                    # The library will attempt to trigger a retry by raising a ConnectionError, if
+                    # (a) bytes_downloaded is less than response header x-goog-stored-content-length, and
+                    # (b) the object is not gzip-compressed when stored in Cloud Storage.
+                    raise ConnectionError(content_length_msg)
+                else:
+                    msg = _CHECKSUM_MISMATCH.format(
+                        self.media_url,
+                        expected_checksum,
+                        actual_checksum,
+                        checksum_type=self.checksum.upper(),
+                    )
+                    msg += content_length_msg
+                    raise common.DataCorruption(response, msg)
 
     def consume(
         self,

packages/google-resumable-media/tests/system/requests/test_download.py

@@ -460,7 +460,7 @@ def test_corrupt_download(self, add_files, corrupting_transport, checksum):
                 info[checksum],
                 checksum_type=checksum.upper(),
             )
-            assert exc_info.value.args == (msg,)
+            assert msg in exc_info.value.args[0]
 
     def test_corrupt_download_no_check(self, add_files, corrupting_transport):
         for info in ALL_FILES:

packages/google-resumable-media/tests/unit/requests/test_download.py

@@ -105,7 +105,11 @@ def test__write_to_stream_with_hash_check_fail(self, checksum):
         msg = download_mod._CHECKSUM_MISMATCH.format(
             EXAMPLE_URL, bad_checksum, good_checksum, checksum_type=checksum.upper()
         )
-        assert error.args[0] == msg
+        assert msg in error.args[0]
+        assert (
+            f"The download request read {download._bytes_downloaded} bytes of data."
+            in error.args[0]
+        )
 
         # Check mocks.
         response.__enter__.assert_called_once_with()
@@ -167,6 +171,29 @@ def test__write_to_stream_with_invalid_checksum_type(self):
         error = exc_info.value
         assert error.args[0] == "checksum must be ``'md5'``, ``'crc32c'`` or ``None``"
 
+    @pytest.mark.parametrize("checksum", ["md5", "crc32c"])
+    def test__write_to_stream_incomplete_read(self, checksum):
+        stream = io.BytesIO()
+        download = download_mod.Download(EXAMPLE_URL, stream=stream, checksum=checksum)
+
+        chunk1 = b"first chunk"
+        mock_full_content_length = len(chunk1) + 123
+        headers = {"x-goog-stored-content-length": mock_full_content_length}
+        bad_checksum = "d3JvbmcgbiBtYWRlIHVwIQ=="
+        header_value = "crc32c={bad},md5={bad}".format(bad=bad_checksum)
+        headers[_helpers._HASH_HEADER] = header_value
+        response = _mock_response(chunks=[chunk1], headers=headers)
+
+        with pytest.raises(ConnectionError) as exc_info:
+            download._write_to_stream(response)
+
+        assert not download.finished
+        error = exc_info.value
+        assert (
+            f"The download request read {download._bytes_downloaded} bytes of data."
+            in error.args[0]
+        )
+
     def _consume_helper(
         self,
         stream=None,
@@ -285,7 +312,11 @@ def test_consume_with_stream_hash_check_fail(self, checksum):
         msg = download_mod._CHECKSUM_MISMATCH.format(
             EXAMPLE_URL, bad_checksum, good_checksum, checksum_type=checksum.upper()
         )
-        assert error.args[0] == msg
+        assert msg in error.args[0]
+        assert (
+            f"The download request read {download._bytes_downloaded} bytes of data."
+            in error.args[0]
+        )
 
         # Check mocks.
         transport.request.assert_called_once_with(
@@ -544,7 +575,11 @@ def test__write_to_stream_with_hash_check_fail(self, checksum):
         msg = download_mod._CHECKSUM_MISMATCH.format(
             EXAMPLE_URL, bad_checksum, good_checksum, checksum_type=checksum.upper()
         )
-        assert error.args[0] == msg
+        assert msg in error.args[0]
+        assert (
+            f"The download request read {download._bytes_downloaded} bytes of data."
+            in error.args[0]
+        )
 
         # Check mocks.
         response.__enter__.assert_called_once_with()
@@ -577,6 +612,31 @@ def test__write_to_stream_with_invalid_checksum_type(self):
         error = exc_info.value
         assert error.args[0] == "checksum must be ``'md5'``, ``'crc32c'`` or ``None``"
 
+    @pytest.mark.parametrize("checksum", ["md5", "crc32c"])
+    def test__write_to_stream_incomplete_read(self, checksum):
+        stream = io.BytesIO()
+        download = download_mod.RawDownload(
+            EXAMPLE_URL, stream=stream, checksum=checksum
+        )
+
+        chunk1 = b"first chunk"
+        mock_full_content_length = len(chunk1) + 123
+        headers = {"x-goog-stored-content-length": mock_full_content_length}
+        bad_checksum = "d3JvbmcgbiBtYWRlIHVwIQ=="
+        header_value = "crc32c={bad},md5={bad}".format(bad=bad_checksum)
+        headers[_helpers._HASH_HEADER] = header_value
+        response = _mock_raw_response(chunks=[chunk1], headers=headers)
+
+        with pytest.raises(ConnectionError) as exc_info:
+            download._write_to_stream(response)
+
+        assert not download.finished
+        error = exc_info.value
+        assert (
+            f"The download request read {download._bytes_downloaded} bytes of data."
+            in error.args[0]
+        )
+
     def _consume_helper(
         self,
         stream=None,
@@ -699,7 +759,11 @@ def test_consume_with_stream_hash_check_fail(self, checksum):
         msg = download_mod._CHECKSUM_MISMATCH.format(
             EXAMPLE_URL, bad_checksum, good_checksum, checksum_type=checksum.upper()
         )
-        assert error.args[0] == msg
+        assert msg in error.args[0]
+        assert (
+            f"The download request read {download._bytes_downloaded} bytes of data."
+            in error.args[0]
+        )
 
         # Check mocks.
         transport.request.assert_called_once_with(

packages/google-resumable-media/tests_async/unit/requests/test_download.py

@@ -98,7 +98,7 @@ async def test__write_to_stream_with_hash_check_fail(self, checksum):
             good_checksum,
             checksum_type=checksum.upper(),
         )
-        assert error.args[0] == msg
+        assert msg in error.args[0]
 
     @pytest.mark.asyncio
     @pytest.mark.parametrize("checksum", ["md5", "crc32c"])
@@ -263,7 +263,7 @@ async def test_consume_with_stream_hash_check_fail(self, checksum):
             good_checksum,
             checksum_type=checksum.upper(),
         )
-        assert error.args[0] == msg
+        assert msg in error.args[0]
 
         # Check mocks.
         transport.request.assert_called_once_with(
@@ -353,7 +353,7 @@ async def test__write_to_stream_with_hash_check_fail(self, checksum):
             good_checksum,
             checksum_type=checksum.upper(),
         )
-        assert error.args[0] == msg
+        assert msg in error.args[0]
 
     @pytest.mark.asyncio
     async def test__write_to_stream_with_invalid_checksum_type(self):
@@ -482,7 +482,7 @@ async def test_consume_with_stream_hash_check_fail(self, checksum):
             good_checksum,
             checksum_type=checksum.upper(),
         )
-        assert error.args[0] == msg
+        assert msg in error.args[0]
 
         # Check mocks.
         transport.request.assert_called_once_with(