diff --git a/sdk/storage/azure-storage-blob/azure/storage/blob/_blob_client.py b/sdk/storage/azure-storage-blob/azure/storage/blob/_blob_client.py index 311c8b727348..5263ea854bab 100644 --- a/sdk/storage/azure-storage-blob/azure/storage/blob/_blob_client.py +++ b/sdk/storage/azure-storage-blob/azure/storage/blob/_blob_client.py @@ -94,7 +94,7 @@ class BlobClient(StorageAccountHostsMixin): # pylint: disable=too-many-public-m The hostname of the secondary endpoint. :keyword int max_block_size: The maximum chunk size for uploading a block blob in chunks. Defaults to 4*1024*1024, or 4MB. - :keyword int max_single_put_size: If the blob size is less than max_single_put_size, then the blob will be + :keyword int max_single_put_size: If the blob size is less than or equal max_single_put_size, then the blob will be uploaded with only one http PUT request. If the blob size is larger than max_single_put_size, the blob will be uploaded in chunks. Defaults to 64*1024*1024, or 64MB. :keyword int min_large_block_upload_threshold: The minimum chunk size required to use the memory efficient diff --git a/sdk/storage/azure-storage-blob/azure/storage/blob/_blob_service_client.py b/sdk/storage/azure-storage-blob/azure/storage/blob/_blob_service_client.py index 2eb6eeddd478..7a97c82fcdeb 100644 --- a/sdk/storage/azure-storage-blob/azure/storage/blob/_blob_service_client.py +++ b/sdk/storage/azure-storage-blob/azure/storage/blob/_blob_service_client.py @@ -77,7 +77,7 @@ class BlobServiceClient(StorageAccountHostsMixin): The hostname of the secondary endpoint. :keyword int max_block_size: The maximum chunk size for uploading a block blob in chunks. Defaults to 4*1024*1024, or 4MB. - :keyword int max_single_put_size: If the blob size is less than max_single_put_size, then the blob will be + :keyword int max_single_put_size: If the blob size is less than or equal max_single_put_size, then the blob will be uploaded with only one http PUT request. If the blob size is larger than max_single_put_size, the blob will be uploaded in chunks. Defaults to 64*1024*1024, or 64MB. :keyword int min_large_block_upload_threshold: The minimum chunk size required to use the memory efficient diff --git a/sdk/storage/azure-storage-blob/azure/storage/blob/_container_client.py b/sdk/storage/azure-storage-blob/azure/storage/blob/_container_client.py index 9521b9839766..ed893793aa24 100644 --- a/sdk/storage/azure-storage-blob/azure/storage/blob/_container_client.py +++ b/sdk/storage/azure-storage-blob/azure/storage/blob/_container_client.py @@ -97,7 +97,7 @@ class ContainerClient(StorageAccountHostsMixin): The hostname of the secondary endpoint. :keyword int max_block_size: The maximum chunk size for uploading a block blob in chunks. Defaults to 4*1024*1024, or 4MB. - :keyword int max_single_put_size: If the blob size is less than max_single_put_size, then the blob will be + :keyword int max_single_put_size: If the blob size is less than or equal max_single_put_size, then the blob will be uploaded with only one http PUT request. If the blob size is larger than max_single_put_size, the blob will be uploaded in chunks. Defaults to 64*1024*1024, or 64MB. :keyword int min_large_block_upload_threshold: The minimum chunk size required to use the memory efficient diff --git a/sdk/storage/azure-storage-blob/azure/storage/blob/_generated/version.py b/sdk/storage/azure-storage-blob/azure/storage/blob/_generated/version.py index 0b20a9ed22a5..be045899fa00 100644 --- a/sdk/storage/azure-storage-blob/azure/storage/blob/_generated/version.py +++ b/sdk/storage/azure-storage-blob/azure/storage/blob/_generated/version.py @@ -9,5 +9,5 @@ # regenerated. # -------------------------------------------------------------------------- -VERSION = "2019-10-10" +VERSION = "2019-12-12" diff --git a/sdk/storage/azure-storage-blob/azure/storage/blob/_shared/base_client.py b/sdk/storage/azure-storage-blob/azure/storage/blob/_shared/base_client.py index f566f4ba1291..e91228fabfd1 100644 --- a/sdk/storage/azure-storage-blob/azure/storage/blob/_shared/base_client.py +++ b/sdk/storage/azure-storage-blob/azure/storage/blob/_shared/base_client.py @@ -246,6 +246,8 @@ def _create_pipeline(self, credential, **kwargs): DistributedTracingPolicy(**kwargs), HttpLoggingPolicy(**kwargs) ] + if kwargs.get("_additional_pipeline_policies"): + policies = policies + kwargs.get("_additional_pipeline_policies") return config, Pipeline(config.transport, policies=policies) def _batch_send( diff --git a/sdk/storage/azure-storage-blob/azure/storage/blob/_shared/base_client_async.py b/sdk/storage/azure-storage-blob/azure/storage/blob/_shared/base_client_async.py index 3c806d796864..1b13658c9824 100644 --- a/sdk/storage/azure-storage-blob/azure/storage/blob/_shared/base_client_async.py +++ b/sdk/storage/azure-storage-blob/azure/storage/blob/_shared/base_client_async.py @@ -102,6 +102,8 @@ def _create_pipeline(self, credential, **kwargs): DistributedTracingPolicy(**kwargs), HttpLoggingPolicy(**kwargs), ] + if kwargs.get("_additional_pipeline_policies"): + policies = policies + kwargs.get("_additional_pipeline_policies") return config, AsyncPipeline(config.transport, policies=policies) async def _batch_send( diff --git a/sdk/storage/azure-storage-blob/azure/storage/blob/_shared/constants.py b/sdk/storage/azure-storage-blob/azure/storage/blob/_shared/constants.py index 7fb05b559850..f67ea29cc137 100644 --- a/sdk/storage/azure-storage-blob/azure/storage/blob/_shared/constants.py +++ b/sdk/storage/azure-storage-blob/azure/storage/blob/_shared/constants.py @@ -17,9 +17,10 @@ # for python 3.5+, there was a change to the definition of the socket timeout (as far as socket.sendall is concerned) # The socket timeout is now the maximum total duration to send all data. if sys.version_info >= (3, 5): - # the timeout to connect is 20 seconds, and the read timeout is 2000 seconds - # the 2000 seconds was calculated with: 100MB (max block size)/ 50KB/s (an arbitrarily chosen minimum upload speed) - READ_TIMEOUT = 2000 + # the timeout to connect is 20 seconds, and the read timeout is 80000 seconds + # the 80000 seconds was calculated with: + # 4000MB (max block size)/ 50KB/s (an arbitrarily chosen minimum upload speed) + READ_TIMEOUT = 80000 STORAGE_OAUTH_SCOPE = "https://storage.azure.com/.default" diff --git a/sdk/storage/azure-storage-blob/azure/storage/blob/_upload_helpers.py b/sdk/storage/azure-storage-blob/azure/storage/blob/_upload_helpers.py index 85d40698864f..d3862eb72646 100644 --- a/sdk/storage/azure-storage-blob/azure/storage/blob/_upload_helpers.py +++ b/sdk/storage/azure-storage-blob/azure/storage/blob/_upload_helpers.py @@ -80,8 +80,8 @@ def upload_block_blob( # pylint: disable=too-many-locals blob_headers = kwargs.pop('blob_headers', None) tier = kwargs.pop('standard_blob_tier', None) - # Do single put if the size is smaller than config.max_single_put_size - if adjusted_count is not None and (adjusted_count < blob_settings.max_single_put_size): + # Do single put if the size is smaller than or equal config.max_single_put_size + if adjusted_count is not None and (adjusted_count <= blob_settings.max_single_put_size): try: data = data.read(length) if not isinstance(data, six.binary_type): diff --git a/sdk/storage/azure-storage-blob/azure/storage/blob/aio/_blob_client_async.py b/sdk/storage/azure-storage-blob/azure/storage/blob/aio/_blob_client_async.py index 5861ca03fd97..2b626f16e85d 100644 --- a/sdk/storage/azure-storage-blob/azure/storage/blob/aio/_blob_client_async.py +++ b/sdk/storage/azure-storage-blob/azure/storage/blob/aio/_blob_client_async.py @@ -74,7 +74,7 @@ class BlobClient(AsyncStorageAccountHostsMixin, BlobClientBase): # pylint: disa The hostname of the secondary endpoint. :keyword int max_block_size: The maximum chunk size for uploading a block blob in chunks. Defaults to 4*1024*1024, or 4MB. - :keyword int max_single_put_size: If the blob size is less than max_single_put_size, then the blob will be + :keyword int max_single_put_size: If the blob size is less than or equal max_single_put_size, then the blob will be uploaded with only one http PUT request. If the blob size is larger than max_single_put_size, the blob will be uploaded in chunks. Defaults to 64*1024*1024, or 64MB. :keyword int min_large_block_upload_threshold: The minimum chunk size required to use the memory efficient diff --git a/sdk/storage/azure-storage-blob/azure/storage/blob/aio/_blob_service_client_async.py b/sdk/storage/azure-storage-blob/azure/storage/blob/aio/_blob_service_client_async.py index 7a834c6d1915..4d45dc336e81 100644 --- a/sdk/storage/azure-storage-blob/azure/storage/blob/aio/_blob_service_client_async.py +++ b/sdk/storage/azure-storage-blob/azure/storage/blob/aio/_blob_service_client_async.py @@ -76,7 +76,7 @@ class BlobServiceClient(AsyncStorageAccountHostsMixin, BlobServiceClientBase): The hostname of the secondary endpoint. :keyword int max_block_size: The maximum chunk size for uploading a block blob in chunks. Defaults to 4*1024*1024, or 4MB. - :keyword int max_single_put_size: If the blob size is less than max_single_put_size, then the blob will be + :keyword int max_single_put_size: If the blob size is less than or equal max_single_put_size, then the blob will be uploaded with only one http PUT request. If the blob size is larger than max_single_put_size, the blob will be uploaded in chunks. Defaults to 64*1024*1024, or 64MB. :keyword int min_large_block_upload_threshold: The minimum chunk size required to use the memory efficient diff --git a/sdk/storage/azure-storage-blob/azure/storage/blob/aio/_container_client_async.py b/sdk/storage/azure-storage-blob/azure/storage/blob/aio/_container_client_async.py index a0133468b2a0..c881e68bd286 100644 --- a/sdk/storage/azure-storage-blob/azure/storage/blob/aio/_container_client_async.py +++ b/sdk/storage/azure-storage-blob/azure/storage/blob/aio/_container_client_async.py @@ -79,7 +79,7 @@ class ContainerClient(AsyncStorageAccountHostsMixin, ContainerClientBase): The hostname of the secondary endpoint. :keyword int max_block_size: The maximum chunk size for uploading a block blob in chunks. Defaults to 4*1024*1024, or 4MB. - :keyword int max_single_put_size: If the blob size is less than max_single_put_size, then the blob will be + :keyword int max_single_put_size: If the blob size is less than or equal max_single_put_size, then the blob will be uploaded with only one http PUT request. If the blob size is larger than max_single_put_size, the blob will be uploaded in chunks. Defaults to 64*1024*1024, or 64MB. :keyword int min_large_block_upload_threshold: The minimum chunk size required to use the memory efficient diff --git a/sdk/storage/azure-storage-blob/azure/storage/blob/aio/_upload_helpers.py b/sdk/storage/azure-storage-blob/azure/storage/blob/aio/_upload_helpers.py index b936ee6cfc59..7842e03f9d50 100644 --- a/sdk/storage/azure-storage-blob/azure/storage/blob/aio/_upload_helpers.py +++ b/sdk/storage/azure-storage-blob/azure/storage/blob/aio/_upload_helpers.py @@ -56,7 +56,7 @@ async def upload_block_blob( # pylint: disable=too-many-locals tier = kwargs.pop('standard_blob_tier', None) # Do single put if the size is smaller than config.max_single_put_size - if adjusted_count is not None and (adjusted_count < blob_settings.max_single_put_size): + if adjusted_count is not None and (adjusted_count <= blob_settings.max_single_put_size): try: data = data.read(length) if not isinstance(data, six.binary_type): diff --git a/sdk/storage/azure-storage-blob/tests/test_largest_block_blob.py b/sdk/storage/azure-storage-blob/tests/test_largest_block_blob.py new file mode 100644 index 000000000000..c409f29c4a7d --- /dev/null +++ b/sdk/storage/azure-storage-blob/tests/test_largest_block_blob.py @@ -0,0 +1,348 @@ +# coding: utf-8 + +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for +# license information. +# -------------------------------------------------------------------------- + +import pytest + +from os import path, remove, sys, urandom +import platform +import uuid + +from azure.core.pipeline.policies import HTTPPolicy +from azure.storage.blob import ( + BlobServiceClient, + BlobBlock +) +from azure.storage.blob._shared.base_client import format_shared_key_credential + +from _shared.testcase import StorageTestCase, GlobalStorageAccountPreparer + +# ------------------------------------------------------------------------------ +TEST_BLOB_PREFIX = 'largestblob' +LARGEST_BLOCK_SIZE = 4000 * 1024 * 1024 +LARGEST_SINGLE_UPLOAD_SIZE = 5000 * 1024 * 1024 + +# ------------------------------------------------------------------------------ +if platform.python_implementation() == 'PyPy': + pytest.skip("Skip tests for Pypy", allow_module_level=True) + +class StorageLargestBlockBlobTest(StorageTestCase): + def _setup(self, storage_account, key, additional_policies=None, min_large_block_upload_threshold=1 * 1024 * 1024, + max_single_put_size=32 * 1024): + self.bsc = BlobServiceClient( + self.account_url(storage_account, "blob"), + credential=key, + max_single_put_size=max_single_put_size, + max_block_size=LARGEST_BLOCK_SIZE, + min_large_block_upload_threshold=min_large_block_upload_threshold, + _additional_pipeline_policies=additional_policies) + self.config = self.bsc._config + self.container_name = self.get_resource_name('utcontainer') + self.container_name = self.container_name + str(uuid.uuid4()) + + if self.is_live: + self.bsc.create_container(self.container_name) + + def _teardown(self, file_name): + if path.isfile(file_name): + try: + remove(file_name) + except: + pass + + # --Helpers----------------------------------------------------------------- + def _get_blob_reference(self): + return self.get_resource_name(TEST_BLOB_PREFIX) + + def _create_blob(self): + blob_name = self._get_blob_reference() + blob = self.bsc.get_blob_client(self.container_name, blob_name) + blob.upload_blob(b'') + return blob + + # --Test cases for block blobs -------------------------------------------- + @pytest.mark.live_test_only + @pytest.mark.skip(reason="This takes really long time") + @GlobalStorageAccountPreparer() + def test_put_block_bytes_largest(self, resource_group, location, storage_account, storage_account_key): + self._setup(storage_account, storage_account_key) + blob = self._create_blob() + + # Act + data = urandom(LARGEST_BLOCK_SIZE) + blockId = str(uuid.uuid4()).encode('utf-8') + resp = blob.stage_block( + blockId, + data, + length=LARGEST_BLOCK_SIZE) + blob.commit_block_list([BlobBlock(blockId)]) + block_list = blob.get_block_list() + + # Assert + self.assertIsNotNone(resp) + assert 'content_md5' in resp + assert 'content_crc64' in resp + assert 'request_id' in resp + self.assertIsNotNone(block_list) + self.assertEqual(len(block_list), 2) + self.assertEqual(len(block_list[1]), 0) + self.assertEqual(len(block_list[0]), 1) + self.assertEqual(block_list[0][0].size, LARGEST_BLOCK_SIZE) + + @pytest.mark.live_test_only + @GlobalStorageAccountPreparer() + def test_put_block_bytes_largest_without_network(self, resource_group, location, storage_account, storage_account_key): + payload_dropping_policy = PayloadDroppingPolicy() + credential_policy = format_shared_key_credential([storage_account.name, "dummy"], storage_account_key) + self._setup(storage_account, storage_account_key, [payload_dropping_policy, credential_policy]) + blob = self._create_blob() + + # Act + data = urandom(LARGEST_BLOCK_SIZE) + blockId = str(uuid.uuid4()).encode('utf-8') + resp = blob.stage_block( + blockId, + data, + length=LARGEST_BLOCK_SIZE) + blob.commit_block_list([BlobBlock(blockId)]) + block_list = blob.get_block_list() + + # Assert + self.assertIsNotNone(resp) + assert 'content_md5' in resp + assert 'content_crc64' in resp + assert 'request_id' in resp + self.assertIsNotNone(block_list) + self.assertEqual(len(block_list), 2) + self.assertEqual(len(block_list[1]), 0) + self.assertEqual(len(block_list[0]), 1) + self.assertEqual(payload_dropping_policy.put_block_counter, 1) + self.assertEqual(payload_dropping_policy.put_block_sizes[0], LARGEST_BLOCK_SIZE) + + @pytest.mark.live_test_only + @pytest.mark.skip(reason="This takes really long time") + @GlobalStorageAccountPreparer() + def test_put_block_stream_largest(self, resource_group, location, storage_account, storage_account_key): + self._setup(storage_account, storage_account_key) + blob = self._create_blob() + + # Act + stream = LargeStream(LARGEST_BLOCK_SIZE) + blockId = str(uuid.uuid4()) + requestId = str(uuid.uuid4()) + resp = blob.stage_block( + blockId, + stream, + length=LARGEST_BLOCK_SIZE, + client_request_id=requestId) + blob.commit_block_list([BlobBlock(blockId)]) + block_list = blob.get_block_list() + + # Assert + self.assertIsNotNone(resp) + assert 'content_md5' in resp + assert 'content_crc64' in resp + assert 'request_id' in resp + self.assertIsNotNone(block_list) + self.assertEqual(len(block_list), 2) + self.assertEqual(len(block_list[1]), 0) + self.assertEqual(len(block_list[0]), 1) + self.assertEqual(block_list[0][0].size, LARGEST_BLOCK_SIZE) + + @pytest.mark.live_test_only + @GlobalStorageAccountPreparer() + def test_put_block_stream_largest_without_network(self, resource_group, location, storage_account, storage_account_key): + payload_dropping_policy = PayloadDroppingPolicy() + credential_policy = format_shared_key_credential([storage_account.name, "dummy"], storage_account_key) + self._setup(storage_account, storage_account_key, [payload_dropping_policy, credential_policy]) + blob = self._create_blob() + + # Act + stream = LargeStream(LARGEST_BLOCK_SIZE) + blockId = str(uuid.uuid4()) + requestId = str(uuid.uuid4()) + resp = blob.stage_block( + blockId, + stream, + length=LARGEST_BLOCK_SIZE, + client_request_id=requestId) + blob.commit_block_list([BlobBlock(blockId)]) + block_list = blob.get_block_list() + + # Assert + self.assertIsNotNone(resp) + assert 'content_md5' in resp + assert 'content_crc64' in resp + assert 'request_id' in resp + self.assertIsNotNone(block_list) + self.assertEqual(len(block_list), 2) + self.assertEqual(len(block_list[1]), 0) + self.assertEqual(len(block_list[0]), 1) + self.assertEqual(payload_dropping_policy.put_block_counter, 1) + self.assertEqual(payload_dropping_policy.put_block_sizes[0], LARGEST_BLOCK_SIZE) + + @pytest.mark.live_test_only + @pytest.mark.skip(reason="This takes really long time") + @GlobalStorageAccountPreparer() + def test_create_largest_blob_from_path(self, resource_group, location, storage_account, storage_account_key): + self._setup(storage_account, storage_account_key) + blob_name = self._get_blob_reference() + blob = self.bsc.get_blob_client(self.container_name, blob_name) + FILE_PATH = 'largest_blob_from_path.temp.{}.dat'.format(str(uuid.uuid4())) + with open(FILE_PATH, 'wb') as stream: + largeStream = LargeStream(LARGEST_BLOCK_SIZE, 100 * 1024 * 1024) + chunk = largeStream.read() + while chunk: + stream.write(chunk) + chunk = largeStream.read() + + # Act + with open(FILE_PATH, 'rb') as stream: + blob.upload_blob(stream, max_concurrency=2) + + # Assert + self._teardown(FILE_PATH) + + @pytest.mark.live_test_only + @GlobalStorageAccountPreparer() + def test_create_largest_blob_from_path_without_network(self, resource_group, location, storage_account, storage_account_key): + payload_dropping_policy = PayloadDroppingPolicy() + credential_policy = format_shared_key_credential([storage_account.name, "dummy"], storage_account_key) + self._setup(storage_account, storage_account_key, [payload_dropping_policy, credential_policy]) + blob_name = self._get_blob_reference() + blob = self.bsc.get_blob_client(self.container_name, blob_name) + FILE_PATH = 'largest_blob_from_path.temp.{}.dat'.format(str(uuid.uuid4())) + with open(FILE_PATH, 'wb') as stream: + largeStream = LargeStream(LARGEST_BLOCK_SIZE, 100 * 1024 * 1024) + chunk = largeStream.read() + while chunk: + stream.write(chunk) + chunk = largeStream.read() + + # Act + with open(FILE_PATH, 'rb') as stream: + blob.upload_blob(stream, max_concurrency=2) + + # Assert + self._teardown(FILE_PATH) + self.assertEqual(payload_dropping_policy.put_block_counter, 1) + self.assertEqual(payload_dropping_policy.put_block_sizes[0], LARGEST_BLOCK_SIZE) + + @pytest.mark.skip(reason="This takes really long time") + @pytest.mark.live_test_only + @GlobalStorageAccountPreparer() + def test_create_largest_blob_from_stream_without_network(self, resource_group, location, storage_account, storage_account_key): + payload_dropping_policy = PayloadDroppingPolicy() + credential_policy = format_shared_key_credential([storage_account.name, "dummy"], storage_account_key) + self._setup(storage_account, storage_account_key, [payload_dropping_policy, credential_policy]) + blob_name = self._get_blob_reference() + blob = self.bsc.get_blob_client(self.container_name, blob_name) + + number_of_blocks = 50000 + + stream = LargeStream(LARGEST_BLOCK_SIZE*number_of_blocks) + + # Act + blob.upload_blob(stream, max_concurrency=1) + + # Assert + self.assertEqual(payload_dropping_policy.put_block_counter, number_of_blocks) + self.assertEqual(payload_dropping_policy.put_block_sizes[0], LARGEST_BLOCK_SIZE) + + @pytest.mark.live_test_only + @GlobalStorageAccountPreparer() + def test_create_largest_blob_from_stream_single_upload_without_network(self, resource_group, location, storage_account, storage_account_key): + payload_dropping_policy = PayloadDroppingPolicy() + credential_policy = format_shared_key_credential([storage_account.name, "dummy"], storage_account_key) + self._setup(storage_account, storage_account_key, [payload_dropping_policy, credential_policy], + max_single_put_size=LARGEST_SINGLE_UPLOAD_SIZE) + blob_name = self._get_blob_reference() + blob = self.bsc.get_blob_client(self.container_name, blob_name) + + stream = LargeStream(LARGEST_SINGLE_UPLOAD_SIZE) + + # Act + blob.upload_blob(stream, length=LARGEST_SINGLE_UPLOAD_SIZE, max_concurrency=1) + + # Assert + self.assertEqual(payload_dropping_policy.put_block_counter, 0) + self.assertEqual(payload_dropping_policy.put_blob_counter, 1) + + +class LargeStream: + def __init__(self, length, initial_buffer_length=1024*1024): + self._base_data = urandom(initial_buffer_length) + self._base_data_length = initial_buffer_length + self._position = 0 + self._remaining = length + + def read(self, size=None): + if self._remaining == 0: + return b"" + + if size is None: + e = self._base_data_length + else: + e = size + e = min(e, self._remaining) + if e > self._base_data_length: + self._base_data = urandom(e) + self._base_data_length = e + self._remaining = self._remaining - e + return self._base_data[:e] + + def remaining(self): + return self._remaining + + +class PayloadDroppingPolicy(HTTPPolicy): + def __init__(self): + super().__init__() + self.put_block_counter = 0 + self.put_block_sizes = [] + self.put_blob_counter = 0 + self.put_blob_sizes = [] + + def send(self, request): # type: (PipelineRequest) -> PipelineResponse + if _is_put_block_request(request): + if request.http_request.body: + self.put_block_counter = self.put_block_counter + 1 + self.put_block_sizes.append(_get_body_length(request)) + replacement = "dummy_body" + request.http_request.body = replacement + request.http_request.headers["Content-Length"] = str(len(replacement)) + elif _is_put_blob_request(request): + if request.http_request.body: + self.put_blob_counter = self.put_blob_counter + 1 + self.put_blob_sizes.append(_get_body_length(request)) + replacement = "dummy_body" + request.http_request.body = replacement + request.http_request.headers["Content-Length"] = str(len(replacement)) + return self.next.send(request) + + +def _is_put_block_request(request): + query = request.http_request.query + return query and "comp" in query and query["comp"] == "block" + +def _is_put_blob_request(request): + query = request.http_request.query + return request.http_request.method == "PUT" and not query + +def _get_body_length(request): + body = request.http_request.body + length = 0 + if hasattr(body, "read"): + chunk = body.read(10*1024*1024) + while chunk: + length = length + len(chunk) + chunk = body.read(10 * 1024 * 1024) + else: + length = len(body) + return length + +# ------------------------------------------------------------------------------ diff --git a/sdk/storage/azure-storage-blob/tests/test_largest_block_blob_async.py b/sdk/storage/azure-storage-blob/tests/test_largest_block_blob_async.py new file mode 100644 index 000000000000..6ba7b698c863 --- /dev/null +++ b/sdk/storage/azure-storage-blob/tests/test_largest_block_blob_async.py @@ -0,0 +1,384 @@ +# coding: utf-8 + +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for +# license information. +# -------------------------------------------------------------------------- +from io import BytesIO + +import pytest + +from os import path, remove, urandom +import platform +import uuid + +from azure.core.pipeline.policies import SansIOHTTPPolicy +from azure.core.pipeline.transport import AioHttpTransport +from multidict import CIMultiDict, CIMultiDictProxy + +from azure.storage.blob.aio import ( + BlobServiceClient +) +from azure.storage.blob import ( + BlobBlock +) +from azure.storage.blob._shared.base_client import format_shared_key_credential +from azure.storage.blob._shared.constants import CONNECTION_TIMEOUT, READ_TIMEOUT + +from _shared.asynctestcase import AsyncStorageTestCase +from _shared.testcase import GlobalStorageAccountPreparer + +# ------------------------------------------------------------------------------ +TEST_BLOB_PREFIX = 'largestblob' +LARGEST_BLOCK_SIZE = 4000 * 1024 * 1024 +LARGEST_SINGLE_UPLOAD_SIZE = 5000 * 1024 * 1024 + +# ------------------------------------------------------------------------------ +if platform.python_implementation() == 'PyPy': + pytest.skip("Skip tests for Pypy", allow_module_level=True) + + +class AiohttpTestTransport(AioHttpTransport): + """Workaround to vcrpy bug: https://github.com/kevin1024/vcrpy/pull/461 + """ + async def send(self, request, **config): + response = await super(AiohttpTestTransport, self).send(request, **config) + if not isinstance(response.headers, CIMultiDictProxy): + response.headers = CIMultiDictProxy(CIMultiDict(response.internal_response.headers)) + response.content_type = response.headers.get("content-type") + return response + + +class StorageLargestBlockBlobTestAsync(AsyncStorageTestCase): + async def _setup(self, storage_account, key, additional_policies=None, min_large_block_upload_threshold=1 * 1024 * 1024, + max_single_put_size=32 * 1024): + self.bsc = BlobServiceClient( + self.account_url(storage_account, "blob"), + credential=key, + max_single_put_size=max_single_put_size, + max_block_size=LARGEST_BLOCK_SIZE, + min_large_block_upload_threshold=min_large_block_upload_threshold, + _additional_pipeline_policies=additional_policies, + transport=AiohttpTestTransport( + connection_timeout=CONNECTION_TIMEOUT, + read_timeout=READ_TIMEOUT + )) + self.config = self.bsc._config + self.container_name = self.get_resource_name('utcontainer') + self.container_name = self.container_name + str(uuid.uuid4()) + + if self.is_live: + await self.bsc.create_container(self.container_name) + + def _teardown(self, file_name): + if path.isfile(file_name): + try: + remove(file_name) + except: + pass + + # --Helpers----------------------------------------------------------------- + def _get_blob_reference(self): + return self.get_resource_name(TEST_BLOB_PREFIX) + + async def _create_blob(self): + blob_name = self._get_blob_reference() + blob = self.bsc.get_blob_client(self.container_name, blob_name) + await blob.upload_blob(b'') + return blob + + # --Test cases for block blobs -------------------------------------------- + @pytest.mark.live_test_only + @pytest.mark.skip(reason="This takes really long time") + @GlobalStorageAccountPreparer() + @AsyncStorageTestCase.await_prepared_test + async def test_put_block_bytes_largest(self, resource_group, location, storage_account, storage_account_key): + await self._setup(storage_account, storage_account_key) + blob = await self._create_blob() + + # Act + data = urandom(LARGEST_BLOCK_SIZE) + blockId = str(uuid.uuid4()).encode('utf-8') + resp = await blob.stage_block( + blockId, + data, + length=LARGEST_BLOCK_SIZE) + await blob.commit_block_list([BlobBlock(blockId)]) + block_list = await blob.get_block_list() + + # Assert + self.assertIsNotNone(resp) + assert 'content_md5' in resp + assert 'content_crc64' in resp + assert 'request_id' in resp + self.assertIsNotNone(block_list) + self.assertEqual(len(block_list), 2) + self.assertEqual(len(block_list[1]), 0) + self.assertEqual(len(block_list[0]), 1) + self.assertEqual(block_list[0][0].size, LARGEST_BLOCK_SIZE) + + @pytest.mark.live_test_only + @GlobalStorageAccountPreparer() + @AsyncStorageTestCase.await_prepared_test + async def test_put_block_bytes_largest_without_network(self, resource_group, location, storage_account, storage_account_key): + payload_dropping_policy = PayloadDroppingPolicy() + credential_policy = format_shared_key_credential([storage_account.name, "dummy"], storage_account_key) + await self._setup(storage_account, storage_account_key, [payload_dropping_policy, credential_policy]) + blob = await self._create_blob() + + # Act + data = urandom(LARGEST_BLOCK_SIZE) + blockId = str(uuid.uuid4()).encode('utf-8') + resp = await blob.stage_block( + blockId, + data, + length=LARGEST_BLOCK_SIZE) + await blob.commit_block_list([BlobBlock(blockId)]) + block_list = await blob.get_block_list() + + # Assert + self.assertIsNotNone(resp) + assert 'content_md5' in resp + assert 'content_crc64' in resp + assert 'request_id' in resp + self.assertIsNotNone(block_list) + self.assertEqual(len(block_list), 2) + self.assertEqual(len(block_list[1]), 0) + self.assertEqual(len(block_list[0]), 1) + self.assertEqual(payload_dropping_policy.put_block_counter, 1) + self.assertEqual(payload_dropping_policy.put_block_sizes[0], LARGEST_BLOCK_SIZE) + + @pytest.mark.live_test_only + @pytest.mark.skip(reason="This takes really long time") + @GlobalStorageAccountPreparer() + @AsyncStorageTestCase.await_prepared_test + async def test_put_block_stream_largest(self, resource_group, location, storage_account, storage_account_key): + await self._setup(storage_account, storage_account_key) + blob = await self._create_blob() + + # Act + stream = LargeStream(LARGEST_BLOCK_SIZE) + blockId = str(uuid.uuid4()) + requestId = str(uuid.uuid4()) + resp = await blob.stage_block( + blockId, + stream, + length=LARGEST_BLOCK_SIZE, + client_request_id=requestId) + await blob.commit_block_list([BlobBlock(blockId)]) + block_list = await blob.get_block_list() + + # Assert + self.assertIsNotNone(resp) + assert 'content_md5' in resp + assert 'content_crc64' in resp + assert 'request_id' in resp + self.assertIsNotNone(block_list) + self.assertEqual(len(block_list), 2) + self.assertEqual(len(block_list[1]), 0) + self.assertEqual(len(block_list[0]), 1) + self.assertEqual(block_list[0][0].size, LARGEST_BLOCK_SIZE) + + @pytest.mark.live_test_only + @GlobalStorageAccountPreparer() + @AsyncStorageTestCase.await_prepared_test + async def test_put_block_stream_largest_without_network(self, resource_group, location, storage_account, storage_account_key): + payload_dropping_policy = PayloadDroppingPolicy() + credential_policy = format_shared_key_credential([storage_account.name, "dummy"], storage_account_key) + await self._setup(storage_account, storage_account_key, [payload_dropping_policy, credential_policy]) + blob = await self._create_blob() + + # Act + stream = LargeStream(LARGEST_BLOCK_SIZE) + blockId = str(uuid.uuid4()) + requestId = str(uuid.uuid4()) + resp = await blob.stage_block( + blockId, + stream, + length=LARGEST_BLOCK_SIZE, + client_request_id=requestId) + await blob.commit_block_list([BlobBlock(blockId)]) + block_list = await blob.get_block_list() + + # Assert + self.assertIsNotNone(resp) + assert 'content_md5' in resp + assert 'content_crc64' in resp + assert 'request_id' in resp + self.assertIsNotNone(block_list) + self.assertEqual(len(block_list), 2) + self.assertEqual(len(block_list[1]), 0) + self.assertEqual(len(block_list[0]), 1) + self.assertEqual(payload_dropping_policy.put_block_counter, 1) + self.assertEqual(payload_dropping_policy.put_block_sizes[0], LARGEST_BLOCK_SIZE) + + @pytest.mark.live_test_only + @pytest.mark.skip(reason="This takes really long time") + @GlobalStorageAccountPreparer() + @AsyncStorageTestCase.await_prepared_test + async def test_create_largest_blob_from_path(self, resource_group, location, storage_account, storage_account_key): + await self._setup(storage_account, storage_account_key) + blob_name = self._get_blob_reference() + blob = self.bsc.get_blob_client(self.container_name, blob_name) + FILE_PATH = 'largest_blob_from_path.temp.{}.dat'.format(str(uuid.uuid4())) + with open(FILE_PATH, 'wb') as stream: + largeStream = LargeStream(LARGEST_BLOCK_SIZE, 100 * 1024 * 1024) + chunk = largeStream.read() + while chunk: + stream.write(chunk) + chunk = largeStream.read() + + # Act + with open(FILE_PATH, 'rb') as stream: + await blob.upload_blob(stream, max_concurrency=2) + + # Assert + self._teardown(FILE_PATH) + + @pytest.mark.live_test_only + @GlobalStorageAccountPreparer() + @AsyncStorageTestCase.await_prepared_test + async def test_create_largest_blob_from_path_without_network(self, resource_group, location, storage_account, storage_account_key): + payload_dropping_policy = PayloadDroppingPolicy() + credential_policy = format_shared_key_credential([storage_account.name, "dummy"], storage_account_key) + await self._setup(storage_account, storage_account_key, [payload_dropping_policy, credential_policy]) + blob_name = self._get_blob_reference() + blob = self.bsc.get_blob_client(self.container_name, blob_name) + FILE_PATH = 'largest_blob_from_path.temp.{}.dat'.format(str(uuid.uuid4())) + with open(FILE_PATH, 'wb') as stream: + largeStream = LargeStream(LARGEST_BLOCK_SIZE, 100 * 1024 * 1024) + chunk = largeStream.read() + while chunk: + stream.write(chunk) + chunk = largeStream.read() + + # Act + with open(FILE_PATH, 'rb') as stream: + await blob.upload_blob(stream, max_concurrency=2) + + # Assert + self._teardown(FILE_PATH) + self.assertEqual(payload_dropping_policy.put_block_counter, 1) + self.assertEqual(payload_dropping_policy.put_block_sizes[0], LARGEST_BLOCK_SIZE) + + @pytest.mark.skip(reason="This takes really long time") + @pytest.mark.live_test_only + @GlobalStorageAccountPreparer() + @AsyncStorageTestCase.await_prepared_test + async def test_create_largest_blob_from_stream_without_network(self, resource_group, location, storage_account, storage_account_key): + payload_dropping_policy = PayloadDroppingPolicy() + credential_policy = format_shared_key_credential([storage_account.name, "dummy"], storage_account_key) + await self._setup(storage_account, storage_account_key, [payload_dropping_policy, credential_policy]) + blob_name = self._get_blob_reference() + blob = self.bsc.get_blob_client(self.container_name, blob_name) + + number_of_blocks = 50000 + + stream = LargeStream(LARGEST_BLOCK_SIZE*number_of_blocks) + + # Act + await blob.upload_blob(stream, max_concurrency=1) + + # Assert + self.assertEqual(payload_dropping_policy.put_block_counter, number_of_blocks) + self.assertEqual(payload_dropping_policy.put_block_sizes[0], LARGEST_BLOCK_SIZE) + + @pytest.mark.live_test_only + @GlobalStorageAccountPreparer() + @AsyncStorageTestCase.await_prepared_test + async def test_create_largest_blob_from_stream_single_upload_without_network(self, resource_group, location, storage_account, storage_account_key): + payload_dropping_policy = PayloadDroppingPolicy() + credential_policy = format_shared_key_credential([storage_account.name, "dummy"], storage_account_key) + await self._setup(storage_account, storage_account_key, [payload_dropping_policy, credential_policy], + max_single_put_size=LARGEST_SINGLE_UPLOAD_SIZE) + blob_name = self._get_blob_reference() + blob = self.bsc.get_blob_client(self.container_name, blob_name) + + stream = LargeStream(LARGEST_SINGLE_UPLOAD_SIZE) + + # Act + await blob.upload_blob(stream, length=LARGEST_SINGLE_UPLOAD_SIZE, max_concurrency=1) + + # Assert + self.assertEqual(payload_dropping_policy.put_block_counter, 0) + self.assertEqual(payload_dropping_policy.put_blob_counter, 1) + + +class LargeStream(BytesIO): + def __init__(self, length, initial_buffer_length=1024 * 1024): + super().__init__() + self._base_data = urandom(initial_buffer_length) + self._base_data_length = initial_buffer_length + self._position = 0 + self._remaining = length + self._closed = False + + def read(self, size=None): + if self._remaining == 0: + return b"" + + if size is None: + e = self._base_data_length + else: + e = size + e = min(e, self._remaining) + if e > self._base_data_length: + self._base_data = urandom(e) + self._base_data_length = e + self._remaining = self._remaining - e + return self._base_data[:e] + + def remaining(self): + return self._remaining + + def close(self): + self._closed = True + + +class PayloadDroppingPolicy(SansIOHTTPPolicy): + def __init__(self): + super().__init__() + self.put_block_counter = 0 + self.put_block_sizes = [] + self.put_blob_counter = 0 + self.put_blob_sizes = [] + + def on_request(self, request): # type: (PipelineRequest) -> Union[None, Awaitable[None]] + if _is_put_block_request(request): + if request.http_request.body: + self.put_block_counter = self.put_block_counter + 1 + self.put_block_sizes.append(_get_body_length(request)) + replacement = "dummy_body" + request.http_request.body = replacement + request.http_request.headers["Content-Length"] = str(len(replacement)) + elif _is_put_blob_request(request): + if request.http_request.body: + self.put_blob_counter = self.put_blob_counter + 1 + self.put_blob_sizes.append(_get_body_length(request)) + replacement = "dummy_body" + request.http_request.body = replacement + request.http_request.headers["Content-Length"] = str(len(replacement)) + + +def _is_put_block_request(request): + query = request.http_request.query + return query and "comp" in query and query["comp"] == "block" + +def _is_put_blob_request(request): + query = request.http_request.query + return request.http_request.method == "PUT" and not query + +def _get_body_length(request): + body = request.http_request.body + length = 0 + if hasattr(body, "read"): + chunk = body.read(10*1024*1024) + while chunk: + length = length + len(chunk) + chunk = body.read(10 * 1024 * 1024) + else: + length = len(body) + return length + +# ------------------------------------------------------------------------------ diff --git a/sdk/storage/azure-storage-file-datalake/azure/storage/filedatalake/_data_lake_file_client.py b/sdk/storage/azure-storage-file-datalake/azure/storage/filedatalake/_data_lake_file_client.py index ffea24237d3a..ce6d4b64d8d2 100644 --- a/sdk/storage/azure-storage-file-datalake/azure/storage/filedatalake/_data_lake_file_client.py +++ b/sdk/storage/azure-storage-file-datalake/azure/storage/filedatalake/_data_lake_file_client.py @@ -337,6 +337,9 @@ def upload_data(self, data, # type: Union[AnyStr, Iterable[AnyStr], IO[AnyStr]] The match condition to use upon the etag. :keyword int timeout: The timeout parameter is expressed in seconds. + :keyword int chunk_size: + The maximum chunk size for uploading a file in chunks. + Defaults to 100*1024*1024, or 100MB. :return: response dict (Etag and last modified). """ options = self._upload_options( diff --git a/sdk/storage/azure-storage-file-datalake/azure/storage/filedatalake/_shared/base_client.py b/sdk/storage/azure-storage-file-datalake/azure/storage/filedatalake/_shared/base_client.py index 8b65417e9a9a..27f52dd413fc 100644 --- a/sdk/storage/azure-storage-file-datalake/azure/storage/filedatalake/_shared/base_client.py +++ b/sdk/storage/azure-storage-file-datalake/azure/storage/filedatalake/_shared/base_client.py @@ -246,6 +246,8 @@ def _create_pipeline(self, credential, **kwargs): DistributedTracingPolicy(**kwargs), HttpLoggingPolicy(**kwargs) ] + if kwargs.get("_additional_pipeline_policies"): + policies = policies + kwargs.get("_additional_pipeline_policies") return config, Pipeline(config.transport, policies=policies) def _batch_send( diff --git a/sdk/storage/azure-storage-file-datalake/azure/storage/filedatalake/_shared/base_client_async.py b/sdk/storage/azure-storage-file-datalake/azure/storage/filedatalake/_shared/base_client_async.py index 3c806d796864..1b13658c9824 100644 --- a/sdk/storage/azure-storage-file-datalake/azure/storage/filedatalake/_shared/base_client_async.py +++ b/sdk/storage/azure-storage-file-datalake/azure/storage/filedatalake/_shared/base_client_async.py @@ -102,6 +102,8 @@ def _create_pipeline(self, credential, **kwargs): DistributedTracingPolicy(**kwargs), HttpLoggingPolicy(**kwargs), ] + if kwargs.get("_additional_pipeline_policies"): + policies = policies + kwargs.get("_additional_pipeline_policies") return config, AsyncPipeline(config.transport, policies=policies) async def _batch_send( diff --git a/sdk/storage/azure-storage-file-datalake/azure/storage/filedatalake/_upload_helper.py b/sdk/storage/azure-storage-file-datalake/azure/storage/filedatalake/_upload_helper.py index e1db768eff3b..165c62af9cee 100644 --- a/sdk/storage/azure-storage-file-datalake/azure/storage/filedatalake/_upload_helper.py +++ b/sdk/storage/azure-storage-file-datalake/azure/storage/filedatalake/_upload_helper.py @@ -39,6 +39,7 @@ def upload_datalake_file( # pylint: disable=unused-argument properties = kwargs.pop('properties', None) path_http_headers = kwargs.pop('path_http_headers', None) modified_access_conditions = kwargs.pop('modified_access_conditions', None) + chunk_size = kwargs.pop('chunk_size', 100 * 1024 * 1024) if not overwrite: # if customers didn't specify access conditions, they cannot flush data to existing file @@ -67,7 +68,7 @@ def upload_datalake_file( # pylint: disable=unused-argument service=client, uploader_class=DataLakeFileChunkUploader, total_size=length, - chunk_size=100 * 1024 * 1024, + chunk_size=chunk_size, stream=stream, max_concurrency=max_concurrency, validate_content=validate_content, diff --git a/sdk/storage/azure-storage-file-datalake/azure/storage/filedatalake/aio/_data_lake_file_client_async.py b/sdk/storage/azure-storage-file-datalake/azure/storage/filedatalake/aio/_data_lake_file_client_async.py index 3b6b894d193e..da3aeed1137c 100644 --- a/sdk/storage/azure-storage-file-datalake/azure/storage/filedatalake/aio/_data_lake_file_client_async.py +++ b/sdk/storage/azure-storage-file-datalake/azure/storage/filedatalake/aio/_data_lake_file_client_async.py @@ -255,6 +255,9 @@ async def upload_data(self, data, # type: Union[AnyStr, Iterable[AnyStr], IO[An The match condition to use upon the etag. :keyword int timeout: The timeout parameter is expressed in seconds. + :keyword int chunk_size: + The maximum chunk size for uploading a file in chunks. + Defaults to 100*1024*1024, or 100MB. :return: response dict (Etag and last modified). """ options = self._upload_options( diff --git a/sdk/storage/azure-storage-file-datalake/azure/storage/filedatalake/aio/_upload_helper.py b/sdk/storage/azure-storage-file-datalake/azure/storage/filedatalake/aio/_upload_helper.py index caa678b7db64..c94079579d17 100644 --- a/sdk/storage/azure-storage-file-datalake/azure/storage/filedatalake/aio/_upload_helper.py +++ b/sdk/storage/azure-storage-file-datalake/azure/storage/filedatalake/aio/_upload_helper.py @@ -39,6 +39,7 @@ async def upload_datalake_file( # pylint: disable=unused-argument properties = kwargs.pop('properties', None) path_http_headers = kwargs.pop('path_http_headers', None) modified_access_conditions = kwargs.pop('modified_access_conditions', None) + chunk_size = kwargs.pop('chunk_size', 100 * 1024 * 1024) if not overwrite: # if customers didn't specify access conditions, they cannot flush data to existing file @@ -67,7 +68,7 @@ async def upload_datalake_file( # pylint: disable=unused-argument service=client, uploader_class=DataLakeFileChunkUploader, total_size=length, - chunk_size=100 * 1024 * 1024, + chunk_size=chunk_size, stream=stream, max_concurrency=max_concurrency, validate_content=validate_content, diff --git a/sdk/storage/azure-storage-file-datalake/tests/test_large_file.py b/sdk/storage/azure-storage-file-datalake/tests/test_large_file.py new file mode 100644 index 000000000000..272dd233b814 --- /dev/null +++ b/sdk/storage/azure-storage-file-datalake/tests/test_large_file.py @@ -0,0 +1,180 @@ +# coding: utf-8 + +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for +# license information. +# -------------------------------------------------------------------------- +import unittest +from os import urandom + +import pytest +import re +from azure.core.pipeline.policies import HTTPPolicy + +from azure.core.exceptions import ResourceExistsError +from azure.storage.blob._shared.base_client import format_shared_key_credential +from azure.storage.filedatalake import DataLakeServiceClient +from testcase import ( + StorageTestCase, + record, + TestMode +) + +# ------------------------------------------------------------------------------ +TEST_DIRECTORY_PREFIX = 'directory' +TEST_FILE_PREFIX = 'file' +FILE_PATH = 'file_output.temp.dat' +LARGEST_BLOCK_SIZE = 4000 * 1024 * 1024 +# ------------------------------------------------------------------------------ + + +class LargeFileTest(StorageTestCase): + def setUp(self): + super(LargeFileTest, self).setUp() + url = self._get_account_url() + self.payload_dropping_policy = PayloadDroppingPolicy() + credential_policy = format_shared_key_credential([self.settings.STORAGE_DATA_LAKE_ACCOUNT_NAME, "dummy"], + self.settings.STORAGE_DATA_LAKE_ACCOUNT_KEY) + self.dsc = DataLakeServiceClient(url, + credential=self.settings.STORAGE_DATA_LAKE_ACCOUNT_KEY, + logging_enable=True, + _additional_pipeline_policies=[self.payload_dropping_policy, credential_policy]) + self.config = self.dsc._config + + self.file_system_name = self.get_resource_name('filesystem') + + if not self.is_playback(): + file_system = self.dsc.get_file_system_client(self.file_system_name) + try: + file_system.create_file_system(timeout=5) + except ResourceExistsError: + pass + + def tearDown(self): + if not self.is_playback(): + try: + self.dsc.delete_file_system(self.file_system_name) + except: + pass + + return super(LargeFileTest, self).tearDown() + + @pytest.mark.live_test_only + def test_append_large_stream_without_network(self): + directory_name = self.get_resource_name(TEST_DIRECTORY_PREFIX) + + # Create a directory to put the file under that + directory_client = self.dsc.get_directory_client(self.file_system_name, directory_name) + directory_client.create_directory() + + file_client = directory_client.get_file_client('filename') + file_client.create_file() + + data = LargeStream(LARGEST_BLOCK_SIZE) + + # Act + response = file_client.append_data(data, 0, LARGEST_BLOCK_SIZE) + + self.assertIsNotNone(response) + self.assertEqual(self.payload_dropping_policy.append_counter, 1) + self.assertEqual(self.payload_dropping_policy.append_sizes[0], LARGEST_BLOCK_SIZE) + + @pytest.mark.live_test_only + def test_upload_large_stream_without_network(self): + directory_name = self.get_resource_name(TEST_DIRECTORY_PREFIX) + + # Create a directory to put the file under that + directory_client = self.dsc.get_directory_client(self.file_system_name, directory_name) + directory_client.create_directory() + + file_client = directory_client.get_file_client('filename') + file_client.create_file() + + length = 2*LARGEST_BLOCK_SIZE + data = LargeStream(length) + + # Act + response = file_client.upload_data(data, length, overwrite=True, chunk_size = LARGEST_BLOCK_SIZE) + + self.assertIsNotNone(response) + self.assertEqual(self.payload_dropping_policy.append_counter, 2) + self.assertEqual(self.payload_dropping_policy.append_sizes[0], LARGEST_BLOCK_SIZE) + self.assertEqual(self.payload_dropping_policy.append_sizes[1], LARGEST_BLOCK_SIZE) + + +class LargeStream: + def __init__(self, length, initial_buffer_length=1024*1024): + self._base_data = urandom(initial_buffer_length) + self._base_data_length = initial_buffer_length + self._position = 0 + self._remaining = length + + def read(self, size=None): + if self._remaining == 0: + return b"" + + if size is None: + e = self._base_data_length + else: + e = size + e = min(e, self._remaining) + if e > self._base_data_length: + self._base_data = urandom(e) + self._base_data_length = e + self._remaining = self._remaining - e + return self._base_data[:e] + + def remaining(self): + return self._remaining + + +class PayloadDroppingPolicy(HTTPPolicy): + def __init__(self): + super().__init__() + self.append_counter = 0 + self.append_sizes = [] + self.dummy_body = "dummy_body" + + def send(self, request): # type: (PipelineRequest) -> PipelineResponse + if _is_append_request(request): + if request.http_request.body: + position = self.append_counter*len(self.dummy_body) + request.http_request.url = re.sub(r'position=\d+', "position=" + str(position), request.http_request.url) + self.append_sizes.append(_get_body_length(request)) + replacement = self.dummy_body + request.http_request.body = replacement + request.http_request.headers["Content-Length"] = str(len(replacement)) + self.append_counter = self.append_counter + 1 + if _is_flush_request(request): + position = self.append_counter * len(self.dummy_body) + request.http_request.url = re.sub(r'position=\d+', "position=" + str(position), request.http_request.url) + return self.next.send(request) + + +def _is_append_request(request): + query = request.http_request.query + return query and "action" in query and query["action"] == "append" + + +def _is_flush_request(request): + query = request.http_request.query + return query and "action" in query and query["action"] == "flush" + + +def _get_body_length(request): + body = request.http_request.body + length = 0 + if hasattr(body, "read"): + chunk = body.read(10*1024*1024) + while chunk: + length = length + len(chunk) + chunk = body.read(10 * 1024 * 1024) + else: + length = len(body) + return length + + +# ------------------------------------------------------------------------------ +if __name__ == '__main__': + unittest.main() diff --git a/sdk/storage/azure-storage-file-datalake/tests/test_large_file_async.py b/sdk/storage/azure-storage-file-datalake/tests/test_large_file_async.py new file mode 100644 index 000000000000..2bfdef448ab2 --- /dev/null +++ b/sdk/storage/azure-storage-file-datalake/tests/test_large_file_async.py @@ -0,0 +1,207 @@ +# coding: utf-8 + +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for +# license information. +# -------------------------------------------------------------------------- +import re +import unittest +import asyncio +from io import BytesIO +from os import urandom + +import pytest + +from azure.core.exceptions import ResourceExistsError +from azure.core.pipeline.policies import SansIOHTTPPolicy +from azure.storage.blob._shared.base_client import format_shared_key_credential +from azure.storage.filedatalake.aio import DataLakeServiceClient +from testcase import ( + StorageTestCase, + record, + TestMode +) + +# ------------------------------------------------------------------------------ +TEST_DIRECTORY_PREFIX = 'directory' +TEST_FILE_PREFIX = 'file' +FILE_PATH = 'file_output.temp.dat' +LARGEST_BLOCK_SIZE = 4000 * 1024 * 1024 +# ------------------------------------------------------------------------------ + + +class LargeFileTest(StorageTestCase): + def setUp(self): + super(LargeFileTest, self).setUp() + url = self._get_account_url() + self.payload_dropping_policy = PayloadDroppingPolicy() + credential_policy = format_shared_key_credential([self.settings.STORAGE_DATA_LAKE_ACCOUNT_NAME, "dummy"], + self.settings.STORAGE_DATA_LAKE_ACCOUNT_KEY) + self.dsc = DataLakeServiceClient(url, + credential=self.settings.STORAGE_DATA_LAKE_ACCOUNT_KEY, + _additional_pipeline_policies=[self.payload_dropping_policy, credential_policy]) + loop = asyncio.get_event_loop() + loop.run_until_complete(self.dsc.__aenter__()) + + self.config = self.dsc._config + + self.file_system_name = self.get_resource_name('filesystem') + + if not self.is_playback(): + file_system = self.dsc.get_file_system_client(self.file_system_name) + try: + loop = asyncio.get_event_loop() + loop.run_until_complete(file_system.create_file_system(timeout=5)) + + except ResourceExistsError: + pass + + def tearDown(self): + if not self.is_playback(): + try: + loop = asyncio.get_event_loop() + loop.run_until_complete(self.dsc.delete_file_system(self.file_system_name)) + loop.run_until_complete(self.dsc.__aexit__()) + except: + pass + + return super(LargeFileTest, self).tearDown() + + # --Helpers----------------------------------------------------------------- + def _get_directory_reference(self, prefix=TEST_DIRECTORY_PREFIX): + directory_name = self.get_resource_name(prefix) + return directory_name + + # --Helpers----------------------------------------------------------------- + + async def _test_append_large_stream_without_network(self): + directory_name = self._get_directory_reference() + + # Create a directory to put the file under that + directory_client = self.dsc.get_directory_client(self.file_system_name, directory_name) + await directory_client.create_directory() + + file_client = directory_client.get_file_client('filename') + await file_client.create_file() + + data = LargeStream(LARGEST_BLOCK_SIZE) + + # Act + response = await file_client.append_data(data, 0, LARGEST_BLOCK_SIZE) + + self.assertIsNotNone(response) + self.assertEqual(self.payload_dropping_policy.append_counter, 1) + self.assertEqual(self.payload_dropping_policy.append_sizes[0], LARGEST_BLOCK_SIZE) + + @pytest.mark.live_test_only + def test_append_large_stream_without_network(self): + loop = asyncio.get_event_loop() + loop.run_until_complete(self._test_append_large_stream_without_network()) + + async def _test_upload_large_stream_without_network(self): + directory_name = self.get_resource_name(TEST_DIRECTORY_PREFIX) + + # Create a directory to put the file under that + directory_client = self.dsc.get_directory_client(self.file_system_name, directory_name) + await directory_client.create_directory() + + file_client = directory_client.get_file_client('filename') + await file_client.create_file() + + length = 2*LARGEST_BLOCK_SIZE + data = LargeStream(length) + + # Act + response = await file_client.upload_data(data, length, overwrite=True, chunk_size = LARGEST_BLOCK_SIZE) + + self.assertIsNotNone(response) + self.assertEqual(self.payload_dropping_policy.append_counter, 2) + self.assertEqual(self.payload_dropping_policy.append_sizes[0], LARGEST_BLOCK_SIZE) + self.assertEqual(self.payload_dropping_policy.append_sizes[1], LARGEST_BLOCK_SIZE) + + @pytest.mark.live_test_only + def test_upload_large_stream_without_network(self): + loop = asyncio.get_event_loop() + loop.run_until_complete(self._test_upload_large_stream_without_network()) + + +class LargeStream(BytesIO): + def __init__(self, length, initial_buffer_length=1024 * 1024): + super().__init__() + self._base_data = urandom(initial_buffer_length) + self._base_data_length = initial_buffer_length + self._position = 0 + self._remaining = length + self._closed = False + + def read(self, size=None): + if self._remaining == 0: + return b"" + + if size is None: + e = self._base_data_length + else: + e = size + e = min(e, self._remaining) + if e > self._base_data_length: + self._base_data = urandom(e) + self._base_data_length = e + self._remaining = self._remaining - e + return self._base_data[:e] + + def remaining(self): + return self._remaining + + def close(self): + self._closed = True + + +class PayloadDroppingPolicy(SansIOHTTPPolicy): + def __init__(self): + super().__init__() + self.append_counter = 0 + self.append_sizes = [] + self.dummy_body = "dummy_body" + + def on_request(self, request): # type: (PipelineRequest) -> Union[None, Awaitable[None]] + if _is_append_request(request): + if request.http_request.body: + position = self.append_counter*len(self.dummy_body) + request.http_request.url = re.sub(r'position=\d+', "position=" + str(position), request.http_request.url) + self.append_sizes.append(_get_body_length(request)) + replacement = self.dummy_body + request.http_request.body = replacement + request.http_request.headers["Content-Length"] = str(len(replacement)) + self.append_counter = self.append_counter + 1 + elif _is_flush_request(request): + position = self.append_counter * len(self.dummy_body) + request.http_request.url = re.sub(r'position=\d+', "position=" + str(position), request.http_request.url) + + +def _is_append_request(request): + query = request.http_request.query + return query and "action" in query and query["action"] == "append" + + +def _is_flush_request(request): + query = request.http_request.query + return query and "action" in query and query["action"] == "flush" + + +def _get_body_length(request): + body = request.http_request.body + length = 0 + if hasattr(body, "read"): + chunk = body.read(10*1024*1024) + while chunk: + length = length + len(chunk) + chunk = body.read(10 * 1024 * 1024) + else: + length = len(body) + return length + + +# ------------------------------------------------------------------------------ +if __name__ == '__main__': + unittest.main() diff --git a/sdk/storage/azure-storage-file-share/azure/storage/fileshare/_shared/base_client.py b/sdk/storage/azure-storage-file-share/azure/storage/fileshare/_shared/base_client.py index 8b65417e9a9a..27f52dd413fc 100644 --- a/sdk/storage/azure-storage-file-share/azure/storage/fileshare/_shared/base_client.py +++ b/sdk/storage/azure-storage-file-share/azure/storage/fileshare/_shared/base_client.py @@ -246,6 +246,8 @@ def _create_pipeline(self, credential, **kwargs): DistributedTracingPolicy(**kwargs), HttpLoggingPolicy(**kwargs) ] + if kwargs.get("_additional_pipeline_policies"): + policies = policies + kwargs.get("_additional_pipeline_policies") return config, Pipeline(config.transport, policies=policies) def _batch_send( diff --git a/sdk/storage/azure-storage-file-share/azure/storage/fileshare/_shared/base_client_async.py b/sdk/storage/azure-storage-file-share/azure/storage/fileshare/_shared/base_client_async.py index 3c806d796864..1b13658c9824 100644 --- a/sdk/storage/azure-storage-file-share/azure/storage/fileshare/_shared/base_client_async.py +++ b/sdk/storage/azure-storage-file-share/azure/storage/fileshare/_shared/base_client_async.py @@ -102,6 +102,8 @@ def _create_pipeline(self, credential, **kwargs): DistributedTracingPolicy(**kwargs), HttpLoggingPolicy(**kwargs), ] + if kwargs.get("_additional_pipeline_policies"): + policies = policies + kwargs.get("_additional_pipeline_policies") return config, AsyncPipeline(config.transport, policies=policies) async def _batch_send( diff --git a/sdk/storage/azure-storage-queue/azure/storage/queue/_shared/base_client.py b/sdk/storage/azure-storage-queue/azure/storage/queue/_shared/base_client.py index 8b65417e9a9a..27f52dd413fc 100644 --- a/sdk/storage/azure-storage-queue/azure/storage/queue/_shared/base_client.py +++ b/sdk/storage/azure-storage-queue/azure/storage/queue/_shared/base_client.py @@ -246,6 +246,8 @@ def _create_pipeline(self, credential, **kwargs): DistributedTracingPolicy(**kwargs), HttpLoggingPolicy(**kwargs) ] + if kwargs.get("_additional_pipeline_policies"): + policies = policies + kwargs.get("_additional_pipeline_policies") return config, Pipeline(config.transport, policies=policies) def _batch_send( diff --git a/sdk/storage/azure-storage-queue/azure/storage/queue/_shared/base_client_async.py b/sdk/storage/azure-storage-queue/azure/storage/queue/_shared/base_client_async.py index 3c806d796864..1b13658c9824 100644 --- a/sdk/storage/azure-storage-queue/azure/storage/queue/_shared/base_client_async.py +++ b/sdk/storage/azure-storage-queue/azure/storage/queue/_shared/base_client_async.py @@ -102,6 +102,8 @@ def _create_pipeline(self, credential, **kwargs): DistributedTracingPolicy(**kwargs), HttpLoggingPolicy(**kwargs), ] + if kwargs.get("_additional_pipeline_policies"): + policies = policies + kwargs.get("_additional_pipeline_policies") return config, AsyncPipeline(config.transport, policies=policies) async def _batch_send(