chore(idempotency): refactored expire in-progress logic

Author: rubenfonseca

aws_lambda_powertools/utilities/idempotency/persistence/dynamodb.py

@@ -167,20 +167,36 @@ def _put_record(self, data_record: DataRecord) -> None:
         now = datetime.datetime.now()
         try:
             logger.debug(f"Putting record for idempotency key: {data_record.idempotency_key}")
+
+            condition_expression = "attribute_not_exists(#id) OR #now < :now"
+            expression_attribute_names = {
+                "#id": self.key_attr,
+                "#now": self.expiry_attr,
+                "#status": self.status_attr,
+            }
+            expression_attribute_values = {":now": int(now.timestamp())}
+
+            # When we want to expire_in_progress invocations, we check if the in_progress timestamp exists
+            # and we are past that timestamp. We also make sure the status is INPROGRESS because we don't
+            # want to repeat COMPLETE invocations.
+            #
+            # We put this in an if block because customers might want to disable the feature,
+            # reverting to the old behavior that relies just on the idempotency key.
+            if self.expires_in_progress:
+                condition_expression += (
+                    " OR ("
+                    "attribute_exists(#in_progress_expiry) AND "
+                    "#in_progress_expiry < :now AND #status = :inprogress"
+                    ")"
+                )
+                expression_attribute_names["#in_progress_expiry"] = self.in_progress_expiry_attr
+                expression_attribute_values[":inprogress"] = STATUS_CONSTANTS["INPROGRESS"]
+
             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)"
-                ),
-                ExpressionAttributeNames={
-                    "#id": self.key_attr,
-                    "#now": self.expiry_attr,
-                    "#in_progress_expiry": self.in_progress_expiry_attr,
-                    "#status": self.status_attr,
-                },
-                ExpressionAttributeValues={":now": int(now.timestamp()), ":inprogress": STATUS_CONSTANTS["INPROGRESS"]},
+                ConditionExpression=condition_expression,
+                ExpressionAttributeNames=expression_attribute_names,
+                ExpressionAttributeValues=expression_attribute_values,
             )
         except self.table.meta.client.exceptions.ConditionalCheckFailedException:
             logger.debug(f"Failed to put record for already existing idempotency key: {data_record.idempotency_key}")
@@ -200,16 +216,16 @@ def _update_record(self, data_record: DataRecord):
             "#status": self.status_attr,
         }
 
-        if self.expires_in_progress:
-            update_expression += ", #in_progress_expiry = :in_progress_expiry"
-            expression_attr_values[":in_progress_expiry"] = data_record.in_progress_expiry_timestamp
-            expression_attr_names["#in_progress_expiry"] = self.in_progress_expiry_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
 
+        if self.expires_in_progress:
+            update_expression += ", #in_progress_expiry = :in_progress_expiry"
+            expression_attr_values[":in_progress_expiry"] = data_record.in_progress_expiry_timestamp
+            expression_attr_names["#in_progress_expiry"] = self.in_progress_expiry_attr
+
         kwargs = {
             "Key": self._get_key(data_record.idempotency_key),
             "UpdateExpression": update_expression,