ask_candid/base/lambda_base.py

from typing import Any
from time import sleep
import json

import boto3


class LambdaInvokeBase:
    """Base class for AWS Lambda direct-invocation based classes. Each class which inherits from this only serves a
    single function.

    Parameters
    ----------
    function_name : str
        Name of the Lambda function to invoke
    access_key : Optional[str], optional
        AWS access key, by default None
    secret_key : Optional[str], optional
        AWS secret key, by default None
    """

    errors = frozenset([
        "Unhandled"
    ])

    def __init__(
        self, function_name: str,
        access_key: str | None = None, secret_key: str | None = None,
    ) -> None:
        if access_key is not None and secret_key is not None:
            self._client = boto3.client(
                "lambda",
                aws_access_key_id=access_key,
                aws_secret_access_key=secret_key,
                region_name="us-east-1",
            )
        else:
            self._client = boto3.client("lambda", region_name='us-east-1')

        self.function_name = function_name

    def _submit_request(self, payload: dict[str, Any]) -> dict[str, Any] | list[Any]:
        response = self._client.invoke(
            FunctionName=self.function_name,
            InvocationType="RequestResponse",
            Payload=json.dumps(payload),
        )

        if response.get("FunctionError") in self.errors:
            # could use recursion, but we need to keep track of number of function calls
            sleep(1)
            response = self._client.invoke(
                FunctionName=self.function_name,
                InvocationType="RequestResponse",
                Payload=json.dumps(payload),
            )

        return json.loads(response["Payload"].read())