diff --git a/tracecat/secrets/service.py b/tracecat/secrets/service.py index c361ca00e..2c50b4738 100644 --- a/tracecat/secrets/service.py +++ b/tracecat/secrets/service.py @@ -115,7 +115,18 @@ async def get_secret_by_name( secret_name: str, environment: str | None = None, ) -> Secret: - """Get a workspace secret by name.""" + """Get a workspace secret by name. + + Args: + secret_name: The name of the secret to retrieve + environment: Optional environment to filter by. If provided, only returns secrets for that environment. + + Returns: + The matching Secret object + + Raises: + TracecatNotFoundError: If no secret is found with the given name/environment or if multiple secrets are found + """ statement = select(Secret).where( Secret.owner_id == self.role.workspace_id, diff --git a/tracecat/validation/service.py b/tracecat/validation/service.py index c771f9493..f5dcdad4f 100644 --- a/tracecat/validation/service.py +++ b/tracecat/validation/service.py @@ -3,7 +3,7 @@ from typing import Any from pydantic import ValidationError -from sqlalchemy.exc import MultipleResultsFound, NoResultFound +from sqlalchemy.exc import MultipleResultsFound from sqlmodel.ext.asyncio.session import AsyncSession from tracecat_registry import RegistrySecret @@ -19,7 +19,7 @@ from tracecat.registry.actions.models import RegistryActionInterface from tracecat.registry.actions.service import RegistryActionsService from tracecat.secrets.service import SecretsService -from tracecat.types.exceptions import RegistryValidationError +from tracecat.types.exceptions import RegistryValidationError, TracecatNotFoundError from tracecat.validation.common import json_schema_to_pydantic, secret_validator from tracecat.validation.models import ( ExprValidationResult, @@ -45,18 +45,17 @@ async def validate_single_secret( try: defined_secret = await secrets_service.get_secret_by_name( - registry_secret.name, - environment=environment, - ) - except (NoResultFound, MultipleResultsFound) as e: - secret_repr = f"{registry_secret.name!r} (env: {environment!r})" - msg = ( - f"Secret {secret_repr} is missing in the secrets manager." - if isinstance(e, NoResultFound) - else f"Multiple secrets found when searching for secret {secret_repr}." + registry_secret.name, environment=environment ) + except TracecatNotFoundError as e: # If the secret is required, we fail early if not registry_secret.optional: + secret_repr = f"{registry_secret.name!r} (env: {environment!r})" + match e.__cause__: + case MultipleResultsFound(): + msg = f"Multiple secrets found when searching for secret {secret_repr}." + case _: + msg = f"Secret {secret_repr} is missing in the secrets manager." return [ SecretValidationResult( status="error",