Skip to content

Commit 22f36da

Browse files
authored
fix: propagate timeout in BlobWriter (#1186)
Fixes #1184
1 parent 9e4d1d8 commit 22f36da

File tree

2 files changed

+14
-4
lines changed

2 files changed

+14
-4
lines changed

google/cloud/storage/fileio.py

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -406,9 +406,15 @@ def _upload_chunks_from_buffer(self, num_chunks):
406406

407407
upload, transport = self._upload_and_transport
408408

409+
# Attach timeout if specified in the keyword arguments.
410+
# Otherwise, the default timeout will be used from the media library.
411+
kwargs = {}
412+
if "timeout" in self._upload_kwargs:
413+
kwargs = {"timeout": self._upload_kwargs.get("timeout")}
414+
409415
# Upload chunks. The SlidingBuffer class will manage seek position.
410416
for _ in range(num_chunks):
411-
upload.transmit_next_chunk(transport)
417+
upload.transmit_next_chunk(transport, **kwargs)
412418

413419
# Wipe the buffer of chunks uploaded, preserving any remaining data.
414420
self._buffer.flush()

tests/unit/test_fileio.py

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -346,6 +346,7 @@ def test_write(self, mock_warn):
346346
blob = mock.Mock()
347347
upload = mock.Mock()
348348
transport = mock.Mock()
349+
timeout = 600
349350

350351
blob._initiate_resumable_upload.return_value = (upload, transport)
351352

@@ -354,7 +355,10 @@ def test_write(self, mock_warn):
354355
# arguments are used.
355356
# It would be normal to use a context manager here, but not doing so
356357
# gives us more control over close() for test purposes.
357-
upload_kwargs = {"if_metageneration_match": 1}
358+
upload_kwargs = {
359+
"if_metageneration_match": 1,
360+
"timeout": timeout,
361+
}
358362
chunk_size = 8 # Note: Real upload requires a multiple of 256KiB.
359363
writer = self._make_blob_writer(
360364
blob,
@@ -366,7 +370,7 @@ def test_write(self, mock_warn):
366370

367371
# The transmit_next_chunk method must actually consume bytes from the
368372
# sliding buffer for the flush() feature to work properly.
369-
upload.transmit_next_chunk.side_effect = lambda _: writer._buffer.read(
373+
upload.transmit_next_chunk.side_effect = lambda _, timeout: writer._buffer.read(
370374
chunk_size
371375
)
372376

@@ -388,7 +392,7 @@ def test_write(self, mock_warn):
388392
retry=None,
389393
**upload_kwargs
390394
)
391-
upload.transmit_next_chunk.assert_called_with(transport)
395+
upload.transmit_next_chunk.assert_called_with(transport, timeout=timeout)
392396
self.assertEqual(upload.transmit_next_chunk.call_count, 4)
393397

394398
# Write another byte, finalize and close.

0 commit comments

Comments
 (0)