feat(event_handler): add VPCLatticeResolver (#2601)

leandrodamascena · web-flow · commit 1905e4cb022b · 2023-06-28T14:04:20.000+02:00
diff --git a/aws_lambda_powertools/event_handler/__init__.py b/aws_lambda_powertools/event_handler/__init__.py
@@ -2,16 +2,19 @@
 Event handler decorators for common Lambda events
 """
 
-from .api_gateway import (
+from aws_lambda_powertools.event_handler.api_gateway import (
     ALBResolver,
     APIGatewayHttpResolver,
     ApiGatewayResolver,
     APIGatewayRestResolver,
     CORSConfig,
     Response,
 )
-from .appsync import AppSyncResolver
-from .lambda_function_url import LambdaFunctionUrlResolver
+from aws_lambda_powertools.event_handler.appsync import AppSyncResolver
+from aws_lambda_powertools.event_handler.lambda_function_url import (
+    LambdaFunctionUrlResolver,
+)
+from aws_lambda_powertools.event_handler.vpc_lattice import VPCLatticeResolver
 
 __all__ = [
     "AppSyncResolver",
@@ -22,4 +25,5 @@
     "CORSConfig",
     "LambdaFunctionUrlResolver",
     "Response",
+    "VPCLatticeResolver",
 ]
diff --git a/aws_lambda_powertools/event_handler/api_gateway.py b/aws_lambda_powertools/event_handler/api_gateway.py
@@ -33,6 +33,7 @@
     APIGatewayProxyEvent,
     APIGatewayProxyEventV2,
     LambdaFunctionUrlEvent,
+    VPCLatticeEvent,
 )
 from aws_lambda_powertools.utilities.data_classes.common import BaseProxyEvent
 from aws_lambda_powertools.utilities.typing import LambdaContext
@@ -53,6 +54,7 @@ class ProxyEventType(Enum):
     APIGatewayProxyEvent = "APIGatewayProxyEvent"
     APIGatewayProxyEventV2 = "APIGatewayProxyEventV2"
     ALBEvent = "ALBEvent"
+    VPCLatticeEvent = "VPCLatticeEvent"
     LambdaFunctionUrlEvent = "LambdaFunctionUrlEvent"
 
 
@@ -683,13 +685,17 @@ def _to_proxy_event(self, event: Dict) -> BaseProxyEvent:
         if self._proxy_type == ProxyEventType.LambdaFunctionUrlEvent:
             logger.debug("Converting event to Lambda Function URL contract")
             return LambdaFunctionUrlEvent(event)
+        if self._proxy_type == ProxyEventType.VPCLatticeEvent:
+            logger.debug("Converting event to VPC Lattice contract")
+            return VPCLatticeEvent(event)
         logger.debug("Converting event to ALB contract")
         return ALBEvent(event)
 
     def _resolve(self) -> ResponseBuilder:
         """Resolves the response or return the not found response"""
         method = self.current_event.http_method.upper()
         path = self._remove_prefix(self.current_event.path)
+
         for route in self._static_routes + self._dynamic_routes:
             if method != route.method:
                 continue
diff --git a/aws_lambda_powertools/event_handler/vpc_lattice.py b/aws_lambda_powertools/event_handler/vpc_lattice.py
@@ -0,0 +1,53 @@
+from typing import Callable, Dict, List, Optional
+
+from aws_lambda_powertools.event_handler import CORSConfig
+from aws_lambda_powertools.event_handler.api_gateway import (
+    ApiGatewayResolver,
+    ProxyEventType,
+)
+from aws_lambda_powertools.utilities.data_classes import VPCLatticeEvent
+
+
+class VPCLatticeResolver(ApiGatewayResolver):
+    """VPC Lattice resolver
+
+    Documentation:
+    - https://docs.aws.amazon.com/lambda/latest/dg/services-vpc-lattice.html
+    - https://docs.aws.amazon.com/lambda/latest/dg/services-vpc-lattice.html#vpc-lattice-receiving-events
+
+    Examples
+    --------
+    Simple example integrating with Tracer
+
+    ```python
+    from aws_lambda_powertools import Tracer
+    from aws_lambda_powertools.event_handler import VPCLatticeResolver
+
+    tracer = Tracer()
+    app = VPCLatticeResolver()
+
+    @app.get("/get-call")
+    def simple_get():
+        return {"message": "Foo"}
+
+    @app.post("/post-call")
+    def simple_post():
+        post_data: dict = app.current_event.json_body
+        return {"message": post_data}
+
+    @tracer.capture_lambda_handler
+    def lambda_handler(event, context):
+        return app.resolve(event, context)
+    """
+
+    current_event: VPCLatticeEvent
+
+    def __init__(
+        self,
+        cors: Optional[CORSConfig] = None,
+        debug: Optional[bool] = None,
+        serializer: Optional[Callable[[Dict], str]] = None,
+        strip_prefixes: Optional[List[str]] = None,
+    ):
+        """Amazon VPC Lattice resolver"""
+        super().__init__(ProxyEventType.VPCLatticeEvent, cors, debug, serializer, strip_prefixes)
diff --git a/aws_lambda_powertools/logging/correlation_paths.py b/aws_lambda_powertools/logging/correlation_paths.py
@@ -8,3 +8,4 @@
 EVENT_BRIDGE = "id"
 LAMBDA_FUNCTION_URL = API_GATEWAY_REST
 S3_OBJECT_LAMBDA = "xAmzRequestId"
+VPC_LATTICE = 'headers."x-amzn-trace-id"'
diff --git a/aws_lambda_powertools/utilities/data_classes/vpc_lattice.py b/aws_lambda_powertools/utilities/data_classes/vpc_lattice.py
@@ -1,14 +1,18 @@
 from typing import Any, Dict, Optional
 
-from aws_lambda_powertools.utilities.data_classes.common import DictWrapper
+from aws_lambda_powertools.shared.headers_serializer import (
+    BaseHeadersSerializer,
+    HttpApiHeadersSerializer,
+)
+from aws_lambda_powertools.utilities.data_classes.common import BaseProxyEvent
 from aws_lambda_powertools.utilities.data_classes.shared_functions import (
     base64_decode,
     get_header_value,
     get_query_string_value,
 )
 
 
-class VPCLatticeEvent(DictWrapper):
+class VPCLatticeEvent(BaseProxyEvent):
     @property
     def body(self) -> str:
         """The VPC Lattice body."""
@@ -54,6 +58,19 @@ def raw_path(self) -> str:
         """The raw VPC Lattice request path."""
         return self["raw_path"]
 
+    # VPCLattice event has no path field
+    # Added here for consistency with the BaseProxyEvent class
+    @property
+    def path(self) -> str:
+        return self["raw_path"]
+
+    # VPCLattice event has no http_method field
+    # Added here for consistency with the BaseProxyEvent class
+    @property
+    def http_method(self) -> str:
+        """The HTTP method used. Valid values include: DELETE, GET, HEAD, OPTIONS, PATCH, POST, and PUT."""
+        return self["method"]
+
     def get_query_string_value(self, name: str, default_value: Optional[str] = None) -> Optional[str]:
         """Get query string value by name
 
@@ -101,3 +118,7 @@ def get_header_value(
             default_value=default_value,
             case_sensitive=case_sensitive,
         )
+
+    def header_serializer(self) -> BaseHeadersSerializer:
+        # When using the VPC Lattice integration, we have multiple HTTP Headers.
+        return HttpApiHeadersSerializer()
diff --git a/docs/core/event_handler/api_gateway.md b/docs/core/event_handler/api_gateway.md
@@ -3,7 +3,7 @@ title: REST API
 description: Core utility
 ---
 
-Event handler for Amazon API Gateway REST and HTTP APIs, Application Loader Balancer (ALB), and Lambda Function URLs.
+Event handler for Amazon API Gateway REST and HTTP APIs, Application Loader Balancer (ALB), Lambda Function URLs, and VPC Lattice.
 
 ## Key Features
 
@@ -20,6 +20,8 @@ Event handler for Amazon API Gateway REST and HTTP APIs, Application Loader Bala
 
 If you're using any API Gateway integration, you must have an existing [API Gateway Proxy integration](https://docs.aws.amazon.com/apigateway/latest/developerguide/set-up-lambda-proxy-integrations.html){target="_blank"} or [ALB](https://docs.aws.amazon.com/elasticloadbalancing/latest/application/lambda-functions.html){target="_blank"} configured to invoke your Lambda function.
 
+In case of using [VPC Lattice](https://docs.aws.amazon.com/lambda/latest/dg/services-vpc-lattice.html){target="_blank"}, you must have a service network configured to invoke your Lambda function.
+
 This is the sample infrastructure for API Gateway and Lambda Function URLs we are using for the examples in this documentation.
 
 ???+ info "There is no additional permissions or dependencies required to use this utility."
@@ -42,7 +44,7 @@ Before you decorate your functions to handle a given path and HTTP method(s), yo
 
 A resolver will handle request resolution, including [one or more routers](#split-routes-with-router), and give you access to the current event via typed properties.
 
-For resolvers, we provide: `APIGatewayRestResolver`, `APIGatewayHttpResolver`, `ALBResolver`, and `LambdaFunctionUrlResolver`. From here on, we will default to `APIGatewayRestResolver` across examples.
+For resolvers, we provide: `APIGatewayRestResolver`, `APIGatewayHttpResolver`, `ALBResolver`, `LambdaFunctionUrlResolver`, and `VPCLatticeResolver`. From here on, we will default to `APIGatewayRestResolver` across examples.
 
 ???+ info "Auto-serialization"
     We serialize `Dict` responses as JSON, trim whitespace for compact responses, set content-type to `application/json`, and
@@ -116,6 +118,22 @@ When using [AWS Lambda Function URL](https://docs.aws.amazon.com/lambda/latest/d
     --8<-- "examples/event_handler_lambda_function_url/src/getting_started_lambda_function_url_resolver.json"
     ```
 
+#### VPC Lattice
+
+When using [VPC Lattice with AWS Lambda](https://docs.aws.amazon.com/lambda/latest/dg/services-vpc-lattice.html){target="_blank"}, you can use `VPCLatticeResolver`.
+
+=== "getting_started_vpclattice_resolver.py"
+
+    ```python hl_lines="5 11" title="Using VPC Lattice resolver"
+    --8<-- "examples/event_handler_rest/src/getting_started_vpclattice_resolver.py"
+    ```
+
+=== "getting_started_vpclattice_resolver.json"
+
+    ```json hl_lines="2 3" title="Example payload delivered to the handler"
+    --8<-- "examples/event_handler_rest/src/getting_started_vpclattice_resolver.json"
+    ```
+
 ### Dynamic routes
 
 You can use `/todos/<todo_id>` to configure dynamic URL paths, where `<todo_id>` will be resolved at runtime.
diff --git a/examples/event_handler_rest/src/getting_started_vpclattice_resolver.json b/examples/event_handler_rest/src/getting_started_vpclattice_resolver.json
@@ -0,0 +1,15 @@
+{
+    "raw_path": "/testpath",
+    "method": "GET",
+    "headers": {
+      "user_agent": "curl/7.64.1",
+      "x-forwarded-for": "10.213.229.10",
+      "host": "test-lambda-service-3908sdf9u3u.dkfjd93.vpc-lattice-svcs.us-east-2.on.aws",
+      "accept": "*/*"
+    },
+    "query_string_parameters": {
+      "order-id": "1"
+    },
+    "body": "eyJ0ZXN0IjogImV2ZW50In0=",
+    "is_base64_encoded": true
+  }
diff --git a/examples/event_handler_rest/src/getting_started_vpclattice_resolver.py b/examples/event_handler_rest/src/getting_started_vpclattice_resolver.py
@@ -0,0 +1,28 @@
+import requests
+from requests import Response
+
+from aws_lambda_powertools import Logger, Tracer
+from aws_lambda_powertools.event_handler import VPCLatticeResolver
+from aws_lambda_powertools.logging import correlation_paths
+from aws_lambda_powertools.utilities.typing import LambdaContext
+
+tracer = Tracer()
+logger = Logger()
+app = VPCLatticeResolver()
+
+
+@app.get("/todos")
+@tracer.capture_method
+def get_todos():
+    todos: Response = requests.get("https://jsonplaceholder.typicode.com/todos")
+    todos.raise_for_status()
+
+    # for brevity, we'll limit to the first 10 only
+    return {"todos": todos.json()[:10]}
+
+
+# You can continue to use other utilities just as before
+@logger.inject_lambda_context(correlation_id_path=correlation_paths.APPLICATION_LOAD_BALANCER)
+@tracer.capture_lambda_handler
+def lambda_handler(event: dict, context: LambdaContext) -> dict:
+    return app.resolve(event, context)
diff --git a/tests/events/vpcLatticeEventPathTrailingSlash.json b/tests/events/vpcLatticeEventPathTrailingSlash.json
@@ -0,0 +1,15 @@
+{
+  "raw_path": "/testpath/",
+  "method": "GET",
+  "headers": {
+    "user_agent": "curl/7.64.1",
+    "x-forwarded-for": "10.213.229.10",
+    "host": "test-lambda-service-3908sdf9u3u.dkfjd93.vpc-lattice-svcs.us-east-2.on.aws",
+    "accept": "*/*"
+  },
+  "query_string_parameters": {
+    "order-id": "1"
+  },
+  "body": "eyJ0ZXN0IjogImV2ZW50In0=",
+  "is_base64_encoded": true
+}
diff --git a/tests/functional/event_handler/conftest.py b/tests/functional/event_handler/conftest.py
@@ -0,0 +1,9 @@
+import json
+
+import pytest
+
+
+@pytest.fixture
+def json_dump():
+    # our serializers reduce length to save on costs; fixture to replicate separators
+    return lambda obj: json.dumps(obj, separators=(",", ":"))
diff --git a/tests/functional/event_handler/test_api_gateway.py b/tests/functional/event_handler/test_api_gateway.py
@@ -41,12 +41,6 @@
 from tests.functional.utils import load_event
 
 
-@pytest.fixture
-def json_dump():
-    # our serializers reduce length to save on costs; fixture to replicate separators
-    return lambda obj: json.dumps(obj, separators=(",", ":"))
-
-
 def read_media(file_name: str) -> bytes:
     path = Path(str(Path(__file__).parent.parent.parent.parent) + "/docs/media/" + file_name)
     return path.read_bytes()
diff --git a/tests/functional/event_handler/test_vpc_lattice.py b/tests/functional/event_handler/test_vpc_lattice.py
@@ -0,0 +1,77 @@
+from aws_lambda_powertools.event_handler import (
+    Response,
+    VPCLatticeResolver,
+    content_types,
+)
+from aws_lambda_powertools.event_handler.api_gateway import CORSConfig
+from aws_lambda_powertools.utilities.data_classes import VPCLatticeEvent
+from tests.functional.utils import load_event
+
+
+def test_vpclattice_event():
+    # GIVEN a VPC Lattice event
+    app = VPCLatticeResolver()
+
+    @app.get("/testpath")
+    def foo():
+        assert isinstance(app.current_event, VPCLatticeEvent)
+        assert app.lambda_context == {}
+        return Response(200, content_types.TEXT_HTML, "foo")
+
+    # WHEN calling the event handler
+    result = app(load_event("vpcLatticeEvent.json"), {})
+
+    # THEN process event correctly
+    # AND set the current_event type as VPCLatticeEvent
+    assert result["statusCode"] == 200
+    assert result["headers"]["Content-Type"] == content_types.TEXT_HTML
+    assert result["body"] == "foo"
+
+
+def test_vpclattice_event_path_trailing_slash(json_dump):
+    # GIVEN a VPC Lattice event
+    app = VPCLatticeResolver()
+
+    @app.get("/testpath")
+    def foo():
+        assert isinstance(app.current_event, VPCLatticeEvent)
+        assert app.lambda_context == {}
+        return Response(200, content_types.TEXT_HTML, "foo")
+
+    # WHEN calling the event handler using path with trailing "/"
+    result = app(load_event("vpcLatticeEventPathTrailingSlash.json"), {})
+
+    # THEN
+    assert result["statusCode"] == 404
+    assert result["headers"]["Content-Type"] == content_types.APPLICATION_JSON
+    expected = {"statusCode": 404, "message": "Not found"}
+    assert result["body"] == json_dump(expected)
+
+
+def test_cors_preflight_body_is_empty_not_null():
+    # GIVEN CORS is configured
+    app = VPCLatticeResolver(cors=CORSConfig())
+
+    event = {"raw_path": "/my/request", "method": "OPTIONS", "headers": {}}
+
+    # WHEN calling the event handler
+    result = app(event, {})
+
+    # THEN there body should be empty strings
+    assert result["body"] == ""
+
+
+def test_vpclattice_url_no_matches():
+    # GIVEN a VPC Lattice event
+    app = VPCLatticeResolver()
+
+    @app.post("/no_match")
+    def foo():
+        raise RuntimeError()
+
+    # WHEN calling the event handler
+    result = app(load_event("vpcLatticeEvent.json"), {})
+
+    # THEN process event correctly
+    # AND return 404 because the event doesn't match any known route
+    assert result["statusCode"] == 404
