From 7bc3c94dd3e6cd366d9a550726699322020c2e2c Mon Sep 17 00:00:00 2001 From: Zunli Hu Date: Mon, 28 Dec 2020 17:44:06 +0800 Subject: [PATCH 1/2] fix OOM --- .../storage/operations/fs_file.py | 12 +++--- .../test_storage_adls_upload_scenarios.py | 42 +++++++++++++++++++ 2 files changed, 48 insertions(+), 6 deletions(-) create mode 100644 src/azure-cli/azure/cli/command_modules/storage/tests/latest/test_storage_adls_upload_scenarios.py diff --git a/src/azure-cli/azure/cli/command_modules/storage/operations/fs_file.py b/src/azure-cli/azure/cli/command_modules/storage/operations/fs_file.py index 306c3ad2bf2..ea8eaf3283f 100644 --- a/src/azure-cli/azure/cli/command_modules/storage/operations/fs_file.py +++ b/src/azure-cli/azure/cli/command_modules/storage/operations/fs_file.py @@ -89,9 +89,6 @@ def upload_file(cmd, client, local_path, overwrite=None, content_settings=None, if_match=None, if_none_match=None, if_modified_since=None, if_unmodified_since=None, umask=None, permissions=None): - count = os.path.getsize(local_path) - with open(local_path, 'rb') as stream: - data = stream.read(count) from azure.core import MatchConditions upload_file_args = { 'content_settings': content_settings, @@ -126,7 +123,9 @@ def upload_file(cmd, client, local_path, overwrite=None, content_settings=None, upload_file_args['match_condition'] = MatchConditions.IfPresent try: - return client.upload_data(data=data, length=count, overwrite=overwrite, **upload_file_args) + count = os.path.getsize(local_path) + with open(local_path, 'rb') as stream: + return client.upload_data(data=stream, length=count, overwrite=overwrite, **upload_file_args) except HttpResponseError as ex: StorageErrorCode = cmd.get_models("_shared.models#StorageErrorCode", resource_type=ResourceType.DATA_STORAGE_FILEDATALAKE) @@ -135,5 +134,6 @@ def upload_file(cmd, client, local_path, overwrite=None, content_settings=None, raise CLIError("You cannot upload to an existing non-empty file with overwrite=false. " "Please set --overwrite to overwrite the existing file.") raise ex - - return client.upload_data(data=data, length=count, overwrite=overwrite, **upload_file_args) + count = os.path.getsize(local_path) + with open(local_path, 'rb') as stream: + return client.upload_data(data=stream, length=count, overwrite=overwrite, **upload_file_args) diff --git a/src/azure-cli/azure/cli/command_modules/storage/tests/latest/test_storage_adls_upload_scenarios.py b/src/azure-cli/azure/cli/command_modules/storage/tests/latest/test_storage_adls_upload_scenarios.py new file mode 100644 index 00000000000..5300eb05792 --- /dev/null +++ b/src/azure-cli/azure/cli/command_modules/storage/tests/latest/test_storage_adls_upload_scenarios.py @@ -0,0 +1,42 @@ +# -------------------------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for license information. +# -------------------------------------------------------------------------------------------- + +import os +from datetime import datetime, timedelta +from azure.cli.testsdk import (LiveScenarioTest, ResourceGroupPreparer, StorageAccountPreparer, + JMESPathCheck, JMESPathCheckExists, NoneCheck, api_version_constraint) +from azure.cli.core.profiles import ResourceType +from azure.cli.testsdk.decorators import serial_test +from ..storage_test_util import StorageScenarioMixin + + +@api_version_constraint(ResourceType.MGMT_STORAGE, min_api='2016-12-01') +class StorageADLSUploadLiveTests(LiveScenarioTest): + @ResourceGroupPreparer() + @StorageAccountPreparer(kind='StorageV2', hns=True) + def test_storage_file_upload_2G_file(self, resource_group, storage_account): + file_size_kb = 2 * 1024 * 1024 + file_system = self.create_random_name(prefix='cont', length=24) + local_dir = self.create_temp_dir() + local_file = self.create_temp_file(file_size_kb, full_random=True) + file_name = self.create_random_name(prefix='adls', length=24) + account_key = self.cmd('storage account keys list -n {} -g {} --query "[0].value" -otsv' + .format(storage_account, resource_group)).output + + self.set_env('AZURE_STORAGE_ACCOUNT', storage_account) + self.set_env('AZURE_STORAGE_KEY', account_key) + + self.cmd('storage fs create -n {}'.format(file_system)) + + self.cmd('storage fs file exists -p {} -f {}'.format(file_name, file_system), + checks=JMESPathCheck('exists', False)) + + self.cmd('storage fs file upload -f {} -s "{}" -p {} '.format(file_system, local_file, file_name)) + + self.cmd('storage fs file exists -f {} -p {}'.format(file_system, file_name), + checks=JMESPathCheck('exists', True)) + + self.cmd('storage fs file show -p {} -f {}'.format(file_name, file_system), + checks=[JMESPathCheck('properties.contentLength', file_size_kb * 1024)]) From 379901107e617445cd840ef6690e751b8274eb0f Mon Sep 17 00:00:00 2001 From: Zunli Hu Date: Mon, 28 Dec 2020 20:52:50 +0800 Subject: [PATCH 2/2] fix style and test --- .../azure/cli/command_modules/storage/operations/fs_file.py | 6 ++++-- .../tests/latest/test_storage_adls_upload_scenarios.py | 3 +-- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/src/azure-cli/azure/cli/command_modules/storage/operations/fs_file.py b/src/azure-cli/azure/cli/command_modules/storage/operations/fs_file.py index ea8eaf3283f..35d224ca419 100644 --- a/src/azure-cli/azure/cli/command_modules/storage/operations/fs_file.py +++ b/src/azure-cli/azure/cli/command_modules/storage/operations/fs_file.py @@ -125,7 +125,8 @@ def upload_file(cmd, client, local_path, overwrite=None, content_settings=None, try: count = os.path.getsize(local_path) with open(local_path, 'rb') as stream: - return client.upload_data(data=stream, length=count, overwrite=overwrite, **upload_file_args) + response = client.upload_data(data=stream, length=count, overwrite=overwrite, **upload_file_args) + return response except HttpResponseError as ex: StorageErrorCode = cmd.get_models("_shared.models#StorageErrorCode", resource_type=ResourceType.DATA_STORAGE_FILEDATALAKE) @@ -136,4 +137,5 @@ def upload_file(cmd, client, local_path, overwrite=None, content_settings=None, raise ex count = os.path.getsize(local_path) with open(local_path, 'rb') as stream: - return client.upload_data(data=stream, length=count, overwrite=overwrite, **upload_file_args) + response = client.upload_data(data=stream, length=count, overwrite=overwrite, **upload_file_args) + return response diff --git a/src/azure-cli/azure/cli/command_modules/storage/tests/latest/test_storage_adls_upload_scenarios.py b/src/azure-cli/azure/cli/command_modules/storage/tests/latest/test_storage_adls_upload_scenarios.py index 5300eb05792..ad297e8ecba 100644 --- a/src/azure-cli/azure/cli/command_modules/storage/tests/latest/test_storage_adls_upload_scenarios.py +++ b/src/azure-cli/azure/cli/command_modules/storage/tests/latest/test_storage_adls_upload_scenarios.py @@ -19,7 +19,6 @@ class StorageADLSUploadLiveTests(LiveScenarioTest): def test_storage_file_upload_2G_file(self, resource_group, storage_account): file_size_kb = 2 * 1024 * 1024 file_system = self.create_random_name(prefix='cont', length=24) - local_dir = self.create_temp_dir() local_file = self.create_temp_file(file_size_kb, full_random=True) file_name = self.create_random_name(prefix='adls', length=24) account_key = self.cmd('storage account keys list -n {} -g {} --query "[0].value" -otsv' @@ -39,4 +38,4 @@ def test_storage_file_upload_2G_file(self, resource_group, storage_account): checks=JMESPathCheck('exists', True)) self.cmd('storage fs file show -p {} -f {}'.format(file_name, file_system), - checks=[JMESPathCheck('properties.contentLength', file_size_kb * 1024)]) + checks=[JMESPathCheck('size', file_size_kb * 1024)])