From 1c53296d80685d12ce3da0b14b765299438e1a37 Mon Sep 17 00:00:00 2001 From: Mathieu Cloutier Date: Fri, 16 May 2025 11:47:25 +0200 Subject: [PATCH 1/7] Added test --- .../test_apigateway_integrations.py | 79 +++++++++++++++++++ ...test_apigateway_integrations.snapshot.json | 13 +++ ...st_apigateway_integrations.validation.json | 3 + 3 files changed, 95 insertions(+) diff --git a/tests/aws/services/apigateway/test_apigateway_integrations.py b/tests/aws/services/apigateway/test_apigateway_integrations.py index d9780595186c8..bd87b465baf15 100644 --- a/tests/aws/services/apigateway/test_apigateway_integrations.py +++ b/tests/aws/services/apigateway/test_apigateway_integrations.py @@ -724,6 +724,85 @@ def invoke_api(url) -> requests.Response: snapshot.match("invoke-path-else", response_data_3.json()) +@markers.aws.validated +def test_integration_mock_with_response_override_in_request_template( + create_rest_apigw, aws_client, snapshot +): + expected_status = 444 + api_id, _, root_id = create_rest_apigw( + name=f"test-api-{short_uid()}", + description="this is my api", + ) + + aws_client.apigateway.put_method( + restApiId=api_id, + resourceId=root_id, + httpMethod="GET", + authorizationType="NONE", + ) + + aws_client.apigateway.put_method_response( + restApiId=api_id, resourceId=root_id, httpMethod="GET", statusCode="200" + ) + + request_template = textwrap.dedent(f""" + #set($context.responseOverride.status = {expected_status}) + #set($context.responseOverride.header.foo = "bar") + #set($context.responseOverride.custom = "is also passed around") + {{ + "statusCode": 200 + }} + """) + + aws_client.apigateway.put_integration( + restApiId=api_id, + resourceId=root_id, + httpMethod="GET", + integrationHttpMethod="POST", + type="MOCK", + requestParameters={}, + requestTemplates={"application/json": request_template}, + ) + response_template = textwrap.dedent(""" + #set($statusOverride = $context.responseOverride.status) + #set($fooHeader = $context.responseOverride.header.foo) + #set($custom = $context.responseOverride.custom) + { + "statusOverride": "$statusOverride", + "fooHeader": "$fooHeader", + "custom": "$custom" + } + """) + + aws_client.apigateway.put_integration_response( + restApiId=api_id, + resourceId=root_id, + httpMethod="GET", + statusCode="200", + selectionPattern="2\\d{2}", + responseTemplates={"application/json": response_template}, + ) + stage_name = "dev" + aws_client.apigateway.create_deployment(restApiId=api_id, stageName=stage_name) + + invocation_url = api_invoke_url(api_id=api_id, stage=stage_name) + + def invoke_api(url) -> requests.Response: + _response = requests.get(url, verify=False) + assert _response.status_code == expected_status + return _response + + response_data = retry(invoke_api, sleep=2, retries=10, url=invocation_url) + assert response_data.headers["foo"] == "bar" + snapshot.match( + "response", + { + "body": response_data.json(), + "status_code": response_data.status_code, + }, + ) + + @pytest.fixture def default_vpc(aws_client): vpcs = aws_client.ec2.describe_vpcs() diff --git a/tests/aws/services/apigateway/test_apigateway_integrations.snapshot.json b/tests/aws/services/apigateway/test_apigateway_integrations.snapshot.json index 821a3a98b8c3b..94de18011bc1f 100644 --- a/tests/aws/services/apigateway/test_apigateway_integrations.snapshot.json +++ b/tests/aws/services/apigateway/test_apigateway_integrations.snapshot.json @@ -1078,5 +1078,18 @@ } } } + }, + "tests/aws/services/apigateway/test_apigateway_integrations.py::test_integration_mock_with_response_override_in_request_template": { + "recorded-date": "16-05-2025, 09:46:05", + "recorded-content": { + "response": { + "body": { + "custom": "is also passed around", + "fooHeader": "bar", + "statusOverride": "444" + }, + "status_code": 444 + } + } } } diff --git a/tests/aws/services/apigateway/test_apigateway_integrations.validation.json b/tests/aws/services/apigateway/test_apigateway_integrations.validation.json index 9a6dd24061fb0..2780b17c042f4 100644 --- a/tests/aws/services/apigateway/test_apigateway_integrations.validation.json +++ b/tests/aws/services/apigateway/test_apigateway_integrations.validation.json @@ -20,6 +20,9 @@ "tests/aws/services/apigateway/test_apigateway_integrations.py::test_integration_mock_with_request_overrides_in_response_template": { "last_validated_date": "2024-11-06T23:09:04+00:00" }, + "tests/aws/services/apigateway/test_apigateway_integrations.py::test_integration_mock_with_response_override_in_request_template": { + "last_validated_date": "2025-05-16T09:46:05+00:00" + }, "tests/aws/services/apigateway/test_apigateway_integrations.py::test_put_integration_response_with_response_template": { "last_validated_date": "2024-05-30T16:15:58+00:00" }, From af65f66ef1b9e3c67dae03a1b422f0379f05e92d Mon Sep 17 00:00:00 2001 From: Mathieu Cloutier Date: Fri, 16 May 2025 11:50:15 +0200 Subject: [PATCH 2/7] update handlers --- .../execute_api/handlers/integration_request.py | 13 ++++++++----- .../next_gen/execute_api/template_mapping.py | 10 +++++----- 2 files changed, 13 insertions(+), 10 deletions(-) diff --git a/localstack-core/localstack/services/apigateway/next_gen/execute_api/handlers/integration_request.py b/localstack-core/localstack/services/apigateway/next_gen/execute_api/handlers/integration_request.py index 8f0be6a7a6236..21d20f2a07833 100644 --- a/localstack-core/localstack/services/apigateway/next_gen/execute_api/handlers/integration_request.py +++ b/localstack-core/localstack/services/apigateway/next_gen/execute_api/handlers/integration_request.py @@ -22,7 +22,7 @@ MappingTemplateParams, MappingTemplateVariables, ) -from ..variables import ContextVarsRequestOverride +from ..variables import ContextVariables, ContextVarsRequestOverride LOG = logging.getLogger(__name__) @@ -119,13 +119,16 @@ def __call__( converted_body = self.convert_body(context) - body, request_override = self.render_request_template_mapping( + body, context_variables = self.render_request_template_mapping( context=context, body=converted_body, template=request_template ) # mutate the ContextVariables with the requestOverride result, as we copy the context when rendering the # template to avoid mutation on other fields # the VTL responseTemplate can access the requestOverride + request_override: ContextVarsRequestOverride = context_variables["requestOverride"] context.context_variables["requestOverride"] = request_override + # Response override attributes set in the request template will be available in the response template + context.context_variables["responseOverride"] = context_variables["responseOverride"] # TODO: log every override that happens afterwards (in a loop on `request_override`) merge_recursive(request_override, request_data_mapping, overwrite=True) @@ -180,7 +183,7 @@ def render_request_template_mapping( context: RestApiInvocationContext, body: str | bytes, template: str, - ) -> tuple[bytes, ContextVarsRequestOverride]: + ) -> tuple[bytes, ContextVariables]: request: InvocationRequest = context.invocation_request if not template: @@ -191,7 +194,7 @@ def render_request_template_mapping( except UnicodeError: raise InternalServerError("Internal server error") - body, request_override = self._vtl_template.render_request( + body, context_variables = self._vtl_template.render_request( template=template, variables=MappingTemplateVariables( context=context.context_variables, @@ -206,7 +209,7 @@ def render_request_template_mapping( ), ), ) - return to_bytes(body), request_override + return to_bytes(body), context_variables @staticmethod def get_request_template(integration: Integration, request: InvocationRequest) -> str: diff --git a/localstack-core/localstack/services/apigateway/next_gen/execute_api/template_mapping.py b/localstack-core/localstack/services/apigateway/next_gen/execute_api/template_mapping.py index ac47dd3465107..92665e1f8507d 100644 --- a/localstack-core/localstack/services/apigateway/next_gen/execute_api/template_mapping.py +++ b/localstack-core/localstack/services/apigateway/next_gen/execute_api/template_mapping.py @@ -262,21 +262,21 @@ def prepare_namespace(self, variables, source: str = APIGW_SOURCE) -> dict[str, def render_request( self, template: str, variables: MappingTemplateVariables - ) -> tuple[str, ContextVarsRequestOverride]: + ) -> tuple[str, ContextVariables]: variables_copy: MappingTemplateVariables = copy.deepcopy(variables) variables_copy["context"]["requestOverride"] = ContextVarsRequestOverride( querystring={}, header={}, path={} ) + variables_copy["context"]["responseOverride"] = ContextVarsResponseOverride( + header={}, status=0 + ) result = self.render_vtl(template=template.strip(), variables=variables_copy) - return result, variables_copy["context"]["requestOverride"] + return result, variables_copy["context"] def render_response( self, template: str, variables: MappingTemplateVariables ) -> tuple[str, ContextVarsResponseOverride]: variables_copy: MappingTemplateVariables = copy.deepcopy(variables) - variables_copy["context"]["responseOverride"] = ContextVarsResponseOverride( - header={}, status=0 - ) result = self.render_vtl(template=template.strip(), variables=variables_copy) return result, variables_copy["context"]["responseOverride"] From 676f9bcce5458e3d11ffb47919d82781e44f1549 Mon Sep 17 00:00:00 2001 From: Mathieu Cloutier Date: Fri, 16 May 2025 12:04:53 +0200 Subject: [PATCH 3/7] Handle case where there is no mappings --- .../next_gen/execute_api/handlers/integration_request.py | 8 ++++++-- .../apigateway/next_gen/execute_api/template_mapping.py | 4 ++++ tests/unit/services/apigateway/test_template_mapping.py | 3 ++- 3 files changed, 12 insertions(+), 3 deletions(-) diff --git a/localstack-core/localstack/services/apigateway/next_gen/execute_api/handlers/integration_request.py b/localstack-core/localstack/services/apigateway/next_gen/execute_api/handlers/integration_request.py index 21d20f2a07833..087fe7c600b82 100644 --- a/localstack-core/localstack/services/apigateway/next_gen/execute_api/handlers/integration_request.py +++ b/localstack-core/localstack/services/apigateway/next_gen/execute_api/handlers/integration_request.py @@ -125,10 +125,14 @@ def __call__( # mutate the ContextVariables with the requestOverride result, as we copy the context when rendering the # template to avoid mutation on other fields # the VTL responseTemplate can access the requestOverride - request_override: ContextVarsRequestOverride = context_variables["requestOverride"] + request_override: ContextVarsRequestOverride = context_variables.get( + "requestOverride", {} + ) context.context_variables["requestOverride"] = request_override # Response override attributes set in the request template will be available in the response template - context.context_variables["responseOverride"] = context_variables["responseOverride"] + context.context_variables["responseOverride"] = context_variables.get( + "responseOverride", {} + ) # TODO: log every override that happens afterwards (in a loop on `request_override`) merge_recursive(request_override, request_data_mapping, overwrite=True) diff --git a/localstack-core/localstack/services/apigateway/next_gen/execute_api/template_mapping.py b/localstack-core/localstack/services/apigateway/next_gen/execute_api/template_mapping.py index 92665e1f8507d..8332e01425069 100644 --- a/localstack-core/localstack/services/apigateway/next_gen/execute_api/template_mapping.py +++ b/localstack-core/localstack/services/apigateway/next_gen/execute_api/template_mapping.py @@ -277,6 +277,10 @@ def render_response( self, template: str, variables: MappingTemplateVariables ) -> tuple[str, ContextVarsResponseOverride]: variables_copy: MappingTemplateVariables = copy.deepcopy(variables) + if not variables_copy["context"].get("responseOverride"): + variables_copy["context"]["responseOverride"] = ContextVarsResponseOverride( + header={}, status=0 + ) result = self.render_vtl(template=template.strip(), variables=variables_copy) return result, variables_copy["context"]["responseOverride"] diff --git a/tests/unit/services/apigateway/test_template_mapping.py b/tests/unit/services/apigateway/test_template_mapping.py index 25d711a04b54b..12db254c91370 100644 --- a/tests/unit/services/apigateway/test_template_mapping.py +++ b/tests/unit/services/apigateway/test_template_mapping.py @@ -122,9 +122,10 @@ def test_render_custom_template(self, format): template = TEMPLATE_JSON if format == APPLICATION_JSON else TEMPLATE_XML template += REQUEST_OVERRIDE - rendered_request, request_override = ApiGatewayVtlTemplate().render_request( + rendered_request, context_variable = ApiGatewayVtlTemplate().render_request( template=template, variables=variables ) + request_override = context_variable["requestOverride"] if format == APPLICATION_JSON: rendered_request = json.loads(rendered_request) assert rendered_request.get("body") == {"spam": "eggs"} From b12fdc5af9f48f4b04f9f4ca755712b5319ea094 Mon Sep 17 00:00:00 2001 From: Mathieu Cloutier Date: Fri, 16 May 2025 12:43:13 +0200 Subject: [PATCH 4/7] Swiotch responsibility of generating override context variable to the parser to prevent if block everywhere --- .../execute_api/handlers/integration_request.py | 12 ++++++------ .../execute_api/handlers/integration_response.py | 2 +- .../next_gen/execute_api/handlers/parse.py | 9 ++++++++- .../next_gen/execute_api/template_mapping.py | 11 ----------- .../apigateway/test_apigateway_integrations.py | 9 ++++++--- .../test_apigateway_integrations.snapshot.json | 13 +++++++++++-- .../test_apigateway_integrations.validation.json | 7 +++++-- .../apigateway/test_handler_integration_request.py | 4 ++++ .../apigateway/test_handler_integration_response.py | 9 +++++++-- .../services/apigateway/test_template_mapping.py | 4 ++++ 10 files changed, 52 insertions(+), 28 deletions(-) diff --git a/localstack-core/localstack/services/apigateway/next_gen/execute_api/handlers/integration_request.py b/localstack-core/localstack/services/apigateway/next_gen/execute_api/handlers/integration_request.py index 087fe7c600b82..0a2587d6e8605 100644 --- a/localstack-core/localstack/services/apigateway/next_gen/execute_api/handlers/integration_request.py +++ b/localstack-core/localstack/services/apigateway/next_gen/execute_api/handlers/integration_request.py @@ -119,18 +119,18 @@ def __call__( converted_body = self.convert_body(context) - body, context_variables = self.render_request_template_mapping( + body, mapped_context_variables = self.render_request_template_mapping( context=context, body=converted_body, template=request_template ) # mutate the ContextVariables with the requestOverride result, as we copy the context when rendering the # template to avoid mutation on other fields # the VTL responseTemplate can access the requestOverride - request_override: ContextVarsRequestOverride = context_variables.get( + request_override: ContextVarsRequestOverride = mapped_context_variables.get( "requestOverride", {} ) context.context_variables["requestOverride"] = request_override # Response override attributes set in the request template will be available in the response template - context.context_variables["responseOverride"] = context_variables.get( + context.context_variables["responseOverride"] = mapped_context_variables.get( "responseOverride", {} ) # TODO: log every override that happens afterwards (in a loop on `request_override`) @@ -191,14 +191,14 @@ def render_request_template_mapping( request: InvocationRequest = context.invocation_request if not template: - return to_bytes(body), {} + return to_bytes(body), context.context_variables try: body_utf8 = to_str(body) except UnicodeError: raise InternalServerError("Internal server error") - body, context_variables = self._vtl_template.render_request( + body, mapped_context_variables = self._vtl_template.render_request( template=template, variables=MappingTemplateVariables( context=context.context_variables, @@ -213,7 +213,7 @@ def render_request_template_mapping( ), ), ) - return to_bytes(body), context_variables + return to_bytes(body), mapped_context_variables @staticmethod def get_request_template(integration: Integration, request: InvocationRequest) -> str: diff --git a/localstack-core/localstack/services/apigateway/next_gen/execute_api/handlers/integration_response.py b/localstack-core/localstack/services/apigateway/next_gen/execute_api/handlers/integration_response.py index 02d09db8332c1..34816eb689b1e 100644 --- a/localstack-core/localstack/services/apigateway/next_gen/execute_api/handlers/integration_response.py +++ b/localstack-core/localstack/services/apigateway/next_gen/execute_api/handlers/integration_response.py @@ -263,7 +263,7 @@ def render_response_template_mapping( self, context: RestApiInvocationContext, template: str, body: bytes | str ) -> tuple[bytes, ContextVarsResponseOverride]: if not template: - return to_bytes(body), ContextVarsResponseOverride(status=0, header={}) + return to_bytes(body), context.context_variables["responseOverride"] # if there are no template, we can pass binary data through if not isinstance(body, str): diff --git a/localstack-core/localstack/services/apigateway/next_gen/execute_api/handlers/parse.py b/localstack-core/localstack/services/apigateway/next_gen/execute_api/handlers/parse.py index 5971e7872ebd7..da787af601e6e 100644 --- a/localstack-core/localstack/services/apigateway/next_gen/execute_api/handlers/parse.py +++ b/localstack-core/localstack/services/apigateway/next_gen/execute_api/handlers/parse.py @@ -18,7 +18,12 @@ from ..header_utils import should_drop_header_from_invocation from ..helpers import generate_trace_id, generate_trace_parent, parse_trace_id from ..moto_helpers import get_stage_variables -from ..variables import ContextVariables, ContextVarsIdentity +from ..variables import ( + ContextVariables, + ContextVarsIdentity, + ContextVarsRequestOverride, + ContextVarsResponseOverride, +) LOG = logging.getLogger(__name__) @@ -159,8 +164,10 @@ def create_context_variables(context: RestApiInvocationContext) -> ContextVariab path=f"/{context.stage}{invocation_request['raw_path']}", protocol="HTTP/1.1", requestId=long_uid(), + requestOverride=ContextVarsRequestOverride(querystring={}, header={}, path={}), requestTime=timestamp(time=now, format=REQUEST_TIME_DATE_FORMAT), requestTimeEpoch=int(now.timestamp() * 1000), + responseOverride=ContextVarsResponseOverride(header={}, status=0), stage=context.stage, ) return context_variables diff --git a/localstack-core/localstack/services/apigateway/next_gen/execute_api/template_mapping.py b/localstack-core/localstack/services/apigateway/next_gen/execute_api/template_mapping.py index 8332e01425069..6dc9753370560 100644 --- a/localstack-core/localstack/services/apigateway/next_gen/execute_api/template_mapping.py +++ b/localstack-core/localstack/services/apigateway/next_gen/execute_api/template_mapping.py @@ -27,7 +27,6 @@ from localstack import config from localstack.services.apigateway.next_gen.execute_api.variables import ( ContextVariables, - ContextVarsRequestOverride, ContextVarsResponseOverride, ) from localstack.utils.aws.templating import APIGW_SOURCE, VelocityUtil, VtlTemplate @@ -264,12 +263,6 @@ def render_request( self, template: str, variables: MappingTemplateVariables ) -> tuple[str, ContextVariables]: variables_copy: MappingTemplateVariables = copy.deepcopy(variables) - variables_copy["context"]["requestOverride"] = ContextVarsRequestOverride( - querystring={}, header={}, path={} - ) - variables_copy["context"]["responseOverride"] = ContextVarsResponseOverride( - header={}, status=0 - ) result = self.render_vtl(template=template.strip(), variables=variables_copy) return result, variables_copy["context"] @@ -277,10 +270,6 @@ def render_response( self, template: str, variables: MappingTemplateVariables ) -> tuple[str, ContextVarsResponseOverride]: variables_copy: MappingTemplateVariables = copy.deepcopy(variables) - if not variables_copy["context"].get("responseOverride"): - variables_copy["context"]["responseOverride"] = ContextVarsResponseOverride( - header={}, status=0 - ) result = self.render_vtl(template=template.strip(), variables=variables_copy) return result, variables_copy["context"]["responseOverride"] diff --git a/tests/aws/services/apigateway/test_apigateway_integrations.py b/tests/aws/services/apigateway/test_apigateway_integrations.py index bd87b465baf15..d3e3198a2d86a 100644 --- a/tests/aws/services/apigateway/test_apigateway_integrations.py +++ b/tests/aws/services/apigateway/test_apigateway_integrations.py @@ -725,8 +725,9 @@ def invoke_api(url) -> requests.Response: @markers.aws.validated +@pytest.mark.parametrize("create_response_template", [True, False]) def test_integration_mock_with_response_override_in_request_template( - create_rest_apigw, aws_client, snapshot + create_rest_apigw, aws_client, snapshot, create_response_template ): expected_status = 444 api_id, _, root_id = create_rest_apigw( @@ -780,7 +781,9 @@ def test_integration_mock_with_response_override_in_request_template( httpMethod="GET", statusCode="200", selectionPattern="2\\d{2}", - responseTemplates={"application/json": response_template}, + responseTemplates={"application/json": response_template} + if create_response_template + else {}, ) stage_name = "dev" aws_client.apigateway.create_deployment(restApiId=api_id, stageName=stage_name) @@ -797,7 +800,7 @@ def invoke_api(url) -> requests.Response: snapshot.match( "response", { - "body": response_data.json(), + "body": response_data.json() if create_response_template else response_data.content, "status_code": response_data.status_code, }, ) diff --git a/tests/aws/services/apigateway/test_apigateway_integrations.snapshot.json b/tests/aws/services/apigateway/test_apigateway_integrations.snapshot.json index 94de18011bc1f..d0e0d59455823 100644 --- a/tests/aws/services/apigateway/test_apigateway_integrations.snapshot.json +++ b/tests/aws/services/apigateway/test_apigateway_integrations.snapshot.json @@ -1079,8 +1079,8 @@ } } }, - "tests/aws/services/apigateway/test_apigateway_integrations.py::test_integration_mock_with_response_override_in_request_template": { - "recorded-date": "16-05-2025, 09:46:05", + "tests/aws/services/apigateway/test_apigateway_integrations.py::test_integration_mock_with_response_override_in_request_template[True]": { + "recorded-date": "16-05-2025, 10:22:21", "recorded-content": { "response": { "body": { @@ -1091,5 +1091,14 @@ "status_code": 444 } } + }, + "tests/aws/services/apigateway/test_apigateway_integrations.py::test_integration_mock_with_response_override_in_request_template[False]": { + "recorded-date": "16-05-2025, 10:22:27", + "recorded-content": { + "response": { + "body": "b''", + "status_code": 444 + } + } } } diff --git a/tests/aws/services/apigateway/test_apigateway_integrations.validation.json b/tests/aws/services/apigateway/test_apigateway_integrations.validation.json index 2780b17c042f4..883298cf6153e 100644 --- a/tests/aws/services/apigateway/test_apigateway_integrations.validation.json +++ b/tests/aws/services/apigateway/test_apigateway_integrations.validation.json @@ -20,8 +20,11 @@ "tests/aws/services/apigateway/test_apigateway_integrations.py::test_integration_mock_with_request_overrides_in_response_template": { "last_validated_date": "2024-11-06T23:09:04+00:00" }, - "tests/aws/services/apigateway/test_apigateway_integrations.py::test_integration_mock_with_response_override_in_request_template": { - "last_validated_date": "2025-05-16T09:46:05+00:00" + "tests/aws/services/apigateway/test_apigateway_integrations.py::test_integration_mock_with_response_override_in_request_template[False]": { + "last_validated_date": "2025-05-16T10:22:27+00:00" + }, + "tests/aws/services/apigateway/test_apigateway_integrations.py::test_integration_mock_with_response_override_in_request_template[True]": { + "last_validated_date": "2025-05-16T10:22:21+00:00" }, "tests/aws/services/apigateway/test_apigateway_integrations.py::test_put_integration_response_with_response_template": { "last_validated_date": "2024-05-30T16:15:58+00:00" diff --git a/tests/unit/services/apigateway/test_handler_integration_request.py b/tests/unit/services/apigateway/test_handler_integration_request.py index 1f08d04547261..df2237b752b1b 100644 --- a/tests/unit/services/apigateway/test_handler_integration_request.py +++ b/tests/unit/services/apigateway/test_handler_integration_request.py @@ -21,6 +21,8 @@ ) from localstack.services.apigateway.next_gen.execute_api.variables import ( ContextVariables, + ContextVarsRequestOverride, + ContextVarsResponseOverride, ) from localstack.testing.config import TEST_AWS_ACCOUNT_ID, TEST_AWS_REGION_NAME @@ -80,6 +82,8 @@ def default_context(): path=f"{TEST_API_STAGE}/resource/{{proxy}}", resourcePath="/resource/{proxy}", stage=TEST_API_STAGE, + requestOverride=ContextVarsRequestOverride(header={}, path={}, querystring={}), + responseOverride=ContextVarsResponseOverride(header={}, status=0), ) return context diff --git a/tests/unit/services/apigateway/test_handler_integration_response.py b/tests/unit/services/apigateway/test_handler_integration_response.py index 8ec1a96fe2a4f..0b8429e704a0a 100644 --- a/tests/unit/services/apigateway/test_handler_integration_response.py +++ b/tests/unit/services/apigateway/test_handler_integration_response.py @@ -16,7 +16,10 @@ IntegrationResponseHandler, InvocationRequestParser, ) -from localstack.services.apigateway.next_gen.execute_api.variables import ContextVariables +from localstack.services.apigateway.next_gen.execute_api.variables import ( + ContextVariables, + ContextVarsResponseOverride, +) from localstack.testing.config import TEST_AWS_ACCOUNT_ID, TEST_AWS_REGION_NAME TEST_API_ID = "test-api" @@ -141,7 +144,9 @@ def ctx(): context.invocation_request = request context.integration = Integration(type=IntegrationType.HTTP) - context.context_variables = ContextVariables() + context.context_variables = ContextVariables( + responseOverride=ContextVarsResponseOverride(header={}, status=0) + ) context.endpoint_response = EndpointResponse( body=b'{"foo":"bar"}', status_code=200, diff --git a/tests/unit/services/apigateway/test_template_mapping.py b/tests/unit/services/apigateway/test_template_mapping.py index 12db254c91370..7b18d7fc17790 100644 --- a/tests/unit/services/apigateway/test_template_mapping.py +++ b/tests/unit/services/apigateway/test_template_mapping.py @@ -16,6 +16,8 @@ ContextVariables, ContextVarsAuthorizer, ContextVarsIdentity, + ContextVarsRequestOverride, + ContextVarsResponseOverride, ) @@ -115,6 +117,7 @@ def test_render_custom_template(self, format): authorizer=ContextVarsAuthorizer(principalId="12233"), identity=ContextVarsIdentity(accountId="00000", apiKey="11111"), resourcePath="/{proxy}", + requestOverride=ContextVarsRequestOverride(header={}, path={}, querystring={}), ), stageVariables={"stageVariable1": "value1", "stageVariable2": "value2"}, ) @@ -194,6 +197,7 @@ def test_render_response_template(self, format): authorizer=ContextVarsAuthorizer(principalId="12233"), identity=ContextVarsIdentity(accountId="00000", apiKey="11111"), resourcePath="/{proxy}", + responseOverride=ContextVarsResponseOverride(header={}, status=0), ), stageVariables={"stageVariable1": "value1", "stageVariable2": "value2"}, ) From 18e80420dca79e83c4238964f721cb95c342fd09 Mon Sep 17 00:00:00 2001 From: Mathieu Cloutier Date: Fri, 16 May 2025 13:55:57 +0200 Subject: [PATCH 5/7] remove overrides from main context --- .../next_gen/execute_api/context.py | 5 +++- .../handlers/integration_request.py | 23 ++++++++----------- .../handlers/integration_response.py | 3 ++- .../next_gen/execute_api/handlers/parse.py | 7 ++++-- .../next_gen/execute_api/template_mapping.py | 20 ++++++++++++---- .../next_gen/execute_api/variables.py | 5 ++++ .../test_handler_integration_request.py | 3 +++ .../test_handler_integration_response.py | 9 +++++--- .../apigateway/test_template_mapping.py | 16 +++++++++---- 9 files changed, 62 insertions(+), 29 deletions(-) diff --git a/localstack-core/localstack/services/apigateway/next_gen/execute_api/context.py b/localstack-core/localstack/services/apigateway/next_gen/execute_api/context.py index 03632d0829aaa..95eeb989f6917 100644 --- a/localstack-core/localstack/services/apigateway/next_gen/execute_api/context.py +++ b/localstack-core/localstack/services/apigateway/next_gen/execute_api/context.py @@ -8,7 +8,7 @@ from localstack.aws.api.apigateway import Integration, Method, Resource from localstack.services.apigateway.models import RestApiDeployment -from .variables import ContextVariables, LoggingContextVariables +from .variables import ContextVariableOverrides, ContextVariables, LoggingContextVariables class InvocationRequest(TypedDict, total=False): @@ -98,6 +98,9 @@ class RestApiInvocationContext(RequestContext): """The Stage variables, also used in parameters mapping and mapping templates""" context_variables: Optional[ContextVariables] """The $context used in data models, authorizers, mapping templates, and CloudWatch access logging""" + context_variable_overrides: Optional[ContextVariableOverrides] + """requestOverrides and responseOverrides are passed from request templates to response templates but are + not in the integration context""" logging_context_variables: Optional[LoggingContextVariables] """Additional $context variables available only for access logging, not yet implemented""" invocation_request: Optional[InvocationRequest] diff --git a/localstack-core/localstack/services/apigateway/next_gen/execute_api/handlers/integration_request.py b/localstack-core/localstack/services/apigateway/next_gen/execute_api/handlers/integration_request.py index 0a2587d6e8605..b9cf68b1ab006 100644 --- a/localstack-core/localstack/services/apigateway/next_gen/execute_api/handlers/integration_request.py +++ b/localstack-core/localstack/services/apigateway/next_gen/execute_api/handlers/integration_request.py @@ -22,7 +22,7 @@ MappingTemplateParams, MappingTemplateVariables, ) -from ..variables import ContextVariables, ContextVarsRequestOverride +from ..variables import ContextVariableOverrides, ContextVarsRequestOverride LOG = logging.getLogger(__name__) @@ -119,20 +119,16 @@ def __call__( converted_body = self.convert_body(context) - body, mapped_context_variables = self.render_request_template_mapping( + body, mapped_overrides = self.render_request_template_mapping( context=context, body=converted_body, template=request_template ) + # Update the context with the returned mapped overrides + context.context_variable_overrides = mapped_overrides # mutate the ContextVariables with the requestOverride result, as we copy the context when rendering the # template to avoid mutation on other fields - # the VTL responseTemplate can access the requestOverride - request_override: ContextVarsRequestOverride = mapped_context_variables.get( + request_override: ContextVarsRequestOverride = mapped_overrides.get( "requestOverride", {} ) - context.context_variables["requestOverride"] = request_override - # Response override attributes set in the request template will be available in the response template - context.context_variables["responseOverride"] = mapped_context_variables.get( - "responseOverride", {} - ) # TODO: log every override that happens afterwards (in a loop on `request_override`) merge_recursive(request_override, request_data_mapping, overwrite=True) @@ -187,18 +183,18 @@ def render_request_template_mapping( context: RestApiInvocationContext, body: str | bytes, template: str, - ) -> tuple[bytes, ContextVariables]: + ) -> tuple[bytes, ContextVariableOverrides]: request: InvocationRequest = context.invocation_request if not template: - return to_bytes(body), context.context_variables + return to_bytes(body), context.context_variable_overrides try: body_utf8 = to_str(body) except UnicodeError: raise InternalServerError("Internal server error") - body, mapped_context_variables = self._vtl_template.render_request( + body, mapped_overrides = self._vtl_template.render_request( template=template, variables=MappingTemplateVariables( context=context.context_variables, @@ -212,8 +208,9 @@ def render_request_template_mapping( ), ), ), + context_overrides=context.context_variable_overrides, ) - return to_bytes(body), mapped_context_variables + return to_bytes(body), mapped_overrides @staticmethod def get_request_template(integration: Integration, request: InvocationRequest) -> str: diff --git a/localstack-core/localstack/services/apigateway/next_gen/execute_api/handlers/integration_response.py b/localstack-core/localstack/services/apigateway/next_gen/execute_api/handlers/integration_response.py index 34816eb689b1e..2dccb39c74a6b 100644 --- a/localstack-core/localstack/services/apigateway/next_gen/execute_api/handlers/integration_response.py +++ b/localstack-core/localstack/services/apigateway/next_gen/execute_api/handlers/integration_response.py @@ -263,7 +263,7 @@ def render_response_template_mapping( self, context: RestApiInvocationContext, template: str, body: bytes | str ) -> tuple[bytes, ContextVarsResponseOverride]: if not template: - return to_bytes(body), context.context_variables["responseOverride"] + return to_bytes(body), context.context_variable_overrides["responseOverride"] # if there are no template, we can pass binary data through if not isinstance(body, str): @@ -284,6 +284,7 @@ def render_response_template_mapping( ), ), ), + context_overrides=context.context_variable_overrides, ) # AWS ignores the status if the override isn't an integer between 100 and 599 diff --git a/localstack-core/localstack/services/apigateway/next_gen/execute_api/handlers/parse.py b/localstack-core/localstack/services/apigateway/next_gen/execute_api/handlers/parse.py index da787af601e6e..829314807752d 100644 --- a/localstack-core/localstack/services/apigateway/next_gen/execute_api/handlers/parse.py +++ b/localstack-core/localstack/services/apigateway/next_gen/execute_api/handlers/parse.py @@ -19,6 +19,7 @@ from ..helpers import generate_trace_id, generate_trace_parent, parse_trace_id from ..moto_helpers import get_stage_variables from ..variables import ( + ContextVariableOverrides, ContextVariables, ContextVarsIdentity, ContextVarsRequestOverride, @@ -45,6 +46,10 @@ def parse_and_enrich(self, context: RestApiInvocationContext): # then we can create the ContextVariables, used throughout the invocation as payload and to render authorizer # payload, mapping templates and such. context.context_variables = self.create_context_variables(context) + context.context_variable_overrides = ContextVariableOverrides( + requestOverride=ContextVarsRequestOverride(header={}, querystring={}, path={}), + responseOverride=ContextVarsResponseOverride(header={}, status=0), + ) # TODO: maybe adjust the logging LOG.debug("Initializing $context='%s'", context.context_variables) # then populate the stage variables @@ -164,10 +169,8 @@ def create_context_variables(context: RestApiInvocationContext) -> ContextVariab path=f"/{context.stage}{invocation_request['raw_path']}", protocol="HTTP/1.1", requestId=long_uid(), - requestOverride=ContextVarsRequestOverride(querystring={}, header={}, path={}), requestTime=timestamp(time=now, format=REQUEST_TIME_DATE_FORMAT), requestTimeEpoch=int(now.timestamp() * 1000), - responseOverride=ContextVarsResponseOverride(header={}, status=0), stage=context.stage, ) return context_variables diff --git a/localstack-core/localstack/services/apigateway/next_gen/execute_api/template_mapping.py b/localstack-core/localstack/services/apigateway/next_gen/execute_api/template_mapping.py index 6dc9753370560..01beb0114f598 100644 --- a/localstack-core/localstack/services/apigateway/next_gen/execute_api/template_mapping.py +++ b/localstack-core/localstack/services/apigateway/next_gen/execute_api/template_mapping.py @@ -26,6 +26,7 @@ from localstack import config from localstack.services.apigateway.next_gen.execute_api.variables import ( + ContextVariableOverrides, ContextVariables, ContextVarsResponseOverride, ) @@ -260,16 +261,27 @@ def prepare_namespace(self, variables, source: str = APIGW_SOURCE) -> dict[str, return namespace def render_request( - self, template: str, variables: MappingTemplateVariables - ) -> tuple[str, ContextVariables]: + self, + template: str, + variables: MappingTemplateVariables, + context_overrides: ContextVariableOverrides, + ) -> tuple[str, ContextVariableOverrides]: variables_copy: MappingTemplateVariables = copy.deepcopy(variables) + variables_copy["context"].update(copy.deepcopy(context_overrides)) result = self.render_vtl(template=template.strip(), variables=variables_copy) - return result, variables_copy["context"] + return result, ContextVariableOverrides( + requestOverride=variables_copy["context"]["requestOverride"], + responseOverride=variables_copy["context"]["responseOverride"], + ) def render_response( - self, template: str, variables: MappingTemplateVariables + self, + template: str, + variables: MappingTemplateVariables, + context_overrides: ContextVariableOverrides, ) -> tuple[str, ContextVarsResponseOverride]: variables_copy: MappingTemplateVariables = copy.deepcopy(variables) + variables_copy["context"].update(copy.deepcopy(context_overrides)) result = self.render_vtl(template=template.strip(), variables=variables_copy) return result, variables_copy["context"]["responseOverride"] diff --git a/localstack-core/localstack/services/apigateway/next_gen/execute_api/variables.py b/localstack-core/localstack/services/apigateway/next_gen/execute_api/variables.py index 6403f01852752..76d5a40b18710 100644 --- a/localstack-core/localstack/services/apigateway/next_gen/execute_api/variables.py +++ b/localstack-core/localstack/services/apigateway/next_gen/execute_api/variables.py @@ -75,6 +75,11 @@ class ContextVarsResponseOverride(TypedDict): status: int +class ContextVariableOverrides(TypedDict): + requestOverride: ContextVarsRequestOverride + responseOverride: ContextVarsResponseOverride + + class GatewayResponseContextVarsError(TypedDict, total=False): # This variable can only be used for simple variable substitution in a GatewayResponse body-mapping template, # which is not processed by the Velocity Template Language engine, and in access logging. diff --git a/tests/unit/services/apigateway/test_handler_integration_request.py b/tests/unit/services/apigateway/test_handler_integration_request.py index df2237b752b1b..72b021e4b2d63 100644 --- a/tests/unit/services/apigateway/test_handler_integration_request.py +++ b/tests/unit/services/apigateway/test_handler_integration_request.py @@ -20,6 +20,7 @@ PassthroughBehavior, ) from localstack.services.apigateway.next_gen.execute_api.variables import ( + ContextVariableOverrides, ContextVariables, ContextVarsRequestOverride, ContextVarsResponseOverride, @@ -82,6 +83,8 @@ def default_context(): path=f"{TEST_API_STAGE}/resource/{{proxy}}", resourcePath="/resource/{proxy}", stage=TEST_API_STAGE, + ) + context.context_variable_overrides = ContextVariableOverrides( requestOverride=ContextVarsRequestOverride(header={}, path={}, querystring={}), responseOverride=ContextVarsResponseOverride(header={}, status=0), ) diff --git a/tests/unit/services/apigateway/test_handler_integration_response.py b/tests/unit/services/apigateway/test_handler_integration_response.py index 0b8429e704a0a..122af7c5bbc13 100644 --- a/tests/unit/services/apigateway/test_handler_integration_response.py +++ b/tests/unit/services/apigateway/test_handler_integration_response.py @@ -17,7 +17,8 @@ InvocationRequestParser, ) from localstack.services.apigateway.next_gen.execute_api.variables import ( - ContextVariables, + ContextVariableOverrides, + ContextVarsRequestOverride, ContextVarsResponseOverride, ) from localstack.testing.config import TEST_AWS_ACCOUNT_ID, TEST_AWS_REGION_NAME @@ -144,8 +145,10 @@ def ctx(): context.invocation_request = request context.integration = Integration(type=IntegrationType.HTTP) - context.context_variables = ContextVariables( - responseOverride=ContextVarsResponseOverride(header={}, status=0) + context.context_variables = {} + context.context_variable_overrides = ContextVariableOverrides( + requestOverride=ContextVarsRequestOverride(header={}, path={}, querystring={}), + responseOverride=ContextVarsResponseOverride(header={}, status=0), ) context.endpoint_response = EndpointResponse( body=b'{"foo":"bar"}', diff --git a/tests/unit/services/apigateway/test_template_mapping.py b/tests/unit/services/apigateway/test_template_mapping.py index 7b18d7fc17790..9b58d85225736 100644 --- a/tests/unit/services/apigateway/test_template_mapping.py +++ b/tests/unit/services/apigateway/test_template_mapping.py @@ -13,6 +13,7 @@ VelocityUtilApiGateway, ) from localstack.services.apigateway.next_gen.execute_api.variables import ( + ContextVariableOverrides, ContextVariables, ContextVarsAuthorizer, ContextVarsIdentity, @@ -117,16 +118,19 @@ def test_render_custom_template(self, format): authorizer=ContextVarsAuthorizer(principalId="12233"), identity=ContextVarsIdentity(accountId="00000", apiKey="11111"), resourcePath="/{proxy}", - requestOverride=ContextVarsRequestOverride(header={}, path={}, querystring={}), ), stageVariables={"stageVariable1": "value1", "stageVariable2": "value2"}, ) + context_overrides = ContextVariableOverrides( + requestOverride=ContextVarsRequestOverride(header={}, path={}, querystring={}), + responseOverride=ContextVarsResponseOverride(header={}, status=0), + ) template = TEMPLATE_JSON if format == APPLICATION_JSON else TEMPLATE_XML template += REQUEST_OVERRIDE rendered_request, context_variable = ApiGatewayVtlTemplate().render_request( - template=template, variables=variables + template=template, variables=variables, context_overrides=context_overrides ) request_override = context_variable["requestOverride"] if format == APPLICATION_JSON: @@ -197,16 +201,18 @@ def test_render_response_template(self, format): authorizer=ContextVarsAuthorizer(principalId="12233"), identity=ContextVarsIdentity(accountId="00000", apiKey="11111"), resourcePath="/{proxy}", - responseOverride=ContextVarsResponseOverride(header={}, status=0), ), stageVariables={"stageVariable1": "value1", "stageVariable2": "value2"}, ) - + context_overrides = ContextVariableOverrides( + requestOverride=ContextVarsRequestOverride(header={}, path={}, querystring={}), + responseOverride=ContextVarsResponseOverride(header={}, status=0), + ) template = TEMPLATE_JSON if format == APPLICATION_JSON else TEMPLATE_XML template += RESPONSE_OVERRIDE rendered_response, response_override = ApiGatewayVtlTemplate().render_response( - template=template, variables=variables + template=template, variables=variables, context_overrides=context_overrides ) if format == APPLICATION_JSON: rendered_response = json.loads(rendered_response) From bb5dbe1c81d73dcc24345f8b7990f5f27336a0a2 Mon Sep 17 00:00:00 2001 From: Mathieu Cloutier Date: Fri, 16 May 2025 16:00:41 +0200 Subject: [PATCH 6/7] fix test --- .../apigateway/next_gen/execute_api/context.py | 1 + .../apigateway/next_gen/execute_api/test_invoke.py | 11 +++++++++-- 2 files changed, 10 insertions(+), 2 deletions(-) diff --git a/localstack-core/localstack/services/apigateway/next_gen/execute_api/context.py b/localstack-core/localstack/services/apigateway/next_gen/execute_api/context.py index 95eeb989f6917..932eacee71048 100644 --- a/localstack-core/localstack/services/apigateway/next_gen/execute_api/context.py +++ b/localstack-core/localstack/services/apigateway/next_gen/execute_api/context.py @@ -132,3 +132,4 @@ def __init__(self, request: Request): self.endpoint_response = None self.invocation_response = None self.trace_id = None + self.context_variable_overrides = None diff --git a/localstack-core/localstack/services/apigateway/next_gen/execute_api/test_invoke.py b/localstack-core/localstack/services/apigateway/next_gen/execute_api/test_invoke.py index 4ed1a4c0db845..7d40a6ef14379 100644 --- a/localstack-core/localstack/services/apigateway/next_gen/execute_api/test_invoke.py +++ b/localstack-core/localstack/services/apigateway/next_gen/execute_api/test_invoke.py @@ -16,6 +16,11 @@ from .handlers.resource_router import RestAPIResourceRouter from .header_utils import build_multi_value_headers from .template_mapping import dict_to_string +from .variables import ( + ContextVariableOverrides, + ContextVarsRequestOverride, + ContextVarsResponseOverride, +) # TODO: we probably need to write and populate those logs as part of the handler chain itself # and store it in the InvocationContext. That way, we could also retrieve in when calling TestInvoke @@ -150,8 +155,10 @@ def create_test_invocation_context( invocation_context.context_variables = parse_handler.create_context_variables( invocation_context ) - invocation_context.trace_id = parse_handler.populate_trace_id({}) - + invocation_context.context_variable_overrides = ContextVariableOverrides( + requestOverride=ContextVarsRequestOverride(header={}, path={}, querystring={}), + responseOverride=ContextVarsResponseOverride(header={}, status=0), + ) resource = deployment.rest_api.resources[test_request["resourceId"]] resource_method = resource["resourceMethods"][http_method] invocation_context.resource = resource From 38d05a2c9ce30afe7cbd09ee105156cff85401b3 Mon Sep 17 00:00:00 2001 From: Benjamin Simon Date: Fri, 16 May 2025 16:35:17 +0200 Subject: [PATCH 7/7] fix trace id --- .../services/apigateway/next_gen/execute_api/test_invoke.py | 1 + 1 file changed, 1 insertion(+) diff --git a/localstack-core/localstack/services/apigateway/next_gen/execute_api/test_invoke.py b/localstack-core/localstack/services/apigateway/next_gen/execute_api/test_invoke.py index 7d40a6ef14379..0d871077aa707 100644 --- a/localstack-core/localstack/services/apigateway/next_gen/execute_api/test_invoke.py +++ b/localstack-core/localstack/services/apigateway/next_gen/execute_api/test_invoke.py @@ -159,6 +159,7 @@ def create_test_invocation_context( requestOverride=ContextVarsRequestOverride(header={}, path={}, querystring={}), responseOverride=ContextVarsResponseOverride(header={}, status=0), ) + invocation_context.trace_id = parse_handler.populate_trace_id({}) resource = deployment.rest_api.resources[test_request["resourceId"]] resource_method = resource["resourceMethods"][http_method] invocation_context.resource = resource