|
| 1 | +import datetime |
| 2 | +from urllib.parse import parse_qs |
| 3 | + |
| 4 | +from rolo import Request |
| 5 | +from rolo.gateway.chain import HandlerChain |
| 6 | +from werkzeug.datastructures import Headers |
| 7 | + |
| 8 | +from localstack.aws.api.apigateway import TestInvokeMethodRequest, TestInvokeMethodResponse |
| 9 | +from localstack.constants import APPLICATION_JSON |
| 10 | +from localstack.http import Response |
| 11 | +from localstack.utils.strings import to_bytes, to_str |
| 12 | + |
| 13 | +from ...models import RestApiDeployment |
| 14 | +from . import handlers |
| 15 | +from .context import InvocationRequest, RestApiInvocationContext |
| 16 | +from .handlers.resource_router import RestAPIResourceRouter |
| 17 | +from .header_utils import build_multi_value_headers |
| 18 | +from .template_mapping import dict_to_string |
| 19 | + |
| 20 | +# TODO: we probably need to write and populate those logs as part of the handler chain itself |
| 21 | +# and store it in the InvocationContext. That way, we could also retrieve in when calling TestInvoke |
| 22 | + |
| 23 | +TEST_INVOKE_TEMPLATE = """Execution log for request {request_id} |
| 24 | +{formatted_date} : Starting execution for request: {request_id} |
| 25 | +{formatted_date} : HTTP Method: {request_method}, Resource Path: {resource_path} |
| 26 | +{formatted_date} : Method request path: {method_request_path_parameters} |
| 27 | +{formatted_date} : Method request query string: {method_request_query_string} |
| 28 | +{formatted_date} : Method request headers: {method_request_headers} |
| 29 | +{formatted_date} : Method request body before transformations: {method_request_body} |
| 30 | +{formatted_date} : Endpoint request URI: {endpoint_uri} |
| 31 | +{formatted_date} : Endpoint request headers: {endpoint_request_headers} |
| 32 | +{formatted_date} : Endpoint request body after transformations: {endpoint_request_body} |
| 33 | +{formatted_date} : Sending request to {endpoint_uri} |
| 34 | +{formatted_date} : Received response. Status: {endpoint_response_status_code}, Integration latency: {endpoint_response_latency} ms |
| 35 | +{formatted_date} : Endpoint response headers: {endpoint_response_headers} |
| 36 | +{formatted_date} : Endpoint response body before transformations: {endpoint_response_body} |
| 37 | +{formatted_date} : Method response body after transformations: {method_response_body} |
| 38 | +{formatted_date} : Method response headers: {method_response_headers} |
| 39 | +{formatted_date} : Successfully completed execution |
| 40 | +{formatted_date} : Method completed with status: {method_response_status} |
| 41 | +""" |
| 42 | + |
| 43 | + |
| 44 | +def _dump_headers(headers: Headers) -> str: |
| 45 | + if not headers: |
| 46 | + return "{}" |
| 47 | + multi_headers = {key: ",".join(headers.getlist(key)) for key in headers.keys()} |
| 48 | + string_headers = dict_to_string(multi_headers) |
| 49 | + if len(string_headers) > 998: |
| 50 | + return f"{string_headers[:998]} [TRUNCATED]" |
| 51 | + |
| 52 | + return string_headers |
| 53 | + |
| 54 | + |
| 55 | +def log_template(invocation_context: RestApiInvocationContext, response_headers: Headers) -> str: |
| 56 | + # TODO: funny enough, in AWS for the `endpoint_response_headers` in AWS_PROXY, they log the response headers from |
| 57 | + # lambda HTTP Invoke call even though we use the headers from the lambda response itself |
| 58 | + formatted_date = datetime.datetime.now(tz=datetime.UTC).strftime("%a %b %d %H:%M:%S %Z %Y") |
| 59 | + request = invocation_context.invocation_request |
| 60 | + context_var = invocation_context.context_variables |
| 61 | + integration_req = invocation_context.integration_request |
| 62 | + endpoint_resp = invocation_context.endpoint_response |
| 63 | + method_resp = invocation_context.invocation_response |
| 64 | + # TODO: if endpoint_uri is an ARN, it means it's an AWS_PROXY integration |
| 65 | + # this should be transformed to the true URL of a lambda invoke call |
| 66 | + endpoint_uri = integration_req.get("uri", "") |
| 67 | + |
| 68 | + return TEST_INVOKE_TEMPLATE.format( |
| 69 | + formatted_date=formatted_date, |
| 70 | + request_id=context_var["requestId"], |
| 71 | + resource_path=request["path"], |
| 72 | + request_method=request["http_method"], |
| 73 | + method_request_path_parameters=dict_to_string(request["path_parameters"]), |
| 74 | + method_request_query_string=dict_to_string(request["query_string_parameters"]), |
| 75 | + method_request_headers=_dump_headers(request.get("headers")), |
| 76 | + method_request_body=to_str(request.get("body", "")), |
| 77 | + endpoint_uri=endpoint_uri, |
| 78 | + endpoint_request_headers=_dump_headers(integration_req.get("headers")), |
| 79 | + endpoint_request_body=to_str(integration_req.get("body", "")), |
| 80 | + # TODO: measure integration latency |
| 81 | + endpoint_response_latency=150, |
| 82 | + endpoint_response_status_code=endpoint_resp.get("status_code"), |
| 83 | + endpoint_response_body=to_str(endpoint_resp.get("body", "")), |
| 84 | + endpoint_response_headers=_dump_headers(endpoint_resp.get("headers")), |
| 85 | + method_response_status=method_resp.get("status_code"), |
| 86 | + method_response_body=to_str(method_resp.get("body", "")), |
| 87 | + method_response_headers=_dump_headers(response_headers), |
| 88 | + ) |
| 89 | + |
| 90 | + |
| 91 | +def create_test_chain() -> HandlerChain[RestApiInvocationContext]: |
| 92 | + return HandlerChain( |
| 93 | + request_handlers=[ |
| 94 | + handlers.method_request_handler, |
| 95 | + handlers.integration_request_handler, |
| 96 | + handlers.integration_handler, |
| 97 | + handlers.integration_response_handler, |
| 98 | + handlers.method_response_handler, |
| 99 | + ], |
| 100 | + exception_handlers=[ |
| 101 | + handlers.gateway_exception_handler, |
| 102 | + ], |
| 103 | + ) |
| 104 | + |
| 105 | + |
| 106 | +def create_test_invocation_context( |
| 107 | + test_request: TestInvokeMethodRequest, |
| 108 | + deployment: RestApiDeployment, |
| 109 | +) -> RestApiInvocationContext: |
| 110 | + parse_handler = handlers.parse_request |
| 111 | + http_method = test_request["httpMethod"] |
| 112 | + |
| 113 | + # we do not need a true HTTP request for the context, as we are skipping all the parsing steps and using the |
| 114 | + # provider data |
| 115 | + invocation_context = RestApiInvocationContext( |
| 116 | + request=Request(method=http_method), |
| 117 | + ) |
| 118 | + path_query = test_request.get("pathWithQueryString", "/").split("?") |
| 119 | + path = path_query[0] |
| 120 | + multi_query_args: dict[str, list[str]] = {} |
| 121 | + |
| 122 | + if len(path_query) > 1: |
| 123 | + multi_query_args = parse_qs(path_query[1]) |
| 124 | + |
| 125 | + # for the single value parameters, AWS only keeps the last value of the list |
| 126 | + single_query_args = {k: v[-1] for k, v in multi_query_args.items()} |
| 127 | + |
| 128 | + invocation_request = InvocationRequest( |
| 129 | + http_method=http_method, |
| 130 | + path=path, |
| 131 | + raw_path=path, |
| 132 | + query_string_parameters=single_query_args, |
| 133 | + multi_value_query_string_parameters=multi_query_args, |
| 134 | + headers=Headers(test_request.get("headers")), |
| 135 | + # TODO: handle multiValueHeaders |
| 136 | + body=to_bytes(test_request.get("body") or ""), |
| 137 | + ) |
| 138 | + invocation_context.invocation_request = invocation_request |
| 139 | + |
| 140 | + _, path_parameters = RestAPIResourceRouter(deployment).match(invocation_context) |
| 141 | + invocation_request["path_parameters"] = path_parameters |
| 142 | + |
| 143 | + invocation_context.deployment = deployment |
| 144 | + invocation_context.api_id = test_request["restApiId"] |
| 145 | + invocation_context.stage = None |
| 146 | + invocation_context.deployment_id = "" |
| 147 | + invocation_context.account_id = deployment.account_id |
| 148 | + invocation_context.region = deployment.region |
| 149 | + invocation_context.stage_variables = test_request.get("stageVariables", {}) |
| 150 | + invocation_context.context_variables = parse_handler.create_context_variables( |
| 151 | + invocation_context |
| 152 | + ) |
| 153 | + invocation_context.trace_id = parse_handler.populate_trace_id({}) |
| 154 | + |
| 155 | + resource = deployment.rest_api.resources[test_request["resourceId"]] |
| 156 | + resource_method = resource["resourceMethods"][http_method] |
| 157 | + invocation_context.resource = resource |
| 158 | + invocation_context.resource_method = resource_method |
| 159 | + invocation_context.integration = resource_method["methodIntegration"] |
| 160 | + handlers.route_request.update_context_variables_with_resource( |
| 161 | + invocation_context.context_variables, resource |
| 162 | + ) |
| 163 | + |
| 164 | + return invocation_context |
| 165 | + |
| 166 | + |
| 167 | +def run_test_invocation( |
| 168 | + test_request: TestInvokeMethodRequest, deployment: RestApiDeployment |
| 169 | +) -> TestInvokeMethodResponse: |
| 170 | + # validate resource exists in deployment |
| 171 | + invocation_context = create_test_invocation_context(test_request, deployment) |
| 172 | + |
| 173 | + test_chain = create_test_chain() |
| 174 | + # header order is important |
| 175 | + if invocation_context.integration["type"] == "MOCK": |
| 176 | + base_headers = {"Content-Type": APPLICATION_JSON} |
| 177 | + else: |
| 178 | + # we manually add the trace-id, as it is normally added by handlers.response_enricher which adds to much data |
| 179 | + # for the TestInvoke. It needs to be first |
| 180 | + base_headers = { |
| 181 | + "X-Amzn-Trace-Id": invocation_context.trace_id, |
| 182 | + "Content-Type": APPLICATION_JSON, |
| 183 | + } |
| 184 | + |
| 185 | + test_response = Response(headers=base_headers) |
| 186 | + start_time = datetime.datetime.now() |
| 187 | + test_chain.handle(context=invocation_context, response=test_response) |
| 188 | + end_time = datetime.datetime.now() |
| 189 | + |
| 190 | + response_headers = test_response.headers.copy() |
| 191 | + # AWS does not return the Content-Length for TestInvokeMethod |
| 192 | + response_headers.remove("Content-Length") |
| 193 | + |
| 194 | + log = log_template(invocation_context, response_headers) |
| 195 | + |
| 196 | + headers = dict(response_headers) |
| 197 | + multi_value_headers = build_multi_value_headers(response_headers) |
| 198 | + |
| 199 | + return TestInvokeMethodResponse( |
| 200 | + log=log, |
| 201 | + status=test_response.status_code, |
| 202 | + body=test_response.get_data(as_text=True), |
| 203 | + headers=headers, |
| 204 | + multiValueHeaders=multi_value_headers, |
| 205 | + latency=int((end_time - start_time).total_seconds()), |
| 206 | + ) |
0 commit comments