diff --git a/templates/agent/handler.py b/templates/agent/handler.py index f990c50..9d8f1d8 100644 --- a/templates/agent/handler.py +++ b/templates/agent/handler.py @@ -2,9 +2,11 @@ from aws_lambda_powertools.event_handler import BedrockAgentFunctionResolver from aws_lambda_powertools.utilities.data_classes import BedrockAgentEvent from aws_lambda_powertools.utilities.typing import LambdaContext +from pydantic import ValidationError from templates.agent.models import Item from templates.agent.settings import Settings +from templates.models import Entity from templates.repository import Repository settings = Settings() # type: ignore @@ -29,6 +31,11 @@ def get_item(item_id: str) -> dict: The item details or an error message. """ logger.info("Retrieving item", extra={"itemId": item_id}) + try: + Entity(id=item_id) # Validate ID format before querying repository + except ValidationError: + return {"error": "Invalid item ID format"} + try: item = repository.get_item(item_id) if not item: diff --git a/templates/graphql/handler.py b/templates/graphql/handler.py index e000537..5f504fa 100644 --- a/templates/graphql/handler.py +++ b/templates/graphql/handler.py @@ -6,6 +6,7 @@ from templates.graphql.models import Item from templates.graphql.settings import Settings +from templates.models import Entity from templates.repository import Repository settings = Settings() # type: ignore @@ -29,6 +30,11 @@ def get_item(id: str) -> dict | None: Returns: The item if found, or None. """ + try: + Entity(id=id) # Validate ID format before querying repository + except ValidationError: + raise RuntimeError("Invalid item ID format") from None + try: if (item := repository.get_item(id)) is None: return None diff --git a/templates/models.py b/templates/models.py index 5769ace..e2d4ad6 100644 --- a/templates/models.py +++ b/templates/models.py @@ -29,4 +29,5 @@ class Entity(Object): default_factory=lambda: str(uuid4()), min_length=1, max_length=50, + pattern=r"^[a-zA-Z0-9-_]+$", )