30 lines
850 B
Python
30 lines
850 B
Python
from aws_lambda_powertools.event_handler.api_gateway import (
|
|
APIGatewayHttpResolver,
|
|
Response,
|
|
)
|
|
from aws_lambda_powertools.event_handler.middlewares import (
|
|
BaseMiddlewareHandler,
|
|
NextMiddleware,
|
|
)
|
|
|
|
|
|
class CorrelationIdMiddleware(BaseMiddlewareHandler):
|
|
def __init__(self, header: str):
|
|
super().__init__()
|
|
self.header = header
|
|
|
|
def handler(
|
|
self, app: APIGatewayHttpResolver, next_middleware: NextMiddleware
|
|
) -> Response:
|
|
# BEFORE logic
|
|
request_id = app.current_event.request_context.request_id
|
|
correlation_id = app.current_event.headers.get(self.header, request_id)
|
|
|
|
# Call next middleware or route handler ('/todos')
|
|
response = next_middleware(app)
|
|
|
|
# AFTER logic
|
|
response.headers[self.header] = correlation_id
|
|
|
|
return response
|