From e5995e99e3970b0d78f33f4fc4eab7e87eb8d8ca Mon Sep 17 00:00:00 2001
From: Rui Marinho <rui.marinho@uphold.com>
Date: Fri, 23 Aug 2024 10:22:22 +0100
Subject: [PATCH] Fix tags not being applied to AWS EBS Snapshots

---
 barman/cloud_providers/__init__.py     |  2 ++
 barman/cloud_providers/aws_s3.py       | 16 ++++++---
 tests/test_cloud_snapshot_interface.py | 45 ++++++++++++++++++++++++++
 3 files changed, 59 insertions(+), 4 deletions(-)

diff --git a/barman/cloud_providers/__init__.py b/barman/cloud_providers/__init__.py
index c65317bfe..3a96ca249 100644
--- a/barman/cloud_providers/__init__.py
+++ b/barman/cloud_providers/__init__.py
@@ -201,6 +201,7 @@ def get_snapshot_interface(config):
             config.aws_profile,
             config.aws_region,
             config.aws_await_snapshots_timeout,
+            config.tags,
         ]
         return AwsCloudSnapshotInterface(*args)
     else:
@@ -253,6 +254,7 @@ def get_snapshot_interface_from_server_config(server_config):
             server_config.aws_profile,
             server_config.aws_region,
             server_config.aws_await_snapshots_timeout,
+            server_config.tags,
         )
     else:
         raise CloudProviderUnsupported(
diff --git a/barman/cloud_providers/aws_s3.py b/barman/cloud_providers/aws_s3.py
index 7cf4cb330..1742f99c8 100644
--- a/barman/cloud_providers/aws_s3.py
+++ b/barman/cloud_providers/aws_s3.py
@@ -463,7 +463,7 @@ class AwsCloudSnapshotInterface(CloudSnapshotInterface):
         https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/ebs-creating-snapshot.html
     """
 
-    def __init__(self, profile_name=None, region=None, await_snapshots_timeout=3600):
+    def __init__(self, profile_name=None, region=None, await_snapshots_timeout=3600, tags=None):
         """
         Creates the client necessary for creating and managing snapshots.
 
@@ -478,6 +478,7 @@ def __init__(self, profile_name=None, region=None, await_snapshots_timeout=3600)
         self.region = region or self.session.region_name
         self.ec2_client = self.session.client("ec2", region_name=self.region)
         self.await_snapshots_timeout = await_snapshots_timeout
+        self.tags = tags
 
     def _get_waiter_config(self):
         delay = 15
@@ -710,13 +711,20 @@ def _create_snapshot(self, backup_info, volume_name, volume_id):
             volume_name,
             volume_id,
         )
+        tags = [
+            {"Key": "Name", "Value": snapshot_name},
+        ]
+
+        if self.tags is not None:
+            for tag in self.tags:
+                key, value = tag
+                tags.append({"Key": key, "Value": value})
+
         resp = self.ec2_client.create_snapshot(
             TagSpecifications=[
                 {
                     "ResourceType": "snapshot",
-                    "Tags": [
-                        {"Key": "Name", "Value": snapshot_name},
-                    ],
+                    "Tags": tags,
                 }
             ],
             VolumeId=volume_id,
diff --git a/tests/test_cloud_snapshot_interface.py b/tests/test_cloud_snapshot_interface.py
index 17993bf06..8daabdaf5 100644
--- a/tests/test_cloud_snapshot_interface.py
+++ b/tests/test_cloud_snapshot_interface.py
@@ -2740,6 +2740,51 @@ def test_create_snapshot(self, caplog):
             in caplog.text
         )
 
+    def test_create_snapshot_with_tags(self, caplog):
+        """
+        Verify that _create_snapshot calls boto3 with defined tags and returns the expected values.
+        """
+        # GIVEN a new AwsCloudInterface with tags defined
+        snapshot_interface = AwsCloudSnapshotInterface(tags=[
+            ("environment", "production"),
+            ("project", "my-project"),
+            ("service", "barman")
+        ])
+        # AND a backup_info for a given server name and backup ID
+        backup_info = mock.Mock(backup_id=self.backup_id, server_name=self.server_name)
+        # AND a mock create_snapshot function which returns a successful response
+        mock_ec2_client = self._mock_boto3.Session.return_value.client.return_value
+        mock_resp = mock_ec2_client.create_snapshot.return_value
+        mock_resp["State"] = "pending"
+        # AND log level is INFO
+        caplog.set_level(logging.INFO)
+
+        # WHEN _create_snapshot is called
+        volume_name = "my-pgdata-volume"
+        volume_id = "vol-0123456789abcdef01"
+        snapshot_name, _ = snapshot_interface._create_snapshot(
+            backup_info,
+            volume_name,
+            volume_id,
+        )
+
+        # THEN create_snapshot is called on the EC2 client with the expected args
+        mock_ec2_client.create_snapshot.assert_called_once()
+        mock_ec2_client.create_snapshot.assert_called_once_with(
+            TagSpecifications=[
+                {
+                    "ResourceType": "snapshot",
+                    "Tags": [
+                        {"Key": "Name", "Value": snapshot_name},
+                        {"Key": "environment", "Value": "production"},
+                        {"Key": "project", "Value": "my-project"},
+                        {"Key": "service", "Value": "barman"}
+                    ],
+                }
+            ],
+            VolumeId=volume_id,
+        )
+
     def test_create_snapshot_failed(self):
         """
         Verify that _create_snapshot calls boto3 and returns the expected values.