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)