[Storage] az storage blob copy start: Support --tier and --rehydrate-priority

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

@@ -642,7 +642,7 @@
 
 helps['storage blob copy start'] = """
 type: command
-short-summary: Copies a blob asynchronously. Use `az storage blob show` to check the status of the blobs.
+short-summary: Copy a blob asynchronously. Use `az storage blob show` to check the status of the blobs.
 parameters:
   - name: --source-uri -u
     type: string
@@ -656,15 +656,49 @@
         `https://myaccount.blob.core.windows.net/mycontainer/myblob`,
         `https://myaccount.blob.core.windows.net/mycontainer/myblob?snapshot=<DateTime>`,
         `https://otheraccount.blob.core.windows.net/mycontainer/myblob?sastoken`
+  - name: --destination-if-modified-since
+    type: string
+    short-summary: >
+        A DateTime value. Azure expects the date value passed in to be UTC.
+        If timezone is included, any non-UTC datetimes will be converted to UTC.
+        If a date is passed in without timezone info, it is assumed to be UTC.
+        Specify this conditional header to copy the blob only
+        if the destination blob has been modified since the specified date/time.
+        If the destination blob has not been modified, the Blob service returns
+        status code 412 (Precondition Failed).
+  - name: --destination-if-unmodified-since
+    type: string
+    short-summary: >
+        A DateTime value. Azure expects the date value passed in to be UTC.
+        If timezone is included, any non-UTC datetimes will be converted to UTC.
+        If a date is passed in without timezone info, it is assumed to be UTC.
+        Specify this conditional header to copy the blob only
+        if the destination blob has not been modified since the specified
+        date/time. If the destination blob has been modified, the Blob service
+        returns status code 412 (Precondition Failed).
+  - name: --source-if-modified-since
+    type: string
+    short-summary: >
+        A DateTime value. Azure expects the date value passed in to be UTC.
+        If timezone is included, any non-UTC datetimes will be converted to UTC.
+        If a date is passed in without timezone info, it is assumed to be UTC.
+        Specify this conditional header to copy the blob only if the source
+        blob has been modified since the specified date/time.
+  - name: --source-if-unmodified-since
+    type: string
+    short-summary: >
+        A DateTime value. Azure expects the date value passed in to be UTC.
+        If timezone is included, any non-UTC datetimes will be converted to UTC.
+        If a date is passed in without timezone info, it is assumed to be UTC.
+        Specify this conditional header to copy the blob only if the source blob
+        has not been modified since the specified date/time.
 examples:
-  - name: Copies a blob asynchronously. Use `az storage blob show` to check the status of the blobs. (autogenerated)
+  - name: Copy a blob asynchronously. Use `az storage blob show` to check the status of the blobs.
     text: |
         az storage blob copy start --account-key 00000000 --account-name MyAccount --destination-blob MyDestinationBlob --destination-container MyDestinationContainer --source-uri https://storage.blob.core.windows.net/photos
-    crafted: true
-  - name: Copies a blob asynchronously. Use `az storage blob show` to check the status of the blobs (autogenerated)
+  - name: Copy a blob asynchronously. Use `az storage blob show` to check the status of the blobs.
     text: |
         az storage blob copy start --account-name MyAccount --destination-blob MyDestinationBlob --destination-container MyDestinationContainer --sas-token $sas --source-uri https://storage.blob.core.windows.net/photos
-    crafted: true
 """
 
 helps['storage blob copy start-batch'] = """

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

@@ -186,9 +186,6 @@ def load_arguments(self, _):  # pylint: disable=too-many-locals, too-many-statem
         help='Expiration period in days of the Key Policy assigned to the storage account'
     )
 
-    t_blob_tier = self.get_sdk('_generated.models._azure_blob_storage_enums#AccessTierOptional',
-                               resource_type=ResourceType.DATA_STORAGE_BLOB)
-
     allow_cross_tenant_replication_type = CLIArgumentType(
         arg_type=get_three_state_flag(), options_list=['--allow-cross-tenant-replication', '-r'], min_api='2021-04-01',
         help='Allow or disallow cross AAD tenant object replication. The default interpretation is true for this '
@@ -201,6 +198,22 @@ def load_arguments(self, _):  # pylint: disable=too-many-locals, too-many-statem
         arg_group='Azure Files Identity Based Authentication',
         help='Default share permission for users using Kerberos authentication if RBAC role is not assigned.')
 
+    t_blob_tier = self.get_sdk('_generated.models._azure_blob_storage_enums#AccessTierOptional',
+                               resource_type=ResourceType.DATA_STORAGE_BLOB)
+    t_rehydrate_priority = self.get_sdk('_generated.models._azure_blob_storage_enums#RehydratePriority',
+                                        resource_type=ResourceType.DATA_STORAGE_BLOB)
+    tier_type = CLIArgumentType(
+        arg_type=get_enum_type(t_blob_tier), min_api='2019-02-02',
+        help='The tier value to set the blob to. For page blob, the tier correlates to the size of the blob '
+             'and number of allowed IOPS. Possible values are P10, P15, P20, P30, P4, P40, P50, P6, P60, P70, P80 '
+             'and this is only applicable to page blobs on premium storage accounts; For block blob, possible '
+             'values are Archive, Cool and Hot. This is only applicable to block blobs on standard storage accounts.'
+    )
+    rehydrate_priority_type = CLIArgumentType(
+        arg_type=get_enum_type(t_rehydrate_priority), options_list=('--rehydrate-priority', '-r'),
+        min_api='2019-02-02',
+        help='Indicate the priority with which to rehydrate an archived blob.')
+
     with self.argument_context('storage') as c:
         c.argument('container_name', container_name_type)
         c.argument('directory_name', directory_type)
@@ -964,13 +977,37 @@ def load_arguments(self, _):  # pylint: disable=too-many-locals, too-many-statem
                    help='Name of the destination blob. If the exists, it will be overwritten.')
         c.argument('source_lease_id', arg_group='Copy Source')
 
-    with self.argument_context('storage blob copy start') as c:
-        from azure.cli.command_modules.storage._validators import validate_source_uri
+    with self.argument_context('storage blob copy start', resource_type=ResourceType.DATA_STORAGE_BLOB) as c:
+        from ._validators import validate_source_url
 
-        c.register_source_uri_arguments(validator=validate_source_uri)
-        c.argument('requires_sync', arg_type=get_three_state_flag(),
-                   help='Enforce that the service will not return a response until the copy is complete.'
-                   'Not support for standard page blob.')
+        c.register_blob_arguments()
+        c.register_precondition_options()
+        c.register_precondition_options(prefix='source_')
+        c.register_source_uri_arguments(validator=validate_source_url)
+
+        c.ignore('incremental_copy')
+        c.argument('if_match', options_list=['--destination-if-match'])
+        c.argument('if_modified_since', options_list=['--destination-if-modified-since'])
+        c.argument('if_none_match', options_list=['--destination-if-none-match'])
+        c.argument('if_unmodified_since', options_list=['--destination-if-unmodified-since'])
+        c.argument('if_tags_match_condition', options_list=['--destination-tags-condition'])
+
+        c.argument('blob_name', options_list=['--destination-blob', '-b'], required=True,
+                   help='Name of the destination blob. If the exists, it will be overwritten.')
+        c.argument('container_name', options_list=['--destination-container', '-c'], required=True,
+                   help='The container name.')
+        c.extra('destination_lease', options_list='--destination-lease-id',
+                help='The lease ID specified for this header must match the lease ID of the estination blob. '
+                'If the request does not include the lease ID or it is not valid, the operation fails with status '
+                'code 412 (Precondition Failed).')
+        c.extra('source_lease', options_list='--source-lease-id', arg_group='Copy Source',
+                help='Specify this to perform the Copy Blob operation only if the lease ID given matches the '
+                'active lease ID of the source blob.')
+        c.extra('rehydrate_priority', rehydrate_priority_type)
+        c.extra('requires_sync', arg_type=get_three_state_flag(),
+                help='Enforce that the service will not return a response until the copy is complete.')
+        c.extra('tier', tier_type)
+        c.extra('tags', tags_type)
 
     with self.argument_context('storage blob copy start-batch', arg_group='Copy Source') as c:
         from azure.cli.command_modules.storage._validators import get_source_file_or_blob_service_client

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

@@ -387,6 +387,111 @@ def validate_source_uri(cmd, namespace):  # pylint: disable=too-many-statements
     namespace.copy_source = uri
 
 
+def validate_source_url(cmd, namespace):  # pylint: disable=too-many-statements
+    from .util import create_short_lived_blob_sas, create_short_lived_file_sas
+    usage_string = \
+        'Invalid usage: {}. Supply only one of the following argument sets to specify source:' \
+        '\n\t   --source-uri [--source-sas]' \
+        '\n\tOR --source-container --source-blob [--source-account-name & sas] [--source-snapshot]' \
+        '\n\tOR --source-container --source-blob [--source-account-name & key] [--source-snapshot]' \
+        '\n\tOR --source-share --source-path' \
+        '\n\tOR --source-share --source-path [--source-account-name & sas]' \
+        '\n\tOR --source-share --source-path [--source-account-name & key]'
+
+    ns = vars(namespace)
+
+    # source as blob
+    container = ns.pop('source_container', None)
+    blob = ns.pop('source_blob', None)
+    snapshot = ns.pop('source_snapshot', None)
+
+    # source as file
+    share = ns.pop('source_share', None)
+    path = ns.pop('source_path', None)
+    file_snapshot = ns.pop('file_snapshot', None)
+
+    # source credential clues
+    source_account_name = ns.pop('source_account_name', None)
+    source_account_key = ns.pop('source_account_key', None)
+    source_sas = ns.pop('source_sas', None)
+
+    # source in the form of an uri
+    uri = ns.get('source_url', None)
+    if uri:
+        if any([container, blob, snapshot, share, path, file_snapshot, source_account_name,
+                source_account_key]):
+            raise ValueError(usage_string.format('Unused parameters are given in addition to the '
+                                                 'source URI'))
+        if source_sas:
+            source_sas = source_sas.lstrip('?')
+            uri = '{}{}{}'.format(uri, '?', source_sas)
+            namespace.copy_source = uri
+        return
+
+    # ensure either a file or blob source is specified
+    valid_blob_source = container and blob and not share and not path and not file_snapshot
+    valid_file_source = share and path and not container and not blob and not snapshot
+
+    if not valid_blob_source and not valid_file_source:
+        raise ValueError(usage_string.format('Neither a valid blob or file source is specified'))
+    if valid_blob_source and valid_file_source:
+        raise ValueError(usage_string.format('Ambiguous parameters, both blob and file sources are '
+                                             'specified'))
+
+    validate_client_parameters(cmd, namespace)  # must run first to resolve storage account
+
+    if not source_account_name:
+        if source_account_key:
+            raise ValueError(usage_string.format('Source account key is given but account name is not'))
+        # assume that user intends to copy blob in the same account
+        source_account_name = ns.get('account_name', None)
+
+    # determine if the copy will happen in the same storage account
+    same_account = False
+
+    if not source_account_key and not source_sas:
+        if source_account_name == ns.get('account_name', None):
+            same_account = True
+            source_account_key = ns.get('account_key', None)
+            source_sas = ns.get('sas_token', None)
+        else:
+            # the source account is different from destination account but the key is missing try to query one.
+            try:
+                source_account_key = _query_account_key(cmd.cli_ctx, source_account_name)
+            except ValueError:
+                raise ValueError('Source storage account {} not found.'.format(source_account_name))
+
+    # Both source account name and either key or sas (or both) are now available
+    if not source_sas:
+        # generate a sas token even in the same account when the source and destination are not the same kind.
+        if valid_file_source and (ns.get('container_name', None) or not same_account):
+            import os
+            dir_name, file_name = os.path.split(path) if path else (None, '')
+            source_sas = create_short_lived_file_sas(cmd, source_account_name, source_account_key, share,
+                                                     dir_name, file_name)
+        elif valid_blob_source and (ns.get('share_name', None) or not same_account):
+            source_sas = create_short_lived_blob_sas(cmd, source_account_name, source_account_key, container, blob)
+
+    query_params = []
+    if source_sas:
+        query_params.append(source_sas.lstrip('?'))
+    if snapshot:
+        query_params.append('snapshot={}'.format(snapshot))
+    if file_snapshot:
+        query_params.append('sharesnapshot={}'.format(file_snapshot))
+
+    uri = 'https://{0}.{1}.{6}/{2}/{3}{4}{5}'.format(
+        source_account_name,
+        'blob' if valid_blob_source else 'file',
+        container if valid_blob_source else share,
+        encode_for_url(blob if valid_blob_source else path),
+        '?' if query_params else '',
+        '&'.join(query_params),
+        cmd.cli_ctx.cloud.suffixes.storage_endpoint)
+
+    namespace.source_url = uri
+
+
 def validate_blob_type(namespace):
     if not namespace.blob_type:
         namespace.blob_type = 'page' if namespace.file_path.endswith('.vhd') else 'block'

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

@@ -312,6 +312,7 @@ def get_custom_sdk(custom_module, client_factory, resource_type=ResourceType.DAT
                                                                resource_type=ResourceType.DATA_STORAGE_BLOB)) as g:
         from ._transformers import transform_blob_list_output, transform_blob_json_output
         from ._format import transform_blob_output
+        g.storage_custom_command_oauth('copy start', 'copy_blob')
         g.storage_custom_command_oauth('show', 'show_blob_v2', transform=transform_blob_json_output,
                                        table_transformer=transform_blob_output,
                                        exception_handler=show_exception_handler)
@@ -376,7 +377,6 @@ def get_custom_sdk(custom_module, client_factory, resource_type=ResourceType.DAT
         g.storage_command_oauth(
             'metadata show', 'get_blob_metadata', exception_handler=show_exception_handler)
         g.storage_command_oauth('metadata update', 'set_blob_metadata')
-        g.storage_command_oauth('copy start', 'copy_blob')
         g.storage_command_oauth('copy cancel', 'abort_copy_blob')
         g.storage_custom_command_oauth(
             'copy start-batch', 'storage_blob_copy_batch')

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

@@ -837,3 +837,9 @@ def query_blob(client, query_expression, input_config=None, output_config=None,
         return None
 
     return reader.readall().decode("utf-8")
+
+
+def copy_blob(client, source_url, metadata=None, **kwargs):
+    if not kwargs['requires_sync']:
+        kwargs.pop('requires_sync')
+    return client.start_copy_from_url(source_url=source_url, metadata=metadata, incremental_copy=False, **kwargs)

src/azure-cli/azure/cli/command_modules/storage/tests/latest/test_storage_blob_copy_scenarios.py

@@ -86,7 +86,7 @@ def test_storage_blob_copy_same_account_sas(self, resource_group, storage_accoun
                                expiry).output.strip()
 
         self.storage_cmd('storage blob copy start -b dst -c {} --source-blob src --sas-token {} --source-container {} '
-                         '--source-if-unmodified-since "2020-06-29T06:32Z" --destination-if-modified-since '
+                         '--source-if-unmodified-since "2021-06-29T06:32Z" --destination-if-modified-since '
                          '"2020-06-29T06:32Z" ', account_info, target_container, sas, source_container)
 
         from time import sleep, time

src/azure-cli/azure/cli/command_modules/storage/tests/latest/test_storage_blob_scenarios.py

@@ -755,5 +755,31 @@ def test_storage_blob_restore(self, resource_group, storage_account):
             time_to_restore, storage_account, resource_group))
 
 
+class StorageBlobCopyTestScenario(StorageScenarioMixin, ScenarioTest):
+    @ResourceGroupPreparer()
+    @StorageAccountPreparer(kind='storageV2')
+    def test_storage_blob_copy_rehydrate_priority(self, resource_group, storage_account):
+        source_file = self.create_temp_file(16)
+        account_info = self.get_account_info(resource_group, storage_account)
+
+        source_container = self.create_container(account_info)
+        target_container = self.create_container(account_info)
+
+        self.storage_cmd('storage blob upload -c {} -f "{}" -n src ', account_info,
+                         source_container, source_file)
+        self.storage_cmd('storage blob set-tier -c {} -n {} --tier Archive', account_info,
+                         source_container, 'src')
+        self.storage_cmd('az storage blob show -c {} -n {} ', account_info, source_container, 'src') \
+            .assert_with_checks(JMESPathCheck('properties.blobTier', 'Archive'))
+
+        source_uri = self.storage_cmd('storage blob url -c {} -n src', account_info, source_container).output
+
+        self.storage_cmd('storage blob copy start -b dst -c {} --source-uri {} --tier Cool -r High', account_info,
+                         target_container, source_uri)
+        self.storage_cmd('storage blob show -c {} -n {} ', account_info, target_container, 'dst') \
+            .assert_with_checks(JMESPathCheck('properties.blobTier', 'Archive'),
+                                JMESPathCheck('properties.rehydrationStatus', 'rehydrate-pending-to-cool'))
+
+
 if __name__ == '__main__':
     unittest.main()