lambda_function.py

"""AWS Lambda function to enrich data with user preferences from DynamoDB."""

import json
import logging
import os
from typing import Any

import boto3
from botocore.exceptions import ClientError

# Configure logging
logger = logging.getLogger()
logger.setLevel(logging.INFO)

# Initialize DynamoDB resource outside handler for connection reuse
TABLE_NAME = os.environ.get("TABLE_NAME", "")
dynamodb_resource = boto3.resource("dynamodb")
users_table = dynamodb_resource.Table(TABLE_NAME)

# Default user preferences
DEFAULT_PREFERENCES = {
    "language": "en",
    "notifications_enabled": False,
}


def get_user_preferences(user_id: str) -> dict[str, Any] | None:
    """Fetch user preferences from DynamoDB.

    Args:
        user_id: The unique identifier for the user.

    Returns:
        User preferences dict if found, None otherwise.

    Raises:
        ClientError: If DynamoDB request fails.
    """
    try:
        response = users_table.get_item(Key={"id": user_id})
        return response.get("Item")
    except ClientError as e:
        logger.error("Failed to fetch user %s: %s", user_id, e.response["Error"]["Message"])
        raise


def enrich_payload(payload: dict[str, Any]) -> dict[str, Any]:
    """Enrich payload with user preferences.

    Args:
        payload: The original request payload containing user id.

    Returns:
        Enriched payload with user preferences or defaults.
    """
    user_id = payload.get("id")
    if not user_id:
        logger.warning("No user id provided, using default preferences")
        return {**payload, **DEFAULT_PREFERENCES}

    user_prefs = get_user_preferences(user_id)
    if user_prefs:
        logger.info("Found preferences for user %s", user_id)
        return {**payload, **user_prefs}

    logger.info("No preferences found for user %s, using defaults", user_id)
    return {**payload, **DEFAULT_PREFERENCES}


def lambda_handler(event: dict[str, Any], context: Any) -> dict[str, Any]:
    """AWS Lambda handler to enrich incoming data with user preferences.

    Args:
        event: Lambda event containing request data.
        context: Lambda context object.

    Returns:
        Response dict with status code and enriched payload.
    """
    logger.info("Processing event: %s", json.dumps(event))

    try:
        # Handle both direct invocation and API Gateway events
        if isinstance(event, list) and event:
            body = event[0].get("body", "{}")
        elif isinstance(event, dict):
            body = event.get("body", "{}")
        else:
            body = "{}"

        # Parse JSON body if it's a string
        payload = json.loads(body) if isinstance(body, str) else body

        enriched_payload = enrich_payload(payload)

        return {
            "statusCode": 200,
            "body": json.dumps(enriched_payload),
            "headers": {
                "Content-Type": "application/json",
            },
        }

    except json.JSONDecodeError as e:
        logger.error("Invalid JSON in request body: %s", e)
        return {
            "statusCode": 400,
            "body": json.dumps({"error": "Invalid JSON in request body"}),
            "headers": {"Content-Type": "application/json"},
        }
    except ClientError as e:
        logger.error("DynamoDB error: %s", e)
        return {
            "statusCode": 500,
            "body": json.dumps({"error": "Internal server error"}),
            "headers": {"Content-Type": "application/json"},
        }
    except Exception as e:
        logger.error("Unexpected error: %s", e)
        return {
            "statusCode": 500,
            "body": json.dumps({"error": "Internal server error"}),
            "headers": {"Content-Type": "application/json"},
        }