[Storage] Fix #21781 az storage blob upload/download: Progress fix

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

@@ -23,8 +23,7 @@
                           validate_fs_public_access, validate_logging_version, validate_or_policy, validate_policy,
                           get_api_version_type, blob_download_file_path_validator, blob_tier_validator, validate_subnet,
                           validate_immutability_arguments, validate_blob_name_for_upload, validate_share_close_handle,
-                          add_upload_progress_callback, blob_tier_validator_track2, add_download_progress_callback,
-                          services_type_v2, resource_type_type_v2)
+                          blob_tier_validator_track2, services_type_v2, resource_type_type_v2)
 
 
 def load_arguments(self, _):  # pylint: disable=too-many-locals, too-many-statements, too-many-lines, too-many-branches, line-too-long
@@ -968,7 +967,7 @@ def load_arguments(self, _):  # pylint: disable=too-many-locals, too-many-statem
         c.argument('blob_type', options_list=('--type', '-t'), validator=validate_blob_type,
                    arg_type=get_enum_type(get_blob_types()), arg_group="Additional Flags")
         c.argument('validate_content', action='store_true', min_api='2016-05-31', arg_group="Content Control")
-        c.extra('no_progress', progress_type, validator=add_upload_progress_callback, arg_group="Additional Flags")
+        c.extra('no_progress', progress_type, validator=add_progress_callback, arg_group="Additional Flags")
         c.extra('tier', tier_type, validator=blob_tier_validator_track2, arg_group="Additional Flags")
         c.argument('encryption_scope', validator=validate_encryption_scope_client_params,
                    help='A predefined encryption scope used to encrypt the data on the service.',
@@ -994,7 +993,7 @@ def load_arguments(self, _):  # pylint: disable=too-many-locals, too-many-statem
         c.argument('maxsize_condition', arg_group='Content Control')
         c.argument('validate_content', action='store_true', min_api='2016-05-31', arg_group='Content Control')
         c.argument('blob_type', options_list=('--type', '-t'), arg_type=get_enum_type(get_blob_types()))
-        c.extra('no_progress', progress_type, validator=add_upload_progress_callback)
+        c.extra('no_progress', progress_type, validator=add_progress_callback)
         c.extra('tier', tier_type, is_preview=True, validator=blob_tier_validator_track2)
         c.extra('overwrite', arg_type=get_three_state_flag(), is_preview=True,
                 help='Whether the blob to be uploaded should overwrite the current data. If True, blob upload '
@@ -1018,7 +1017,7 @@ def load_arguments(self, _):  # pylint: disable=too-many-locals, too-many-statem
                    help='End of byte range to use for downloading a section of the blob. If end_range is given, '
                         'start_range must be provided. The start_range and end_range params are inclusive. '
                         'Ex: start_range=0, end_range=511 will download first 512 bytes of blob.')
-        c.extra('no_progress', progress_type, validator=add_download_progress_callback)
+        c.extra('no_progress', progress_type, validator=add_progress_callback)
         c.extra('snapshot', help='The snapshot parameter is an opaque DateTime value that, when present, '
                                  'specifies the blob snapshot to retrieve.')
         c.extra('lease', options_list=['--lease-id'], help='Required if the blob has an active lease.')

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

@@ -1055,7 +1055,7 @@ def process_blob_download_batch_parameters(cmd, namespace):
     _process_blob_batch_container_parameters(cmd, namespace)
 
     # 3. Call validators
-    add_download_progress_callback(cmd, namespace)
+    add_progress_callback(cmd, namespace)
 
 
 def process_blob_upload_batch_parameters(cmd, namespace):
@@ -1094,7 +1094,7 @@ def process_blob_upload_batch_parameters(cmd, namespace):
     validate_metadata(namespace)
     t_blob_content_settings = get_sdk(cmd.cli_ctx, ResourceType.DATA_STORAGE_BLOB, '_models#ContentSettings')
     get_content_setting_validator(t_blob_content_settings, update=False)(cmd, namespace)
-    add_upload_progress_callback(cmd, namespace)
+    add_progress_callback(cmd, namespace)
     blob_tier_validator_track2(cmd, namespace)
 
 
@@ -2046,52 +2046,6 @@ def validate_upload_blob(namespace):
         raise InvalidArgumentValueError("usage error: please specify one of --file and --data to upload.")
 
 
-def add_upload_progress_callback(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:
-                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_download_progress_callback(cmd, namespace):
-    def _update_progress(response):
-        if response.http_response.status_code not in [200, 201, 206]:
-            return
-
-        message = getattr(_update_progress, 'message', 'Alive')
-        reuse = getattr(_update_progress, 'reuse', False)
-        current = response.context['download_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:
-                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 validate_blob_arguments(namespace):
     from azure.cli.core.azclierror import RequiredArgumentMissingError
     if not namespace.blob_url and not all([namespace.blob_name, namespace.container_name]):

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

@@ -568,7 +568,7 @@ def upload_blob(cmd, client, file_path=None, container_name=None, blob_name=None
         upload_args['validate_content'] = validate_content
 
     if progress_callback:
-        upload_args['raw_response_hook'] = progress_callback
+        upload_args['progress_hook'] = progress_callback
 
     check_blob_args = {
         'if_modified_since': if_modified_since,
@@ -631,7 +631,7 @@ def download_blob(client, file_path=None, open_mode='wb', start_range=None, end_
         offset = start_range
         length = end_range - start_range + 1
     if progress_callback:
-        kwargs['raw_response_hook'] = progress_callback
+        kwargs['progress_hook'] = progress_callback
     if not file_path:
         kwargs['max_concurrency'] = 1
     download_stream = client.download_blob(offset=offset, length=length, **kwargs)