fix(idempotency): make idempotent_function decorator thread safe (aws-powertools#1899)

* update changelog with latest changes * Add low level dynamodb client * Fix black and mypy warnings * Fix test error + add missing pytest groupings * Simplify parallel function handler * Rename e2e tests * fix(idempotency): updating documentation --------- Co-authored-by: Release bot <aws-devax-open-source@amazon.com> Co-authored-by: Michal Ploski <mploski@amazon.com> Co-authored-by: Leandro Damascena <leandro.damascena@gmail.com>

Author: 4 people

aws_lambda_powertools/event_handler/api_gateway.py

@@ -836,7 +836,6 @@ def route(
  # Override _compile_regex to exclude trailing slashes for route resolution
  @staticmethod
  def _compile_regex(rule: str, base_regex: str = _ROUTE_REGEX):
-
  return super(APIGatewayRestResolver, APIGatewayRestResolver)._compile_regex(rule, "^{}/*$")
 
 

aws_lambda_powertools/logging/utils.py

@@ -12,7 +12,6 @@ def copy_config_to_registered_loggers(
  exclude: Optional[Set[str]] = None,
  include: Optional[Set[str]] = None,
 ) -> None:
-
  """Copies source Logger level and handler to all registered loggers for consistent formatting.
 
  Parameters

aws_lambda_powertools/utilities/feature_flags/schema.py

@@ -272,7 +272,6 @@ def __init__(self, rule: Dict[str, Any], rule_name: str, logger: Optional[Union[
  self.logger = logger or logging.getLogger(__name__)
 
  def validate(self):
-
  if not self.conditions or not isinstance(self.conditions, list):
  self.logger.debug(f"Condition is empty or invalid for rule={self.rule_name}")
  raise SchemaValidationError(f"Invalid condition, rule={self.rule_name}")

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,

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
+ if self.payload_validation_enabled and 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)})

docs/utilities/idempotency.md

@@ -132,13 +132,9 @@ When using `idempotent_function`, you must tell us which keyword parameter in yo
 
 !!! info "We support JSON serializable data, [Python Dataclasses](https://docs.python.org/3.7/library/dataclasses.html){target="_blank"}, [Parser/Pydantic Models](parser.md){target="_blank"}, and our [Event Source Data Classes](./data_classes.md){target="_blank"}."
 
-???+ warning "Limitations"
+???+ warning "Limitation"
  Make sure to call your decorated function using keyword arguments.
 
- Decorated functions with `idempotent_function` are not thread-safe, if the caller uses threading, not the function computation itself.
-
- DynamoDB Persistency layer uses a Resource client [which is not thread-safe](https://boto3.amazonaws.com/v1/documentation/api/latest/guide/resources.html?highlight=multithreading#multithreading-or-multiprocessing-with-resources){target="_blank"}.
-
 === "dataclass_sample.py"
 
  ```python hl_lines="3-4 23 33"

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
+
+

tests/e2e/idempotency/handlers/function_thread_safety_handler.py

@@ -0,0 +1,29 @@
+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)]
+
+ return [
+ {"state": future._state, "exception": future.exception(), "output": future.result()}
+ for future in as_completed(futures)
+ ]

tests/e2e/idempotency/handlers/parallel_execution_handler.py

@@ -12,7 +12,6 @@
 
 @idempotent(persistence_store=persistence_layer)
 def lambda_handler(event, context):
-
  time.sleep(5)
 
  return event

tests/e2e/idempotency/handlers/ttl_cache_expiration_handler.py

@@ -14,7 +14,6 @@
 
 @idempotent(config=config, persistence_store=persistence_layer)
 def lambda_handler(event, context):
-
  time_now = time.time()
 
  return {"time": str(time_now)}