05 - Typing

app-type-hinting.py

import json
import os
import time
from typing import Any, Dict

from aws_lambda_powertools import Logger, Metrics, Tracer
from aws_lambda_powertools.logging import correlation_paths
from aws_lambda_powertools.metrics import MetricUnit
from aws_lambda_powertools.utilities.data_classes import APIGatewayProxyEvent
from aws_lambda_powertools.utilities.typing import LambdaContext

logger = Logger()
tracer = Tracer()
metrics = Metrics()


@tracer.capture_method
def create_payment(payment: Dict[str, str]) -> Dict[str, str]:
    order_key = payment["order_key"]
    logger.info(f"Order key: %s", order_key)

    # Simulate failure
    if "FAILURE" in order_key:
        tracer.put_annotation("PAYMENT_STATUS", "ERROR")
        metrics.add_metric(name="PaymentFailure", unit=MetricUnit.Count, value=1)
        raise RuntimeError("Failed to make payment")

    # Simulate a slow transaction
    time.sleep(2)
    tracer.put_annotation(key="ORDER_KEY", value=order_key)
    tracer.put_annotation(key="PAYMENT_STATUS", value="SUCCESS")
    metrics.add_metric(name="PaymentSuccessful", unit=MetricUnit.Count, value=1)

    return {"order_id": "1111111", "status": "SUCCESS"}


@metrics.log_metrics(capture_cold_start_metric=True)
@tracer.capture_lambda_handler
@logger.inject_lambda_context(log_event=True, correlation_id_path=correlation_paths.API_GATEWAY_REST)
def lambda_handler(_event: Dict[str, Any], context: LambdaContext) -> Dict[str, Any]:
    logger.debug("Function: %s#%s", context.function_name, context.function_version)
    event = APIGatewayProxyEvent(_event)
    assert event.get_header_value("x-api-key") == os.environ["SOME_STATIC_KEY"]
    payment = create_payment(event.json_body)
    return {
        "statusCode": 200,
        "headers": {"Content-Type": "application/json"},
        "body": json.dumps(payment),
    }