Merge pull request #77 from rpocase/20250109-attempt-all-regions

feat: allow partial image registration

Author: toabctl

awspub/exceptions.py

@@ -10,6 +10,10 @@ class MultipleImagesException(Exception):
     pass
 
 
+class IncompleteImageSetException(Exception):
+    pass
+
+
 class BucketDoesNotExistException(Exception):
     def __init__(self, bucket_name: str, *args, **kwargs):
         msg = f"The bucket named '{bucket_name}' does not exist. You will need to create the bucket before proceeding."

awspub/image.py

@@ -5,6 +5,7 @@
 from typing import Any, Dict, List, Optional
 
 import boto3
+import botocore.exceptions
 from mypy_boto3_ec2.client import EC2Client
 from mypy_boto3_ssm import SSMClient
 
@@ -418,6 +419,7 @@ def create(self) -> Dict[str, _ImageInfo]:
         )
 
         images: Dict[str, _ImageInfo] = dict()
+        missing_regions: List[str] = []
         for region in self.image_regions:
             ec2client_region: EC2Client = boto3.client("ec2", region_name=region)
             image_info: Optional[_ImageInfo] = self._get(ec2client_region)
@@ -435,53 +437,10 @@ def create(self) -> Dict[str, _ImageInfo]:
                     )
                 images[region] = image_info
             else:
-                logger.info(
-                    f"creating image with name '{self.image_name}' in "
-                    f"region {ec2client_region.meta.region_name} ..."
-                )
-
-                register_image_kwargs = dict(
-                    Name=self.image_name,
-                    Description=self.conf.get("description", ""),
-                    Architecture=self._ctx.conf["source"]["architecture"],
-                    RootDeviceName=self.conf["root_device_name"],
-                    BlockDeviceMappings=[
-                        {
-                            "Ebs": {
-                                "SnapshotId": snapshot_ids[region],
-                                "VolumeType": self.conf["root_device_volume_type"],
-                                "VolumeSize": self.conf["root_device_volume_size"],
-                            },
-                            "DeviceName": self.conf["root_device_name"],
-                        },
-                        # TODO: make those ephemeral block device mappings configurable
-                        {"VirtualName": "ephemeral0", "DeviceName": "/dev/sdb"},
-                        {"VirtualName": "ephemeral1", "DeviceName": "/dev/sdc"},
-                    ],
-                    EnaSupport=True,
-                    SriovNetSupport="simple",
-                    VirtualizationType="hvm",
-                    BootMode=self.conf["boot_mode"],
-                )
-
-                if self.conf["tpm_support"]:
-                    register_image_kwargs["TpmSupport"] = self.conf["tpm_support"]
-
-                if self.conf["imds_support"]:
-                    register_image_kwargs["ImdsSupport"] = self.conf["imds_support"]
-
-                if self.conf["uefi_data"]:
-                    with open(self.conf["uefi_data"], "r") as f:
-                        uefi_data = f.read()
-                    register_image_kwargs["UefiData"] = uefi_data
-
-                if self.conf["billing_products"]:
-                    register_image_kwargs["BillingProducts"] = self.conf["billing_products"]
-
-                resp = ec2client_region.register_image(**register_image_kwargs)
-                ec2client_region.create_tags(Resources=[resp["ImageId"]], Tags=self._tags)
-                images[region] = _ImageInfo(resp["ImageId"], snapshot_ids[region])
-
+                if image := self._register_image(snapshot_ids[region], ec2client_region):
+                    images[region] = image
+                else:
+                    missing_regions.append(region)
         # wait for the images
         logger.info(f"Waiting for {len(images)} images to be ready the regions ...")
         for region, image_info in images.items():
@@ -500,8 +459,80 @@ def create(self) -> Dict[str, _ImageInfo]:
         if self.conf["share"]:
             self._share(self.conf["share"], images)
 
+        if missing_regions:
+            logger.error("Failed to publish images to all regions", extra={"missing_regions": missing_regions})
+            raise exceptions.IncompleteImageSetException("Incomplete image set published")
+
         return images
 
+    def _register_image(self, snapshot_id: str, ec2client: EC2Client) -> Optional[_ImageInfo]:
+        """
+        Register snapshot_id in region configured for ec2client_region
+
+        :param snapshot_id: snapshot id to use for image registration
+        :type snapshot_id: str
+        :param ec2client: EC2Client for the region to register image to
+        :type ec2client: EC2Client
+        :return: _ImageInfo containing the ImageId SnapshotId pair
+        :rtype: _ImageInfo
+        """
+        logger.info(f"creating image with name '{self.image_name}' in " f"region {ec2client.meta.region_name} ...")
+
+        register_image_kwargs = dict(
+            Name=self.image_name,
+            Description=self.conf.get("description", ""),
+            Architecture=self._ctx.conf["source"]["architecture"],
+            RootDeviceName=self.conf["root_device_name"],
+            BlockDeviceMappings=[
+                {
+                    "Ebs": {
+                        "SnapshotId": snapshot_id,
+                        "VolumeType": self.conf["root_device_volume_type"],
+                        "VolumeSize": self.conf["root_device_volume_size"],
+                    },
+                    "DeviceName": self.conf["root_device_name"],
+                },
+                # TODO: make those ephemeral block device mappings configurable
+                {"VirtualName": "ephemeral0", "DeviceName": "/dev/sdb"},
+                {"VirtualName": "ephemeral1", "DeviceName": "/dev/sdc"},
+            ],
+            EnaSupport=True,
+            SriovNetSupport="simple",
+            VirtualizationType="hvm",
+            BootMode=self.conf["boot_mode"],
+        )
+
+        if self.conf["tpm_support"]:
+            register_image_kwargs["TpmSupport"] = self.conf["tpm_support"]
+
+        if self.conf["imds_support"]:
+            register_image_kwargs["ImdsSupport"] = self.conf["imds_support"]
+
+        if self.conf["uefi_data"]:
+            with open(self.conf["uefi_data"], "r") as f:
+                uefi_data = f.read()
+            register_image_kwargs["UefiData"] = uefi_data
+
+        if self.conf["billing_products"]:
+            register_image_kwargs["BillingProducts"] = self.conf["billing_products"]
+
+        try:
+            resp = ec2client.register_image(**register_image_kwargs)
+        except botocore.exceptions.ClientError as e:
+            if e.response.get("Error", {}).get("Code", None) == "OperationNotPermitted":
+                logger.exception(
+                    "Unable to register image",
+                    extra={
+                        "registration_options": register_image_kwargs,
+                        "region": ec2client.meta.region_name,
+                    },
+                )
+                return None
+            raise e
+
+        ec2client.create_tags(Resources=[resp["ImageId"]], Tags=self._tags)
+        return _ImageInfo(resp["ImageId"], snapshot_id)
+
     def publish(self) -> None:
         """
         Handle all publication steps

awspub/tests/test_image.py

@@ -1,6 +1,7 @@
 import pathlib
-from unittest.mock import patch
+from unittest.mock import MagicMock, patch
 
+import botocore.exceptions
 import pytest
 
 from awspub import context, exceptions, image
@@ -494,3 +495,62 @@ def test_image__share_list_filtered(partition, imagename, share_list_expected):
         ctx = context.Context(curdir / "fixtures/config1.yaml", None)
         img = image.Image(ctx, imagename)
         assert img._share_list_filtered(img.conf["share"]) == share_list_expected
+
+
+@patch("awspub.s3.S3.bucket_region", return_value="region1")
+def test_create__should_allow_partial_registration(s3_bucket_mock):
+    """
+    Test that the create() method allows a partial upload set
+    """
+    with patch("boto3.client") as bclient_mock:
+        instance = bclient_mock.return_value
+
+        ctx = context.Context(curdir / "fixtures/config1.yaml", None)
+        img = image.Image(ctx, "test-image-6")
+        img._image_regions = ["region1", "region2"]
+        img._image_regions_cached = True
+        with patch.object(img, "_get") as get_mock, patch.object(img._snapshot, "copy") as copy_mock:
+            copy_mock.return_value = {r: f"snapshot{i}" for i, r in enumerate(img.image_regions)}
+            get_mock.return_value = None
+            instance.register_image.side_effect = [
+                botocore.exceptions.ClientError(
+                    {
+                        "Error": {
+                            "Code": "OperationNotPermitted",
+                            "Message": "Intentional permission failure for snapshot0",
+                        }
+                    },
+                    "awspub Testing",
+                ),
+                {"ImageId": "id1"},
+            ]
+            with pytest.raises(exceptions.IncompleteImageSetException):
+                img.create() == {"region2": image._ImageInfo("id1", "snapshot1")}
+        # register and create_tags should be called since at least one snapshot made it
+        assert instance.register_image.called
+        assert instance.create_tags.called
+
+
+def test_register_image__should_return_none_on_permission_failures():
+    instance = MagicMock()
+
+    instance.register_image.side_effect = botocore.exceptions.ClientError(
+        {"Error": {"Code": "OperationNotPermitted", "Message": "Testing"}}, "Testing"
+    )
+    ctx = context.Context(curdir / "fixtures/config1.yaml", None)
+    img = image.Image(ctx, "test-image-6")
+    snapshot_ids = {"eu-central-1": "my-snapshot"}
+    assert img._register_image(snapshot_ids["eu-central-1"], instance) is None
+
+
+def test_register_image__should_raise_on_unhandled_client_error():
+    instance = MagicMock()
+
+    instance.register_image.side_effect = botocore.exceptions.ClientError(
+        {"Error": {"Code": "UnsupportedOperation", "Message": "Testing"}}, "Testing"
+    )
+    ctx = context.Context(curdir / "fixtures/config1.yaml", None)
+    img = image.Image(ctx, "test-image-6")
+    snapshot_ids = {"eu-central-1": "my-snapshot"}
+    with pytest.raises(botocore.exceptions.ClientError):
+        img._register_image(snapshot_ids["eu-central-1"], instance) is None