feat(gce): implement both new methods for enabling file uploading and image creation

Author: a-dubs

examples/gce.py

@@ -2,14 +2,17 @@
 # This file is part of pycloudlib. See LICENSE file for license information.
 """Basic examples of various lifecycle with an GCE instance."""
 
+import argparse
+import datetime
 import logging
 import os
 
 import pycloudlib
-from pycloudlib.cloud import ImageType
+from pycloudlib import GCE
+from pycloudlib.cloud import ImageInfo, ImageType
 
 
-def manage_ssh_key(gce):
+def manage_ssh_key(gce: GCE):
     """Manage ssh keys for gce instances."""
     pub_key_path = "gce-pubkey"
     priv_key_path = "gce-privkey"
@@ -27,7 +30,7 @@ def manage_ssh_key(gce):
     gce.use_key(public_key_path=pub_key_path, private_key_path=priv_key_path)
 
 
-def generic(gce):
+def generic(gce: GCE):
     """Show example of using the GCE library.
 
     Connects to GCE and finds the latest daily image. Then runs
@@ -39,32 +42,108 @@ def generic(gce):
         print(inst.execute("lsb_release -a"))
 
 
-def pro(gce):
+def pro(gce: GCE):
     """Show example of running a GCE PRO machine."""
     daily = gce.daily_image("bionic", image_type=ImageType.PRO)
     with gce.launch(daily) as inst:
         inst.wait()
         print(inst.execute("sudo ua status --wait"))
 
 
-def pro_fips(gce):
+def pro_fips(gce: GCE):
     """Show example of running a GCE PRO FIPS machine."""
     daily = gce.daily_image("bionic", image_type=ImageType.PRO_FIPS)
     with gce.launch(daily) as inst:
         inst.wait()
         print(inst.execute("sudo ua status --wait"))
 
 
-def demo():
+def custom_image(gce: GCE, image_name):
+    """Show example of running a GCE custom image."""
+    image_id = gce.get_image_id_from_name(image_name)
+    print(image_id)
+    with gce.launch(image_id=image_id) as instance:
+        instance.wait()
+        print(instance.execute("hostname"))
+        input("Press Enter to teardown instance")
+
+
+def upload_custom_image(gce: GCE, image_name, local_file_path, bucket_name):
+    """Show example of uploading a custom image to GCE."""
+    new_image: ImageInfo = gce.create_image_from_local_file(
+        local_file_path=local_file_path,
+        image_name=image_name,
+        intermediary_storage_name=bucket_name,
+    )
+    print("created new image:", new_image)
+    return new_image
+
+
+def demo_image_creation(
+    local_file_path: str,
+    bucket_name: str,
+    image_name_template: str = "gce-example-image-{}",
+):
+    """Show example of creating a custom image on GCE from a local image file."""
+    # get short date and time for unique tag
+    time_tag = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
+    tag = f"gce-example-{time_tag}"
+
+    with pycloudlib.GCE(tag=tag) as gce:
+        manage_ssh_key(gce)
+
+        new_image: ImageInfo = upload_custom_image(
+            gce,
+            image_name=image_name_template.format(time_tag),
+            local_file_path=local_file_path,
+            bucket_name=bucket_name,
+        )
+
+        with gce.launch(new_image.id) as instance:
+            instance.wait()
+            print(instance.execute("hostname"))
+            print(instance.execute("lsb_release -a"))
+            input("Press Enter to teardown instance")
+
+
+def demo_instances():
     """Show examples of launching GCP instances."""
-    logging.basicConfig(level=logging.DEBUG)
     with pycloudlib.GCE(tag="examples") as gce:
         manage_ssh_key(gce)
-
         generic(gce)
         pro(gce)
         pro_fips(gce)
 
 
+def main():
+    """Take in cli args and run GCE demo scripts."""
+    parser = argparse.ArgumentParser(description="GCE Demo Script")
+    parser.add_argument(
+        "demo_type",
+        choices=["instance", "image"],
+        help="Type of demo to run: 'instance' for basic instance launch demo, or "
+        "'image' for image creation demo",
+        nargs="?",
+        default="instance",
+    )
+    parser.add_argument(
+        "--local_file_path", type=str, help="Local file path for the image creation demo"
+    )
+    parser.add_argument("--bucket_name", type=str, help="Bucket name for the image creation demo")
+    args = parser.parse_args()
+
+    logging.basicConfig(level=logging.DEBUG)
+
+    if args.demo_type == "instance":
+        demo_instances()
+    elif args.demo_type == "image":
+        if not args.local_file_path or not args.bucket_name:
+            parser.error("Image creation demo requires --local_file_path and --bucket_name")
+        demo_image_creation(
+            local_file_path=args.local_file_path,
+            bucket_name=args.bucket_name,
+        )
+
+
 if __name__ == "__main__":
-    demo()
+    main()

pycloudlib/gce/cloud.py

@@ -14,9 +14,9 @@
 
 from google.api_core.exceptions import GoogleAPICallError
 from google.api_core.extended_operation import ExtendedOperation
-from google.cloud import compute_v1
+from google.cloud import compute_v1, storage
 
-from pycloudlib.cloud import BaseCloud, ImageType
+from pycloudlib.cloud import BaseCloud, ImageInfo, ImageType
 from pycloudlib.config import ConfigFile
 from pycloudlib.errors import (
     CloudSetupError,
@@ -109,6 +109,7 @@ def __init__(
         self._instances_client = compute_v1.InstancesClient(credentials=credentials)
         self._zone_operations_client = compute_v1.ZoneOperationsClient(credentials=credentials)
         self._global_operations_client = compute_v1.GlobalOperationsClient(credentials=credentials)
+        self._storage_client = storage.Client(credentials=credentials)
         region = region or self.config.get("region") or "us-west2"
         zone = zone or self.config.get("zone") or "a"
         self.project = project
@@ -498,3 +499,149 @@ def _wait_for_operation(self, operation, operation_type="global", sleep_seconds=
                 "Expected DONE state, but found {} after waiting {} seconds. "
                 "Check GCE console for more details. \n".format(response.status, sleep_seconds)
             )
+
+    def upload_local_file_to_cloud_storage(
+        self,
+        *,
+        local_file_path: str,
+        storage_name: str,
+        remote_file_name: Optional[str] = None,
+    ):
+        """
+        Upload a file to a storage destination on the Cloud.
+
+        Args:
+            local_file_path: The local file path of the image to upload.
+            storage_name: The name of the storage destination on the Cloud to upload the file to.
+            remote_file_name: The name of the file in the storage destination. If not provided,
+            the base name of the local file path will be used.
+
+        Returns:
+            str: URL of the uploaded file in the storage destination.
+        """
+        if not remote_file_name:
+            remote_file_name = os.path.basename(local_file_path)
+
+        bucket = self._storage_client.bucket(storage_name)
+        blob = bucket.blob(remote_file_name)
+
+        # Check if the file already exists in the destination bucket
+        if blob.exists():
+            self._log.info(
+                f"File '{remote_file_name}' already exists in bucket '{storage_name}', "
+                "skipping upload."
+            )
+        else:
+            self._log.info(
+                f"Uploading {local_file_path} to {remote_file_name} in bucket {storage_name}..."
+            )
+            blob.upload_from_filename(local_file_path)
+            self._log.info(f"File {local_file_path} uploaded successfully to {remote_file_name}.")
+
+        return f"http://storage.googleapis.com/{storage_name}/{remote_file_name}"
+
+    def create_image_from_local_file(
+        self,
+        *,
+        local_file_path: str,
+        image_name: str,
+        intermediary_storage_name: str,
+        description: Optional[str] = None,
+    ) -> ImageInfo:
+        """
+        Upload local image file to storage on the Cloud and then create a custom image from it.
+
+        Args:
+            local_file_path: The local file path of the image to upload.
+            image_name: The name to upload the image as and to register.
+            intermediary_storage_name: The intermediary storage destination on the Cloud to upload
+            the file to before creating the image.
+
+        Returns:
+            ImageInfo: Information about the created image.
+        """
+        # Upload the image to GCS
+        remote_file_name = f"{image_name}.tar.gz"
+        self._log.info(
+            "Uploading image '%s' to '%s' as '%s'",
+            image_name,
+            intermediary_storage_name,
+            remote_file_name,
+        )
+        gcs_image_path = self.upload_local_file_to_cloud_storage(
+            local_file_path=local_file_path,
+            storage_name=intermediary_storage_name,
+            remote_file_name=remote_file_name,
+        )
+        # Register the custom image from GCS
+        return self._create_image_from_cloud_storage(
+            image_name=image_name,
+            remote_image_file_url=gcs_image_path,
+        )
+
+    def _create_image_from_cloud_storage(
+        self,
+        *,
+        image_name: str,
+        remote_image_file_url: str,
+        image_description: Optional[str] = None,
+        image_family: Optional[str] = None,
+    ) -> ImageInfo:
+        """
+        Register a custom image in GCE from a file in GCE's Cloud storage using its url.
+
+        Ideally, this url would be returned from the upload_local_file_to_cloud_storage method.
+
+        Args:
+            image_name: The name the image will be created with.
+            remote_image_file_url: The URL of the image file in the Cloud storage.
+            image_description: (Optional) A description of the image.
+            image_family: (Optional) The family name of the image.
+
+        Returns:
+            ImageInfo: Information about the created image.
+        """
+        image_body = compute_v1.Image(
+            name=image_name,
+            raw_disk=compute_v1.RawDisk(source=remote_image_file_url),
+            description=image_description,
+            architecture="x86_64",
+        )
+        if image_family:
+            image_body.family = image_family
+
+        self._log.info(
+            "Attempting to register custom image '%s' from GCS path '%s'...",
+            image_name,
+            remote_image_file_url,
+        )
+
+        try:
+            operation = self._images_client.insert(
+                project=self.project,
+                image_resource=image_body,
+            )
+            self._wait_for_operation(operation)
+            self._log.info(f"Custom image '{image_name}' registered successfully.")
+
+        except Exception as e:
+            self._log.error(f"Failed to create custom image from GCS: {e}")
+            raise e
+
+        return ImageInfo(
+            id=f"projects/{self.project}/global/images/{image_name}",
+            name=image_name,
+        )
+
+    def _wait_for_operation(self, operation: ExtendedOperation):
+        """
+        Wait for a GCE operation to complete.
+
+        Args:
+            operation: The operation to wait for.
+        """
+        self._log.debug("Waiting for operation to complete...")
+        while not operation.done():
+            self._log.debug("Operation is still in progress...")
+            time.sleep(5)
+        self._log.debug("Operation completed.")