feat!: add ability to retain snapshot after cleanup

This allows for marking snapshots to be kept at time of creation
instead of being cleaned up automatically. This allows for still using
context managers for easy cleanup while retaining snapshots where
desired. This functionality can also be used when not using context
managers and instead the cloud.clean() method is manually called,
like in cloud-init's integration tests.

BREAKING CHANGE: signature of cloud.snapshot() public method has
been changed and now all args except for the first (instance)
must be passed as named args.

Author: a-dubs

pycloudlib/azure/cloud.py

@@ -637,7 +637,7 @@ def delete_image(self, image_id, **kwargs):
         if delete_poller.status() == "Succeeded":
             if image_id in self.registered_images:
                 del self.registered_images[image_id]
-                self._log.debug("Image %s was deleted", image_id)
+                self._record_image_deletion(image_id)
         else:
             self._log.debug(
                 "Error deleting %s. Status: %d",
@@ -1102,12 +1102,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
 
@@ -1139,7 +1140,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

@@ -19,6 +19,7 @@
 )
 from pycloudlib.instance import BaseInstance
 from pycloudlib.key import KeyPair
+from pycloudlib.types import ImageInfo
 from pycloudlib.util import (
     get_timestamped_tag,
     log_exception_list,
@@ -47,7 +48,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.config = self._check_and_get_config(config_file, required_values)
@@ -177,12 +179,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
@@ -204,11 +207,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.image_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.image_name,
+                image_info.image_id,
+            )
         return exceptions
 
     def list_keys(self):
@@ -359,3 +369,72 @@ def _get_ssh_keys(
             private_key_path=private_key_path,
             name=name,
         )
+
+    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 argument.
+
+        This method saves the snapshot information to either `created_images` or `preserved_images` 
+        based on the value of `keep_snapshot`. These lists are used by the `BaseCloud`'s `clean()` 
+        method to manage snapshots during cleanup. The snapshot information is also logged in a 
+        consistent format so that individual clouds do NOT need to worry about logging.
+
+        Args:
+            snapshot_id (str): ID of the snapshot (used later to delete the snapshot).
+            snapshot_name (str): Name of the snapshot (for user reference).
+            keep_snapshot (bool): Whether to keep the snapshot after the cloud instance is cleaned up.
+
+        Returns:
+            ImageInfo: An ImageInfo object containing the snapshot information.
+        """
+        image_info = ImageInfo(
+            image_id=snapshot_id,
+            image_name=snapshot_name,
+        )
+        if not keep_snapshot:
+            self.created_images.append(image_info)
+            self._log.info(
+                "Created temporary snapshot %s",
+                image_info,
+            )
+        else:
+            self.preserved_images.append(image_info)
+            self._log.info(
+                "Created permanent snapshot %s",
+                image_info,
+            )
+        return image_info
+
+    def _record_image_deletion(self, image_id: str):
+        """
+        Record the deletion of an image.
+
+        This method should be called after an image is successfully deleted.
+        It will remove the image from the list of created_images or preserved_images
+        so that the cloud does not attempt to re-clean it up later. It will also log
+        the deletion of the image.
+
+        :param image_id: ID of the image that was deleted
+        """
+        if match := [i for i in self.created_images if i.image_id == image_id]:
+            deleted_image = match[0]
+            self.created_images.remove(deleted_image)
+            self._log.debug(
+                "Snapshot %s has been deleted. Will no longer need to be cleaned up later.",
+                deleted_image,
+            )
+        elif match := [i for i in self.preserved_images if i.image_id == image_id]:
+            deleted_image = match[0]
+            self.preserved_images.remove(deleted_image)
+            self._log.debug(
+                "Snapshot %s has been deleted. This snapshot was taken with keep=True, "
+                "but since it has been manually deleted, it will not be preserved.",
+                deleted_image,
+            )
+        else:
+            self._log.debug("Deleted image %s", image_id)

pycloudlib/ec2/cloud.py

@@ -295,6 +295,8 @@ def delete_image(self, image_id, **kwargs):
         self._log.debug("removing custom snapshot %s", snapshot_id)
         self.client.delete_snapshot(SnapshotId=snapshot_id)
 
+        self._record_image_deletion(image_id)
+
     def delete_key(self, name):
         """Delete an uploaded key.
 
@@ -417,12 +419,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
@@ -441,7 +444,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

@@ -315,6 +315,7 @@ def delete_image(self, image_id, **kwargs):
             raise_on_error(operation)
         except GoogleAPICallError as e:
             raise_on_error(e)
+        self._record_image_deletion(image_id)
 
     def get_instance(
         self,
@@ -428,12 +429,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
@@ -471,7 +473,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

@@ -13,7 +13,7 @@
 
 from pycloudlib.cloud import BaseCloud
 from pycloudlib.config import ConfigFile
-from pycloudlib.errors import InvalidTagNameError
+from pycloudlib.errors import InvalidTagNameError, ResourceNotFoundError, ResourceType
 from pycloudlib.ibm._util import get_first as _get_first
 from pycloudlib.ibm._util import iter_resources as _iter_resources
 from pycloudlib.ibm._util import wait_until as _wait_until
@@ -130,7 +130,9 @@ def delete_image(self, image_id: str, **kwargs):
             self._client.delete_image(image_id).get_result()
         except ApiException as e:
             if "does not exist" not in str(e):
-                raise
+                raise ResourceNotFoundError(ResourceType.IMAGE, image_id) from e
+        else:
+            self._record_image_deletion(image_id)
 
     def released_image(self, release, *, arch: str = "amd64", **kwargs):
         """ID of the latest released image for a particular release.
@@ -312,12 +314,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 +350,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

@@ -81,6 +81,7 @@ def delete_image(self, image_id: str, **kwargs):
             ) from e
         except SoftLayer.SoftLayerAPIError as e:
             raise IBMClassicException(f"Error deleting image {image_id}") from e
+        self._record_image_deletion(image_id)
 
     def released_image(self, release, *, disk_size: str = "25G", **kwargs):
         """ID (globalIdentifier) of the latest released image for a particular release.
@@ -267,7 +268,9 @@ def launch(
     def snapshot(
         self,
         instance,
+        *,
         clean=True,
+        keep=False,
         note: Optional[str] = None,
         **kwargs,
     ):
@@ -276,6 +279,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 +294,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

@@ -394,11 +394,10 @@ def delete_image(self, image_id, **kwargs):
             image_id: string, LXD image fingerprint
         """
         self._log.debug("Deleting image: '%s'", image_id)
-
         subp(["lxc", "image", "delete", image_id])
-        self._log.debug("Deleted %s", image_id)
+        self._record_image_deletion(image_id)
 
-    def snapshot(self, instance, clean=True, name=None):
+    def snapshot(self, instance: LXDInstance, *, clean=True, keep=False, name=None):  # type: ignore
         """Take a snapshot of the passed in instance for use as image.
 
         :param instance: The instance to create an image from
@@ -413,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
@@ -425,13 +428,6 @@ def clean(self) -> List[Exception]:
         """
         exceptions = super().clean()
 
-        for snapshot in self.created_snapshots:
-            try:
-                subp(["lxc", "image", "delete", snapshot])
-            except RuntimeError as e:
-                if "Image not found" not in str(e):
-                    exceptions.append(e)
-
         for profile in self.created_profiles:
             try:
                 subp(["lxc", "profile", "delete", profile])