From ccc12b78d90612b31e3674134bb0eda44daff665 Mon Sep 17 00:00:00 2001 From: abhishekacharya200 Date: Tue, 20 Oct 2026 03:45:12 +1100 Subject: [PATCH] fix(storage): honor timeout in upload_chunks_concurrently upload_chunks_concurrently accepted a timeout but never passed it to the XML multipart requests, so initiate, every part upload, finalize and cancel always ran with the (61, 60) default. Pass it through to all of them, including the _upload_part worker. Callers who don't pass a timeout now get the documented 60 seconds (_DEFAULT_TIMEOUT, as every Blob upload and download already sends) instead of (61, 60), so only the connect timeout moves by one second. Fixes #18764 --- .../google/cloud/storage/transfer_manager.py | 10 +- .../tests/unit/test_transfer_manager.py | 124 ++++++++++++++++-- 2 files changed, 119 insertions(+), 15 deletions(-) diff --git a/packages/google-cloud-storage/google/cloud/storage/transfer_manager.py b/packages/google-cloud-storage/google/cloud/storage/transfer_manager.py index b376c5138d72..0ac4c109002a 100644 --- a/packages/google-cloud-storage/google/cloud/storage/transfer_manager.py +++ b/packages/google-cloud-storage/google/cloud/storage/transfer_manager.py @@ -1198,7 +1198,7 @@ def upload_chunks_concurrently( container = XMLMPUContainer(url, filename, headers=headers, retry=retry) - container.initiate(transport=transport, content_type=content_type) + container.initiate(transport=transport, content_type=content_type, timeout=timeout) upload_id = container.upload_id size = os.path.getsize(filename) @@ -1228,6 +1228,7 @@ def upload_chunks_concurrently( checksum=checksum, headers=headers.copy(), retry=retry, + timeout=timeout, ) ) @@ -1241,9 +1242,9 @@ def upload_chunks_concurrently( part_number, etag = future.result() container.register_part(part_number, etag) - container.finalize(blob._get_transport(client)) + container.finalize(blob._get_transport(client), timeout=timeout) except Exception: - container.cancel(blob._get_transport(client)) + container.cancel(blob._get_transport(client), timeout=timeout) raise @@ -1258,6 +1259,7 @@ def _upload_part( checksum, headers, retry, + timeout, ): """Helper function that runs inside a thread or subprocess to upload a part. @@ -1280,7 +1282,7 @@ def _upload_part( headers=headers, retry=retry, ) - part.upload(client._http) + part.upload(client._http, timeout=timeout) return (part_number, part.etag) diff --git a/packages/google-cloud-storage/tests/unit/test_transfer_manager.py b/packages/google-cloud-storage/tests/unit/test_transfer_manager.py index c370f2eede2a..ec63239bfe8c 100644 --- a/packages/google-cloud-storage/tests/unit/test_transfer_manager.py +++ b/packages/google-cloud-storage/tests/unit/test_transfer_manager.py @@ -21,6 +21,7 @@ from google.api_core import exceptions from google.cloud.storage import Blob, Client, transfer_manager +from google.cloud.storage.constants import _DEFAULT_TIMEOUT from google.cloud.storage.exceptions import DataCorruption from google.cloud.storage.retry import DEFAULT_RETRY @@ -1031,13 +1032,17 @@ def test_upload_chunks_concurrently(): worker_type=transfer_manager.THREAD, ) container_mock.initiate.assert_called_once_with( - transport=transport, content_type=blob.content_type + transport=transport, + content_type=blob.content_type, + timeout=_DEFAULT_TIMEOUT, ) container_mock.register_part.assert_any_call(1, ETAG) container_mock.register_part.assert_any_call(2, ETAG) - container_mock.finalize.assert_called_once_with(bucket.client._http) + container_mock.finalize.assert_called_once_with( + bucket.client._http, timeout=_DEFAULT_TIMEOUT + ) - part_mock.upload.assert_called_with(transport) + part_mock.upload.assert_called_with(transport, timeout=_DEFAULT_TIMEOUT) def test_upload_chunks_concurrently_quotes_urls(): @@ -1079,17 +1084,21 @@ def test_upload_chunks_concurrently_quotes_urls(): ) container_mock.initiate.assert_called_once_with( - transport=transport, content_type=blob.content_type + transport=transport, + content_type=blob.content_type, + timeout=_DEFAULT_TIMEOUT, ) container_mock.register_part.assert_any_call(1, ETAG) container_mock.register_part.assert_any_call(2, ETAG) - container_mock.finalize.assert_called_once_with(bucket.client._http) + container_mock.finalize.assert_called_once_with( + bucket.client._http, timeout=_DEFAULT_TIMEOUT + ) container_cls_mock.assert_called_once_with( quoted_url, FILENAME, headers=mock.ANY, retry=DEFAULT_RETRY ) - part_mock.upload.assert_called_with(transport) + part_mock.upload.assert_called_with(transport, timeout=_DEFAULT_TIMEOUT) def test_upload_chunks_concurrently_passes_concurrency_options(): @@ -1133,12 +1142,100 @@ def test_upload_chunks_concurrently_passes_concurrency_options(): pass # The futures don't actually work, so we expect this to abort. # Conveniently, that gives us a chance to test the auto-delete # exception handling feature. - container_mock.cancel.assert_called_once_with(transport) + container_mock.cancel.assert_called_once_with( + transport, timeout=_DEFAULT_TIMEOUT + ) pool_patch.assert_called_with(max_workers=MAX_WORKERS) wait_patch.assert_called_with(mock.ANY, timeout=DEADLINE, return_when=mock.ANY) +def test_upload_chunks_concurrently_passes_timeout(): + bucket = mock.Mock() + bucket.name = "bucket" + bucket.client = _PickleableMockClient(identify_as_client=True) + transport = bucket.client._http + bucket.user_project = None + + blob = Blob("blob", bucket) + blob.content_type = FAKE_CONTENT_TYPE + + FILENAME = "file_a.txt" + SIZE = 2048 + TIMEOUT = (3, 7) + + container_mock = mock.Mock() + container_mock.upload_id = "abcd" + part_mock = mock.Mock() + part_mock.etag = "efgh" + + with ( + mock.patch("os.path.getsize", return_value=SIZE), + mock.patch( + "google.cloud.storage.transfer_manager.XMLMPUContainer", + return_value=container_mock, + ), + mock.patch( + "google.cloud.storage.transfer_manager.XMLMPUPart", return_value=part_mock + ), + ): + transfer_manager.upload_chunks_concurrently( + FILENAME, + blob, + chunk_size=SIZE // 2, + worker_type=transfer_manager.THREAD, + timeout=TIMEOUT, + ) + container_mock.initiate.assert_called_once_with( + transport=transport, content_type=blob.content_type, timeout=TIMEOUT + ) + assert part_mock.upload.call_args_list == [ + mock.call(transport, timeout=TIMEOUT), + mock.call(transport, timeout=TIMEOUT), + ] + container_mock.finalize.assert_called_once_with(transport, timeout=TIMEOUT) + + +def test_upload_chunks_concurrently_passes_timeout_to_cancel(): + bucket = mock.Mock() + bucket.name = "bucket" + bucket.client = _PickleableMockClient(identify_as_client=True) + transport = bucket.client._http + bucket.user_project = None + + blob = Blob("blob", bucket) + + FILENAME = "file_a.txt" + SIZE = 2048 + TIMEOUT = 15 + + container_mock = mock.Mock() + container_mock.upload_id = "abcd" + part_mock = mock.Mock() + part_mock.upload.side_effect = exceptions.ServiceUnavailable("part failed") + + with ( + mock.patch("os.path.getsize", return_value=SIZE), + mock.patch( + "google.cloud.storage.transfer_manager.XMLMPUContainer", + return_value=container_mock, + ), + mock.patch( + "google.cloud.storage.transfer_manager.XMLMPUPart", return_value=part_mock + ), + ): + with pytest.raises(exceptions.ServiceUnavailable): + transfer_manager.upload_chunks_concurrently( + FILENAME, + blob, + chunk_size=SIZE // 2, + worker_type=transfer_manager.THREAD, + timeout=TIMEOUT, + ) + container_mock.cancel.assert_called_once_with(transport, timeout=TIMEOUT) + container_mock.finalize.assert_not_called() + + def test_upload_chunks_concurrently_with_metadata_and_encryption(): import datetime @@ -1239,12 +1336,16 @@ def test_upload_chunks_concurrently_with_metadata_and_encryption(): URL, FILENAME, headers=expected_headers, retry=DEFAULT_RETRY ) container_mock.initiate.assert_called_once_with( - transport=transport, content_type=blob.content_type + transport=transport, + content_type=blob.content_type, + timeout=_DEFAULT_TIMEOUT, ) container_mock.register_part.assert_any_call(1, ETAG) container_mock.register_part.assert_any_call(2, ETAG) - container_mock.finalize.assert_called_once_with(transport) - part_mock.upload.assert_called_with(blob.client._http) + container_mock.finalize.assert_called_once_with( + transport, timeout=_DEFAULT_TIMEOUT + ) + part_mock.upload.assert_called_with(blob.client._http, timeout=_DEFAULT_TIMEOUT) class _PickleableMockBlob: @@ -1387,8 +1488,9 @@ def test__upload_part(): None, {"key", "value"}, retry=DEFAULT_RETRY, + timeout=(3, 7), ) - part.upload.assert_called_once() + part.upload.assert_called_once_with(mock.ANY, timeout=(3, 7)) assert result == (1, ETAG)