RFC: Event Handler for API Gateway Web Sockets

[walmsles]

Is this related to an existing feature request or issue?

#1165

Which Powertools for AWS Lambda (Python) utility does this relate to?

Event Handler - REST API

Summary

To provide an Event Handler class to enable the implementation of Web Socket routes for API Gateway Web Socket connections.

Background

Web Socket connections are serviced by API routes that behave in particular ways. I recently wrote about this in an article here. The implementation of the Web Socket class should take into account current Web Socket route nuances and enable customers to create Web Socket implementations quickly and with significantly less boilerplate code and with a similar implementation pattern - at the end of they day it is all about defining routes.

Use case

Enabling implementation of API Gateway Web Socket routes using exactly the same style of code as existing API Gateway implementations and enabling familiar handling of Web Socket routes.

Proposal

before:

from aws_lambda_powertools import Logger, Tracer
from aws_lambda_powertools.logging import correlation_paths
from aws_lambda_powertools.utilities.typing import LambdaContext

tracer = Tracer()
logger = Logger()

@tracer.capture_method
def socket_connect():
    // do connection processing
    return {"statusCode": 200}

@tracer.capture_method
def order_notify():

    // do notify processing
    return {
         "statusCode": 200,
         "body": "order# 1234 Created",
    }

@tracer.capture_method
def socket_disconnect():

    // do connection processing

    return {"statusCode": 200}

# You can continue to use other utilities just as before
@logger.inject_lambda_context(correlation_id_path=correlation_paths.API_GATEWAY_REST)
@tracer.capture_lambda_handler
def lambda_handler(event: dict, context: LambdaContext) -> dict:
    if event.requestContext.routeKey == '$connect':
        return socket_connect(event, context)
    elif event.requestContext.routeKey == '$disconnect': 
        return socket_disconnect(event, context)
    else:
        return route_not_found(event, context)

after:

from aws_lambda_powertools import Logger, Tracer
from aws_lambda_powertools.event_handler import APIGatewayWebsocketResolver
from aws_lambda_powertools.logging import correlation_paths
from aws_lambda_powertools.utilities.typing import LambdaContext

tracer = Tracer()
logger = Logger()
app = APIGatewayWebsocketResolver()


@app.route("$connect")
@tracer.capture_method
def socket_connect():
    // do connection processing
    return {"statusCode": 200}

@app.route("order/notify")
@tracer.capture_method
def order_notify():

    // do connection processing
    return {
         "statusCode": 200,
         "body": "order# 1234 Created",
    }

@app.route("$disconnect")
@tracer.capture_method
def socket_disconnect():

    // do connection processing

    return {"statusCode": 200}

@app.route("$disconnect")
@tracer.capture_method
def socket_disconnect():

    // do connection processing

    return {"statusCode": 200}

# You can continue to use other utilities just as before
@logger.inject_lambda_context(correlation_id_path=correlation_paths.API_GATEWAY_REST)
@tracer.capture_lambda_handler
def lambda_handler(event: dict, context: LambdaContext) -> dict:
    return app.resolve(event, context)

Out of scope

The implementation must be specifically focused on route definitions and handlers with utilities for common use-cases only and not touch any cloud infrastructure (as per tenets).

Potential challenges

End-to-End testing is one potential challenge - unsure if this style of interface exists, so async testing with live infrastructure will require careful thought and planning.

We need to carefully navigate route definition and look at what we do when no $connect route exists (which implies NO authorisation/authentication in the API gateway solution). Infrastructure as code tools should cater for this nuance, but we need to decide how to handle this situation specifically, which creates a cloud vulnerability.

We need to consider route definition carefully and decide on where this fits in the class hierarchy. Potentially this should exist separately from APIGatewayResolvers, given there are no HTTP Methods and other characteristics. There are Headers but only inbound within the actual event - never outbound.

Lots to think about and discuss, I would like to see a familiar Resolver style implementation for WebSocket routes - they are analogous to APIgatewayResolver routes in many ways. Custom and $default routes are capable of returning an actual response which is sent back to the calling web socket connection in the same way as any other broadcast message.

Dependencies and Integrations

No response

Alternative solutions

Custom implementation with boilerplate code.

Acknowledgment

• This feature request meets Powertools for AWS Lambda (Python) Tenets
• Should this be considered in other Powertools for AWS Lambda languages? i.e. Java, TypeScript, and .NET

[leandrodamascena]

Hi @walmsles! Thanks for opening this RFC to add support for WebSockets, this has been a long-time wish and I'm happy in getting this resolved now.

I'll add this for review next week, okay? If you have a PoC to share with us, that would be great.

Thanks

[walmsles]

Hi @leandrodamascena, no POC yet. No rush. Keen to hear from others in the community about what would make sense for a WebSocket event handler, use cases, existing problems and friction points.

The solution needs to remove boilerplate code and increase developer velocity.

Other thoughts that come to mind:

• Not defining a $connect route should trigger a warning in debug mode (is generally considered a security issue) - maybe this is always?
• Propose a new class hierarchy in the same fashion as AppSync since it is a completely different protocol and feature set
• Must include middleware capability - this makes sense for WebSocket handlers and is a useful feature. Could look to create a new BaseRoute class and push function handler and middleware processing to the new base (DRY).

[stevrobu]

@leandrodamascena - I'd like to help with this.

[leandrodamascena]

Hi @stevrobu! Yes, help us create this new Resolver in Event Handler! For now, we need help discussing the best developer experience and the possible problems/benefits of integrating WebSocket into our utility. Please feel free to make comments, create a PoC, ask questions, and do whatever you can to contribute to this.

Thanks!

[stevrobu]

@leandrodamascena - I have a POC completed. I will get it checked into my branch and add comments/documentation for it by end of the week.

[walmsles]

It would be nice to separate out the different route types. My example above used a central route definition even for the defined routes. But I think it would be nice to represent the defined WebSocket routes as distinct middleware handlers.

from aws_lambda_powertools import Logger, Tracer
from aws_lambda_powertools.event_handler import APIGatewayWebsocketResolver
from aws_lambda_powertools.logging import correlation_paths
from aws_lambda_powertools.utilities.typing import LambdaContext

tracer = Tracer()
logger = Logger()
app = APIGatewayWebsocketResolver()


@app.connect()
@tracer.capture_method
def socket_connect():
    // do connection processing
    return {"statusCode": 200}

@app.default()
def socket_default():
     // do default route handling (catch-all)
    return {"statusCode": 200}


@tracer.capture_method
def socket_connect():
    // do connection processing
    return {"statusCode": 200}


@app.route("order/notify")
@tracer.capture_method
def order_notify():

    // do connection processing
    return {
         "statusCode": 200,
         "body": "order# 1234 Created",
    }

@app.disconnect()
@tracer.capture_method
def socket_disconnect():

    // do connection processing

    return {"statusCode": 200}


# You can continue to use other utilities just as before
@logger.inject_lambda_context(correlation_id_path=correlation_paths.API_GATEWAY_REST)
@tracer.capture_lambda_handler
def lambda_handler(event: dict, context: LambdaContext) -> dict:
    return app.resolve(event, context)

Common Web socket use cases that come to mind:

1. Async Task processing (executing longer-running tasks)
2. Submitting Events to an EDA of some form
3. Channel notifications - web socket connections subscribe to one or more channels for broadcast event notifications

Is it useful to build adapters for these use cases to wrap up some common tooling developers need to do to accelerate the development of these use cases?

Not sure exactly what these should look like but things that come to mind:

• Use Case 1 - this is really just a different form of Use case 2 IMO - a different EDA mechanism for launching a long running task
• Use case 2 - provide an EDA provider enabling users to inject a event submission provider with a transformer to transform the API data into the event pattern required. The actual implementation can be included for EventBridge, SQS and SNS as immediate starters to accelerate these use cases.
• Use case 3: Provide a storage adapter for storing connection data into a data store (optional addition). We could then create a class that represents a WebSocketChannelBroadcast use-case which automates the saving and deletion of connections from a data store injected via an adapter pattern (like the Idempotency Storage Provider idea).

These are typical use cases and all involve a WHOLE heap of boiler plate code which we could provide to accelerate use cases rather than just creating a simple Web Socket adapter imitating the APIRestResolver patterns. The APIRestResolver style pattern is still needed as the basis for adding these kinds of use cases on top.