Add low level dynamodb client

Michal Ploski · Michal Ploski · commit d687ce95157d · 2023-02-05T23:02:50.000+01:00
diff --git a/aws_lambda_powertools/utilities/idempotency/persistence/base.py b/aws_lambda_powertools/utilities/idempotency/persistence/base.py
@@ -37,7 +37,7 @@ class DataRecord:
 
     def __init__(
         self,
-        idempotency_key,
+        idempotency_key: str,
         status: str = "",
         expiry_timestamp: Optional[int] = None,
         in_progress_expiry_timestamp: Optional[int] = None,
diff --git a/aws_lambda_powertools/utilities/idempotency/persistence/dynamodb.py b/aws_lambda_powertools/utilities/idempotency/persistence/dynamodb.py
@@ -4,7 +4,9 @@
 from typing import Any, Dict, Optional
 
 import boto3
+from boto3.dynamodb.types import TypeDeserializer
 from botocore.config import Config
+from botocore.exceptions import ClientError
 
 from aws_lambda_powertools.shared import constants
 from aws_lambda_powertools.utilities.idempotency import BasePersistenceLayer
@@ -79,13 +81,14 @@ def __init__(
 
         self._boto_config = boto_config or Config()
         self._boto3_session = boto3_session or boto3.session.Session()
+        self._client = self._boto3_session.client("dynamodb", config=self._boto_config)
+
         if sort_key_attr == key_attr:
             raise ValueError(f"key_attr [{key_attr}] and sort_key_attr [{sort_key_attr}] cannot be the same!")
 
         if static_pk_value is None:
             static_pk_value = f"idempotency#{os.getenv(constants.LAMBDA_FUNCTION_NAME_ENV, '')}"
 
-        self._table = None
         self.table_name = table_name
         self.key_attr = key_attr
         self.static_pk_value = static_pk_value
@@ -95,31 +98,15 @@ def __init__(
         self.status_attr = status_attr
         self.data_attr = data_attr
         self.validation_key_attr = validation_key_attr
-        super(DynamoDBPersistenceLayer, self).__init__()
 
-    @property
-    def table(self):
-        """
-        Caching property to store boto3 dynamodb Table resource
+        self._deserializer = TypeDeserializer()
 
-        """
-        if self._table:
-            return self._table
-        ddb_resource = self._boto3_session.resource("dynamodb", config=self._boto_config)
-        self._table = ddb_resource.Table(self.table_name)
-        return self._table
-
-    @table.setter
-    def table(self, table):
-        """
-        Allow table instance variable to be set directly, primarily for use in tests
-        """
-        self._table = table
+        super(DynamoDBPersistenceLayer, self).__init__()
 
     def _get_key(self, idempotency_key: str) -> dict:
         if self.sort_key_attr:
-            return {self.key_attr: self.static_pk_value, self.sort_key_attr: idempotency_key}
-        return {self.key_attr: idempotency_key}
+            return {self.key_attr: {"S": self.static_pk_value}, self.sort_key_attr: {"S": idempotency_key}}
+        return {self.key_attr: {"S": idempotency_key}}
 
     def _item_to_data_record(self, item: Dict[str, Any]) -> DataRecord:
         """
@@ -136,36 +123,39 @@ def _item_to_data_record(self, item: Dict[str, Any]) -> DataRecord:
             representation of item
 
         """
+        data = self._deserializer.deserialize({"M": item})
         return DataRecord(
-            idempotency_key=item[self.key_attr],
-            status=item[self.status_attr],
-            expiry_timestamp=item[self.expiry_attr],
-            in_progress_expiry_timestamp=item.get(self.in_progress_expiry_attr),
-            response_data=item.get(self.data_attr),
-            payload_hash=item.get(self.validation_key_attr),
+            idempotency_key=data[self.key_attr],
+            status=data[self.status_attr],
+            expiry_timestamp=data[self.expiry_attr],
+            in_progress_expiry_timestamp=data.get(self.in_progress_expiry_attr),
+            response_data=data.get(self.data_attr),
+            payload_hash=data.get(self.validation_key_attr),
         )
 
     def _get_record(self, idempotency_key) -> DataRecord:
-        response = self.table.get_item(Key=self._get_key(idempotency_key), ConsistentRead=True)
-
+        response = self._client.get_item(
+            TableName=self.table_name, Key=self._get_key(idempotency_key), ConsistentRead=True
+        )
         try:
             item = response["Item"]
-        except KeyError:
-            raise IdempotencyItemNotFoundError
+        except KeyError as exc:
+            raise IdempotencyItemNotFoundError from exc
         return self._item_to_data_record(item)
 
     def _put_record(self, data_record: DataRecord) -> None:
         item = {
             **self._get_key(data_record.idempotency_key),
-            self.expiry_attr: data_record.expiry_timestamp,
-            self.status_attr: data_record.status,
+            self.key_attr: {"S": data_record.idempotency_key},
+            self.expiry_attr: {"N": str(data_record.expiry_timestamp)},
+            self.status_attr: {"S": data_record.status},
         }
 
         if data_record.in_progress_expiry_timestamp is not None:
-            item[self.in_progress_expiry_attr] = data_record.in_progress_expiry_timestamp
+            item[self.in_progress_expiry_attr] = {"N": str(data_record.in_progress_expiry_timestamp)}
 
         if self.payload_validation_enabled:
-            item[self.validation_key_attr] = data_record.payload_hash
+            item[self.validation_key_attr] = {"S": data_record.payload_hash}
 
         now = datetime.datetime.now()
         try:
@@ -199,8 +189,8 @@ def _put_record(self, data_record: DataRecord) -> None:
             condition_expression = (
                 f"{idempotency_key_not_exist} OR {idempotency_expiry_expired} OR ({inprogress_expiry_expired})"
             )
-
-            self.table.put_item(
+            self._client.put_item(
+                TableName=self.table_name,
                 Item=item,
                 ConditionExpression=condition_expression,
                 ExpressionAttributeNames={
@@ -210,22 +200,28 @@ def _put_record(self, data_record: DataRecord) -> None:
                     "#status": self.status_attr,
                 },
                 ExpressionAttributeValues={
-                    ":now": int(now.timestamp()),
-                    ":now_in_millis": int(now.timestamp() * 1000),
-                    ":inprogress": STATUS_CONSTANTS["INPROGRESS"],
+                    ":now": {"N": str(int(now.timestamp()))},
+                    ":now_in_millis": {"N": str(int(now.timestamp() * 1000))},
+                    ":inprogress": {"S": 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}")
-            raise IdempotencyItemAlreadyExistsError
+        except ClientError as exc:
+            error_code = exc.response.get("Error", {}).get("Code")
+            if error_code == "ConditionalCheckFailedException":
+                logger.debug(
+                    f"Failed to put record for already existing idempotency key: {data_record.idempotency_key}"
+                )
+                raise IdempotencyItemAlreadyExistsError from exc
+            else:
+                raise
 
     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,
+            ":expiry": {"N": str(data_record.expiry_timestamp)},
+            ":response_data": {"S": data_record.response_data},
+            ":status": {"S": data_record.status},
         }
         expression_attr_names = {
             "#expiry": self.expiry_attr,
@@ -235,7 +231,7 @@ def _update_record(self, data_record: DataRecord):
 
         if self.payload_validation_enabled:
             update_expression += ", #validation_key = :validation_key"
-            expression_attr_values[":validation_key"] = data_record.payload_hash
+            expression_attr_values[":validation_key"] = {"S": data_record.payload_hash}
             expression_attr_names["#validation_key"] = self.validation_key_attr
 
         kwargs = {
@@ -245,8 +241,8 @@ def _update_record(self, data_record: DataRecord):
             "ExpressionAttributeNames": expression_attr_names,
         }
 
-        self.table.update_item(**kwargs)
+        self._client.update_item(TableName=self.table_name, **kwargs)
 
     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._get_key(data_record.idempotency_key))
+        self._client.delete_item(TableName=self.table_name, Key={**self._get_key(data_record.idempotency_key)})
diff --git a/mypy.ini b/mypy.ini
@@ -29,6 +29,9 @@ ignore_missing_imports = True
 [mypy-boto3.dynamodb.conditions]
 ignore_missing_imports = True
 
+[mypy-boto3.dynamodb.types]
+ignore_missing_imports = True
+
 [mypy-botocore.config]
 ignore_missing_imports = True
 
@@ -58,3 +61,5 @@ ignore_missing_imports = True
 
 [mypy-ijson]
 ignore_missing_imports = True
+
+
diff --git a/tests/e2e/idempotency/handlers/parallel_functions_handler.py b/tests/e2e/idempotency/handlers/parallel_functions_handler.py
@@ -0,0 +1,30 @@
+import os
+import time
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from threading import current_thread
+
+from aws_lambda_powertools.utilities.idempotency import (
+    DynamoDBPersistenceLayer,
+    idempotent_function,
+)
+
+TABLE_NAME = os.getenv("IdempotencyTable", "")
+persistence_layer = DynamoDBPersistenceLayer(table_name=TABLE_NAME)
+threads_count = 2
+
+
+@idempotent_function(persistence_store=persistence_layer, data_keyword_argument="record")
+def record_handler(record):
+    time_now = time.time()
+    return {"thread_name": current_thread().name, "time": str(time_now)}
+
+
+def lambda_handler(event, context):
+    with ThreadPoolExecutor(max_workers=threads_count) as executor:
+        futures = [executor.submit(record_handler, **{"record": event}) for _ in range(threads_count)]
+
+    output = []
+    for future in as_completed(futures):
+        output.append({"state": future._state, "exception": future.exception(), "output": future.result()})
+
+    return output
diff --git a/tests/e2e/idempotency/infrastructure.py b/tests/e2e/idempotency/infrastructure.py
@@ -15,6 +15,7 @@ def create_resources(self):
         table.grant_read_write_data(functions["TtlCacheExpirationHandler"])
         table.grant_read_write_data(functions["TtlCacheTimeoutHandler"])
         table.grant_read_write_data(functions["ParallelExecutionHandler"])
+        table.grant_read_write_data(functions["ParallelFunctionsHandler"])
 
     def _create_dynamodb_table(self) -> Table:
         table = dynamodb.Table(
diff --git a/tests/e2e/idempotency/test_idempotency_dynamodb.py b/tests/e2e/idempotency/test_idempotency_dynamodb.py
@@ -22,12 +22,16 @@ def parallel_execution_handler_fn_arn(infrastructure: dict) -> str:
     return infrastructure.get("ParallelExecutionHandlerArn", "")
 
 
+@pytest.fixture
+def parallel_functions_handler_fn_arn(infrastructure: dict) -> str:
+    return infrastructure.get("ParallelFunctionsHandlerArn", "")
+
+
 @pytest.fixture
 def idempotency_table_name(infrastructure: dict) -> str:
     return infrastructure.get("DynamoDBTable", "")
 
 
-@pytest.mark.xdist_group(name="idempotency")
 def test_ttl_caching_expiration_idempotency(ttl_cache_expiration_handler_fn_arn: str):
     # GIVEN
     payload = json.dumps({"message": "Lambda Powertools - TTL 5s"})
@@ -57,7 +61,6 @@ def test_ttl_caching_expiration_idempotency(ttl_cache_expiration_handler_fn_arn:
     assert third_execution_response != second_execution_response
 
 
-@pytest.mark.xdist_group(name="idempotency")
 def test_ttl_caching_timeout_idempotency(ttl_cache_timeout_handler_fn_arn: str):
     # GIVEN
     payload_timeout_execution = json.dumps({"sleep": 5, "message": "Lambda Powertools - TTL 1s"})
@@ -81,7 +84,6 @@ def test_ttl_caching_timeout_idempotency(ttl_cache_timeout_handler_fn_arn: str):
     assert payload_working_execution == execution_working_response
 
 
-@pytest.mark.xdist_group(name="idempotency")
 def test_parallel_execution_idempotency(parallel_execution_handler_fn_arn: str):
     # GIVEN
     arguments = json.dumps({"message": "Lambda Powertools - Parallel execution"})
@@ -97,3 +99,29 @@ def test_parallel_execution_idempotency(parallel_execution_handler_fn_arn: str):
     # THEN
     assert "Execution already in progress with idempotency key" in error_idempotency_execution_response
     assert "Task timed out after" in timeout_execution_response
+
+
+def test_parallel_functions_execution_idempotency(parallel_functions_handler_fn_arn: str):
+    # GIVEN
+    payload = json.dumps({"message": "Lambda Powertools - Parallel functions execution"})
+
+    # WHEN
+    # first execution
+    first_execution, _ = data_fetcher.get_lambda_response(lambda_arn=parallel_functions_handler_fn_arn, payload=payload)
+    first_execution_response = first_execution["Payload"].read().decode("utf-8")
+
+    # the second execution should return the same response as the first execution
+    second_execution, _ = data_fetcher.get_lambda_response(
+        lambda_arn=parallel_functions_handler_fn_arn, payload=payload
+    )
+    second_execution_response = second_execution["Payload"].read().decode("utf-8")
+
+    # THEN
+    # Function threads finished without exception AND
+    # first and second execution is the same
+    for function_thread in json.loads(first_execution_response):
+        assert function_thread["state"] == "FINISHED"
+        assert function_thread["exception"] is None
+        assert function_thread["output"] is not None
+
+    assert first_execution_response == second_execution_response
diff --git a/tests/functional/idempotency/conftest.py b/tests/functional/idempotency/conftest.py
@@ -85,11 +85,11 @@ def expected_params_update_item(serialized_lambda_response, hashed_idempotency_k
             "#status": "status",
         },
         "ExpressionAttributeValues": {
-            ":expiry": stub.ANY,
-            ":response_data": serialized_lambda_response,
-            ":status": "COMPLETED",
+            ":expiry": {"N": stub.ANY},
+            ":response_data": {"S": serialized_lambda_response},
+            ":status": {"S": "COMPLETED"},
         },
-        "Key": {"id": hashed_idempotency_key},
+        "Key": {"id": {"S": hashed_idempotency_key}},
         "TableName": "TEST_TABLE",
         "UpdateExpression": "SET #response_data = :response_data, " "#expiry = :expiry, #status = :status",
     }
@@ -107,12 +107,12 @@ def expected_params_update_item_with_validation(
             "#validation_key": "validation",
         },
         "ExpressionAttributeValues": {
-            ":expiry": stub.ANY,
-            ":response_data": serialized_lambda_response,
-            ":status": "COMPLETED",
-            ":validation_key": hashed_validation_key,
+            ":expiry": {"N": stub.ANY},
+            ":response_data": {"S": serialized_lambda_response},
+            ":status": {"S": "COMPLETED"},
+            ":validation_key": {"S": hashed_validation_key},
         },
-        "Key": {"id": hashed_idempotency_key},
+        "Key": {"id": {"S": hashed_idempotency_key}},
         "TableName": "TEST_TABLE",
         "UpdateExpression": (
             "SET #response_data = :response_data, "
@@ -135,12 +135,16 @@ def expected_params_put_item(hashed_idempotency_key):
             "#status": "status",
             "#in_progress_expiry": "in_progress_expiration",
         },
-        "ExpressionAttributeValues": {":now": stub.ANY, ":now_in_millis": stub.ANY, ":inprogress": "INPROGRESS"},
+        "ExpressionAttributeValues": {
+            ":now": {"N": stub.ANY},
+            ":now_in_millis": {"N": stub.ANY},
+            ":inprogress": {"S": "INPROGRESS"},
+        },
         "Item": {
-            "expiration": stub.ANY,
-            "id": hashed_idempotency_key,
-            "status": "INPROGRESS",
-            "in_progress_expiration": stub.ANY,
+            "expiration": {"N": stub.ANY},
+            "in_progress_expiration": {"N": stub.ANY},
+            "id": {"S": hashed_idempotency_key},
+            "status": {"S": "INPROGRESS"},
         },
         "TableName": "TEST_TABLE",
     }
@@ -159,13 +163,17 @@ def expected_params_put_item_with_validation(hashed_idempotency_key, hashed_vali
             "#status": "status",
             "#in_progress_expiry": "in_progress_expiration",
         },
-        "ExpressionAttributeValues": {":now": stub.ANY, ":now_in_millis": stub.ANY, ":inprogress": "INPROGRESS"},
+        "ExpressionAttributeValues": {
+            ":now": {"N": stub.ANY},
+            ":now_in_millis": {"N": stub.ANY},
+            ":inprogress": {"S": "INPROGRESS"},
+        },
         "Item": {
-            "expiration": stub.ANY,
-            "in_progress_expiration": stub.ANY,
-            "id": hashed_idempotency_key,
-            "status": "INPROGRESS",
-            "validation": hashed_validation_key,
+            "expiration": {"N": stub.ANY},
+            "in_progress_expiration": {"N": stub.ANY},
+            "id": {"S": hashed_idempotency_key},
+            "status": {"S": "INPROGRESS"},
+            "validation": {"S": hashed_validation_key},
         },
         "TableName": "TEST_TABLE",
     }
diff --git a/tests/functional/idempotency/test_idempotency.py b/tests/functional/idempotency/test_idempotency.py
diff --git a/tests/functional/idempotency/utils.py b/tests/functional/idempotency/utils.py
