chore(idempotency): addressed review comments

Author: rubenfonseca

aws_lambda_powertools/utilities/idempotency/base.py

@@ -121,8 +121,9 @@ def _get_remaining_time_in_millis(self) -> Optional[int]:
         """
         Tries to determine the remaining time available for the current lambda invocation.
 
-        Currently, it only works if the idempotent handler decorator is used, since we need to acess the lambda context.
-        However, this could be improved if we start storing the lambda context globally during the invocation.
+        This only works if the idempotent handler decorator is used, since we need to access the lambda context.
+        However, this could be improved if we start storing the lambda context globally during the invocation. One
+        way to do this is to register the lambda context when configuring the IdempotencyConfig object.
 
         Returns
         -------

aws_lambda_powertools/utilities/idempotency/config.py

@@ -14,6 +14,7 @@ def __init__(
         use_local_cache: bool = False,
         local_cache_max_items: int = 256,
         hash_function: str = "md5",
+        lambda_context: Optional[LambdaContext] = None,
     ):
         """
         Initialize the base persistence layer
@@ -34,6 +35,8 @@ def __init__(
             Max number of items to store in local cache, by default 1024
         hash_function: str, optional
             Function to use for calculating hashes, by default md5.
+        lambda_context: LambdaContext, optional
+            Lambda Context containing information about the invocation, function and execution environment.
         """
         self.event_key_jmespath = event_key_jmespath
         self.payload_validation_jmespath = payload_validation_jmespath
@@ -43,7 +46,7 @@ def __init__(
         self.use_local_cache = use_local_cache
         self.local_cache_max_items = local_cache_max_items
         self.hash_function = hash_function
-        self.lambda_context: Optional[LambdaContext] = None
+        self.lambda_context: Optional[LambdaContext] = lambda_context
 
     def register_lambda_context(self, lambda_context: LambdaContext):
         self.lambda_context = lambda_context

aws_lambda_powertools/utilities/idempotency/persistence/base.py

@@ -8,7 +8,6 @@
 import os
 import warnings
 from abc import ABC, abstractmethod
-from math import ceil
 from types import MappingProxyType
 from typing import Any, Dict, Optional
 
@@ -355,12 +354,7 @@ def save_inprogress(self, data: Dict[str, Any], remaining_time_in_millis: Option
             now = datetime.datetime.now()
             period = datetime.timedelta(milliseconds=remaining_time_in_millis)
 
-            # It's very important to use math.ceil here. Otherwise, we might return an integer that will be smaller
-            # than the current time in milliseconds, due to rounding. This will create a scenario where the record
-            # looks already expired in the store, but the invocation is still running.
-            timestamp = ceil((now + period).timestamp())
-
-            data_record.in_progress_expiry_timestamp = timestamp
+            data_record.in_progress_expiry_timestamp = int((now + period).timestamp() * 1000)
         else:
             warnings.warn("Expires in progress is enabled but we couldn't determine the remaining time left")
 

aws_lambda_powertools/utilities/idempotency/persistence/dynamodb.py

@@ -168,19 +168,50 @@ def _put_record(self, data_record: DataRecord) -> None:
         try:
             logger.debug(f"Putting record for idempotency key: {data_record.idempotency_key}")
 
+            # |     LOCKED     |         RETRY if status = "INPROGRESS"                |     RETRY
+            # |----------------|-------------------------------------------------------|-------------> .... (time)
+            # |             Lambda                                              Idempotency Record
+            # |             Timeout                                                 Timeout
+            # |       (in_progress_expiry)                                          (expiry)
+
+            # Conditions to successfully save a record:
+
+            # The idempotency key does not exist:
+            #    - first time that this invocation key is used
+            #    - previous invocation with the same key was deleted due to TTL
+            idempotency_key_not_exist = "attribute_not_exists(#id)"
+
+            # The idempotency record exists but it's expired:
+            idempotency_expiry_expired = "#expiry < :now"
+
+            # The status of the record is "INPROGRESS", there is an in-progress expiry timestamp, but it's expired
+            inprogress_expiry_expired = " AND ".join(
+                [
+                    "#status = :inprogress",
+                    "attribute_exists(#in_progress_expiry)",
+                    "#in_progress_expiry < :now_in_millis",
+                ]
+            )
+            inprogress_expiry_expired = f"({inprogress_expiry_expired})"
+
+            condition_expression = " OR ".join(
+                [idempotency_key_not_exist, idempotency_expiry_expired, inprogress_expiry_expired]
+            )
+
             self.table.put_item(
                 Item=item,
-                ConditionExpression=(
-                    "attribute_not_exists(#id) OR #now < :now OR "
-                    "(attribute_exists(#in_progress_expiry) AND #in_progress_expiry < :now AND #status = :inprogress)"
-                ),
+                ConditionExpression=condition_expression,
                 ExpressionAttributeNames={
                     "#id": self.key_attr,
-                    "#now": self.expiry_attr,
+                    "#expiry": self.expiry_attr,
                     "#in_progress_expiry": self.in_progress_expiry_attr,
                     "#status": self.status_attr,
                 },
-                ExpressionAttributeValues={":now": int(now.timestamp()), ":inprogress": STATUS_CONSTANTS["INPROGRESS"]},
+                ExpressionAttributeValues={
+                    ":now": int(now.timestamp()),
+                    ":now_in_millis": int(now.timestamp() * 1000),
+                    ":inprogress": STATUS_CONSTANTS["INPROGRESS"],
+                },
             )
         except self.table.meta.client.exceptions.ConditionalCheckFailedException:
             logger.debug(f"Failed to put record for already existing idempotency key: {data_record.idempotency_key}")

tests/functional/idempotency/conftest.py

@@ -134,16 +134,16 @@ def expected_params_update_item_with_validation(
 def expected_params_put_item(hashed_idempotency_key):
     return {
         "ConditionExpression": (
-            "attribute_not_exists(#id) OR #now < :now OR "
-            "(attribute_exists(#in_progress_expiry) AND #in_progress_expiry < :now AND #status = :inprogress)"
+            "attribute_not_exists(#id) OR #expiry < :now OR "
+            "(#status = :inprogress AND attribute_exists(#in_progress_expiry) AND #in_progress_expiry < :now_in_millis)"
         ),
         "ExpressionAttributeNames": {
             "#id": "id",
-            "#now": "expiration",
+            "#expiry": "expiration",
             "#status": "status",
             "#in_progress_expiry": "in_progress_expiration",
         },
-        "ExpressionAttributeValues": {":now": stub.ANY, ":inprogress": "INPROGRESS"},
+        "ExpressionAttributeValues": {":now": stub.ANY, ":now_in_millis": stub.ANY, ":inprogress": "INPROGRESS"},
         "Item": {
             "expiration": stub.ANY,
             "id": hashed_idempotency_key,
@@ -158,16 +158,16 @@ def expected_params_put_item(hashed_idempotency_key):
 def expected_params_put_item_with_validation(hashed_idempotency_key, hashed_validation_key):
     return {
         "ConditionExpression": (
-            "attribute_not_exists(#id) OR #now < :now OR "
-            "(attribute_exists(#in_progress_expiry) AND #in_progress_expiry < :now AND #status = :inprogress)"
+            "attribute_not_exists(#id) OR #expiry < :now OR "
+            "(#status = :inprogress AND attribute_exists(#in_progress_expiry) AND #in_progress_expiry < :now_in_millis)"
         ),
         "ExpressionAttributeNames": {
             "#id": "id",
-            "#now": "expiration",
+            "#expiry": "expiration",
             "#status": "status",
             "#in_progress_expiry": "in_progress_expiration",
         },
-        "ExpressionAttributeValues": {":now": stub.ANY, ":inprogress": "INPROGRESS"},
+        "ExpressionAttributeValues": {":now": stub.ANY, ":now_in_millis": stub.ANY, ":inprogress": "INPROGRESS"},
         "Item": {
             "expiration": stub.ANY,
             "in_progress_expiration": stub.ANY,

tests/functional/idempotency/utils.py

@@ -19,16 +19,16 @@ def build_idempotency_put_item_stub(
     idempotency_key_hash = f"{function_name}.{handler_name}#{hash_idempotency_key(data)}"
     return {
         "ConditionExpression": (
-            "attribute_not_exists(#id) OR #now < :now OR "
-            "(attribute_exists(#in_progress_expiry) AND #in_progress_expiry < :now AND #status = :inprogress)"
+            "attribute_not_exists(#id) OR #expiry < :now OR "
+            "(#status = :inprogress AND attribute_exists(#in_progress_expiry) AND #in_progress_expiry < :now_in_millis)"
         ),
         "ExpressionAttributeNames": {
             "#id": "id",
-            "#now": "expiration",
+            "#expiry": "expiration",
             "#status": "status",
             "#in_progress_expiry": "in_progress_expiration",
         },
-        "ExpressionAttributeValues": {":now": stub.ANY, ":inprogress": "INPROGRESS"},
+        "ExpressionAttributeValues": {":now": stub.ANY, ":now_in_millis": stub.ANY, ":inprogress": "INPROGRESS"},
         "Item": {
             "expiration": stub.ANY,
             "id": idempotency_key_hash,