diff --git a/chalice/app.py b/chalice/app.py index 77962f439..18fcb7e03 100644 --- a/chalice/app.py +++ b/chalice/app.py @@ -1436,7 +1436,8 @@ class AuthResponse(object): 'PATCH', 'POST', 'PUT', 'GET'] def __init__(self, routes: List[Union[str, 'AuthRoute']], - principal_id: str, context: Optional[Dict[str, str]] = None): + principal_id: str, context: Optional[Dict[str, str]] = None, + usage_identifier_key: Optional[str] = None): self.routes: List[Union[str, 'AuthRoute']] = routes self.principal_id: str = principal_id # The request is used to generate full qualified ARNs @@ -1445,13 +1446,17 @@ def __init__(self, routes: List[Union[str, 'AuthRoute']], if context is None: context = {} self.context: Dict[str, str] = context + self.usage_identifier_key: Optional[str] = usage_identifier_key def to_dict(self, request: AuthRequest) -> Dict[str, Any]: - return { + response = { 'context': self.context, 'principalId': self.principal_id, 'policyDocument': self._generate_policy(request), } + if self.usage_identifier_key: + response['usageIdentifierKey'] = self.usage_identifier_key + return response def _generate_policy(self, request: AuthRequest) -> Dict[str, Any]: allowed_resources = self._generate_allowed_resources(request) diff --git a/docs/source/api.rst b/docs/source/api.rst index b9ca7613b..e271cce09 100644 --- a/docs/source/api.rst +++ b/docs/source/api.rst @@ -812,6 +812,11 @@ These classes are used when defining built-in authorizers in Chalice. will be accessible in the ``app.current_request.context`` in all subsequent authorized requests for this user. + .. attribute:: usage_identifier_key + + An optional string value that represents a usage plan's api + key if the ``apiKeySource`` for that plan is set to ``AUTHORIZER``. + .. class:: AuthRoute(path, methods) This class be used in the ``routes`` attribute of a diff --git a/tests/unit/test_app.py b/tests/unit/test_app.py index ea33f3ce3..d681b9a76 100644 --- a/tests/unit/test_app.py +++ b/tests/unit/test_app.py @@ -1507,6 +1507,39 @@ def test_root_resource(auth_request): } +def test_auth_response_with_usage_identifier_key(auth_request): + event = { + 'type': 'TOKEN', + 'authorizationToken': 'authtoken', + 'methodArn': 'arn:aws:execute-api:us-west-2:1:id/dev/GET/a', + } + auth_app = app.Chalice('builtin-auth') + + response = { + 'context': {}, + 'principalId': 'principal', + 'policyDocument': { + 'Version': '2012-10-17', + 'Statement': [ + {'Action': 'execute-api:Invoke', + 'Effect': 'Allow', + 'Resource': [ + 'arn:aws:execute-api:us-west-2:1:id/dev/*/a' + ]} + ] + }, + 'usageIdentifierKey': 'api-key' + } + + @auth_app.authorizer() + def builtin_auth(auth_request): + return app.AuthResponse(['/a'], 'principal', + usage_identifier_key='api-key') + + actual = builtin_auth(event, None) + assert actual == response + + def test_can_register_scheduled_event_with_str(sample_app): @sample_app.schedule('rate(1 minute)') def foo(event):