def test_chunkify() -> None:
data = [1, 2, 3, 4, 5]
assert list(chunkify(data, size=2)) == [[1, 2], [3, 4], [5]]
assert list(chunkify([], size=2)) == []
generator = chunkify(data, size=2)
assert next(generator) == [1, 2]
assert next(generator) == [3, 4]
assert next(generator) == [5]
with pytest.raises(StopIteration):
next(generator)
def batch_delete_records(self, records: Iterable[_RecordType]) -> None:
"""
Delete records from DB.
See `DynamoTable.batch_delete`.
Arguments:
records -- Full or partial records to delete.
"""
for records_chunk in chunkify(records, self.max_batch_size):
delete_data_table = DataTable[_RecordType](
record_class=self.record_class).add_record(*records_chunk)
self.batch_delete(delete_data_table)
def _execute_method_batch_update_item(self,
data_table: DataTable) -> DataTable:
self._validate_data_table_has_table_keys(data_table)
record_chunks = chunkify(data_table.get_records(), self.MAX_BATCH_SIZE)
table_name = self.table_resource.name
for record_chunk in record_chunks:
request_list = []
for record in record_chunk:
request_list.append({"PutRequest": {"Item": dict(record)}})
request_items = {table_name: request_list}
self._batch_write_item(
RequestItems=request_items,
**self._extra_params,
)
return data_table
def batch_upsert_records(
self,
records: Iterable[_RecordType],
set_if_not_exists_keys: Iterable[str] = (),
) -> None:
"""
Upsert records to DB.
See `DynamoTable.batch_upsert`.
Arguments:
records -- Full or partial records data.
set_if_not_exists_keys -- List of keys to set only if they no do exist in DB.
"""
for records_chunk in chunkify(records, self.max_batch_size):
upsert_data_table = DataTable(
record_class=self.record_class).add_record(*records_chunk)
self.batch_upsert(upsert_data_table,
set_if_not_exists_keys=set_if_not_exists_keys)
def _execute_method_batch_get_item(self,
data_table: DataTable) -> DataTable:
self._validate_data_table_has_table_keys(data_table)
record_chunks = chunkify(data_table.get_records(), self.MAX_BATCH_SIZE)
table_name = self.table_resource.name
response_table = DataTable[Dict[str, Any]]()
for record_chunk in record_chunks:
key_data_list = []
for record in record_chunk:
key_data = {
k: v
for k, v in record.items() if k in self.table_keys
}
key_data_list.append(key_data)
request_items = {
table_name: {
"Keys": key_data_list,
"ConsistentRead": self._consistent_read
}
}
response = self._batch_get_item(
RequestItems=request_items,
**self._extra_params,
)
if response.get("Responses", {}).get(table_name):
response_table.add_record(*response["Responses"][table_name])
result = DataTable[Dict[str, Any]]()
for record in data_table.get_records():
key_data = {
k: v
for k, v in record.items() if k in self.table_keys
}
response_records = response_table.filter_records(
key_data).get_records()
for response_record in response_records:
record.update(response_record)
result.add_record(record)
return result
def _execute_method_batch_delete_item(self,
data_table: DataTable) -> DataTable:
self._validate_data_table_has_table_keys(data_table)
record_chunks = chunkify(data_table.get_records(), self.MAX_BATCH_SIZE)
table_name = self.table_resource.name
for record_chunk in record_chunks:
request_list = []
for record in record_chunk:
key_data = {
k: v
for k, v in record.items() if k in self.table_keys
}
request_item = {"DeleteRequest": {"Key": key_data}}
if request_item not in request_list:
request_list.append(request_item)
request_items = {table_name: request_list}
self._batch_write_item(
RequestItems=request_items,
**self._extra_params,
)
return data_table
def batch_get_records(
self,
records: Iterable[_RecordType],
consistent_read: bool = False) -> Iterator[_RecordType]:
"""
Get records as an iterator from DB.
See `DynamoTable.batch_get`.
Arguments:
records -- Full or partial records data.
consistent_read -- `ConsistentRead` boto3 parameter.
Yields:
Found or not found record data.
"""
for records_chunk in chunkify(records, self.max_batch_size):
get_data_table = DataTable(
record_class=self.record_class).add_record(*records_chunk)
result_data_table = self.batch_get(get_data_table,
consistent_read=consistent_read)
for record in result_data_table.get_records():
yield self._convert_record(record)
def clear_table(
self,
partition_key: Optional[str] = None,
partition_key_prefix: Optional[str] = None,
sort_key: Optional[str] = None,
sort_key_prefix: Optional[str] = None,
index: Optional[DynamoTableIndex] = None,
filter_expression: Optional[ConditionExpressionType] = None,
limit: Optional[int] = None,
) -> None:
"""
Remove records from DB.
If `partition_key` and `partition_key_prefix` are None - deletes all records.
Arguments:
partition_key -- Partition key value.
sort_key -- Sort key value.
partition_key_prefix -- Partition key prefix value.
sort_key_prefix -- Sort key prefix value.
index -- DynamoTableIndex instance, primary index is used if not provided.
filter_expression -- Query filter expression.
limit -- Max number of results.
"""
if partition_key is not None:
records = self.query(
partition_key=partition_key,
index=index,
sort_key=sort_key,
sort_key_prefix=sort_key_prefix,
filter_expression=filter_expression,
limit=limit,
projection=self._get_keys_projection(),
)
else:
filter_expressions: List[ConditionExpression] = []
data = {}
if partition_key_prefix:
filter_expressions.append(
ConditionExpression(self.partition_key_name,
operator="begins_with"))
data[self.partition_key_name] = partition_key_prefix
if sort_key and self.sort_key_name:
filter_expressions.append(
ConditionExpression(self.sort_key_name))
data[self.sort_key_name] = sort_key
if sort_key_prefix and self.sort_key_name:
filter_expressions.append(
ConditionExpression(self.sort_key_name,
operator="begins_with"))
data[self.sort_key_name] = sort_key_prefix
if not filter_expressions:
filter_expression = None
else:
filter_expression = filter_expressions[0]
for part in filter_expressions[1:]:
filter_expression = filter_expression & part
records = self.scan(
filter_expression=filter_expression,
data=data,
projection=self._get_keys_projection(),
)
for records_chunk in chunkify(records, self.max_batch_size):
existing_records = DataTable(
record_class=self.record_class).add_record(*records_chunk)
self.dynamo_query_class.build_batch_delete_item(
logger=self._logger).table(
table_keys=self.table_keys,
table=self.table,
).execute(existing_records)