feat(idempotency): Add raise_on_no_idempotency_key flag (#297)

* feat(idempotencty): Raise error when event data is None

GIVEN a persistence_store with event_key_jmespath = `body`
WHEN getting the hashed idempotency key with an event without a `body` key
THEN raise IdempotencyValidationError

* feat(idempotency): Add raise_on_no_idempotency_key

Add `raise_on_no_idempotency_key` to enforce a idempotency key value is
passed into the lambda.

* fix: Include empty data structures

Take in account for empty data structures, including non-lists iterables like tuples, dictionaries

Co-authored-by: Heitor Lessa <[email protected]>

* tests: Include tests for empty data structures

Co-authored-by: Heitor Lessa <[email protected]>

* fix: tests

* docs: Make some doc changes

Co-authored-by: Heitor Lessa <[email protected]>

Author: Michael Brewer

Diff for: aws_lambda_powertools/utilities/idempotency/exceptions.py

@@ -43,3 +43,9 @@ class IdempotencyPersistenceLayerError(Exception):
     """
     Unrecoverable error from the data store
     """
+
+
+class IdempotencyKeyError(Exception):
+    """
+    Payload does not contain a idempotent key
+    """

Diff for: aws_lambda_powertools/utilities/idempotency/persistence/base.py

@@ -6,6 +6,7 @@
 import hashlib
 import json
 import logging
+import warnings
 from abc import ABC, abstractmethod
 from types import MappingProxyType
 from typing import Any, Dict
@@ -17,6 +18,7 @@
 from aws_lambda_powertools.utilities.idempotency.exceptions import (
     IdempotencyInvalidStatusError,
     IdempotencyItemAlreadyExistsError,
+    IdempotencyKeyError,
     IdempotencyValidationError,
 )
 
@@ -112,6 +114,7 @@ def __init__(
         use_local_cache: bool = False,
         local_cache_max_items: int = 256,
         hash_function: str = "md5",
+        raise_on_no_idempotency_key: bool = False,
     ) -> None:
         """
         Initialize the base persistence layer
@@ -130,6 +133,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.
+        raise_on_no_idempotency_key: bool, optional
+            Raise exception if no idempotency key was found in the request, by default False
         """
         self.event_key_jmespath = event_key_jmespath
         if self.event_key_jmespath:
@@ -143,6 +148,7 @@ def __init__(
             self.validation_key_jmespath = jmespath.compile(payload_validation_jmespath)
             self.payload_validation_enabled = True
         self.hash_function = getattr(hashlib, hash_function)
+        self.raise_on_no_idempotency_key = raise_on_no_idempotency_key
 
     def _get_hashed_idempotency_key(self, lambda_event: Dict[str, Any]) -> str:
         """
@@ -162,8 +168,18 @@ def _get_hashed_idempotency_key(self, lambda_event: Dict[str, Any]) -> str:
         data = lambda_event
         if self.event_key_jmespath:
             data = self.event_key_compiled_jmespath.search(lambda_event)
+
+        if self.is_missing_idempotency_key(data):
+            warnings.warn(f"No value found for idempotency_key. jmespath: {self.event_key_jmespath}")
+            if self.raise_on_no_idempotency_key:
+                raise IdempotencyKeyError("No data found to create a hashed idempotency_key")
+
         return self._generate_hash(data)
 
+    @staticmethod
+    def is_missing_idempotency_key(data) -> bool:
+        return data is None or not data or all(x is None for x in data)
+
     def _get_hashed_payload(self, lambda_event: Dict[str, Any]) -> str:
         """
         Extract data from lambda event using validation key jmespath, and return a hashed representation

Diff for: docs/utilities/idempotency.md

@@ -254,11 +254,24 @@ idempotent invocations.
     ```
 
 In this example, the "userDetail" and "productId" keys are used as the payload to generate the idempotency key. If
-we try to send the same request but with a different amount, Lambda will raise `IdempotencyValidationError`. Without
+we try to send the same request but with a different amount, we will raise `IdempotencyValidationError`. Without
 payload validation, we would have returned the same result as we did for the initial request. Since we're also
 returning an amount in the response, this could be quite confusing for the client. By using payload validation on the
 amount field, we prevent this potentially confusing behaviour and instead raise an Exception.
 
+### Making idempotency key required
+
+If you want to enforce that an idempotency key is required, you can set `raise_on_no_idempotency_key` to `True`,
+and we will raise `IdempotencyKeyError` if none was found.
+
+```python hl_lines="4"
+DynamoDBPersistenceLayer(
+    event_key_jmespath="body",
+    table_name="IdempotencyTable",
+    raise_on_no_idempotency_key=True
+    )
+```
+
 ### Changing dynamoDB attribute names
 If you want to use an existing DynamoDB table, or wish to change the name of the attributes used to store items in the
 table, you can do so when you construct the `DynamoDBPersistenceLayer` instance.
@@ -278,7 +291,7 @@ This example demonstrates changing the attribute names to custom values:
     ```python hl_lines="5-10"
     persistence_layer = DynamoDBPersistenceLayer(
         event_key_jmespath="[userDetail, productId]",
-        table_name="IdempotencyTable",)
+        table_name="IdempotencyTable",
         key_attr="idempotency_key",
         expiry_attr="expires_at",
         status_attr="current_status",

Diff for: tests/functional/idempotency/conftest.py

@@ -152,7 +152,7 @@ def hashed_validation_key(lambda_apigw_event):
 @pytest.fixture
 def persistence_store(config, request, default_jmespath):
     persistence_store = DynamoDBPersistenceLayer(
-        event_key_jmespath=default_jmespath,
+        event_key_jmespath=request.param.get("event_key_jmespath") or default_jmespath,
         table_name=TABLE_NAME,
         boto_config=config,
         use_local_cache=request.param["use_local_cache"],

Diff for: tests/functional/idempotency/test_idempotency.py

@@ -1,5 +1,7 @@
 import copy
+import json
 import sys
+from hashlib import md5
 
 import pytest
 from botocore import stub
@@ -8,11 +10,12 @@
     IdempotencyAlreadyInProgressError,
     IdempotencyInconsistentStateError,
     IdempotencyInvalidStatusError,
+    IdempotencyKeyError,
     IdempotencyPersistenceLayerError,
     IdempotencyValidationError,
 )
 from aws_lambda_powertools.utilities.idempotency.idempotency import idempotent
-from aws_lambda_powertools.utilities.idempotency.persistence.base import DataRecord
+from aws_lambda_powertools.utilities.idempotency.persistence.base import BasePersistenceLayer, DataRecord
 from aws_lambda_powertools.utilities.validation import envelopes, validator
 
 TABLE_NAME = "TEST_TABLE"
@@ -638,3 +641,52 @@ def test_delete_from_cache_when_empty(persistence_store):
     except KeyError:
         # THEN we should not get a KeyError
         pytest.fail("KeyError should not happen")
+
+
+def test_is_missing_idempotency_key():
+    # GIVEN None THEN is_missing_idempotency_key is True
+    assert BasePersistenceLayer.is_missing_idempotency_key(None)
+    # GIVEN a list of Nones THEN is_missing_idempotency_key is True
+    assert BasePersistenceLayer.is_missing_idempotency_key([None, None])
+    # GIVEN a list of all not None THEN is_missing_idempotency_key is false
+    assert BasePersistenceLayer.is_missing_idempotency_key([None, "Value"]) is False
+    # GIVEN a str THEN is_missing_idempotency_key is false
+    assert BasePersistenceLayer.is_missing_idempotency_key("Value") is False
+    # GIVEN an empty tuple THEN is_missing_idempotency_key is false
+    assert BasePersistenceLayer.is_missing_idempotency_key(())
+    # GIVEN an empty list THEN is_missing_idempotency_key is false
+    assert BasePersistenceLayer.is_missing_idempotency_key([])
+    # GIVEN an empty dictionary THEN is_missing_idempotency_key is false
+    assert BasePersistenceLayer.is_missing_idempotency_key({})
+    # GIVEN an empty str THEN is_missing_idempotency_key is false
+    assert BasePersistenceLayer.is_missing_idempotency_key("")
+
+
+@pytest.mark.parametrize("persistence_store", [{"use_local_cache": False, "event_key_jmespath": "body"}], indirect=True)
+def test_default_no_raise_on_missing_idempotency_key(persistence_store):
+    # GIVEN a persistence_store with use_local_cache = False and event_key_jmespath = "body"
+    assert persistence_store.use_local_cache is False
+    assert "body" in persistence_store.event_key_jmespath
+
+    # WHEN getting the hashed idempotency key for an event with no `body` key
+    hashed_key = persistence_store._get_hashed_idempotency_key({})
+
+    # THEN return the hash of None
+    assert md5(json.dumps(None).encode()).hexdigest() == hashed_key
+
+
+@pytest.mark.parametrize(
+    "persistence_store", [{"use_local_cache": False, "event_key_jmespath": "[body, x]"}], indirect=True
+)
+def test_raise_on_no_idempotency_key(persistence_store):
+    # GIVEN a persistence_store with raise_on_no_idempotency_key and no idempotency key in the request
+    persistence_store.raise_on_no_idempotency_key = True
+    assert persistence_store.use_local_cache is False
+    assert "body" in persistence_store.event_key_jmespath
+
+    # WHEN getting the hashed idempotency key for an event with no `body` key
+    with pytest.raises(IdempotencyKeyError) as excinfo:
+        persistence_store._get_hashed_idempotency_key({})
+
+    # THEN raise IdempotencyKeyError error
+    assert "No data found to create a hashed idempotency_key" in str(excinfo.value)