Merge pull request #37 from AllenInstitute/feature/add-update_s3_path_tags-function

Add update_path_tags function for managing S3 object tags

Author: rpmcginty

src/aibs_informatics_aws_utils/s3.py

@@ -17,6 +17,7 @@
     Callable,
     Dict,
     List,
+    Literal,
     Mapping,
     Optional,
     Pattern,
@@ -575,9 +576,72 @@ def get_s3_path_stats(s3_path: S3URI, **kwargs) -> S3PathStats:
     )
 
 
-# TODO: Two things
-# 1. allow for a way to specify `size_only` for this fn when transfering large number of files
-# 2. add flag for failing if no source data exists.
+def update_path_tags(
+    s3_path: S3URI,
+    tags: Dict[str, str],
+    mode: Literal["replace", "append", "delete"] = "append",
+    recursive: bool = True,
+    **kwargs,
+):
+    """Update tags for all objects at or prefixed by the specified path
+
+    There are three modes for updating tags:
+    - replace: Replace all existing tags with new tags
+    - append: Merge new tags with existing tags
+    - delete: Delete specified tags from existing tags (values do not matter)
+
+    If recursive is True and s3_path is an object prefix, all objects under the prefix
+    will have their tags updated. If there is an object at s3_path, it will also have
+    its tags updated.
+
+    Args:
+        s3_path (S3URI): S3 path or prefix to update tags for
+        tags (Dict[str, str]): Tags to update
+        mode (Literal["replace", "append", "delete"]): Tag update mode.
+            Options:
+                - replace: Replace all existing tags with new tags
+                - append: Merge new tags with existing tags
+                - delete: Delete specified tags from existing tags (values do not matter)
+            Defaults to "append".
+        recursive (bool): Whether to update tags recursively for all objects under prefix.
+            Defaults to True.
+    """
+    if recursive and is_object_prefix(s3_path, **kwargs):
+        s3_paths = list_s3_paths(s3_path=s3_path, **kwargs)
+        logger.info(f"Updating tags for {len(s3_paths)} objects under prefix {s3_path}")
+        for nested_s3_path in s3_paths:
+            update_path_tags(nested_s3_path, tags, mode, recursive=False, **kwargs)
+    if is_object(s3_path, **kwargs):
+        s3 = get_s3_client(**kwargs)
+
+        current_tags = {
+            tag["Key"]: tag["Value"]
+            for tag in s3.get_object_tagging(
+                Bucket=s3_path.bucket,
+                Key=s3_path.key,
+            ).get("TagSet", [])
+        }
+
+        if mode == "replace":
+            new_tags = tags
+        elif mode in ["append", "delete"]:
+            new_tags = current_tags.copy()
+            if mode == "append":
+                new_tags.update(tags)
+            else:
+                for tag_key in tags.keys():
+                    if tag_key in new_tags:
+                        del new_tags[tag_key]
+        else:
+            raise ValueError(f"Unknown tag update mode: {mode}")  # pragma: no cover
+
+        s3.put_object_tagging(
+            Bucket=s3_path.bucket,
+            Key=s3_path.key,
+            Tagging={"TagSet": [{"Key": k, "Value": v} for k, v in new_tags.items()]},
+        )
+
+
 def sync_paths(
     source_path: Union[Path, S3URI],
     destination_path: Union[Path, S3URI],

test/aibs_informatics_aws_utils/test_s3.py

@@ -3,7 +3,7 @@
 import re
 from pathlib import Path
 from time import sleep
-from typing import Optional
+from typing import Dict, Optional
 
 import moto
 import requests
@@ -49,6 +49,7 @@
     process_transfer_requests,
     should_sync,
     sync_paths,
+    update_path_tags,
     update_s3_storage_class,
     upload_file,
     upload_json,
@@ -105,6 +106,17 @@ def client__list_objects_v2(self, **kwargs):
             kwargs["Bucket"] = self.DEFAULT_BUCKET_NAME
         return self.s3_client.list_objects_v2(**kwargs)
 
+    def _get_tag_dict(self, s3_path: S3URI) -> Dict[str, str]:
+        response = self.s3_client.get_object_tagging(Bucket=s3_path.bucket, Key=s3_path.key)
+        return {tag["Key"]: tag["Value"] for tag in response.get("TagSet", [])}
+
+    def _put_tags(self, s3_path: S3URI, tags: Dict[str, str]):
+        self.s3_client.put_object_tagging(
+            Bucket=s3_path.bucket,
+            Key=s3_path.key,
+            Tagging={"TagSet": [{"Key": key, "Value": value} for key, value in tags.items()]},
+        )
+
     def test__get_presigned_urls__default_generates_READ_ONLY_urls(self):
         ## Setup
         s3_path_a = self.get_s3_path("path/to/object_a")
@@ -737,6 +749,72 @@ def test__delete_s3_path__handles_folder(self):
         delete_s3_path(s3_path=s3_path)
         self.assertEqual(0, len(list_s3_paths(s3_path)))
 
+    def test__update_path_tags__replace_mode_overwrites_existing_tags(self):
+        s3_path = self.put_object("path/to/tagged.txt", "content")
+        self._put_tags(s3_path, {"old": "tag"})
+
+        update_path_tags(
+            s3_path,
+            {"new": "value"},
+            mode="replace",
+            region=self.DEFAULT_REGION,
+        )
+
+        self.assertDictEqual(self._get_tag_dict(s3_path), {"new": "value"})
+
+    def test__update_path_tags__append_mode_merges_and_overrides(self):
+        s3_path = self.put_object("path/to/append.txt", "content")
+        self._put_tags(s3_path, {"keep": "1", "override": "old"})
+
+        update_path_tags(
+            s3_path,
+            {"override": "new", "fresh": "2"},
+            mode="append",
+            region=self.DEFAULT_REGION,
+        )
+
+        self.assertDictEqual(
+            self._get_tag_dict(s3_path),
+            {"keep": "1", "override": "new", "fresh": "2"},
+        )
+
+    def test__update_path_tags__delete_mode_removes_requested_keys(self):
+        s3_path = self.put_object("path/to/delete.txt", "content")
+        self._put_tags(s3_path, {"keep": "1", "drop": "2"})
+
+        update_path_tags(
+            s3_path,
+            {"drop": "ignored", "missing": "value"},
+            mode="delete",
+            region=self.DEFAULT_REGION,
+        )
+
+        self.assertDictEqual(self._get_tag_dict(s3_path), {"keep": "1"})
+
+    def test__update_path_tags__recursive_prefix_updates_all_children(self):
+        prefix_path = self.get_s3_path("path/to/prefix")
+        child_one = self.put_object("path/to/prefix/a.txt", "first")
+        child_two = self.put_object("path/to/prefix/nested/b.txt", "second")
+        outside = self.put_object("path/to/other.txt", "outside")
+
+        self._put_tags(child_one, {"existing": "one"})
+        self._put_tags(child_two, {"existing": "two"})
+        self._put_tags(outside, {"existing": "outside"})
+
+        update_path_tags(
+            prefix_path,
+            {"batch": "true"},
+            mode="append",
+            region=self.DEFAULT_REGION,
+        )
+
+        for s3_path, label in ((child_one, "one"), (child_two, "two")):
+            tags = self._get_tag_dict(s3_path)
+            self.assertEqual(tags["existing"], label)
+            self.assertEqual(tags["batch"], "true")
+
+        self.assertNotIn("batch", self._get_tag_dict(outside))
+
     def test__update_s3_storage_class__handles_shallow_to_GLACIER(self):
         s3_root = self.get_s3_path("source/path/")
 
@@ -880,7 +958,7 @@ def test__should_sync__s3_to_local__multipart_upload_with_custom_chunk_size_work
 
         s3.upload_file(
             Filename=str(orig_file),
-            **source_path.as_dict(),
+            **source_path.as_dict(),  # type: ignore
             Config=TransferConfig(multipart_threshold=1024, multipart_chunksize=1024),
         )
         destination_path.write_text(orig_file.read_text())
@@ -897,7 +975,7 @@ def test__should_sync__s3_to_local__multipart_upload_chunksize__gt__default(self
         # This does not upload as multipart for custom chunk size
         self.s3_client.upload_file(
             Filename=str(orig_file),
-            **source_path1.as_dict(),
+            **source_path1.as_dict(),  # type: ignore
             Config=TransferConfig(
                 multipart_threshold=AWS_S3_DEFAULT_CHUNK_SIZE_BYTES * 2,
                 multipart_chunksize=AWS_S3_DEFAULT_CHUNK_SIZE_BYTES * 2,
@@ -906,7 +984,7 @@ def test__should_sync__s3_to_local__multipart_upload_chunksize__gt__default(self
         # This one uploads as multipart even though there is only one part
         self.s3_client.upload_file(
             Filename=str(orig_file),
-            **source_path2.as_dict(),
+            **source_path2.as_dict(),  # type: ignore
             Config=TransferConfig(
                 multipart_threshold=AWS_S3_DEFAULT_CHUNK_SIZE_BYTES,
                 multipart_chunksize=AWS_S3_DEFAULT_CHUNK_SIZE_BYTES * 2,
@@ -928,7 +1006,7 @@ def test__should_sync__handles_multipart_upload_chunksize(self):
         # This does not upload as multipart for custom chunk size
         self.s3_client.upload_file(
             Filename=str(orig_file),
-            **source_path.as_dict(),
+            **source_path.as_dict(),  # type: ignore
             Config=TransferConfig(multipart_threshold=MB, multipart_chunksize=MB),
         )
 
@@ -945,7 +1023,7 @@ def test__should_sync__handles_multipart_upload_chunksize__single_part(self):
         # This does not upload as multipart for custom chunk size
         self.s3_client.upload_file(
             Filename=str(orig_file),
-            **source_path.as_dict(),
+            **source_path.as_dict(),  # type: ignore
             Config=TransferConfig(multipart_threshold=MB, multipart_chunksize=2 * MB),
         )
 
@@ -962,7 +1040,7 @@ def test__should_sync__handles_multipart_upload__threshold_not_passed(self):
         # This does not upload as multipart for custom chunk size
         self.s3_client.upload_file(
             Filename=str(orig_file),
-            **source_path.as_dict(),
+            **source_path.as_dict(),  # type: ignore
             Config=TransferConfig(multipart_threshold=3 * MB, multipart_chunksize=MB),
         )
 
@@ -979,7 +1057,7 @@ def test__should_sync__handles_multipart_upload__gt_threshold_passed(self):
         # This does not upload as multipart for custom chunk size
         self.s3_client.upload_file(
             Filename=str(orig_file),
-            **source_path.as_dict(),
+            **source_path.as_dict(),  # type: ignore
             Config=TransferConfig(multipart_threshold=2 * MB, multipart_chunksize=MB),
         )
         destination_path.write_text(orig_file.read_text())