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.7.5
+1!11.0.0

examples/az.py

@@ -5,7 +5,7 @@
 import logging
 
 import pycloudlib
-from pycloudlib.cloud import ImageType
+from pycloudlib.types import ImageType
 
 cloud_config = """#cloud-config
 runcmd:

examples/ec2.py

@@ -6,7 +6,7 @@
 import os
 
 import pycloudlib
-from pycloudlib.cloud import ImageType
+from pycloudlib.types import ImageType
 
 
 def hot_add(ec2, daily):

examples/gce.py

@@ -6,7 +6,7 @@
 import os
 
 import pycloudlib
-from pycloudlib.cloud import ImageType
+from pycloudlib.types import ImageType
 
 
 def manage_ssh_key(gce):

examples/lxd.py

@@ -6,7 +6,7 @@
 import textwrap
 
 import pycloudlib
-from pycloudlib.cloud import ImageType
+from pycloudlib.types import ImageType
 
 RELEASE = "noble"
 

pycloudlib/azure/cloud.py

@@ -15,7 +15,8 @@
 
 from pycloudlib.azure import security_types, util
 from pycloudlib.azure.instance import AzureInstance, VMInstanceStatus
-from pycloudlib.cloud import BaseCloud, ImageType
+from pycloudlib.cloud import BaseCloud
+from pycloudlib.types import ImageType
 from pycloudlib.config import ConfigFile
 from pycloudlib.errors import (
     InstanceNotFoundError,
@@ -636,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",
@@ -1101,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
 
@@ -1138,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

@@ -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
@@ -20,24 +21,14 @@
 )
 from pycloudlib.instance import BaseInstance
 from pycloudlib.key import KeyPair
+from pycloudlib.types import ImageInfo, ImageType
 from pycloudlib.util import (
     get_timestamped_tag,
     log_exception_list,
 )
 
 _RequiredValues = Optional[Sequence[Optional[Any]]]
 
-
-@enum.unique
-class ImageType(enum.Enum):
-    """Allowed image types when launching cloud images."""
-
-    GENERIC = "generic"
-    MINIMAL = "minimal"
-    PRO = "Pro"
-    PRO_FIPS = "Pro FIPS"
-
-
 class BaseCloud(ABC):
     """Base Cloud Class."""
 
@@ -58,7 +49,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)
@@ -189,12 +181,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
@@ -216,11 +209,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):
@@ -371,3 +371,70 @@ 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 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

@@ -6,7 +6,8 @@
 
 import botocore
 
-from pycloudlib.cloud import BaseCloud, ImageType
+from pycloudlib.cloud import BaseCloud
+from pycloudlib.types import ImageType
 from pycloudlib.config import ConfigFile
 from pycloudlib.ec2.instance import EC2Instance
 from pycloudlib.ec2.util import _get_session, _tag_resource
@@ -294,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.
 
@@ -416,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
@@ -440,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

@@ -16,7 +16,8 @@
 from google.api_core.extended_operation import ExtendedOperation
 from google.cloud import compute_v1
 
-from pycloudlib.cloud import BaseCloud, ImageType
+from pycloudlib.cloud import BaseCloud
+from pycloudlib.types import ImageType
 from pycloudlib.config import ConfigFile
 from pycloudlib.errors import (
     CloudSetupError,
@@ -314,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,
@@ -427,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
@@ -470,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]: