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

VERSION

@@ -1 +1 @@
-1!10.2.0
+1!11.0.0

pycloudlib/azure/cloud.py

@@ -627,7 +627,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",
@@ -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,36 @@ class ImageType(enum.Enum):
     PRO_FIPS = "Pro FIPS"
 
 
+@dataclasses.dataclass
+class ImageInfo:
+    """Dataclass that represents an image on any given cloud."""
+
+    image_id: str
+    image_name: str
+
+    def __str__(self):
+        """Return a human readable string representation of the image."""
+        return f"{self.image_name} [id: {self.image_id}]"
+
+    def __repr__(self):
+        """Return a string representation of the image."""
+        return f"ImageInfo(id={self.image_id}, name={self.image_name})"
+
+    def __eq__(self, other):
+        """
+        Check if two ImageInfo objects represent the same image.
+
+        Only the id is used for comparison since this should be the unique identifier for an image.
+        """
+        if not isinstance(other, ImageInfo):
+            return False
+        return self.image_id == other.image_id
+
+    def __dict__(self):
+        """Return a dictionary representation of the image."""
+        return {"image_id": self.image_id, "image_name": self.image_name}
+
+
 class BaseCloud(ABC):
     """Base Cloud Class."""
 
@@ -58,7 +89,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 +217,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 +245,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):
@@ -311,6 +351,18 @@ def _validate_tag(tag: str):
             raise InvalidTagNameError(tag=tag, rules_failed=rules_failed)
 
     def _get_ssh_keys(self) -> KeyPair:
+        """
+        Get the ssh key pair to use for the cloud instance.
+
+        If no key pair is provided in the config file, the default key pair
+        will be used. The default key pair is the id_rsa or id_ed25519 key
+        in the user's .ssh directory.
+
+        :raises PycloudlibError: if no public key path is provided and no default key is found
+        :raises PycloudlibError: if the public key path provided in the config does not exist
+
+        :return: KeyPair object with the public and private key paths
+        """
         user = getpass.getuser()
         # check if id_rsa or id_ed25519 keys exist in the user's .ssh directory
         possible_default_keys = [
@@ -340,3 +392,70 @@ 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(
+            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

@@ -291,6 +291,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.
 
@@ -413,12 +415,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 +440,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

@@ -314,6 +314,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,
@@ -427,12 +428,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 +472,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,8 @@ 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
+        else:
+            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 +269,9 @@ def launch(
     def snapshot(
         self,
         instance,
+        *,
         clean=True,
+        keep=False,
         note: Optional[str] = None,
         **kwargs,
     ):
@@ -276,6 +280,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 +295,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"]