KMS: Move Tagging Functionality into methods on the Provider (#13595)

aidehn · web-flow · commit 0eded68ce31d · 2026-02-05T11:45:46.000+08:00
diff --git a/localstack-core/localstack/services/kms/models.py b/localstack-core/localstack/services/kms/models.py
@@ -49,8 +49,8 @@
     UnsupportedOperationException,
 )
 from localstack.constants import TAG_KEY_CUSTOM_ID
-from localstack.services.kms.exceptions import TagException, ValidationException
-from localstack.services.kms.utils import is_valid_key_arn, validate_tag
+from localstack.services.kms.exceptions import ValidationException
+from localstack.services.kms.utils import is_valid_key_arn, validate_tag, validate_tag_list
 from localstack.services.stores import AccountRegionBundle, BaseStore, LocalAttribute
 from localstack.utils.aws.arns import get_partition, kms_alias_arn, kms_key_arn
 from localstack.utils.crypto import decrypt, encrypt
@@ -630,12 +630,7 @@ def add_tags(self, tags: TagList) -> None:
         if not tags:
             return
 
-        unique_tag_keys = {tag["TagKey"] for tag in tags}
-        if len(unique_tag_keys) < len(tags):
-            raise TagException("Duplicate tag keys")
-
-        if len(tags) > 50:
-            raise TagException("Too many tags")
+        validate_tag_list(tags)
 
         # Do not care if we overwrite an existing tag:
         # https://docs.aws.amazon.com/kms/latest/APIReference/API_TagResource.html
diff --git a/localstack-core/localstack/services/kms/provider.py b/localstack-core/localstack/services/kms/provider.py
@@ -106,6 +106,9 @@
     ScheduleKeyDeletionResponse,
     SignRequest,
     SignResponse,
+    Tag,
+    TagKeyList,
+    TagList,
     TagResourceRequest,
     UnsupportedOperationException,
     UntagResourceRequest,
@@ -389,6 +392,20 @@ def _parse_key_id(key_id_or_arn: str, context: RequestContext) -> tuple[str, str
     def _is_rsa_spec(key_spec: str) -> bool:
         return key_spec in [KeySpec.RSA_2048, KeySpec.RSA_3072, KeySpec.RSA_4096]
 
+    # These tagging methods are overwritten in the Pro implementation.
+    def _get_key_tags(self, key: KmsKey, account_id: str, region: str) -> TagList:
+        return [Tag(TagKey=key, TagValue=value) for key, value in key.tags.items()]
+
+    def _set_key_tags(self, key: KmsKey, account_id: str, region: str, tags: TagList) -> None:
+        return
+
+    def _remove_key_tags(
+        self, key: KmsKey, account_id: str, region: str, tag_keys: TagKeyList
+    ) -> None:
+        # AWS doesn't seem to mind removal of a non-existent tag, so we do not raise any exception.
+        for tag_key in tag_keys:
+            key.tags.pop(tag_key, None)
+
     #
     # Operation Handlers
     #
@@ -400,6 +417,7 @@ def create_key(
         request: CreateKeyRequest = None,
     ) -> CreateKeyResponse:
         key = self._create_kms_key(context.account_id, context.region, request)
+        self._set_key_tags(key, context.account_id, context.region, request.get("Tags", []))
         return CreateKeyResponse(KeyMetadata=key.metadata)
 
     @handler("ScheduleKeyDeletion", expand=False)
@@ -1466,15 +1484,15 @@ def list_resource_tags(
         key = self._get_kms_key(
             context.account_id, context.region, request.get("KeyId"), any_key_state_allowed=True
         )
-        keys_list = PaginatedList(
-            [{"TagKey": tag_key, "TagValue": tag_value} for tag_key, tag_value in key.tags.items()]
-        )
+        keys_list = PaginatedList(self._get_key_tags(key, context.account_id, context.region))
         page, next_token = keys_list.get_page(
             lambda tag: tag.get("TagKey"),
             next_token=request.get("Marker"),
             page_size=request.get("Limit", 50),
         )
-        kwargs = {"NextMarker": next_token, "Truncated": True} if next_token else {}
+        kwargs = (
+            {"NextMarker": next_token, "Truncated": True} if next_token else {"Truncated": False}
+        )
         return ListResourceTagsResponse(Tags=page, **kwargs)
 
     @handler("RotateKeyOnDemand", expand=False)
@@ -1498,29 +1516,30 @@ def rotate_key_on_demand(
 
     @handler("TagResource", expand=False)
     def tag_resource(self, context: RequestContext, request: TagResourceRequest) -> None:
+        tags = request["Tags"]
         key = self._get_kms_key(
             context.account_id,
             context.region,
             request.get("KeyId"),
             enabled_key_allowed=True,
             disabled_key_allowed=True,
         )
-        key.add_tags(request.get("Tags"))
+        key.add_tags(tags)
+        self._set_key_tags(key, context.account_id, context.region, tags)
 
     @handler("UntagResource", expand=False)
     def untag_resource(self, context: RequestContext, request: UntagResourceRequest) -> None:
+        if not (tag_keys := request.get("TagKeys", [])):
+            return
+
         key = self._get_kms_key(
             context.account_id,
             context.region,
             request.get("KeyId"),
             enabled_key_allowed=True,
             disabled_key_allowed=True,
         )
-        if not request.get("TagKeys"):
-            return
-        for tag_key in request.get("TagKeys"):
-            # AWS doesn't seem to mind removal of a non-existent tag, so we do not raise any exception.
-            key.tags.pop(tag_key, None)
+        self._remove_key_tags(key, context.account_id, context.region, tag_keys)
 
     def derive_shared_secret(
         self,
diff --git a/localstack-core/localstack/services/kms/utils.py b/localstack-core/localstack/services/kms/utils.py
@@ -1,7 +1,7 @@
 import re
 from collections.abc import Callable
 
-from localstack.aws.api.kms import DryRunOperationException, Tag, TagException
+from localstack.aws.api.kms import DryRunOperationException, Tag, TagException, TagList
 from localstack.services.kms.exceptions import ValidationException
 from localstack.utils.aws.arns import ARN_PARTITION_REGEX
 
@@ -60,6 +60,14 @@ def validate_tag(tag_position: int, tag: Tag) -> None:
         raise TagException("Tags beginning with aws: are reserved")
 
 
+def validate_tag_list(tag_list: TagList) -> None:
+    unique_tag_keys = {tag["TagKey"] for tag in tag_list}
+    if len(unique_tag_keys) < len(tag_list):
+        raise TagException("Duplicate tag keys")
+    if len(tag_list) > 50:
+        raise TagException("Too many tags")
+
+
 def execute_dry_run_capable[T](func: Callable[..., T], dry_run: bool, *args, **kwargs) -> T:
     """
     Executes a function unless dry run mode is enabled.
