"""Reusable connector to DynamoDB.""" from typing import TYPE_CHECKING, Any, Dict, List, Optional from boto3 import session from pdp import config from pdp.constants.dynamo import DYNAMODB_BATCH_WRITE_ITEM_SIZE from pdp.utils.dynamo import ( batch_write_item, decode_pagination_cursor, deserialize_dynamo_item_to_dict, encode_pagination_cursor, encode_pagination_shorthand, get_batch_write_item_keys, get_opts, serialize_dict_to_dynamo_item, ) if TYPE_CHECKING: from mypy_boto3_dynamodb import DynamoDBClient class DynamoDbConnector: """Connect to DynamoDB.""" def __init__(self, table_name: str, hash_key: str, range_key: Optional[str] = None): """Init method.""" boto3_session = session.Session() self._client = boto3_session.client("dynamodb", **get_opts(config)) self._table_name = table_name self._hash_key = hash_key self._range_key = range_key @property def client(self) -> "DynamoDBClient": """Return a typed DynamoDB client.""" return self._client def query_by_hash_key( self, hash_key: str, range_key: Optional[str] = None, cursor: Optional[str] = None, ) -> Dict[str, Any]: """Query by hash key and optional range key. Args: hash_key (str): Key to match. range_key (Optional[str] = None): Key to match. Optional. cursor (str): Encoded string, Optional. Describes where to continue pagination from. Returns: Dictionary of matching items, and pagination information """ expression_attribute_values = {":hash_key": {"S": hash_key}} key_condition_expression = f"{self._hash_key} = :hash_key" if range_key and not self._range_key: raise ValueError("A range key was given when not expected.") if range_key and self._range_key: expression_attribute_values[":range_key"] = {"S": range_key} key_condition_expression = ( f"{key_condition_expression} AND {self._range_key} = :range_key" ) query_args: Dict[str, Any] = { "ExpressionAttributeValues": expression_attribute_values, "KeyConditionExpression": key_condition_expression, "TableName": self._table_name, } if cursor: cursor = decode_pagination_cursor(cursor) query_args["ExclusiveStartKey"] = cursor response = self.client.query(**query_args) last_evaluated_key: Any = response.get("LastEvaluatedKey") shorthand_key = None if last_evaluated_key: shorthand_key = encode_pagination_shorthand( last_evaluated_key, self._hash_key, range_key=self._range_key ) last_evaluated_key = encode_pagination_cursor(last_evaluated_key) return { "items": response.get("Items", []), "cursor": {"cursor": last_evaluated_key, "shorthand": shorthand_key}, } def update_item( self, hash_key: str, item: Dict[str, Any], range_key: Optional[str] = None ) -> Dict[str, Any]: """Upsert an item using Dynamo's UpdateItem method. Args: hash_key (str): Hash key to match on range_key (Optional[str] = None): Range key to match on (optional) item (dict): Dictionary of DynamoDB attribute, value pairs. For example: { "tenant_type": {"S": "account"}, "roles": {"L": [{"M": {"role": {"S": "settings_admin"}}}]}, } Returns: Response: Dict of updated attributes """ expression_attribute_names = {f"#{attr}": attr for attr in item} expression_attribute_values = {f":{attr}": item[attr] for attr in item} update_expression_list = [] for attr in item: update_expression_list.append(f"#{attr} = :{attr}") update_expression = "SET " + ", ".join(update_expression_list) key = {self._hash_key: {"S": hash_key}} if range_key and self._range_key: key = {self._hash_key: {"S": hash_key}, self._range_key: {"S": range_key}} response = self.client.update_item( TableName=self._table_name, Key=key, ExpressionAttributeNames=expression_attribute_names, ExpressionAttributeValues=expression_attribute_values, UpdateExpression=update_expression, ReturnValues="UPDATED_NEW", # returns only updated attributes ) return response.get("Attributes", {}) async def delete_items(self, keys: List[Dict[str, Any]]) -> List[Dict[str, Any]]: """Delete a list of keys, and return any unprocessed keys. Each `key` must contain the full primary key representing an item. If a DynamoDB table uses both hash and range keys, both must be provided as the primary key in order to delete an item. """ return await self._write_items(keys, is_delete=True) async def put_items(self, items: List[Dict[str, Any]]) -> List[Dict[str, Any]]: """Put items, and return any unprocessed items. Each `item` must contain the full primary key representing the item to be created/updated. If a DynamoDB table uses both hash and range keys, both must be provided as the primary key in order to successfully create/update an item. Generally, this should not be used to update existing items as it will entirely overwrite the item. """ return await self._write_items(items, is_delete=False) async def _write_items( self, items: List[Dict[str, Any]], is_delete: bool, ) -> List[Dict[str, Any]]: """Create batches of DYNAMODB_BATCH_WRITE_ITEM_SIZE from the list of items, transform items for DynamoDB's convention, and call utils batch_write_item. Use `is_delete=True` when `items` is a list of primary keys for bulk deletion. Use `is_delete=False` when `items` is a list of items for bulk creation. """ if not len(items): return [] unprocessed_items: List[Dict[str, Any]] = [] for offset in range(0, len(items), DYNAMODB_BATCH_WRITE_ITEM_SIZE): batch = items[offset : offset + DYNAMODB_BATCH_WRITE_ITEM_SIZE] batch_request_items = self._transform_items_for_batch_write_item_requests( batch, is_delete=is_delete, ) batch_unprocessed_items = await batch_write_item( self._client, batch_request_items, ) unprocessed_items.extend(batch_unprocessed_items.get(self._table_name, [])) return self._transform_batch_write_item_requests_to_items( unprocessed_items, is_delete=is_delete, ) def _transform_items_for_batch_write_item_requests( self, items: List[Dict[str, Any]], is_delete: bool, ) -> Dict[str, List[Dict[str, Any]]]: """Transform a list of dicts (which can be primary keys or items) for a DynamoDB BatchWriteItem request. Use `is_delete=True` when `items` is a list of primary keys for bulk deletion. Use `is_delete=False` when `items` is a list of items for bulk creation. ``` """ payload_keys = get_batch_write_item_keys(is_delete) return { self._table_name: [ { payload_keys.request_type: { payload_keys.subelement: serialize_dict_to_dynamo_item(item), } } for item in items ] } def _transform_batch_write_item_requests_to_items( self, request_items: List[Dict[str, Any]], is_delete: bool, ) -> List[Dict[str, Any]]: """Return list of keys/items from a list of BatchWriteItem RequestItems. Use `is_delete=True` when `items` is a list of primary keys for bulk deletion. Use `is_delete=False` when `items` is a list of items for bulk creation. """ payload_keys = get_batch_write_item_keys(is_delete) return [ deserialize_dynamo_item_to_dict( item[payload_keys.request_type][payload_keys.subelement] ) for item in request_items ]