v11 - feat!: add ability to retain snapshot after cleanup

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])