tracer-cloud--opensre/tests/shared/infrastructure_sdk/resources/lambda_.py

"""Lambda function management."""

import io
import os
import shutil
import subprocess
import tempfile
import time
import zipfile
from pathlib import Path
from typing import Any

from botocore.exceptions import ClientError

from tests.shared.infrastructure_sdk.deployer import (
    DEFAULT_REGION,
    get_boto3_client,
    get_standard_tags,
)


def bundle_code(source_dir: Path, requirements_file: Path | None = None) -> bytes:
    """Create deployment zip with dependencies.

    Args:
        source_dir: Directory containing Lambda code.
        requirements_file: Optional path to requirements.txt.

    Returns:
        Zip file contents as bytes.
    """
    with tempfile.TemporaryDirectory() as tmp_dir:
        tmp_path = Path(tmp_dir)
        package_dir = tmp_path / "package"
        package_dir.mkdir()

        # Install dependencies if requirements file provided
        if requirements_file and requirements_file.exists():
            subprocess.run(
                [
                    "python3",
                    "-m",
                    "pip",
                    "install",
                    "-q",
                    "-r",
                    str(requirements_file),
                    "-t",
                    str(package_dir),
                    "--platform",
                    "manylinux2014_x86_64",
                    "--only-binary=:all:",
                ],
                check=True,
                capture_output=True,
            )

        # Copy source files
        for item in source_dir.iterdir():
            if item.is_file() and item.suffix == ".py":
                shutil.copy2(item, package_dir)
            elif item.is_dir() and not item.name.startswith("."):
                shutil.copytree(item, package_dir / item.name)

        # Create zip
        zip_buffer = io.BytesIO()
        with zipfile.ZipFile(zip_buffer, "w", zipfile.ZIP_DEFLATED) as zf:
            for root, _dirs, files in os.walk(package_dir):
                for file in files:
                    file_path = Path(root) / file
                    arcname = file_path.relative_to(package_dir)
                    zf.write(file_path, arcname)

        return zip_buffer.getvalue()


def bundle_single_file(handler_file: Path, requirements_file: Path | None = None) -> bytes:
    """Create deployment zip from a single handler file.

    Args:
        handler_file: Path to the handler Python file.
        requirements_file: Optional path to requirements.txt.

    Returns:
        Zip file contents as bytes.
    """
    with tempfile.TemporaryDirectory() as tmp_dir:
        tmp_path = Path(tmp_dir)
        package_dir = tmp_path / "package"
        package_dir.mkdir()

        # Install dependencies if requirements file provided
        if requirements_file and requirements_file.exists():
            subprocess.run(
                [
                    "python3",
                    "-m",
                    "pip",
                    "install",
                    "-q",
                    "-r",
                    str(requirements_file),
                    "-t",
                    str(package_dir),
                    "--platform",
                    "manylinux2014_x86_64",
                    "--only-binary=:all:",
                ],
                check=True,
                capture_output=True,
            )

        # Copy handler file
        shutil.copy2(handler_file, package_dir)

        # Create zip
        zip_buffer = io.BytesIO()
        with zipfile.ZipFile(zip_buffer, "w", zipfile.ZIP_DEFLATED) as zf:
            for root, _dirs, files in os.walk(package_dir):
                for file in files:
                    file_path = Path(root) / file
                    arcname = file_path.relative_to(package_dir)
                    zf.write(file_path, arcname)

        return zip_buffer.getvalue()


def create_function(
    name: str,
    role_arn: str,
    handler: str,
    code_zip: bytes,
    runtime: str = "python3.11",
    timeout: int = 30,
    memory: int = 128,
    environment: dict[str, str] | None = None,
    stack_name: str | None = None,
    region: str = DEFAULT_REGION,
    layers: list[str] | None = None,
) -> dict[str, Any]:
    """Create Lambda function.

    Args:
        name: Function name.
        role_arn: ARN of the execution role.
        handler: Handler specification (e.g., "handler.lambda_handler").
        code_zip: Zip file bytes.
        runtime: Python runtime version.
        timeout: Function timeout in seconds.
        memory: Memory allocation in MB.
        environment: Environment variables.
        stack_name: Stack name for tagging.
        region: AWS region.
        layers: Optional list of layer ARNs.

    Returns:
        Dictionary with function info: arn, name, version.
    """
    lambda_client = get_boto3_client("lambda", region)

    # Prepare configuration
    config: dict[str, Any] = {
        "FunctionName": name,
        "Runtime": runtime,
        "Role": role_arn,
        "Handler": handler,
        "Code": {"ZipFile": code_zip},
        "Timeout": timeout,
        "MemorySize": memory,
    }

    if environment:
        config["Environment"] = {"Variables": environment}

    if stack_name:
        config["Tags"] = {t["Key"]: t["Value"] for t in get_standard_tags(stack_name)}

    if layers:
        config["Layers"] = layers

    # Retry logic for IAM propagation delays
    max_retries = 5
    retry_delay = 10  # seconds
    response: dict[str, Any] = {}

    for attempt in range(max_retries):
        try:
            response = lambda_client.create_function(**config)
            break
        except ClientError as e:
            error_code = e.response["Error"]["Code"]
            if error_code == "ResourceConflictException":
                # Function exists, update it
                update_function_code(name, code_zip, region)
                return update_function_configuration(
                    name,
                    environment=environment,
                    timeout=timeout,
                    memory=memory,
                    region=region,
                )
            elif (
                error_code == "InvalidParameterValueException"
                and "cannot be assumed" in str(e)
                and attempt < max_retries - 1
            ):
                # IAM role not ready yet, retry
                time.sleep(retry_delay)
                continue
            raise

    # Wait for function to be active
    _wait_for_function_active(name, lambda_client)

    return {
        "arn": response["FunctionArn"],
        "name": response["FunctionName"],
        "version": response.get("Version", "$LATEST"),
    }


def update_function_code(
    name: str, code_zip: bytes, region: str = DEFAULT_REGION
) -> dict[str, Any]:
    """Update existing function code.

    Args:
        name: Function name.
        code_zip: Zip file bytes.
        region: AWS region.

    Returns:
        Dictionary with function info: arn, name, version.
    """
    lambda_client = get_boto3_client("lambda", region)

    response = lambda_client.update_function_code(
        FunctionName=name,
        ZipFile=code_zip,
    )

    # Wait for update to complete
    _wait_for_function_active(name, lambda_client)

    return {
        "arn": response["FunctionArn"],
        "name": response["FunctionName"],
        "version": response.get("Version", "$LATEST"),
    }


def update_function_configuration(
    name: str,
    environment: dict[str, str] | None = None,
    timeout: int | None = None,
    memory: int | None = None,
    region: str = DEFAULT_REGION,
) -> dict[str, Any]:
    """Update function configuration.

    Args:
        name: Function name.
        environment: New environment variables (replaces existing).
        timeout: New timeout in seconds.
        memory: New memory in MB.
        region: AWS region.

    Returns:
        Updated function info.
    """
    lambda_client = get_boto3_client("lambda", region)

    config: dict[str, Any] = {"FunctionName": name}

    if environment is not None:
        try:
            current_config = lambda_client.get_function_configuration(FunctionName=name)
            current_env = current_config.get("Environment", {}).get("Variables", {})
            merged_env = {**current_env, **environment}
        except ClientError:
            merged_env = dict(environment)
        config["Environment"] = {"Variables": merged_env}
    if timeout is not None:
        config["Timeout"] = timeout
    if memory is not None:
        config["MemorySize"] = memory

    response = lambda_client.update_function_configuration(**config)

    _wait_for_function_active(name, lambda_client)

    return {
        "arn": response["FunctionArn"],
        "name": response["FunctionName"],
        "version": response.get("Version", "$LATEST"),
    }


def delete_function(name: str, region: str = DEFAULT_REGION) -> None:
    """Delete Lambda function.

    Args:
        name: Function name.
        region: AWS region.
    """
    lambda_client = get_boto3_client("lambda", region)

    try:
        lambda_client.delete_function(FunctionName=name)
    except ClientError as e:
        if e.response["Error"]["Code"] != "ResourceNotFoundException":
            raise


def add_permission(
    function_name: str,
    statement_id: str,
    principal: str,
    source_arn: str | None = None,
    region: str = DEFAULT_REGION,
) -> None:
    """Add invoke permission to Lambda.

    Args:
        function_name: Function name.
        statement_id: Unique statement ID.
        principal: Service principal (e.g., "apigateway.amazonaws.com").
        source_arn: Optional source ARN for condition.
        region: AWS region.
    """
    lambda_client = get_boto3_client("lambda", region)

    params: dict[str, Any] = {
        "FunctionName": function_name,
        "StatementId": statement_id,
        "Action": "lambda:InvokeFunction",
        "Principal": principal,
    }

    if source_arn:
        params["SourceArn"] = source_arn

    try:
        lambda_client.add_permission(**params)
    except ClientError as e:
        if e.response["Error"]["Code"] == "ResourceConflictException":
            pass  # Permission already exists
        else:
            raise


def invoke_function(
    function_name: str,
    payload: dict[str, Any] | None = None,
    region: str = DEFAULT_REGION,
) -> dict[str, Any]:
    """Invoke a Lambda function synchronously.

    Args:
        function_name: Function name.
        payload: Request payload.
        region: AWS region.

    Returns:
        Response payload.
    """
    import json

    lambda_client = get_boto3_client("lambda", region)

    params: dict[str, Any] = {
        "FunctionName": function_name,
        "InvocationType": "RequestResponse",
    }

    if payload:
        params["Payload"] = json.dumps(payload)

    response = lambda_client.invoke(**params)

    response_payload = response["Payload"].read()
    if response_payload:
        result: dict[str, Any] = json.loads(response_payload)
        return result
    return {}


def get_function(name: str, region: str = DEFAULT_REGION) -> dict[str, Any] | None:
    """Get function details.

    Args:
        name: Function name.
        region: AWS region.

    Returns:
        Function configuration or None if not found.
    """
    lambda_client = get_boto3_client("lambda", region)

    try:
        response = lambda_client.get_function(FunctionName=name)
        config: dict[str, Any] = response["Configuration"]
        return config
    except ClientError as e:
        if e.response["Error"]["Code"] == "ResourceNotFoundException":
            return None
        raise


def _wait_for_function_active(name: str, client: Any, max_attempts: int = 30) -> None:
    """Wait for function to be in Active state."""
    for _ in range(max_attempts):
        response = client.get_function(FunctionName=name)
        state = response["Configuration"].get("State", "Active")
        last_update_status = response["Configuration"].get("LastUpdateStatus", "Successful")

        if state == "Active" and last_update_status in ("Successful", None):
            return

        time.sleep(2)

    raise TimeoutError(f"Lambda function {name} did not become active")