07 - Exception handling

app-exception-handling.py

import json
import os
import time
from typing import Any, Callable, 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.middleware_factory import lambda_handler_decorator
from aws_lambda_powertools.utilities.data_classes import APIGatewayProxyEvent
from aws_lambda_powertools.utilities.idempotency import DynamoDBPersistenceLayer, IdempotencyConfig, idempotent
from aws_lambda_powertools.utilities.idempotency.exceptions import (
    IdempotencyAlreadyInProgressError,
    IdempotencyKeyError,
    IdempotencyPersistenceLayerError,
)
from aws_lambda_powertools.utilities.typing import LambdaContext

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


class PaymentServiceError(Exception):
    """Payment service error"""


def build_response(status_code: int = 200, body: Dict[str, Any] = None) -> Dict[str, Any]:
    """Build an api gateway response"""
    return {"statusCode": status_code, "headers": {"Content-Type": "application/json"}, "body": json.dumps(body)}


@lambda_handler_decorator(trace_execution=True)
def exception_handler(handler: Callable[[Dict, LambdaContext], Dict], event: Dict, context: LambdaContext) -> Dict:
    """Utility to convert errors to api gateway response"""
    try:
        return handler(event, context)
    except Exception as ex:
        tracer.put_annotation("ERROR_TYPE", type(ex).__name__)
        logger.exception("Failed to perform handler")
        if isinstance(ex, IdempotencyAlreadyInProgressError):
            return build_response(400, {"message": "Transaction in progress, please try again later"})
        if isinstance(ex, IdempotencyKeyError):
            return build_response(400, {"message": "Request is missing an idempotent key"})
        if isinstance(ex, IdempotencyPersistenceLayerError):
            return build_response(503, {"message": "Failure with idempotency persistent layer"})
        if isinstance(ex, PaymentServiceError):
            return build_response(402, {"message": "Payment failed"})
        return build_response(503, {"message": "Something went wrong"})


@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 PaymentServiceError("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"}


@exception_handler
@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)
@idempotent(
    config=IdempotencyConfig(
        event_key_jmespath="powertools_json(body).order_key",
        raise_on_no_idempotency_key=True,
        use_local_cache=True,
        local_cache_max_items=512,
        expires_after_seconds=24 * 60 * 60,
    ),
    persistence_store=DynamoDBPersistenceLayer(table_name=os.environ["IDEMPOTENCY_TABLE_NAME"]),
)
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 build_response(body=payment)