From 64033bc3dcb34a82ca5fe8cbcdfcc0a780207665 Mon Sep 17 00:00:00 2001 From: Mathieu Cloutier Date: Thu, 10 Oct 2024 16:27:33 -0600 Subject: [PATCH] implement retrieving from _custom_id_ tag --- moto/apigateway/utils.py | 11 ++- moto/secretsmanager/utils.py | 12 ++- moto/utilities/id_generator.py | 98 ++++++++++++++----- tests/test_utilities/test_id_generator.py | 112 ++++++++++++++++++++++ 4 files changed, 203 insertions(+), 30 deletions(-) create mode 100644 tests/test_utilities/test_id_generator.py diff --git a/moto/apigateway/utils.py b/moto/apigateway/utils.py index fa7cf6630495..3fcc25f2828e 100644 --- a/moto/apigateway/utils.py +++ b/moto/apigateway/utils.py @@ -5,7 +5,7 @@ import yaml from moto.moto_api._internal import mock_random as random -from moto.utilities.id_generator import ResourceIdentifier, generate_str_id +from moto.utilities.id_generator import ResourceIdentifier, Tags, generate_str_id class ApigwIdentifier(ResourceIdentifier): @@ -14,10 +14,13 @@ class ApigwIdentifier(ResourceIdentifier): def __init__(self, account_id: str, region: str, name: str): super().__init__(account_id, region, name) - def generate(self, existing_ids: Union[List[str], None] = None) -> str: + def generate( + self, existing_ids: Union[List[str], None] = None, tags: Tags = None + ) -> str: return generate_str_id( - self, - existing_ids, + resource_identifier=self, + existing_ids=existing_ids, + tags=tags, length=10, include_digits=True, lower_case=True, diff --git a/moto/secretsmanager/utils.py b/moto/secretsmanager/utils.py index 84d34d4a80dd..e4b55622b176 100644 --- a/moto/secretsmanager/utils.py +++ b/moto/secretsmanager/utils.py @@ -2,7 +2,12 @@ import string from moto.moto_api._internal import mock_random as random -from moto.utilities.id_generator import ExistingIds, ResourceIdentifier, generate_str_id +from moto.utilities.id_generator import ( + ExistingIds, + ResourceIdentifier, + Tags, + generate_str_id, +) from moto.utilities.utils import ARN_PARTITION_REGEX, get_partition @@ -104,10 +109,11 @@ class SecretsManagerSecretIdentifier(ResourceIdentifier): def __init__(self, account_id: str, region: str, secret_id: str): super().__init__(account_id, region, name=secret_id) - def generate(self, existing_ids: ExistingIds = None) -> str: + def generate(self, existing_ids: ExistingIds = None, tags: Tags = None) -> str: id_string = generate_str_id( - existing_ids=existing_ids, resource_identifier=self, + existing_ids=existing_ids, + tags=tags, length=6, include_digits=False, ) diff --git a/moto/utilities/id_generator.py b/moto/utilities/id_generator.py index 5d8f9cd3c080..53abaa8787c2 100644 --- a/moto/utilities/id_generator.py +++ b/moto/utilities/id_generator.py @@ -1,10 +1,23 @@ import abc +import logging import threading -from typing import Any, Callable, Dict, List, Union +from typing import Any, Callable, Dict, List, TypedDict, Union from moto.moto_api._internal import mock_random +log = logging.getLogger(__name__) + ExistingIds = Union[List[str], None] +Tags = Union[Dict[str, str], None] + +# Custom resource tag to override the generated resource ID. +TAG_KEY_CUSTOM_ID = "_custom_id_" + + +class IdSourceContext(TypedDict, total=False): + resource_identifier: "ResourceIdentifier" + tags: Tags + existing_ids: ExistingIds class ResourceIdentifier(abc.ABC): @@ -25,9 +38,8 @@ def __init__(self, account_id: str, region: str, name: str): self.name = name or "" @abc.abstractmethod - def generate(self, existing_ids: ExistingIds = None) -> str: ... - - """ If `existing_ids` is provided, we will not return a custom id if it is already on the list""" + def generate(self, existing_ids: ExistingIds = None, tags: Tags = None) -> str: + """Method to generate a resource id""" @property def unique_identifier(self) -> str: @@ -35,13 +47,16 @@ def unique_identifier(self) -> str: [self.account_id, self.region, self.service, self.resource, self.name] ) + def __str__(self): + return self.unique_identifier + class MotoIdManager: """class to manage custom ids. Do not create instance and instead use the `id_manager` instance created below.""" _custom_ids: Dict[str, str] - _id_sources: List[Callable[[ResourceIdentifier], Union[str, None]]] + _id_sources: List[Callable[[IdSourceContext], Union[str, None]]] _lock: threading.RLock @@ -50,13 +65,14 @@ def __init__(self) -> None: self._lock = threading.RLock() self._id_sources = [] + self.add_id_source(self.get_id_from_tags) self.add_id_source(self.get_custom_id) - def get_custom_id( - self, resource_identifier: ResourceIdentifier - ) -> Union[str, None]: + def get_custom_id(self, id_source_context: IdSourceContext) -> Union[str, None]: # retrieves a custom_id for a resource. Returns None - return self._custom_ids.get(resource_identifier.unique_identifier) + if resource_identifier := id_source_context.get("resource_identifier"): + return self._custom_ids.get(resource_identifier.unique_identifier) + return None def set_custom_id( self, resource_identifier: ResourceIdentifier, custom_id: str @@ -73,18 +89,29 @@ def unset_custom_id(self, resource_identifier: ResourceIdentifier) -> None: self._custom_ids.pop(resource_identifier.unique_identifier, None) def add_id_source( - self, id_source: Callable[[ResourceIdentifier], Union[str, None]] + self, id_source: Callable[[IdSourceContext], Union[str, None]] ) -> None: self._id_sources.append(id_source) - def find_id_from_sources( - self, resource_identifier: ResourceIdentifier, existing_ids: List[str] - ) -> Union[str, None]: + @staticmethod + def get_id_from_tags(id_source_context: IdSourceContext) -> Union[str, None]: + if tags := id_source_context.get("tags"): + return tags.get(TAG_KEY_CUSTOM_ID) + + return None + + def find_id_from_sources(self, id_source_context: IdSourceContext) -> Union[str, None]: + existing_ids = id_source_context.get("existing_ids") or [] for id_source in self._id_sources: - if ( - found_id := id_source(resource_identifier) - ) and found_id not in existing_ids: - return found_id + if found_id := id_source(id_source_context): + if found_id in existing_ids: + log.debug( + f"Found id {found_id} for resource {id_source_context.get('resource_identifier')}, " + "but a resource already exists with this id." + ) + else: + return found_id + return None @@ -92,17 +119,41 @@ def find_id_from_sources( def moto_id(fn: Callable[..., str]) -> Callable[..., str]: - # Decorator for helping in creation of static ids within Moto. + """ + Decorator for helping in creation of static ids. + + The decorated function should accept the following parameters + + :param resource_identifier + :param existing_ids + If provided, we will omit returning a custom id if it is already on the list + :param tags + If provided will look for a tag named `_custom_id_`. This will take precedence over registered custom ids + """ + def _wrapper( resource_identifier: ResourceIdentifier, - existing_ids: ExistingIds, + existing_ids: ExistingIds = None, + tags: Tags = None, **kwargs: Dict[str, Any], ) -> str: - if found_id := moto_id_manager.find_id_from_sources( - resource_identifier, existing_ids or [] + if resource_identifier and ( + found_id := moto_id_manager.find_id_from_sources( + IdSourceContext( + resource_identifier=resource_identifier, + existing_ids=existing_ids, + tags=tags, + ) + ) ): return found_id - return fn(resource_identifier, existing_ids, **kwargs) + + return fn( + resource_identifier=resource_identifier, + existing_ids=existing_ids, + tags=tags, + **kwargs, + ) return _wrapper @@ -110,7 +161,8 @@ def _wrapper( @moto_id def generate_str_id( # type: ignore resource_identifier: ResourceIdentifier, - existing_ids: ExistingIds, + existing_ids: ExistingIds = None, + tags: Tags = None, length: int = 20, include_digits: bool = True, lower_case: bool = False, diff --git a/tests/test_utilities/test_id_generator.py b/tests/test_utilities/test_id_generator.py new file mode 100644 index 000000000000..014f076a41d9 --- /dev/null +++ b/tests/test_utilities/test_id_generator.py @@ -0,0 +1,112 @@ +from moto.utilities.id_generator import ( + TAG_KEY_CUSTOM_ID, + ExistingIds, + ResourceIdentifier, + Tags, + moto_id, +) + +ACCOUNT = "account" +REGION = "us-east-1" +RESOURCE_NAME = "my-resource" + +CUSTOM_ID = "custom" +GENERATED_ID = "generated" +TAG_ID = "fromTag" +SERVICE = "test-service" +RESOURCE = "test-resource" + + +@moto_id +def generate_test_id( + resource_identifier: ResourceIdentifier, + existing_ids: ExistingIds = None, + tags: Tags = None, +): + return GENERATED_ID + + +class TestResourceIdentifier(ResourceIdentifier): + service = SERVICE + resource = RESOURCE + + def generate(self, existing_ids: ExistingIds = None, tags: Tags = None) -> str: + return generate_test_id( + resource_identifier=self, existing_ids=existing_ids, tags=tags + ) + + +def test_generate_with_no_resource_identifier(): + generated_id = generate_test_id(None) + assert generated_id == GENERATED_ID + + +def test_generate_with_matching_resource_identifier(set_custom_id): + resource_identifier = TestResourceIdentifier(ACCOUNT, REGION, RESOURCE_NAME) + + set_custom_id(resource_identifier, CUSTOM_ID) + + generated_id = generate_test_id(resource_identifier=resource_identifier) + assert generated_id == CUSTOM_ID + + +def test_generate_with_non_matching_resource_identifier(set_custom_id): + resource_identifier = TestResourceIdentifier(ACCOUNT, REGION, RESOURCE_NAME) + resource_identifier_2 = TestResourceIdentifier(ACCOUNT, REGION, "non-matching") + + set_custom_id(resource_identifier, CUSTOM_ID) + + generated_id = generate_test_id(resource_identifier=resource_identifier_2) + assert generated_id == GENERATED_ID + + +def test_generate_with_custom_id_tag(): + resource_identifier = TestResourceIdentifier(ACCOUNT, REGION, RESOURCE_NAME) + + generated_id = generate_test_id( + resource_identifier=resource_identifier, tags={TAG_KEY_CUSTOM_ID: TAG_ID} + ) + assert generated_id == TAG_ID + + +def test_generate_with_custom_id_tag_has_priority(set_custom_id): + resource_identifier = TestResourceIdentifier(ACCOUNT, REGION, RESOURCE_NAME) + + set_custom_id(resource_identifier, CUSTOM_ID) + generated_id = generate_test_id( + resource_identifier=resource_identifier, tags={TAG_KEY_CUSTOM_ID: TAG_ID} + ) + assert generated_id == TAG_ID + + +def test_generate_with_existing_id(set_custom_id): + resource_identifier = TestResourceIdentifier(ACCOUNT, REGION, RESOURCE_NAME) + + set_custom_id(resource_identifier, CUSTOM_ID) + generated_id = generate_test_id( + resource_identifier=resource_identifier, existing_ids=[CUSTOM_ID] + ) + assert generated_id == GENERATED_ID + + +def test_generate_with_tags_and_existing_id(set_custom_id): + resource_identifier = TestResourceIdentifier(ACCOUNT, REGION, RESOURCE_NAME) + + generated_id = generate_test_id( + resource_identifier=resource_identifier, + existing_ids=[TAG_ID], + tags={TAG_KEY_CUSTOM_ID: TAG_ID}, + ) + assert generated_id == GENERATED_ID + + +def test_generate_with_tags_fallback(set_custom_id): + resource_identifier = TestResourceIdentifier(ACCOUNT, REGION, RESOURCE_NAME) + + set_custom_id(resource_identifier, CUSTOM_ID) + generated_id = generate_test_id( + resource_identifier=resource_identifier, + existing_ids=[TAG_ID], + tags={TAG_KEY_CUSTOM_ID: TAG_ID}, + ) + assert generated_id == CUSTOM_ID