{Storage} Migrate az storage file upload/upload-batch to track2 sdk (#22803)

* migrate_storage_share_track2

* fix conflict

* fix conflict

* fix checks

* fix checks

* fix checks

* Update src/azure-cli/azure/cli/command_modules/storage/_help.py

Co-authored-by: Yishi Wang <[email protected]>

* Update _transformers.py

* fix conflict

* fix changes

* fix conflicts

* fix conflicts

* fix conflicts

* fix conflicts

* fix conflicts

Co-authored-by: Yishi Wang <[email protected]>

Author: jonie001

src/azure-cli/azure/cli/command_modules/storage/_help.py

@@ -1895,6 +1895,12 @@
     crafted: true
 """
 
+helps['storage file download'] = """
+type: command
+short-summary: Download a file to a file path, with automatic chunking and progress notifications.
+long-summary: Return an instance of File with properties and metadata.
+"""
+
 helps['storage file download-batch'] = """
 type: command
 short-summary: Download files from an Azure Storage File Share to a local directory in a batch operation.

src/azure-cli/azure/cli/command_modules/storage/_params.py

@@ -1952,27 +1952,35 @@ def load_arguments(self, _):  # pylint: disable=too-many-locals, too-many-statem
         c.extra('timeout', help='Request timeout in seconds. Applies to each call to the service.', type=int)
 
     with self.argument_context('storage file upload') as c:
+        from ._validators import add_progress_callback_v2
         t_file_content_settings = self.get_sdk('file.models#ContentSettings')
 
         c.register_path_argument(default_file_param='local_file_path')
         c.register_content_settings_argument(t_file_content_settings, update=False, guess_from_file='local_file_path')
-        c.argument('local_file_path', options_list='--source', type=file_type, completer=FilesCompleter())
-        c.extra('no_progress', progress_type)
-        c.argument('max_connections', type=int)
+        c.argument('local_file_path', options_list='--source', type=file_type, completer=FilesCompleter(),
+                   help='Path of the local file to upload as the file content.')
+        c.extra('no_progress', progress_type, validator=add_progress_callback_v2)
+        c.argument('max_connections', type=int, help='Maximum number of parallel connections to use.')
+        c.extra('share_name', share_name_type, required=True)
+        c.argument('validate_content', action='store_true', min_api='2016-05-31',
+                   help='If true, calculates an MD5 hash for each range of the file. The storage service checks the '
+                        'hash of the content that has arrived with the hash that was sent. This is primarily valuable '
+                        'for detecting bitflips on the wire if using http instead of https as https (the default) will '
+                        'already validate. Note that this MD5 hash is not stored with the file.')
 
     with self.argument_context('storage file url') as c:
         c.register_path_argument(fileshare=True)
         c.extra('share_name', share_name_type, required=True)
         c.argument('protocol', arg_type=get_enum_type(['http', 'https'], 'https'), help='Protocol to use.')
 
     with self.argument_context('storage file upload-batch') as c:
-        from ._validators import process_file_upload_batch_parameters
+        from ._validators import process_file_upload_batch_parameters, add_progress_callback_v2
         c.argument('source', options_list=('--source', '-s'), validator=process_file_upload_batch_parameters)
         c.argument('destination', options_list=('--destination', '-d'))
         c.argument('max_connections', arg_group='Download Control', type=int)
         c.argument('validate_content', action='store_true', min_api='2016-05-31')
         c.register_content_settings_argument(t_file_content_settings, update=False, arg_group='Content Settings')
-        c.extra('no_progress', progress_type)
+        c.extra('no_progress', progress_type, validator=add_progress_callback_v2)
 
     with self.argument_context('storage file download-batch') as c:
         from ._validators import process_file_download_batch_parameters

src/azure-cli/azure/cli/command_modules/storage/_validators.py

@@ -1141,6 +1141,28 @@ def _update_progress(current, total):
         message = getattr(_update_progress, 'message', 'Alive')
         reuse = getattr(_update_progress, 'reuse', False)
 
+        if total:
+            hook.add(message=message, value=current, total_val=total)
+            if total == current and not reuse:
+                hook.end()
+    hook = cmd.cli_ctx.get_progress_controller(det=True)
+    _update_progress.hook = hook
+
+    if not namespace.no_progress:
+        namespace.progress_callback = _update_progress
+    del namespace.no_progress
+
+
+def add_progress_callback_v2(cmd, namespace):
+    def _update_progress(response):
+        if response.http_response.status_code not in [200, 201]:
+            return
+
+        message = getattr(_update_progress, 'message', 'Alive')
+        reuse = getattr(_update_progress, 'reuse', False)
+        current = response.context['upload_stream_current']
+        total = response.context['data_stream_total']
+
         if total:
             hook.add(message=message, value=current, total_val=total)
             if total == current and not reuse:
@@ -1281,6 +1303,7 @@ def process_file_upload_batch_parameters(cmd, namespace):
             namespace.account_name = identifier.account_name
 
     namespace.source = os.path.realpath(namespace.source)
+    namespace.share_name = namespace.destination
 
 
 def process_file_download_batch_parameters(cmd, namespace):

src/azure-cli/azure/cli/command_modules/storage/commands.py

@@ -672,8 +672,6 @@ def get_custom_sdk(custom_module, client_factory, resource_type=ResourceType.DAT
         from ._format import transform_boolean_for_table, transform_file_output
         from ._exception_handler import file_related_exception_handler
         g.storage_command('download', 'get_file_to_path', exception_handler=file_related_exception_handler)
-        g.storage_command('upload', 'create_file_from_path', exception_handler=file_related_exception_handler)
-        g.storage_custom_command('upload-batch', 'storage_file_upload_batch')
         g.storage_custom_command(
             'download-batch', 'storage_file_download_batch')
         g.storage_custom_command('delete-batch', 'storage_file_delete_batch')
@@ -703,6 +701,10 @@ def get_custom_sdk(custom_module, client_factory, resource_type=ResourceType.DAT
         g.storage_command('copy cancel', 'abort_copy')
         g.storage_custom_command('copy start-batch', 'storage_file_copy_batch', client_factory=cf_share_client)
 
+        g.storage_custom_command('upload', 'storage_file_upload', exception_handler=file_related_exception_handler)
+        g.storage_custom_command('upload-batch', 'storage_file_upload_batch',
+                                 custom_command_type=get_custom_sdk('file', client_factory=cf_share_client))
+
     with self.command_group('storage cors', get_custom_sdk('cors', multi_service_properties_factory)) as g:
         from ._transformers import transform_cors_list_output
 

src/azure-cli/azure/cli/command_modules/storage/operations/file.py

@@ -137,6 +137,30 @@ def list_share_files(cmd, client, directory_name=None, timeout=None, exclude_dir
     return results
 
 
+def storage_file_upload(client, local_file_path, content_settings=None,
+                        metadata=None, validate_content=False, progress_callback=None, max_connections=2, timeout=None):
+    upload_args = {
+        'content_settings': content_settings,
+        'metadata': metadata,
+        'validate_content': validate_content,
+        'max_concurrency': max_connections,
+        'timeout': timeout
+    }
+    if progress_callback:
+        upload_args['raw_response_hook'] = progress_callback
+    # Because the contents of the uploaded file may be too large, it should be passed into the a stream object,
+    # upload_file() read file data in batches to avoid OOM problems
+    count = os.path.getsize(local_file_path)
+    with open(local_file_path, 'rb') as stream:
+        response = client.upload_file(data=stream, length=count, **upload_args)
+
+    if 'content_md5' in response:
+        if isinstance(response['content_md5'], bytearray):
+            response['content_md5'] = ''.join(hex(x) for x in response['content_md5'])
+
+    return response
+
+
 def storage_file_upload_batch(cmd, client, destination, source, destination_path=None, pattern=None, dryrun=False,
                               validate_content=False, content_settings=None, max_connections=1, metadata=None,
                               progress_callback=None):
@@ -145,37 +169,41 @@ def storage_file_upload_batch(cmd, client, destination, source, destination_path
     from azure.cli.command_modules.storage.util import glob_files_locally, normalize_blob_file_path
 
     source_files = list(glob_files_locally(source, pattern))
-    settings_class = cmd.get_models('file.models#ContentSettings')
+    settings_class = cmd.get_models('_models#ContentSettings')
 
     if dryrun:
         logger.info('upload files to file share')
         logger.info('    account %s', client.account_name)
         logger.info('      share %s', destination)
         logger.info('      total %d', len(source_files))
-        return [{'File': client.make_file_url(destination, os.path.dirname(dst) or None, os.path.basename(dst)),
+        dst = None
+        kwargs = {
+            'dir_name': os.path.dirname(dst),
+            'file_name': os.path.basename(dst)
+        }
+
+        return [{'File': create_file_url(client, **kwargs),
                  'Type': guess_content_type(src, content_settings, settings_class).content_type} for src, dst in
                 source_files]
 
     # TODO: Performance improvement
     # 1. Upload files in parallel
-    def _upload_action(src, dst):
-        dst = normalize_blob_file_path(destination_path, dst)
-        dir_name = os.path.dirname(dst)
-        file_name = os.path.basename(dst)
+    def _upload_action(src, dst2):
+        dst2 = normalize_blob_file_path(destination_path, dst2)
+        dir_name = os.path.dirname(dst2)
+        file_name = os.path.basename(dst2)
 
-        _make_directory_in_files_share(client, destination, dir_name)
-        create_file_args = {'share_name': destination, 'directory_name': dir_name, 'file_name': file_name,
-                            'local_file_path': src, 'progress_callback': progress_callback,
-                            'content_settings': guess_content_type(src, content_settings, settings_class),
-                            'metadata': metadata, 'max_connections': max_connections}
-
-        if cmd.supported_api_version(min_api='2016-05-31'):
-            create_file_args['validate_content'] = validate_content
+        _make_directory_in_files_share(client, destination, dir_name, V2=True)
 
         logger.warning('uploading %s', src)
-        client.create_file_from_path(**create_file_args)
+        storage_file_upload(client.get_file_client(dst2), src, content_settings, metadata, validate_content,
+                            progress_callback, max_connections)
 
-        return client.make_file_url(destination, dir_name, file_name)
+        args = {
+            'dir_name': dir_name,
+            'file_name': file_name
+        }
+        return create_file_url(client, **args)
 
     return list(_upload_action(src, dst) for src, dst in source_files)
 
@@ -384,7 +412,6 @@ def _create_file_and_directory_from_file(cmd, file_service, source_file_service,
 def _make_directory_in_files_share(file_service, file_share, directory_path, existing_dirs=None, V2=False):
     """
     Create directories recursively.
-
     This method accept a existing_dirs set which serves as the cache of existing directory. If the
     parameter is given, the method will search the set first to avoid repeatedly create directory
     which already exists.