Spaces:
Running
Running
| from typing import List, Dict, Union, Optional, 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: Optional[str] = None, secret_key: Optional[str] = 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]) -> Union[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()) | |