feat(idempotench): add mechanism to register lambda context

Author: rubenfonseca

aws_lambda_powertools/utilities/idempotency/base.py

@@ -74,6 +74,7 @@ def __init__(
         self.data = deepcopy(_prepare_data(function_payload))
         self.fn_args = function_args
         self.fn_kwargs = function_kwargs
+        self.config = config
 
         persistence_store.configure(config, self.function.__name__)
         self.persistence_store = persistence_store
@@ -129,6 +130,10 @@ def _get_remaining_time_in_millis(self) -> Optional[int]:
             Remaining time in millis, or None if the remaining time cannot be determined.
         """
 
+        # Look to see if we have stored a Lambda Context
+        if self.config.lambda_context is not None:
+            self.config.lambda_context.get_remainig_time_in_millis()
+
         # Look into fn_args to see if we have a lambda context
         if self.fn_args and len(self.fn_args) == 2 and getattr(self.fn_args[1], "get_remaining_time_in_millis", None):
             return self.fn_args[1].get_remaining_time_in_millis()

aws_lambda_powertools/utilities/idempotency/config.py

@@ -1,5 +1,7 @@
 from typing import Dict, Optional
 
+from aws_lambda_powertools.utilities.typing import LambdaContext
+
 
 class IdempotencyConfig:
     def __init__(
@@ -41,3 +43,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 = None
+
+    def register_lambda_context(self, lambda_context: LambdaContext):
+        self.lambda_context = lambda_context

docs/utilities/idempotency.md

@@ -400,11 +400,6 @@ after this timestamp, and the record is still marked `INPROGRESS`, we execute th
 already expired. This means that if an invocation expired during execution, it will be quickly executed again on the
 next retry.
 
-???+ info "Info: Calculating the remaining available time"
-    For now this only works with the `idempotent` decorator. At the moment we
-    don't have access to the Lambda context when using the
-    `idempotent_function` so enabling this option is a no-op in that scenario.
-
 <center>
 ```mermaid
 sequenceDiagram
@@ -439,6 +434,32 @@ sequenceDiagram
 <i>Idempotent sequence for Lambda timeouts</i>
 </center>
 
+???+ info "Info: Calculating the remaining available time"
+    When using the `idempotent` decorator, we captura and calculate the remaining available time for you.
+    However, when using the `idempotent_function`, the functionality doesn't work out of the box. You'll
+    need to register the Lambda context on your handler:
+
+```python hl_lines="8 16" title="Registering the Lambda context"
+from aws_lambda_powertools.utilities.data_classes.sqs_event import SQSRecord
+from aws_lambda_powertools.utilities.idempotency import (
+    IdempotencyConfig, idempotent_function
+)
+
+persistence_layer = DynamoDBPersistenceLayer(table_name="...")
+
+config = IdempotencyConfig()
+
+@idempotent_function(data_keyword_argument="record", persistence_store=persistence_layer, config=config)
+def record_handler(record: SQSRecord):
+    return {"message": record["body"]}
+
+
+def lambda_handler(event, context):
+    config.register_lambda_context(context)
+
+    return record_handler(event)
+```
+
 ### Handling exceptions
 
 If you are using the `idempotent` decorator on your Lambda handler, any unhandled exceptions that are raised during the code execution will cause **the record in the persistence layer to be deleted**.