feat: optimization to reduce number of database calls, reorganize persistence layer modules

Author: Tom McCarthy

aws_lambda_powertools/utilities/idempotency/__init__.py

@@ -2,7 +2,9 @@
 Utility for adding idempotency to lambda functions
 """
 
+from aws_lambda_powertools.utilities.idempotency.persistence.base import BasePersistenceLayer
+from aws_lambda_powertools.utilities.idempotency.persistence.dynamodb import DynamoDBPersistenceLayer
+
 from .idempotency import idempotent
-from .persistence import BasePersistenceLayer, DynamoDBPersistenceLayer
 
 __all__ = ("DynamoDBPersistenceLayer", "BasePersistenceLayer", "idempotent")

aws_lambda_powertools/utilities/idempotency/idempotency.py

@@ -5,10 +5,10 @@
 from typing import Any, Callable, Dict
 
 from aws_lambda_powertools.middleware_factory import lambda_handler_decorator
+from aws_lambda_powertools.utilities.idempotency.persistence.base import STATUS_CONSTANTS, BasePersistenceLayer
 
 from ..typing import LambdaContext
-from .exceptions import AlreadyInProgressError, ItemNotFoundError
-from .persistence import STATUS_CONSTANTS, BasePersistenceLayer
+from .exceptions import AlreadyInProgressError, ItemAlreadyExistsError, ItemNotFoundError
 
 logger = logging.getLogger(__name__)
 
@@ -22,7 +22,7 @@ def idempotent(
     handler: Callable[[Any, LambdaContext], Any],
     event: Dict[str, Any],
     context: LambdaContext,
-    persistence: BasePersistenceLayer,
+    persistence_store: BasePersistenceLayer,
 ) -> Any:
     """
     Middleware to handle idempotency
@@ -35,7 +35,7 @@ def idempotent(
         Lambda's Event
     context: Dict
         Lambda's Context
-    persistence: BasePersistenceLayer
+    persistence_store: BasePersistenceLayer
         Instance of BasePersistenceLayer to store data
 
     Examples
@@ -45,41 +45,46 @@ def idempotent(
         >>>
         >>> persistence_store = DynamoDBPersistenceLayer(event_key="body", table_name="idempotency_store")
         >>>
-        >>> @idempotent(persistence=persistence_store)
+        >>> @idempotent(persistence_store=persistence_store)
         >>> def handler(event, context):
         >>>     return {"StatusCode": 200}
     """
 
-    persistence_instance = persistence
     try:
-        event_record = persistence_instance.get_record(event)
-    except ItemNotFoundError:
-        persistence_instance.save_inprogress(event=event)
-        return _call_lambda(handler=handler, persistence_instance=persistence_instance, event=event, context=context)
+        # We call save_inprogress first as an optimization for the most common case where no idempotent record already
+        # exists. If it succeeds, there's no need to call get_record.
+        persistence_store.save_inprogress(event=event)
+    except ItemAlreadyExistsError:
+        try:
+            event_record = persistence_store.get_record(event)
+        except ItemNotFoundError:
+            return _call_lambda(handler=handler, persistence_store=persistence_store, event=event, context=context)
 
-    if event_record.status == STATUS_CONSTANTS["EXPIRED"]:
-        return _call_lambda(handler=handler, persistence_instance=persistence_instance, event=event, context=context)
+        if event_record.status == STATUS_CONSTANTS["EXPIRED"]:
+            return _call_lambda(handler=handler, persistence_store=persistence_store, event=event, context=context)
 
-    if event_record.status == STATUS_CONSTANTS["INPROGRESS"]:
-        raise AlreadyInProgressError(
-            f"Execution already in progress with idempotency key: "
-            f"{persistence_instance.event_key}={event_record.idempotency_key}"
-        )
+        if event_record.status == STATUS_CONSTANTS["INPROGRESS"]:
+            raise AlreadyInProgressError(
+                f"Execution already in progress with idempotency key: "
+                f"{persistence_store.event_key}={event_record.idempotency_key}"
+            )
 
-    if event_record.status == STATUS_CONSTANTS["COMPLETED"]:
-        return event_record.response_json_as_dict()
+        if event_record.status == STATUS_CONSTANTS["COMPLETED"]:
+            return event_record.response_json_as_dict()
+
+    return _call_lambda(handler=handler, persistence_store=persistence_store, event=event, context=context)
 
 
 def _call_lambda(
-    handler: Callable, persistence_instance: BasePersistenceLayer, event: Dict[str, Any], context: LambdaContext
+    handler: Callable, persistence_store: BasePersistenceLayer, event: Dict[str, Any], context: LambdaContext
 ) -> Any:
     """
 
     Parameters
     ----------
     handler: Callable
         Lambda handler
-    persistence_instance: BasePersistenceLayer
+    persistence_store: BasePersistenceLayer
         Instance of persistence layer
     event
         Lambda event
@@ -89,8 +94,8 @@ def _call_lambda(
     try:
         handler_response = handler(event, context)
     except Exception as ex:
-        persistence_instance.save_error(event=event, exception=ex)
+        persistence_store.save_error(event=event, exception=ex)
         raise
     else:
-        persistence_instance.save_success(event=event, result=handler_response)
+        persistence_store.save_success(event=event, result=handler_response)
     return handler_response

aws_lambda_powertools/utilities/idempotency/persistence/__init__.py

@@ -7,14 +7,16 @@
 import json
 import logging
 from abc import ABC, abstractmethod
-from typing import Any, Dict, Optional, Union
+from typing import Any, Dict
 
-import boto3
 import jmespath
-from botocore.config import Config
 
-from .cache_dict import LRUDict
-from .exceptions import IdempotencyValidationerror, InvalidStatusError, ItemAlreadyExistsError, ItemNotFoundError
+from aws_lambda_powertools.utilities.idempotency.cache_dict import LRUDict
+from aws_lambda_powertools.utilities.idempotency.exceptions import (
+    IdempotencyValidationerror,
+    InvalidStatusError,
+    ItemAlreadyExistsError,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -297,8 +299,16 @@ def save_inprogress(self, event: Dict[str, Any]) -> None:
         )
 
         logger.debug(f"Saving in progress record for idempotency key: {data_record.idempotency_key}")
+
+        if self.use_local_cache:
+            record = self._retrieve_from_cache(idempotency_key=data_record.idempotency_key)
+            if record:
+                raise ItemAlreadyExistsError
+
         self._put_record(data_record)
 
+        # This has to come after _put_record. If _put_record call raises ItemAlreadyExists we shouldn't populate the
+        # cache with an "INPROGRESS" record as we don't know the status in the data store at this point.
         if self.use_local_cache:
             self._save_to_cache(data_record)
 
@@ -354,6 +364,10 @@ def get_record(self, lambda_event) -> DataRecord:
                 return cached_record
 
         record = self._get_record(idempotency_key)
+
+        if self.use_local_cache:
+            self._save_to_cache(data_record=record)
+
         self._validate_payload(lambda_event, record)
         return record
 
@@ -416,145 +430,3 @@ def _delete_record(self, data_record: DataRecord) -> None:
         """
 
         raise NotImplementedError
-
-
-class DynamoDBPersistenceLayer(BasePersistenceLayer):
-    def __init__(
-        self,
-        table_name: str,  # Can we use the lambda function name?
-        key_attr: Optional[str] = "id",
-        expiry_attr: Optional[str] = "expiration",
-        status_attr: Optional[str] = "status",
-        data_attr: Optional[str] = "data",
-        validation_key_attr: Optional[str] = "validation",
-        boto_config: Optional[Config] = None,
-        *args,
-        **kwargs,
-    ):
-        """
-        Initialize the DynamoDB client
-
-        Parameters
-        ----------
-        table_name: str
-            Name of the table to use for storing execution records
-        key_attr: str, optional
-            DynamoDB attribute name for key, by default "id"
-        expiry_attr: str, optional
-            DynamoDB attribute name for expiry timestamp, by default "expiration"
-        status_attr: str, optional
-            DynamoDB attribute name for status, by default "status"
-        data_attr: str, optional
-            DynamoDB attribute name for response data, by default "data"
-        boto_config: botocore.config.Config, optional
-            Botocore configuration to pass during client initialization
-        args
-        kwargs
-
-        Examples
-        --------
-        **Create a DynamoDB persistence layer with custom settings**
-            >>> from aws_lambda_powertools.utilities.idempotency import idempotent, DynamoDBPersistenceLayer
-            >>>
-            >>> persistence_store = DynamoDBPersistenceLayer(event_key="body", table_name="idempotency_store")
-            >>>
-            >>> @idempotent(persistence=persistence_store)
-            >>> def handler(event, context):
-            >>>     return {"StatusCode": 200}
-        """
-
-        boto_config = boto_config or Config()
-        self._ddb_resource = boto3.resource("dynamodb", config=boto_config)
-        self.table_name = table_name
-        self.table = self._ddb_resource.Table(self.table_name)
-        self.key_attr = key_attr
-        self.expiry_attr = expiry_attr
-        self.status_attr = status_attr
-        self.data_attr = data_attr
-        self.validation_key_attr = validation_key_attr
-        super(DynamoDBPersistenceLayer, self).__init__(*args, **kwargs)
-
-    def _item_to_data_record(self, item: Dict[str, Union[str, int]]) -> DataRecord:
-        """
-        Translate raw item records from DynamoDB to DataRecord
-
-        Parameters
-        ----------
-        item: Dict[str, Union[str, int]]
-            Item format from dynamodb response
-
-        Returns
-        -------
-        DataRecord
-            representation of item
-
-        """
-        return DataRecord(
-            idempotency_key=item[self.key_attr],
-            status=item[self.status_attr],
-            expiry_timestamp=item[self.expiry_attr],
-            response_data=item.get(self.data_attr),
-            payload_hash=item.get(self.validation_key_attr),
-        )
-
-    def _get_record(self, idempotency_key) -> DataRecord:
-        response = self.table.get_item(Key={self.key_attr: idempotency_key}, ConsistentRead=True)
-
-        try:
-            item = response["Item"]
-        except KeyError:
-            raise ItemNotFoundError
-        return self._item_to_data_record(item)
-
-    def _put_record(self, data_record: DataRecord) -> None:
-
-        item = {
-            self.key_attr: data_record.idempotency_key,
-            self.expiry_attr: data_record.expiry_timestamp,
-            self.status_attr: STATUS_CONSTANTS["INPROGRESS"],
-        }
-
-        if self.payload_validation_enabled:
-            item[self.validation_key_attr] = data_record.payload_hash
-
-        now = datetime.datetime.now()
-        try:
-            logger.debug(f"Putting record for idempotency key: {data_record.idempotency_key}")
-            self.table.put_item(
-                Item=item,
-                ConditionExpression=f"attribute_not_exists({self.key_attr}) OR expiration < :now",
-                ExpressionAttributeValues={":now": int(now.timestamp())},
-            )
-        except self._ddb_resource.meta.client.exceptions.ConditionalCheckFailedException:
-            logger.debug(f"Failed to put record for already existing idempotency key: {data_record.idempotency_key}")
-            raise ItemAlreadyExistsError
-
-    def _update_record(self, data_record: DataRecord):
-        logger.debug(f"Updating record for idempotency key: {data_record.idempotency_key}")
-        update_expression = "SET #response_data = :response_data, #expiry = :expiry, #status = :status"
-        expression_attr_values = {
-            ":expiry": data_record.expiry_timestamp,
-            ":response_data": data_record.response_data,
-            ":status": data_record.status,
-        }
-        expression_attr_names = {
-            "#response_data": self.data_attr,
-            "#expiry": self.expiry_attr,
-            "#status": self.status_attr,
-        }
-
-        if self.payload_validation_enabled:
-            update_expression += ", #validation_key = :validation_key"
-            expression_attr_values[":validation_key"] = data_record.payload_hash
-            expression_attr_names["#validation_key"] = self.validation_key_attr
-
-        self.table.update_item(
-            Key={self.key_attr: data_record.idempotency_key},
-            UpdateExpression=update_expression,
-            ExpressionAttributeValues=expression_attr_values,
-            ExpressionAttributeNames=expression_attr_names,
-        )
-
-    def _delete_record(self, data_record: DataRecord) -> None:
-        logger.debug(f"Deleting record for idempotency key: {data_record.idempotency_key}")
-        self.table.delete_item(Key={self.key_attr: data_record.idempotency_key},)