implement retrieving from _custom_id_ tag

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,

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,
         )

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,23 +38,25 @@ 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:
         return ".".join(
             [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,44 +89,80 @@ 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
 
 
 moto_id_manager = MotoIdManager()
 
 
 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
 
 
 @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,

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