Migrate code

Author: william051200

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

@@ -24,7 +24,7 @@
     validate_asg_names_or_ids, validate_keyvault, _validate_proximity_placement_group,
     validate_vm_name_for_monitor_metrics)
 
-from azure.cli.command_modules.vm._vm_utils import MSI_LOCAL_ID
+from azure.cli.command_modules.vm._vm_utils import MSI_LOCAL_ID, CachingTypes
 from azure.cli.command_modules.vm._image_builder import ScriptType
 
 from azure.cli.command_modules.monitor.validators import validate_metric_dimension
@@ -37,7 +37,7 @@ def load_arguments(self, _):
     # Model imports
     DiskStorageAccountTypes = self.get_models('DiskStorageAccountTypes', operation_group='disks')
     SnapshotStorageAccountTypes = self.get_models('SnapshotStorageAccountTypes', operation_group='snapshots')
-    UpgradeMode, CachingTypes, OperatingSystemTypes = self.get_models('UpgradeMode', 'CachingTypes', 'OperatingSystemTypes')
+    UpgradeMode, OperatingSystemTypes = self.get_models('UpgradeMode', 'OperatingSystemTypes')
     HyperVGenerationTypes = self.get_models('HyperVGenerationTypes')
     DedicatedHostLicenseTypes = self.get_models('DedicatedHostLicenseTypes')
     OrchestrationServiceNames, OrchestrationServiceStateAction = self.get_models('OrchestrationServiceNames', 'OrchestrationServiceStateAction', operation_group='virtual_machine_scale_sets')

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

@@ -2161,27 +2161,33 @@ def process_image_create_namespace(cmd, namespace):
                                   'virtualMachines', 'Microsoft.Compute')
         res = parse_resource_id(res_id)
         if res['type'] == 'virtualMachines':
-            compute_client = _compute_client_factory(cmd.cli_ctx, subscription_id=res['subscription'])
-            vm_info = compute_client.virtual_machines.get(res['resource_group'], res['name'])
+            from .operations.vm import VMShow
+            command_args = {
+                'subscription': res['subscription'],
+                'resource_group': res['resource_group'],
+                'vm_name': res['name']
+            }
+            vm_info = VMShow(cli_ctx=cmd.cli_ctx)(command_args=command_args)
             source_from_vm = True
     except ResourceNotFoundError:
         pass
 
     if source_from_vm:
         # pylint: disable=no-member
-        namespace.os_type = vm_info.storage_profile.os_disk.os_type
+        namespace.os_type = vm_info.get('storageProfile', {}).get('osDisk', {}).get('osType', '')
         namespace.source_virtual_machine = res_id
         if namespace.data_disk_sources:
             raise CLIError("'--data-disk-sources' is not allowed when capturing "
                            "images from virtual machines")
     else:
-        namespace.os_blob_uri, namespace.os_disk, namespace.os_snapshot, _, _ = _figure_out_storage_source(cmd.cli_ctx, namespace.resource_group_name, namespace.source)  # pylint: disable=line-too-long
+        namespace.os_blob_uri, namespace.os_disk, namespace.os_snapshot, _, _ = \
+            _figure_out_storage_source_by_aaz(cmd.cli_ctx, namespace.resource_group_name, namespace.source)
         namespace.data_blob_uris = []
         namespace.data_disks = []
         namespace.data_snapshots = []
         if namespace.data_disk_sources:
             for data_disk_source in namespace.data_disk_sources:
-                source_blob_uri, source_disk, source_snapshot, _, _ = _figure_out_storage_source(
+                source_blob_uri, source_disk, source_snapshot, _, _ = _figure_out_storage_source_by_aaz(
                     cmd.cli_ctx, namespace.resource_group_name, data_disk_source)
                 if source_blob_uri:
                     namespace.data_blob_uris.append(source_blob_uri)
@@ -2218,6 +2224,30 @@ def _figure_out_storage_source(cli_ctx, resource_group_name, source):
     return (source_blob_uri, source_disk, source_snapshot, source_restore_point, source_info)
 
 
+def _figure_out_storage_source_by_aaz(cli_ctx, resource_group_name, source):
+    source_blob_uri = None
+    source_disk = None
+    source_snapshot = None
+    source_info = None
+    source_restore_point = None
+    if urlparse(source).scheme:  # a uri?
+        source_blob_uri = source
+    elif '/disks/' in source.lower():
+        source_disk = source
+    elif '/snapshots/' in source.lower():
+        source_snapshot = source
+    elif '/restorepoints/' in source.lower():
+        source_restore_point = source
+    else:
+        source_info, is_snapshot = _get_disk_or_snapshot_info_by_aaz(cli_ctx, resource_group_name, source)
+        if is_snapshot:
+            source_snapshot = source_info.get('id')
+        else:
+            source_disk = source_info.get('id')
+
+    return (source_blob_uri, source_disk, source_snapshot, source_restore_point, source_info)
+
+
 def _get_disk_or_snapshot_info(cli_ctx, resource_group_name, source):
     compute_client = _compute_client_factory(cli_ctx)
     is_snapshot = True
@@ -2231,6 +2261,28 @@ def _get_disk_or_snapshot_info(cli_ctx, resource_group_name, source):
     return info, is_snapshot
 
 
+def _get_disk_or_snapshot_info_by_aaz(cli_ctx, resource_group_name, source):
+    from .aaz.latest.snapshot import Show as SnapshotShow
+    from .aaz.latest.disk import Show as DiskShow
+    is_snapshot = True
+
+    try:
+        command_args = {
+            'resource_group': resource_group_name,
+            'snapshot_name': source
+        }
+        info = SnapshotShow(cli_ctx=cli_ctx)(command_args=command_args)
+    except ResourceNotFoundError:
+        command_args = {
+            'resource_group': resource_group_name,
+            'disk_name': source
+        }
+        is_snapshot = False
+        info = DiskShow(cli_ctx=cli_ctx)(command_args=command_args)
+
+    return info, is_snapshot
+
+
 def process_disk_encryption_namespace(cmd, namespace):
     namespace.disk_encryption_keyvault = _get_resource_id(cmd.cli_ctx, namespace.disk_encryption_keyvault,
                                                           namespace.resource_group_name,

src/azure-cli/azure/cli/command_modules/vm/_vm_utils.py

@@ -805,6 +805,12 @@ class VMGuestPatchClassificationLinux(Enum):
     SECURITY = 'Security'
 
 
+class CachingTypes(Enum):
+    NONE = 'None'
+    READ_ONLY = 'ReadOnly'
+    READ_WRITE = 'ReadWrite'
+
+
 class DiskCreateOptionTypes(Enum):
     ATTACH = 'Attach'
     COPY = 'Copy'

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

@@ -221,7 +221,7 @@ def load_command_table(self, _):
         self.command_table['disk-encryption-set identity remove'] = DiskEncryptionSetIdentityRemove(loader=self)
         g.custom_show_command('show', 'show_disk_encryption_set_identity')
 
-    with self.command_group('image', compute_image_sdk) as g:
+    with self.command_group('image') as g:
         g.custom_command('create', 'create_image', validator=process_image_create_namespace)
 
     with self.command_group('image builder', image_builder_image_templates_sdk, custom_command_type=image_builder_custom) as g: