feat: allow partial image registration

this changes the `awspub create` path to allow for partial snapshot
registration when OperationNotPermitted errors occur. in the event
of an OperationNotPermitted ClientError all remaining regions will
be attempted.

Author: rpocase

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,8 +437,10 @@ def create(self) -> Dict[str, _ImageInfo]:
                     )
                 images[region] = image_info
             else:
-                images[region] = self._register_image(snapshot_ids[region], region, ec2client_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():
@@ -455,9 +459,13 @@ 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) -> _ImageInfo:
+    def _register_image(self, snapshot_id: str, ec2client: EC2Client) -> Optional[_ImageInfo]:
         """
         Register snapshot_id in region configured for ec2client_region
 
@@ -508,7 +516,20 @@ def _register_image(self, snapshot_id: str, ec2client: EC2Client) -> _ImageInfo:
         if self.conf["billing_products"]:
             register_image_kwargs["BillingProducts"] = self.conf["billing_products"]
 
-        resp = ec2client.register_image(**register_image_kwargs)
+        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)
 

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