From c49b004c7f88c46979d366c5767577fc0c40ac6e Mon Sep 17 00:00:00 2001 From: Jonathan Gangi Date: Fri, 8 Nov 2024 15:10:04 -0300 Subject: [PATCH] AWS: Allow setting tags to snapshots or AMIs only This commit introduces two new optional properties on `AWSPublishingMetadata` named `snapshot_tags` and `ami_tags`, which when set allows to add tags just for snapshot imports or the registered image respectively. Rationale: Currently our team is facing sporadic issues on community AMIs workflow when uploading the image to different billing types (`access` or `hourly`) after a retry, as the second attempt will look for the image name, fail to find, search for tags and find an AMI which concerns the other billing type. This was initially fixed by release-engineering/pubtools-marketplacesvm#74 by adding the billing type as tags. However, the fix brings another potential issue: we will tag everything with the billing tags, including the S3 object, which could being unecessarily uploaded twice, as well as having a tag which doesn't concern the S3 RAW object, but the snapshot/image only. With this change on `cloudimg` we aim to provide the billing type tags just for the snapshot and AMI, leaving the S3 object with the previous ones as it shouldn't receive the billing tag. This change also allows other teams to have a more granular control on which tags should go where. Refers to SPSTRAT-451 --- cloudimg/aws.py | 25 +++++++++-- tests/test_aws.py | 106 ++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 127 insertions(+), 4 deletions(-) diff --git a/cloudimg/aws.py b/cloudimg/aws.py index 45fb009..7fccf53 100644 --- a/cloudimg/aws.py +++ b/cloudimg/aws.py @@ -83,6 +83,12 @@ class AWSPublishingMetadata(PublishingMetadata): billing_products (list, optional): Billing product identifiers boot_mode (str, optional): The boot mode for booting up the AMI on EC2. + + snapshot_tags (dict, optional): Tags to be applied to the snapshot + import only. + + ami_tags (dict, optional): Tags to be applied to the registered AMI + only. """ def __init__(self, *args, **kwargs): @@ -91,6 +97,8 @@ def __init__(self, *args, **kwargs): self.billing_products = kwargs.pop('billing_products', None) bmode_str = kwargs.pop('boot_mode', None) or "not_set" self.boot_mode = AWSBootMode[bmode_str] + self.snapshot_tags = kwargs.pop('snapshot_tags', None) + self.ami_tags = kwargs.pop('ami_tags', None) super(AWSPublishingMetadata, self).__init__(*args, **kwargs) @@ -558,6 +566,13 @@ def publish(self, metadata): Returns: An EC2 Image """ + def add_tags(tag_parameter_name, extra_kwargs): + new_tags = getattr(metadata, tag_parameter_name, None) + if new_tags: + tags = extra_kwargs.get("tags") or {} + new_tags.update(tags) + extra_kwargs.update({"tags": new_tags}) + log.info('Searching for image: %s', metadata.image_name) image = ( self.get_image_by_name(metadata.image_name) or @@ -578,8 +593,7 @@ def publish(self, metadata): # Set tags when they're provided extra_kwargs = {} - if metadata.tags: - extra_kwargs.update({"tags": metadata.tags}) + add_tags("tags", extra_kwargs) if not obj: log.info('Object does not exist: %s', metadata.object_name) @@ -590,6 +604,7 @@ def publish(self, metadata): else: log.info('Object already exists') + add_tags("snapshot_tags", extra_kwargs) snapshot = self.import_snapshot(obj, metadata.snapshot_name, **extra_kwargs) @@ -765,8 +780,10 @@ def register_image(self, snapshot, metadata): BillingProducts=metadata.billing_products, **optional_kwargs, ) - if metadata.tags: - self.tag_image(image, metadata.tags) + if metadata.tags or metadata.ami_tags: + tags = metadata.tags or {} + tags.update(metadata.ami_tags or {}) + self.tag_image(image, tags) return image def share_image(self, image, accounts=[], groups=[]): diff --git a/tests/test_aws.py b/tests/test_aws.py index f551f4e..7d98cba 100644 --- a/tests/test_aws.py +++ b/tests/test_aws.py @@ -700,6 +700,85 @@ def test_publish_tags(self, self.md.object_name, tags=tags) + @patch('cloudimg.aws.AWSService.upload_to_container') + @patch('cloudimg.aws.AWSService.import_snapshot') + @patch('cloudimg.aws.AWSService.register_image') + @patch('cloudimg.aws.AWSService.share_image') + @patch('cloudimg.aws.AWSService.get_image_by_tags') + @patch('cloudimg.aws.AWSService.get_image_by_name') + @patch('cloudimg.aws.AWSService.get_snapshot_by_name') + @patch('cloudimg.aws.AWSService.get_object_by_name') + def test_publish_snapshot_tags(self, + get_object_by_name, + get_snapshot_by_name, + get_image_by_name, + get_image_by_tags, + share_image, + register_image, + import_snapshot, + upload_to_container): + get_image_by_name.return_value = None + get_image_by_tags.return_value = None + get_snapshot_by_name.return_value = None + get_object_by_name.return_value = None + self.md.tags = None + self.md.snapshot_tags = {"snapshot": "tag"} + published = self.svc.publish(self.md) + + share_image.assert_called_once_with(published, + accounts=[], + groups=[]) + self.assertEqual(register_image.call_count, 1) + import_snapshot.assert_called_once_with( + ANY, + self.md.snapshot_name, + tags={"snapshot": "tag"}, + ) + upload_to_container.assert_called_once_with(self.md.image_path, + self.md.container, + self.md.object_name) + + @patch('cloudimg.aws.AWSService.upload_to_container') + @patch('cloudimg.aws.AWSService.import_snapshot') + @patch('cloudimg.aws.AWSService.register_image') + @patch('cloudimg.aws.AWSService.share_image') + @patch('cloudimg.aws.AWSService.get_image_by_tags') + @patch('cloudimg.aws.AWSService.get_image_by_name') + @patch('cloudimg.aws.AWSService.get_snapshot_by_name') + @patch('cloudimg.aws.AWSService.get_object_by_name') + def test_publish_merged_snapshot_tags( + self, + get_object_by_name, + get_snapshot_by_name, + get_image_by_name, + get_image_by_tags, + share_image, + register_image, + import_snapshot, + upload_to_container, + ): + get_image_by_name.return_value = None + get_image_by_tags.return_value = None + get_snapshot_by_name.return_value = None + get_object_by_name.return_value = None + self.md.tags = {"foo": "bar"} + self.md.snapshot_tags = {"snapshot": "tag"} + published = self.svc.publish(self.md) + + share_image.assert_called_once_with(published, + accounts=[], + groups=[]) + self.assertEqual(register_image.call_count, 1) + import_snapshot.assert_called_once_with( + ANY, + self.md.snapshot_name, + tags={"foo": "bar", "snapshot": "tag"}, + ) + upload_to_container.assert_called_once_with(self.md.image_path, + self.md.container, + self.md.object_name, + tags={"foo": "bar"}) + @patch('cloudimg.aws.AWSService.tag_image') def test_register_image_no_tags(self, tag_image): self.mock_register_image.return_value = "fakeimg" @@ -741,6 +820,33 @@ def test_register_image_tags(self, tag_image): tag_image.assert_called_once_with("fakeimg", self.md.tags) self.assertEqual(res, "fakeimg") + @patch('cloudimg.aws.AWSService.tag_image') + def test_register_image_ami_tags(self, tag_image): + self.md.tags = None + self.md.ami_tags = {"ami": "tag"} + self.mock_register_image.return_value = "fakeimg" + + res = self.svc.register_image(MagicMock(), self.md) + + self.mock_register_image.assert_called_once() + tag_image.assert_called_once_with("fakeimg", {"ami": "tag"}) + self.assertEqual(res, "fakeimg") + + @patch('cloudimg.aws.AWSService.tag_image') + def test_register_image_merged_ami_tags(self, tag_image): + self.md.tags = {"tag": "tag"} + self.md.ami_tags = {"ami": "tag"} + self.mock_register_image.return_value = "fakeimg" + + res = self.svc.register_image(MagicMock(), self.md) + + self.mock_register_image.assert_called_once() + tag_image.assert_called_once_with( + "fakeimg", + {"tag": "tag", "ami": "tag"}, + ) + self.assertEqual(res, "fakeimg") + @patch('cloudimg.aws.AWSService.tag_image') def test_register_image_boot_mode(self, tag_image): self.mock_register_image.return_value = "fakeimg"