Adding feature to export snapshots instead of images. (#583)

* Adding feature to export snapshots instead of images. This allows for non-generalized osStates to be copied and deployed across regions. * minor linting fixes * Update create_target.py hacking things * Adding new-line. More Linting issues.
Azure · Apr 2, 2019 · 5a25c81 · 5a25c81
1 parent 208df3d
commit 5a25c81
Show file tree

Hide file tree

Showing 3 changed files with 47 additions and 30 deletions.
diff --git a/src/image-copy/azext_imagecopy/__init__.py b/src/image-copy/azext_imagecopy/__init__.py
@@ -49,6 +49,8 @@ def load_arguments(self, _):
             c.argument('temporary_resource_group_name', options_list=['--temporary_resource_group_name'],
                        default='image-copy-rg',
                        help='Resource Group name where temporary storage account will be created.')
+            c.argument('export_as_snapshot', options_list=['--export-as-snapshot'], action='store_true', default=False,
+                       help='Include this switch to export the copies as snapshots instead of images.')
 
 
 COMMAND_LOADER_CLS = ImageCopyCommandsLoader
diff --git a/src/image-copy/azext_imagecopy/create_target.py b/src/image-copy/azext_imagecopy/create_target.py
@@ -14,13 +14,15 @@
 STORAGE_ACCOUNT_NAME_LENGTH = 24
 
 
+# pylint: disable=too-many-statements
 # pylint: disable=too-many-locals
 def create_target_image(location, transient_resource_group_name, source_type, source_object_name,
                         source_os_disk_snapshot_name, source_os_disk_snapshot_url, source_os_type,
                         target_resource_group_name, azure_pool_frequency, tags, target_name, target_subscription,
-                        timeout):
+                        export_as_snapshot, timeout):
 
-    random_string = get_random_string(STORAGE_ACCOUNT_NAME_LENGTH - len(location))
+    random_string = get_random_string(
+        STORAGE_ACCOUNT_NAME_LENGTH - len(location))
 
     # create the target storage account. storage account name must be lowercase.
     logger.warn(
@@ -49,7 +51,8 @@ def create_target_image(location, transient_resource_group_name, source_type, so
 
     expiry_format = "%Y-%m-%dT%H:%MZ"
     expiry = datetime.datetime.utcnow() + datetime.timedelta(seconds=timeout)
-    logger.debug("create target storage sas using timeout seconds: %d", timeout)
+    logger.debug(
+        "create target storage sas using timeout seconds: %d", timeout)
 
     cli_cmd = prepare_cli_command(['storage', 'account', 'generate-sas',
                                    '--account-name', target_storage_account_name,
@@ -103,8 +106,13 @@ def create_target_image(location, transient_resource_group_name, source_type, so
     target_blob_path = target_blob_endpoint + \
         target_container_name + '/' + blob_name
     target_snapshot_name = source_os_disk_snapshot_name + '-' + location
+    if export_as_snapshot:
+        snapshot_resource_group_name = target_resource_group_name
+    else:
+        snapshot_resource_group_name = transient_resource_group_name
+
     cli_cmd = prepare_cli_command(['snapshot', 'create',
-                                   '--resource-group', transient_resource_group_name,
+                                   '--resource-group', snapshot_resource_group_name,
                                    '--name', target_snapshot_name,
                                    '--location', location,
                                    '--source', target_blob_path],
@@ -113,27 +121,30 @@ def create_target_image(location, transient_resource_group_name, source_type, so
     json_output = run_cli_command(cli_cmd, return_as_json=True)
     target_snapshot_id = json_output['id']
 
-    # Create the final image
-    logger.warn("%s - Creating final image", location)
-    if target_name is None:
-        target_image_name = source_object_name
-        if source_type != 'image':
-            target_image_name += '-image'
-        target_image_name += '-' + location
+    # Optionally create the final image
+    if export_as_snapshot:
+        logger.warn("%s - Skipping image creation", location)
     else:
-        target_image_name = target_name
-
-    cli_cmd = prepare_cli_command(['image', 'create',
-                                   '--resource-group', target_resource_group_name,
-                                   '--name', target_image_name,
-                                   '--location', location,
-                                   '--source', target_blob_path,
-                                   '--os-type', source_os_type,
-                                   '--source', target_snapshot_id],
-                                  tags=tags,
-                                  subscription=target_subscription)
-
-    run_cli_command(cli_cmd)
+        logger.warn("%s - Creating final image", location)
+        if target_name is None:
+            target_image_name = source_object_name
+            if source_type != 'image':
+                target_image_name += '-image'
+            target_image_name += '-' + location
+        else:
+            target_image_name = target_name
+
+        cli_cmd = prepare_cli_command(['image', 'create',
+                                       '--resource-group', target_resource_group_name,
+                                       '--name', target_image_name,
+                                       '--location', location,
+                                       '--source', target_blob_path,
+                                       '--os-type', source_os_type,
+                                       '--source', target_snapshot_id],
+                                      tags=tags,
+                                      subscription=target_subscription)
+
+        run_cli_command(cli_cmd)
 
 
 def wait_for_blob_copy_operation(blob_name, target_container_name, target_storage_account_name,
@@ -167,7 +178,8 @@ def wait_for_blob_copy_operation(blob_name, target_container_name, target_storag
             return
 
     if copy_status != 'success':
-        logger.error("The copy operation didn't succeed. Last status: %s", copy_status)
+        logger.error(
+            "The copy operation didn't succeed. Last status: %s", copy_status)
         logger.error("Command run: %s", cli_cmd)
         logger.error("Command output: %s", json_output)
 

diff --git a/src/image-copy/azext_imagecopy/custom.py b/src/image-copy/azext_imagecopy/custom.py
@@ -18,7 +18,7 @@
 def imagecopy(source_resource_group_name, source_object_name, target_location,
               target_resource_group_name, temporary_resource_group_name, source_type='image',
               cleanup='false', parallel_degree=-1, tags=None, target_name=None,
-              target_subscription=None, timeout=3600):
+              target_subscription=None, export_as_snapshot='false', timeout=3600):
 
     # get the os disk id from source vm/image
     logger.warn("Getting os disk id of the source vm/image")
@@ -47,14 +47,16 @@ def imagecopy(source_resource_group_name, source_object_name, target_location,
             if source_os_disk_id is None:
                 raise TypeError
             source_os_disk_type = "BLOB"
-            logger.debug("found %s: %s", source_os_disk_type, source_os_disk_id)
+            logger.debug("found %s: %s", source_os_disk_type,
+                         source_os_disk_id)
         except TypeError:
             try:  # images created by e.g. image-copy extension
                 source_os_disk_id = json_cmd_output['storageProfile']['osDisk']['snapshot']['id']
                 if source_os_disk_id is None:
                     raise TypeError
                 source_os_disk_type = "SNAPSHOT"
-                logger.debug("found %s: %s", source_os_disk_type, source_os_disk_id)
+                logger.debug("found %s: %s", source_os_disk_type,
+                             source_os_disk_id)
             except TypeError:
                 pass
 
@@ -79,7 +81,8 @@ def imagecopy(source_resource_group_name, source_object_name, target_location,
     run_cli_command(cli_cmd)
 
     # Get SAS URL for the snapshotName
-    logger.warn("Getting sas url for the source snapshot with timeout seconds: %d", timeout)
+    logger.warn(
+        "Getting sas url for the source snapshot with timeout seconds: %d", timeout)
     if timeout < 3600:
         logger.warn("Timeout should be greater than 3600")
         raise CLIError('Inavlid Timeout')
@@ -131,7 +134,7 @@ def imagecopy(source_resource_group_name, source_object_name, target_location,
             tasks.append((location, transient_resource_group_name, source_type,
                           source_object_name, source_os_disk_snapshot_name, source_os_disk_snapshot_url,
                           source_os_type, target_resource_group_name, azure_pool_frequency,
-                          tags, target_name, target_subscription, timeout))
+                          tags, target_name, target_subscription, export_as_snapshot, timeout))
 
         logger.warn("Starting async process for all locations")