Skip to content

Commit 20c0b74

Browse files
ran-isenbergRan IsenberganafalcaoAna Falcaoleandrodamascena
authored
feat(parser): add models for API GW Websockets events (#5597)
* feature(parser): Parser models for API GW Websockets Events * code review fixes * fix typo in the doc. add optional model * fix optional field * change names to snake case --------- Co-authored-by: Ran Isenberg <[email protected]> Co-authored-by: Ana Falcão <[email protected]> Co-authored-by: Ana Falcao <[email protected]> Co-authored-by: Leandro Damascena <[email protected]>
1 parent d1a58cd commit 20c0b74

File tree

10 files changed

+347
-1
lines changed

10 files changed

+347
-1
lines changed

‎aws_lambda_powertools/utilities/parser/envelopes/__init__.py

+2
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
from .apigw import ApiGatewayEnvelope
2+
from .apigw_websocket import ApiGatewayWebSocketEnvelope
23
from .apigwv2 import ApiGatewayV2Envelope
34
from .base import BaseEnvelope
45
from .bedrock_agent import BedrockAgentEnvelope
@@ -17,6 +18,7 @@
1718
__all__ = [
1819
"ApiGatewayEnvelope",
1920
"ApiGatewayV2Envelope",
21+
"ApiGatewayWebSocketEnvelope",
2022
"BedrockAgentEnvelope",
2123
"CloudWatchLogsEnvelope",
2224
"DynamoDBStreamEnvelope",

‎aws_lambda_powertools/utilities/parser/envelopes/apigw_websocket.py

+41
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,41 @@
1+
from __future__ import annotations
2+
3+
import logging
4+
from typing import TYPE_CHECKING, Any
5+
6+
from aws_lambda_powertools.utilities.parser.envelopes.base import BaseEnvelope
7+
from aws_lambda_powertools.utilities.parser.models import APIGatewayWebSocketMessageEventModel
8+
9+
if TYPE_CHECKING:
10+
from aws_lambda_powertools.utilities.parser.types import Model
11+
12+
logger = logging.getLogger(__name__)
13+
14+
15+
class ApiGatewayWebSocketEnvelope(BaseEnvelope):
16+
"""API Gateway WebSockets envelope to extract data within body key of messages routes
17+
(not disconnect or connect)"""
18+
19+
def parse(self, data: dict[str, Any] | Any | None, model: type[Model]) -> Model | None:
20+
"""Parses data found with model provided
21+
22+
Parameters
23+
----------
24+
data : dict
25+
Lambda event to be parsed
26+
model : type[Model]
27+
Data model provided to parse after extracting data using envelope
28+
29+
Returns
30+
-------
31+
Any
32+
Parsed detail payload with model provided
33+
"""
34+
logger.debug(
35+
f"Parsing incoming data with Api Gateway WebSockets model {APIGatewayWebSocketMessageEventModel}",
36+
)
37+
parsed_envelope: APIGatewayWebSocketMessageEventModel = APIGatewayWebSocketMessageEventModel.model_validate(
38+
data,
39+
)
40+
logger.debug(f"Parsing event payload in `detail` with {model}")
41+
return self._parse(data=parsed_envelope.body, model=model)

‎aws_lambda_powertools/utilities/parser/models/__init__.py

+18
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,16 @@
77
APIGatewayEventRequestContext,
88
APIGatewayProxyEventModel,
99
)
10+
from .apigw_websocket import (
11+
APIGatewayWebSocketConnectEventModel,
12+
APIGatewayWebSocketConnectEventRequestContext,
13+
APIGatewayWebSocketDisconnectEventModel,
14+
APIGatewayWebSocketDisconnectEventRequestContext,
15+
APIGatewayWebSocketEventIdentity,
16+
APIGatewayWebSocketEventRequestContextBase,
17+
APIGatewayWebSocketMessageEventModel,
18+
APIGatewayWebSocketMessageEventRequestContext,
19+
)
1020
from .apigwv2 import (
1121
ApiGatewayAuthorizerRequestV2,
1222
APIGatewayProxyEventV2Model,
@@ -105,6 +115,14 @@
105115
__all__ = [
106116
"APIGatewayProxyEventV2Model",
107117
"ApiGatewayAuthorizerRequestV2",
118+
"APIGatewayWebSocketEventIdentity",
119+
"APIGatewayWebSocketMessageEventModel",
120+
"APIGatewayWebSocketMessageEventRequestContext",
121+
"APIGatewayWebSocketConnectEventModel",
122+
"APIGatewayWebSocketConnectEventRequestContext",
123+
"APIGatewayWebSocketDisconnectEventRequestContext",
124+
"APIGatewayWebSocketDisconnectEventModel",
125+
"APIGatewayWebSocketEventRequestContextBase",
108126
"RequestContextV2",
109127
"RequestContextV2Http",
110128
"RequestContextV2Authorizer",

‎aws_lambda_powertools/utilities/parser/models/apigw_websocket.py

+63
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,63 @@
1+
from datetime import datetime
2+
from typing import Dict, List, Literal, Optional, Type, Union
3+
4+
from pydantic import BaseModel, Field
5+
from pydantic.networks import IPvAnyNetwork
6+
7+
8+
class APIGatewayWebSocketEventIdentity(BaseModel):
9+
source_ip: IPvAnyNetwork = Field(alias="sourceIp")
10+
user_agent: Optional[str] = Field(None, alias="userAgent")
11+
12+
class APIGatewayWebSocketEventRequestContextBase(BaseModel):
13+
extended_request_id: str = Field(alias="extendedRequestId")
14+
request_time: str = Field(alias="requestTime")
15+
stage: str = Field(alias="stage")
16+
connected_at: datetime = Field(alias="connectedAt")
17+
request_time_epoch: datetime = Field(alias="requestTimeEpoch")
18+
identity: APIGatewayWebSocketEventIdentity = Field(alias="identity")
19+
request_id: str = Field(alias="requestId")
20+
domain_name: str = Field(alias="domainName")
21+
connection_id: str = Field(alias="connectionId")
22+
api_id: str = Field(alias="apiId")
23+
24+
25+
class APIGatewayWebSocketMessageEventRequestContext(APIGatewayWebSocketEventRequestContextBase):
26+
route_key: str = Field(alias="routeKey")
27+
message_id: str = Field(alias="messageId")
28+
event_type: Literal["MESSAGE"] = Field(alias="eventType")
29+
message_direction: Literal["IN", "OUT"] = Field(alias="messageDirection")
30+
31+
32+
class APIGatewayWebSocketConnectEventRequestContext(APIGatewayWebSocketEventRequestContextBase):
33+
route_key: Literal["$connect"] = Field(alias="routeKey")
34+
event_type: Literal["CONNECT"] = Field(alias="eventType")
35+
message_direction: Literal["IN"] = Field(alias="messageDirection")
36+
37+
38+
class APIGatewayWebSocketDisconnectEventRequestContext(APIGatewayWebSocketEventRequestContextBase):
39+
route_key: Literal["$disconnect"] = Field(alias="routeKey")
40+
disconnect_status_code: int = Field(alias="disconnectStatusCode")
41+
event_type: Literal["DISCONNECT"] = Field(alias="eventType")
42+
message_direction: Literal["IN"] = Field(alias="messageDirection")
43+
disconnect_reason: str = Field(alias="disconnectReason")
44+
45+
46+
class APIGatewayWebSocketConnectEventModel(BaseModel):
47+
headers: Dict[str, str] = Field(alias="headers")
48+
multi_value_headers: Dict[str, List[str]] = Field(alias="multiValueHeaders")
49+
request_context: APIGatewayWebSocketConnectEventRequestContext = Field(alias="requestContext")
50+
is_base64_encoded: bool = Field(alias="isBase64Encoded")
51+
52+
53+
class APIGatewayWebSocketDisconnectEventModel(BaseModel):
54+
headers: Dict[str, str] = Field(alias="headers")
55+
multi_value_headers: Dict[str, List[str]] = Field(alias="multiValueHeaders")
56+
request_context: APIGatewayWebSocketDisconnectEventRequestContext = Field(alias="requestContext")
57+
is_base64_encoded: bool = Field(alias="isBase64Encoded")
58+
59+
60+
class APIGatewayWebSocketMessageEventModel(BaseModel):
61+
request_context: APIGatewayWebSocketMessageEventRequestContext = Field(alias="requestContext")
62+
is_base64_encoded: bool = Field(alias="isBase64Encoded")
63+
body: Optional[Union[str, Type[BaseModel]]] = Field(None, alias="body")

‎docs/utilities/parser.md

+5-1
Original file line numberDiff line numberDiff line change
@@ -108,6 +108,9 @@ The example above uses `SqsModel`. Other built-in models can be found below.
108108
| **ApiGatewayAuthorizerRequest** | Lambda Event Source payload for Amazon API Gateway Lambda Authorizer with Request |
109109
| **APIGatewayProxyEventV2Model** | Lambda Event Source payload for Amazon API Gateway v2 payload |
110110
| **ApiGatewayAuthorizerRequestV2** | Lambda Event Source payload for Amazon API Gateway v2 Lambda Authorizer |
111+
| **APIGatewayWebSocketMessageEventModel** | Lambda Event Source payload for Amazon API Gateway WebSocket API message body |
112+
| **APIGatewayWebSocketConnectEventModel** | Lambda Event Source payload for Amazon API Gateway WebSocket API $connect message |
113+
| **APIGatewayWebSocketDisconnectEventModel** | Lambda Event Source payload for Amazon API Gateway WebSocket API $disconnect message |
111114
| **BedrockAgentEventModel** | Lambda Event Source payload for Bedrock Agents |
112115
| **CloudFormationCustomResourceCreateModel** | Lambda Event Source payload for AWS CloudFormation `CREATE` operation |
113116
| **CloudFormationCustomResourceUpdateModel** | Lambda Event Source payload for AWS CloudFormation `UPDATE` operation |
@@ -188,8 +191,9 @@ You can use pre-built envelopes provided by the Parser to extract and parse spec
188191
| **KinesisFirehoseEnvelope** | 1. Parses data using `KinesisFirehoseModel` which will base64 decode it. ``2. Parses records in in` Records` key using your model`` and returns them in a list. | `List[Model]` |
189192
| **SnsEnvelope** | 1. Parses data using `SnsModel`. ``2. Parses records in `body` key using your model`` and return them in a list. | `List[Model]` |
190193
| **SnsSqsEnvelope** | 1. Parses data using `SqsModel`. `` 2. Parses SNS records in `body` key using `SnsNotificationModel`. `` 3. Parses data in `Message` key using your model and return them in a list. | `List[Model]` |
191-
| **ApiGatewayEnvelope** | 1. Parses data using `APIGatewayProxyEventModel`. ``2. Parses `body` key using your model`` and returns it. | `Model` |
192194
| **ApiGatewayV2Envelope** | 1. Parses data using `APIGatewayProxyEventV2Model`. ``2. Parses `body` key using your model`` and returns it. | `Model` |
195+
| **ApiGatewayEnvelope** | 1. Parses data using `APIGatewayProxyEventModel`. ``2. Parses `body` key using your model`` and returns it. | `Model` |
196+
| **ApiGatewayWebSocketEnvelope** | 1. Parses data using `APIGatewayWebSocketMessageEventModel`. ``2. Parses `body` key using your model`` and returns it. | `Model` |
193197
| **LambdaFunctionUrlEnvelope** | 1. Parses data using `LambdaFunctionUrlModel`. ``2. Parses `body` key using your model`` and returns it. | `Model` |
194198
| **KafkaEnvelope** | 1. Parses data using `KafkaRecordModel`. ``2. Parses `value` key using your model`` and returns it. | `Model` |
195199
| **VpcLatticeEnvelope** | 1. Parses data using `VpcLatticeModel`. ``2. Parses `value` key using your model`` and returns it. | `Model` |

‎tests/events/apiGatewayWebSocketApiConnect.json

+40
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,40 @@
1+
{
2+
"headers": {
3+
"Host": "fjnq7njcv2.execute-api.us-east-1.amazonaws.com",
4+
"Sec-WebSocket-Extensions": "permessage-deflate; client_max_window_bits",
5+
"Sec-WebSocket-Key": "+W5xw47OHh3OTFsWKjGu9Q==",
6+
"Sec-WebSocket-Version": "13",
7+
"X-Amzn-Trace-Id": "Root=1-6731ebfc-08e1e656421db73c5d2eef31",
8+
"X-Forwarded-For": "166.90.225.1",
9+
"X-Forwarded-Port": "443",
10+
"X-Forwarded-Proto": "https"
11+
},
12+
"multiValueHeaders": {
13+
"Host": ["fjnq7njcv2.execute-api.us-east-1.amazonaws.com"],
14+
"Sec-WebSocket-Extensions": ["permessage-deflate; client_max_window_bits"],
15+
"Sec-WebSocket-Key": ["+W5xw47OHh3OTFsWKjGu9Q=="],
16+
"Sec-WebSocket-Version": ["13"],
17+
"X-Amzn-Trace-Id": ["Root=1-6731ebfc-08e1e656421db73c5d2eef31"],
18+
"X-Forwarded-For": ["166.90.225.1"],
19+
"X-Forwarded-Port": ["443"],
20+
"X-Forwarded-Proto": ["https"]
21+
},
22+
"requestContext": {
23+
"routeKey": "$connect",
24+
"eventType": "CONNECT",
25+
"extendedRequestId": "BFHPhFe3IAMF95g=",
26+
"requestTime": "11/Nov/2024:11:35:24 +0000",
27+
"messageDirection": "IN",
28+
"stage": "prod",
29+
"connectedAt": 1731324924553,
30+
"requestTimeEpoch": 1731324924561,
31+
"identity": {
32+
"sourceIp": "166.90.225.1"
33+
},
34+
"requestId": "BFHPhFe3IAMF95g=",
35+
"domainName": "asasasas.execute-api.us-east-1.amazonaws.com",
36+
"connectionId": "BFHPhfCWIAMCKlQ=",
37+
"apiId": "asasasas"
38+
},
39+
"isBase64Encoded": false
40+
}

‎tests/events/apiGatewayWebSocketApiDisconnect.json

+34
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,34 @@
1+
{
2+
"headers": {
3+
"Host": "asasasas.execute-api.us-east-1.amazonaws.com",
4+
"x-api-key": "",
5+
"X-Forwarded-For": "",
6+
"x-restapi": ""
7+
},
8+
"multiValueHeaders": {
9+
"Host": ["asasasas.execute-api.us-east-1.amazonaws.com"],
10+
"x-api-key": [""],
11+
"X-Forwarded-For": [""],
12+
"x-restapi": [""]
13+
},
14+
"requestContext": {
15+
"routeKey": "$disconnect",
16+
"disconnectStatusCode": 1005,
17+
"eventType": "DISCONNECT",
18+
"extendedRequestId": "BFbOeE87IAMF31w=",
19+
"requestTime": "11/Nov/2024:13:51:49 +0000",
20+
"messageDirection": "IN",
21+
"disconnectReason": "Client-side close frame status not set",
22+
"stage": "prod",
23+
"connectedAt": 1731332735513,
24+
"requestTimeEpoch": 1731333109875,
25+
"identity": {
26+
"sourceIp": "166.90.225.1"
27+
},
28+
"requestId": "BFbOeE87IAMF31w=",
29+
"domainName": "asasasas.execute-api.us-east-1.amazonaws.com",
30+
"connectionId": "BFaT_fALIAMCKug=",
31+
"apiId": "asasasas"
32+
},
33+
"isBase64Encoded": false
34+
}

‎tests/events/apiGatewayWebSocketApiMessage.json

+22
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,22 @@
1+
{
2+
"requestContext": {
3+
"routeKey": "chat",
4+
"messageId": "BFaVtfGSIAMCKug=",
5+
"eventType": "MESSAGE",
6+
"extendedRequestId": "BFaVtH2HoAMFZEQ=",
7+
"requestTime": "11/Nov/2024:13:45:46 +0000",
8+
"messageDirection": "IN",
9+
"stage": "prod",
10+
"connectedAt": 1731332735513,
11+
"requestTimeEpoch": 1731332746514,
12+
"identity": {
13+
"sourceIp": "166.90.225.1"
14+
},
15+
"requestId": "BFaVtH2HoAMFZEQ=",
16+
"domainName": "asasasas.execute-api.us-east-1.amazonaws.com",
17+
"connectionId": "BFaT_fALIAMCKug=",
18+
"apiId": "asasasas"
19+
},
20+
"body": "{\"action\": \"chat\", \"message\": \"Hello from client\"}",
21+
"isBase64Encoded": false
22+
}

‎tests/unit/parser/_pydantic/schemas.py

+5
Original file line numberDiff line numberDiff line change
@@ -87,6 +87,11 @@ class MyApiGatewayBusiness(BaseModel):
8787
username: str
8888

8989

90+
class MyApiGatewayWebSocketBusiness(BaseModel):
91+
message: str
92+
action: str
93+
94+
9095
class MyALambdaFuncUrlBusiness(BaseModel):
9196
message: str
9297
username: str

0 commit comments

Comments
 (0)