feat!: add ability to retain snapshot after cleanup

Author: a-dubs

pycloudlib/azure/cloud.py

@@ -1092,12 +1092,13 @@ class compatibility.
 
         raise InstanceNotFoundError(resource_id=instance_id)
 
-    def snapshot(self, instance, clean=True, delete_provisioned_user=True, **kwargs):
+    def snapshot(self, instance, *, clean=True, keep=False, delete_provisioned_user=True, **kwargs):
         """Snapshot an instance and generate an image from it.
 
         Args:
             instance: Instance to snapshot
             clean: Run instance clean method before taking snapshot
+            keep: keep the snapshot after the cloud instance is cleaned up
             delete_provisioned_user: Deletes the last provisioned user
             kwargs: Other named arguments specific to this implementation
 
@@ -1129,7 +1130,11 @@ def snapshot(self, instance, clean=True, delete_provisioned_user=True, **kwargs)
         image_id = image.id
         image_name = image.name
 
-        self.created_images.append(image_id)
+        self._store_snapshot_info(
+            snapshot_id=image_id,
+            snapshot_name=image_name,
+            keep_snapshot=keep,
+        )
 
         self.registered_images[image_id] = {
             "name": image_name,

pycloudlib/cloud.py

@@ -1,6 +1,7 @@
 # This file is part of pycloudlib. See LICENSE file for license information.
 """Base class for all other clouds to provide consistent set of functions."""
 
+import dataclasses
 import enum
 import getpass
 import io
@@ -38,6 +39,14 @@ class ImageType(enum.Enum):
     PRO_FIPS = "Pro FIPS"
 
 
+@dataclasses.dataclass
+class ImageInfo:
+    """Dataclass to hold image information."""
+
+    id: str
+    name: str
+
+
 class BaseCloud(ABC):
     """Base Cloud Class."""
 
@@ -58,7 +67,8 @@ def __init__(
             config_file: path to pycloudlib configuration file
         """
         self.created_instances: List[BaseInstance] = []
-        self.created_images: List[str] = []
+        self.created_images: List[ImageInfo] = []
+        self.preserved_images: List[ImageInfo] = []  # each dict will hold an id and name
 
         self._log = logging.getLogger("{}.{}".format(__name__, self.__class__.__name__))
         self._check_and_set_config(config_file, required_values)
@@ -185,12 +195,13 @@ def launch(
         raise NotImplementedError
 
     @abstractmethod
-    def snapshot(self, instance, clean=True, **kwargs):
+    def snapshot(self, instance, *, clean=True, keep=False, **kwargs):
         """Snapshot an instance and generate an image from it.
 
         Args:
             instance: Instance to snapshot
             clean: run instance clean method before taking snapshot
+            keep: keep the snapshot after the cloud instance is cleaned up
 
         Returns:
             An image id
@@ -212,11 +223,18 @@ def clean(self) -> List[Exception]:
                 instance.delete()
             except Exception as e:
                 exceptions.append(e)
-        for image_id in self.created_images:
+        for image_info in self.created_images:
             try:
-                self.delete_image(image_id)
+                self.delete_image(image_id=image_info.id)
             except Exception as e:
                 exceptions.append(e)
+        for image_info in self.preserved_images:
+            # noop - just log that we're not cleaning up these images
+            self._log.info(
+                "Preserved image %s [id:%s] is NOT being cleaned up.",
+                image_info.name,
+                image_info.id,
+            )
         return exceptions
 
     def list_keys(self):
@@ -340,3 +358,43 @@ def _get_ssh_keys(self) -> KeyPair:
             private_key_path=private_key_path,
             name=self.config.get("key_name", user),
         )
+
+    def _store_snapshot_info(
+        self,
+        snapshot_id: str,
+        snapshot_name: str,
+        keep_snapshot: bool,
+    ) -> ImageInfo:
+        """
+        Save the snapshot information for later cleanup depending on the keep_snapshot arg.
+
+        Will either save the snapshot information to created_images or preserved_images depending
+        on the keep_snapshot arg. These lists are used by the BaseCloud's clean() method to
+        cleanup the snapshots when the cloud instance is cleaned up. The snapshot information
+        is also logged appropriately and in a consistent format.
+
+        :param snapshot_id: ID of the snapshot (this is used later to delete the snapshot)
+        :param snapshot_name: Name of the snapshot (this is for user reference)
+        :param keep_snapshot: Keep the snapshot after the cloud instance is cleaned up
+
+        :return: ImageInfo object with the snapshot information
+        """
+        image_info = ImageInfo(
+            id=snapshot_id,
+            name=snapshot_name,
+        )
+        if not keep_snapshot:
+            self.created_images.append(image_info)
+            self._log.info(
+                "Created temporary snapshot %s [id:%s]",
+                image_info.name,
+                image_info.id,
+            )
+        else:
+            self.preserved_images.append(image_info)
+            self._log.info(
+                "Created permanent snapshot %s [id:%s]",
+                image_info.name,
+                image_info.id,
+            )
+        return image_info

pycloudlib/ec2/cloud.py

@@ -413,12 +413,13 @@ def list_keys(self):
             keypair_names.append(keypair["KeyName"])
         return keypair_names
 
-    def snapshot(self, instance, clean=True):
+    def snapshot(self, instance, *, clean=True, keep=False, **kwargs):
         """Snapshot an instance and generate an image from it.
 
         Args:
             instance: Instance to snapshot
             clean: run instance clean method before taking snapshot
+            keep: keep the snapshot after the cloud instance is cleaned up
 
         Returns:
             An image id
@@ -437,7 +438,12 @@ def snapshot(self, instance, clean=True):
         )
         image_ami_edited = response["ImageId"]
         image = self.resource.Image(image_ami_edited)
-        self.created_images.append(image.id)
+
+        self._store_snapshot_info(
+            snapshot_id=image.id,
+            snapshot_name=image.name,
+            keep_snapshot=keep,
+        )
 
         self._wait_for_snapshot(image)
         _tag_resource(image, self.tag)

pycloudlib/gce/cloud.py

@@ -427,12 +427,13 @@ def launch(
         self.created_instances.append(instance)
         return instance
 
-    def snapshot(self, instance: GceInstance, clean=True, **kwargs):
+    def snapshot(self, instance: GceInstance, *, clean=True, keep=False, **kwargs):
         """Snapshot an instance and generate an image from it.
 
         Args:
             instance: Instance to snapshot
             clean: run instance clean method before taking snapshot
+            keep: keep the snapshot after the cloud instance is cleaned up
 
         Returns:
             An image id
@@ -470,7 +471,11 @@ def snapshot(self, instance: GceInstance, clean=True, **kwargs):
         self._wait_for_operation(operation)
 
         image_id = "projects/{}/global/images/{}".format(self.project, snapshot_name)
-        self.created_images.append(image_id)
+        self._store_snapshot_info(
+            snapshot_name=snapshot_name,
+            snapshot_id=image_id,
+            keep_snapshot=keep,
+        )
         return image_id
 
     def _wait_for_operation(self, operation, operation_type="global", sleep_seconds=300):

pycloudlib/ibm/cloud.py

@@ -312,12 +312,13 @@ def launch(
 
         return instance
 
-    def snapshot(self, instance: IBMInstance, clean: bool = True, **kwargs) -> str:
+    def snapshot(self, instance: IBMInstance, *, clean=True, keep=False, **kwargs) -> str:
         """Snapshot an instance and generate an image from it.
 
         Args:
             instance: Instance to snapshot
             clean: run instance clean method before taking snapshot
+            keep: keep the snapshot after the cloud instance is cleaned up
 
         Returns:
             An image id
@@ -347,7 +348,11 @@ def snapshot(self, instance: IBMInstance, clean: bool = True, **kwargs) -> str:
                 f"Snapshot not available after {timeout_seconds} seconds. Check IBM VPC console."
             ),
         )
-        self.created_images.append(snapshot_id)
+        self._store_snapshot_info(
+            snapshot_name=str(image_prototype["name"]),
+            snapshot_id=snapshot_id,
+            keep_snapshot=keep,
+        )
         return snapshot_id
 
     def list_keys(self) -> List[str]:

pycloudlib/ibm_classic/cloud.py

@@ -267,7 +267,9 @@ def launch(
     def snapshot(
         self,
         instance,
+        *,
         clean=True,
+        keep=False,
         note: Optional[str] = None,
         **kwargs,
     ):
@@ -276,6 +278,7 @@ def snapshot(
         Args:
             instance: Instance to snapshot
             clean: run instance clean method before taking snapshot
+            keep: keep the snapshot after the cloud instance is cleaned up
             note: optional note to add to the snapshot
 
         Returns:
@@ -290,10 +293,10 @@ def snapshot(
             name=f"{self.tag}-snapshot",
             notes=note,
         )
-        self._log.info(
-            "Successfully created snapshot '%s' with ID: %s",
-            snapshot_result["name"],
-            snapshot_result["id"],
+        self._store_snapshot_info(
+            snapshot_name=snapshot_result["name"],
+            snapshot_id=snapshot_result["id"],
+            keep_snapshot=keep,
         )
         return snapshot_result["id"]
 

pycloudlib/lxd/cloud.py

@@ -397,7 +397,7 @@ def delete_image(self, image_id, **kwargs):
         subp(["lxc", "image", "delete", image_id])
         self._log.debug("Deleted %s", image_id)
 
-    def snapshot(self, instance, clean=True, name=None):
+    def snapshot(self, instance, *, clean=True, keep=False, name=None):
         """Take a snapshot of the passed in instance for use as image.
 
         :param instance: The instance to create an image from
@@ -412,7 +412,11 @@ def snapshot(self, instance, clean=True, name=None):
             instance.clean()
 
         snapshot_name = instance.snapshot(name)
-        self.created_snapshots.append(snapshot_name)
+        self._store_snapshot_info(
+            snapshot_name=snapshot_name,
+            snapshot_id=snapshot_name,
+            keep_snapshot=keep,
+        )
         return snapshot_name
 
     # pylint: disable=broad-except

pycloudlib/oci/cloud.py

@@ -313,15 +313,17 @@ def launch(
         self.created_instances.append(instance)
         return instance
 
-    def snapshot(self, instance, clean=True, name=None):
+    def snapshot(self, instance, *, clean=True, keep=False, name=None):
         """Snapshot an instance and generate an image from it.
 
         Args:
             instance: Instance to snapshot
             clean: run instance clean method before taking snapshot
-            name: (Optional) Name of created image
+            keep: Keep the image after the cloud instance is cleaned up
+            name: Name of created image
+
         Returns:
-            An image object
+            The image id of the snapshot
         """
         if clean:
             instance.clean()
@@ -340,6 +342,10 @@ def snapshot(self, instance, clean=True, name=None):
             desired_state="AVAILABLE",
         )
 
-        self.created_images.append(image_data.id)
+        self._store_snapshot_info(
+            snapshot_name=image_data.display_name,
+            snapshot_id=image_data.id,
+            keep_snapshot=keep,
+        )
 
         return image_data.id

pycloudlib/openstack/cloud.py

@@ -171,12 +171,13 @@ def launch(
         self.created_instances.append(instance)
         return instance
 
-    def snapshot(self, instance, clean=True, **kwargs):
+    def snapshot(self, instance, *, clean=True, keep=False, **kwargs):
         """Snapshot an instance and generate an image from it.
 
         Args:
             instance: Instance to snapshot
             clean: run instance clean method before taking snapshot
+            keep: keep the snapshot after the cloud instance is cleaned up
 
         Returns:
             An image id
@@ -188,7 +189,11 @@ def snapshot(self, instance, clean=True, **kwargs):
         image = self.conn.create_image_snapshot(
             "{}-snapshot".format(self.tag), instance.server.id, wait=True
         )
-        self.created_images.append(image.id)
+        self._store_snapshot_info(
+            snapshot_name=image.name,
+            snapshot_id=image.id,
+            keep_snapshot=keep,
+        )
         return image.id
 
     def use_key(self, public_key_path, private_key_path=None, name=None):

pycloudlib/qemu/cloud.py

@@ -542,12 +542,13 @@ def launch(
 
         return instance
 
-    def snapshot(self, instance: QemuInstance, clean=True, **kwargs) -> str:
+    def snapshot(self, instance: QemuInstance, *, clean=True, keep=False, **kwargs) -> str:
         """Snapshot an instance and generate an image from it.
 
         Args:
             instance: Instance to snapshot
             clean: run instance clean method before taking snapshot
+            keep: keep the snapshot after the cloud instance is cleaned up
 
         Returns:
             An image id
@@ -596,7 +597,11 @@ def snapshot(self, instance: QemuInstance, clean=True, **kwargs) -> str:
             snapshot_path,
             instance.instance_path,
         )
-        self.created_images.append(str(snapshot_path))
+        self._store_snapshot_info(
+            snapshot_name=snapshot_path.stem,
+            snapshot_id=str(snapshot_path),
+            keep_snapshot=keep,
+        )
 
         return str(snapshot_path)