From b7b22717ff44e0fdf70b5ebccb7c4fde0a83723d Mon Sep 17 00:00:00 2001 From: Cathy Ouyang Date: Fri, 17 Nov 2023 14:24:55 -0800 Subject: [PATCH 1/2] fix: propagate timeout in BlobWriter --- google/cloud/storage/fileio.py | 8 +++++++- tests/unit/test_fileio.py | 10 +++++++--- 2 files changed, 14 insertions(+), 4 deletions(-) diff --git a/google/cloud/storage/fileio.py b/google/cloud/storage/fileio.py index 639e731ba..257e455cd 100644 --- a/google/cloud/storage/fileio.py +++ b/google/cloud/storage/fileio.py @@ -406,9 +406,15 @@ def _upload_chunks_from_buffer(self, num_chunks): upload, transport = self._upload_and_transport + # Attach timeout if specified in the keyword arguments. + # Otherwise, the default timeout will be used from the media library. + kwargs = {} + if self._upload_kwargs.get("timeout"): + kwargs = {"timeout": self._upload_kwargs.get("timeout")} + # Upload chunks. The SlidingBuffer class will manage seek position. for _ in range(num_chunks): - upload.transmit_next_chunk(transport) + upload.transmit_next_chunk(transport, **kwargs) # Wipe the buffer of chunks uploaded, preserving any remaining data. self._buffer.flush() diff --git a/tests/unit/test_fileio.py b/tests/unit/test_fileio.py index dab272a90..d3ffd3725 100644 --- a/tests/unit/test_fileio.py +++ b/tests/unit/test_fileio.py @@ -346,6 +346,7 @@ def test_write(self, mock_warn): blob = mock.Mock() upload = mock.Mock() transport = mock.Mock() + timeout = 600 blob._initiate_resumable_upload.return_value = (upload, transport) @@ -354,7 +355,10 @@ def test_write(self, mock_warn): # arguments are used. # It would be normal to use a context manager here, but not doing so # gives us more control over close() for test purposes. - upload_kwargs = {"if_metageneration_match": 1} + upload_kwargs = { + "if_metageneration_match": 1, + "timeout": timeout, + } chunk_size = 8 # Note: Real upload requires a multiple of 256KiB. writer = self._make_blob_writer( blob, @@ -366,7 +370,7 @@ def test_write(self, mock_warn): # The transmit_next_chunk method must actually consume bytes from the # sliding buffer for the flush() feature to work properly. - upload.transmit_next_chunk.side_effect = lambda _: writer._buffer.read( + upload.transmit_next_chunk.side_effect = lambda _, timeout: writer._buffer.read( chunk_size ) @@ -388,7 +392,7 @@ def test_write(self, mock_warn): retry=None, **upload_kwargs ) - upload.transmit_next_chunk.assert_called_with(transport) + upload.transmit_next_chunk.assert_called_with(transport, timeout=timeout) self.assertEqual(upload.transmit_next_chunk.call_count, 4) # Write another byte, finalize and close. From 580ee137e3dc725320443de11045b6e1ddbea5cd Mon Sep 17 00:00:00 2001 From: Cathy Ouyang Date: Mon, 27 Nov 2023 13:49:44 -0800 Subject: [PATCH 2/2] append timeout if explicitly specified --- google/cloud/storage/fileio.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/google/cloud/storage/fileio.py b/google/cloud/storage/fileio.py index 257e455cd..97d234983 100644 --- a/google/cloud/storage/fileio.py +++ b/google/cloud/storage/fileio.py @@ -409,7 +409,7 @@ def _upload_chunks_from_buffer(self, num_chunks): # Attach timeout if specified in the keyword arguments. # Otherwise, the default timeout will be used from the media library. kwargs = {} - if self._upload_kwargs.get("timeout"): + if "timeout" in self._upload_kwargs: kwargs = {"timeout": self._upload_kwargs.get("timeout")} # Upload chunks. The SlidingBuffer class will manage seek position.