@@ -2169,27 +2169,33 @@ def process_image_create_namespace(cmd, namespace):
21692169 'virtualMachines' , 'Microsoft.Compute' )
21702170 res = parse_resource_id (res_id )
21712171 if res ['type' ] == 'virtualMachines' :
2172- compute_client = _compute_client_factory (cmd .cli_ctx , subscription_id = res ['subscription' ])
2173- vm_info = compute_client .virtual_machines .get (res ['resource_group' ], res ['name' ])
2172+ from .operations .vm import VMShow
2173+ command_args = {
2174+ 'subscription' : res ['subscription' ],
2175+ 'resource_group' : res ['resource_group' ],
2176+ 'vm_name' : res ['name' ]
2177+ }
2178+ vm_info = VMShow (cli_ctx = cmd .cli_ctx )(command_args = command_args )
21742179 source_from_vm = True
21752180 except ResourceNotFoundError :
21762181 pass
21772182
21782183 if source_from_vm :
21792184 # pylint: disable=no-member
2180- namespace .os_type = vm_info .storage_profile . os_disk . os_type
2185+ namespace .os_type = vm_info .get ( 'storageProfile' , {}). get ( 'osDisk' , {}). get ( 'osType' )
21812186 namespace .source_virtual_machine = res_id
21822187 if namespace .data_disk_sources :
21832188 raise CLIError ("'--data-disk-sources' is not allowed when capturing "
21842189 "images from virtual machines" )
21852190 else :
2186- 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
2191+ namespace .os_blob_uri , namespace .os_disk , namespace .os_snapshot , _ , _ = \
2192+ _figure_out_storage_source_by_aaz (cmd .cli_ctx , namespace .resource_group_name , namespace .source )
21872193 namespace .data_blob_uris = []
21882194 namespace .data_disks = []
21892195 namespace .data_snapshots = []
21902196 if namespace .data_disk_sources :
21912197 for data_disk_source in namespace .data_disk_sources :
2192- source_blob_uri , source_disk , source_snapshot , _ , _ = _figure_out_storage_source (
2198+ source_blob_uri , source_disk , source_snapshot , _ , _ = _figure_out_storage_source_by_aaz (
21932199 cmd .cli_ctx , namespace .resource_group_name , data_disk_source )
21942200 if source_blob_uri :
21952201 namespace .data_blob_uris .append (source_blob_uri )
@@ -2226,6 +2232,30 @@ def _figure_out_storage_source(cli_ctx, resource_group_name, source):
22262232 return (source_blob_uri , source_disk , source_snapshot , source_restore_point , source_info )
22272233
22282234
2235+ def _figure_out_storage_source_by_aaz (cli_ctx , resource_group_name , source ):
2236+ source_blob_uri = None
2237+ source_disk = None
2238+ source_snapshot = None
2239+ source_info = None
2240+ source_restore_point = None
2241+ if urlparse (source ).scheme : # a uri?
2242+ source_blob_uri = source
2243+ elif '/disks/' in source .lower ():
2244+ source_disk = source
2245+ elif '/snapshots/' in source .lower ():
2246+ source_snapshot = source
2247+ elif '/restorepoints/' in source .lower ():
2248+ source_restore_point = source
2249+ else :
2250+ source_info , is_snapshot = _get_disk_or_snapshot_info_by_aaz (cli_ctx , resource_group_name , source )
2251+ if is_snapshot :
2252+ source_snapshot = source_info .get ('id' )
2253+ else :
2254+ source_disk = source_info .get ('id' )
2255+
2256+ return (source_blob_uri , source_disk , source_snapshot , source_restore_point , source_info )
2257+
2258+
22292259def _get_disk_or_snapshot_info (cli_ctx , resource_group_name , source ):
22302260 compute_client = _compute_client_factory (cli_ctx )
22312261 is_snapshot = True
@@ -2239,6 +2269,28 @@ def _get_disk_or_snapshot_info(cli_ctx, resource_group_name, source):
22392269 return info , is_snapshot
22402270
22412271
2272+ def _get_disk_or_snapshot_info_by_aaz (cli_ctx , resource_group_name , source ):
2273+ from .aaz .latest .snapshot import Show as SnapshotShow
2274+ from .aaz .latest .disk import Show as DiskShow
2275+ is_snapshot = True
2276+
2277+ try :
2278+ command_args = {
2279+ 'resource_group' : resource_group_name ,
2280+ 'snapshot_name' : source
2281+ }
2282+ info = SnapshotShow (cli_ctx = cli_ctx )(command_args = command_args )
2283+ except ResourceNotFoundError :
2284+ command_args = {
2285+ 'resource_group' : resource_group_name ,
2286+ 'disk_name' : source
2287+ }
2288+ is_snapshot = False
2289+ info = DiskShow (cli_ctx = cli_ctx )(command_args = command_args )
2290+
2291+ return info , is_snapshot
2292+
2293+
22422294def process_disk_encryption_namespace (cmd , namespace ):
22432295 namespace .disk_encryption_keyvault = _get_resource_id (cmd .cli_ctx , namespace .disk_encryption_keyvault ,
22442296 namespace .resource_group_name ,
0 commit comments