def test_query(self):
record = UserRecord(
email="*****@*****.**",
company="IBM",
name="John",
age=34,
nested={"test": {
"deep": "value"
}},
)
self.table.upsert_record(record)
assert (len(
list(
self.table.query(
partition_key="my_project",
filter_expression=ConditionExpression("nested.test.deep"),
data={"nested.test.deep": "value"},
))) == 1)
assert (len(
list(
self.table.query(
partition_key="my_project",
filter_expression=ConditionExpression("nested.test.deep"),
data={"nested.test.deep": "none"},
))) == 0)
def test_update_item():
table_resource_mock = MagicMock()
query = (DynamoQuery.build_update_item(
update_expression=UpdateExpression("test2"),
condition_expression=ConditionExpression("pk"),
).table(table=table_resource_mock,
table_keys=("pk", "sk")).update("test"))
result = query.execute_dict({
"pk": "pk_value",
"sk": "sk_value",
"test": "data"
})
table_resource_mock.update_item.assert_called_with(
ConditionExpression="#aaa = :aaa",
ExpressionAttributeNames={
"#aaa": "pk",
"#aab": "test"
},
ExpressionAttributeValues={
":aaa": "pk_value",
":aab": "data"
},
Key={
"pk": "pk_value",
"sk": "sk_value"
},
ReturnConsumedCapacity="NONE",
ReturnItemCollectionMetrics="NONE",
ReturnValues="ALL_NEW",
UpdateExpression="SET #aab = :aab",
)
assert list(result.get_records()) == []
with pytest.raises(DynamoQueryError):
DynamoQuery.build_update_item(
condition_expression=ConditionExpression("pk"), ).table(
table=table_resource_mock,
table_keys=("pk", "sk")).execute_dict({
"pk": "value",
"sk": "value",
"test": "data"
})
with pytest.raises(DynamoQueryError):
DynamoQuery.build_update_item(
condition_expression=ConditionExpression("pk"), ).table(
table=table_resource_mock,
table_keys=("pk", "sk")).update(add=["test"
], ).execute_dict({
"pk":
"value",
"sk":
"value",
"test":
"data"
})
def setup_method(self) -> None:
self.result = ConditionExpressionGroup(
[
ConditionExpression("key", ">", "value"),
ConditionExpression("key2")
],
["AND"],
)
self.other = ConditionExpressionGroup(
[ConditionExpression("key3"),
ConditionExpression("key4")], ["OR"])
self.cexpr = ConditionExpression("key5")
def test_init(self) -> None:
assert self.result.key == "key"
assert self.result.operator == "="
assert self.result.value == "value"
assert self.other.key == "key2"
assert self.other.operator == "attribute_exists"
assert self.other.value is True
assert self.nested.key == "nested.key.test"
assert ConditionExpression("key2", "=").value == "key2"
with pytest.raises(ExpressionError):
ConditionExpression("key2", "BETWEEN", ["value"])
with pytest.raises(ExpressionError):
ConditionExpression("key2", "eq")
def test_render(self) -> None:
assert self.result.render() == "{key} = {value__value}"
assert self.other.render() == "attribute_exists({key2})"
assert self.between.render(
) == "{key3} BETWEEN {value1__value} AND {value2__value}"
assert ConditionExpression(
"key", "IN", "value").render() == "{key} IN ({value__value})"
assert (ConditionExpression(
"key", "begins_with",
"value").render() == "begins_with({key}, {value__value})")
assert ConditionExpression(
"key", "attribute_exists").render() == "attribute_exists({key})"
assert (ConditionExpression(
"key", "attribute_exists",
False).render() == "attribute_not_exists({key})")
assert (ConditionExpression(
"key",
"attribute_not_exists").render() == "attribute_not_exists({key})")
assert (ConditionExpression(
"key", "attribute_not_exists",
False).render() == "attribute_exists({key})")
assert (ConditionExpression(
"key", "contains",
"value").render() == "contains({key}, {value__value})")
assert self.nested.render(
) == "{nested}.{key}.{test} = {nested_key_test__value}"
def test_errors() -> None:
filter_expression_mock = MagicMock()
projection_expression_mock = MagicMock()
table_resource_mock = MagicMock()
query = DynamoQuery.build_query(
key_condition_expression=ConditionExpression("key", "contains"),
index_name="my_index",
filter_expression=filter_expression_mock,
projection_expression=projection_expression_mock,
limit=100,
).table(table=table_resource_mock, table_keys=("pk", "sk"))
with pytest.raises(DynamoQueryError):
query.execute_dict({"key": "value"})
query = DynamoQuery.build_query(
key_condition_expression=ConditionExpression("key"),
index_name="my_index",
filter_expression=filter_expression_mock,
projection_expression=projection_expression_mock,
limit=100,
).table(table=table_resource_mock, table_keys=("pk", "sk"))
with pytest.raises(DynamoQueryError):
query.execute_dict({"key1": "value"})
with pytest.raises(DynamoQueryError):
query.execute(DataTable({"key": [1, 2], "b": [3]}))
with pytest.raises(DynamoQueryError):
query.execute(DataTable({"key": [3, DataTable.NOT_SET]}))
with pytest.raises(DynamoQueryError):
DynamoQuery.build_batch_get_item().table(
table=table_resource_mock, table_keys=("pk", "sk")).execute(
DataTable({
"pk": ["test"],
"sk": [DataTable.NOT_SET]
}))
with pytest.raises(DynamoQueryError):
DynamoQuery.build_batch_get_item().table(
table=table_resource_mock,
table_keys=("pk", "sk")).execute(DataTable({"pk": ["test"]}))
def test_methods() -> None:
table_resource_mock = MagicMock()
query = DynamoQuery.build_query(
key_condition_expression=ConditionExpression("key"),
index_name="my_index",
filter_expression=ConditionExpression("test"),
limit=100,
).projection("return")
with pytest.raises(DynamoQueryError):
_ = query.table_resource
with pytest.raises(DynamoQueryError):
_ = query.table_keys
query.table(table=table_resource_mock)
assert str(query) == "<DynamoQuery type=query>"
assert query.table_resource == table_resource_mock
assert not query.was_executed()
assert query.has_more_results()
query.execute_dict({"key": "value", "test": "data"})
assert query.was_executed()
assert query.has_more_results()
assert query.get_last_evaluated_key() == table_resource_mock.query(
).get()
assert query.get_raw_responses() == [table_resource_mock.query()]
table_resource_mock.key_schema = [
{
"AttributeName": "key"
},
{
"NotKey": "not_key"
},
]
assert query.get_table_keys(table_resource_mock) == {"key"}
assert query.limit(10)
with pytest.raises(DynamoQueryError):
DynamoQuery.build_batch_get_item().limit(10)
def main() -> None:
table = boto3.resource("dynamodb").Table("test_dq_users_table")
user_records = [
{
"pk": "*****@*****.**",
"sk": "IBM",
"email": "*****@*****.**",
"company": "IBM",
"name": "John",
"age": 34,
},
{
"pk": "*****@*****.**",
"sk": "CiscoSystems",
"email": "*****@*****.**",
"company": "CiscoSystems",
"name": "Mary",
"age": 34,
},
]
DynamoQuery.build_batch_update_item().table(
table, table_keys={"pk", "sk"}).execute(user_records)
print("Get all records:")
for record in DynamoQuery.build_scan().table(table,
table_keys={"pk", "sk"
}).execute_dict():
print(record)
print("Get John's record:")
for record in (DynamoQuery.build_get_item().table(
table, table_keys={"pk", "sk"}).execute_dict({
"pk": "*****@*****.**",
"sk": "IBM"
})):
print(record)
print("Query by a specific index:")
for record in (DynamoQuery.build_query(
index_name="gsi_name_age",
key_condition_expression=ConditionExpression("name")).table(
table, table_keys={"pk", "sk"}).execute_dict({
"name": "Mary",
"age": 34
})):
print(record)
def test_scan() -> None:
table_resource_mock = MagicMock()
query = (DynamoQuery.build_scan(
filter_expression=ConditionExpression("test"),
projection_expression=ProjectionExpression("test2"),
limit=100,
).table(table=table_resource_mock,
table_keys=("pk", "sk")).projection("test"))
result = query.execute_dict({"test": "value"})
table_resource_mock.scan.assert_called_with(
ExpressionAttributeNames={"#aaa": "test"},
ExpressionAttributeValues={":aaa": "value"},
FilterExpression="#aaa = :aaa",
Limit=100,
ProjectionExpression="#aaa",
)
assert list(result.get_records()) == []
def test_delete_item() -> None:
table_resource_mock = MagicMock()
query = DynamoQuery.build_delete_item(
condition_expression=ConditionExpression("test"), ).table(
table=table_resource_mock, table_keys=("pk", "sk"))
result = query.execute_dict({
"pk": "pk_value",
"sk": "sk_value",
"test": "data"
})
table_resource_mock.delete_item.assert_called_with(
ConditionExpression="#aaa = :aaa",
ExpressionAttributeNames={"#aaa": "test"},
ExpressionAttributeValues={":aaa": "data"},
Key={
"pk": "pk_value",
"sk": "sk_value"
},
ReturnConsumedCapacity="NONE",
ReturnItemCollectionMetrics="NONE",
ReturnValues="ALL_OLD",
)
assert list(result.get_records()) == []
def setup_method(self) -> None:
self.result = ConditionExpression("key", "=", "value")
self.other = ConditionExpression("key2", "attribute_exists")
self.between = ConditionExpression("key3", "BETWEEN",
["value1", "value2"])
self.nested = ConditionExpression("nested.key.test")
class TestConditionExpression:
result: ConditionExpression
other: ConditionExpression
between: ConditionExpression
def setup_method(self) -> None:
self.result = ConditionExpression("key", "=", "value")
self.other = ConditionExpression("key2", "attribute_exists")
self.between = ConditionExpression("key3", "BETWEEN",
["value1", "value2"])
self.nested = ConditionExpression("nested.key.test")
def test_init(self) -> None:
assert self.result.key == "key"
assert self.result.operator == "="
assert self.result.value == "value"
assert self.other.key == "key2"
assert self.other.operator == "attribute_exists"
assert self.other.value is True
assert self.nested.key == "nested.key.test"
assert ConditionExpression("key2", "=").value == "key2"
with pytest.raises(ExpressionError):
ConditionExpression("key2", "BETWEEN", ["value"])
with pytest.raises(ExpressionError):
ConditionExpression("key2", "eq")
def test_methods(self) -> None:
assert self.result.get_format_keys() == {"key"}
assert self.other.get_format_keys() == {"key2"}
assert self.result.get_format_values() == {"value"}
assert self.other.get_format_values() == set()
assert self.between.get_format_values() == {"value1", "value2"}
assert self.result.get_operators() == {"="}
assert self.other.get_operators() == {"attribute_exists"}
assert self.nested.get_format_keys() == {"nested", "key", "test"}
assert self.nested.get_format_values() == {"nested.key.test"}
def test_render(self) -> None:
assert self.result.render() == "{key} = {value__value}"
assert self.other.render() == "attribute_exists({key2})"
assert self.between.render(
) == "{key3} BETWEEN {value1__value} AND {value2__value}"
assert ConditionExpression(
"key", "IN", "value").render() == "{key} IN ({value__value})"
assert (ConditionExpression(
"key", "begins_with",
"value").render() == "begins_with({key}, {value__value})")
assert ConditionExpression(
"key", "attribute_exists").render() == "attribute_exists({key})"
assert (ConditionExpression(
"key", "attribute_exists",
False).render() == "attribute_not_exists({key})")
assert (ConditionExpression(
"key",
"attribute_not_exists").render() == "attribute_not_exists({key})")
assert (ConditionExpression(
"key", "attribute_not_exists",
False).render() == "attribute_exists({key})")
assert (ConditionExpression(
"key", "contains",
"value").render() == "contains({key}, {value__value})")
assert self.nested.render(
) == "{nested}.{key}.{test} = {nested_key_test__value}"
def test_operators(self) -> None:
and_result = self.result & self.other
assert and_result.render(
) == "{key} = {value__value} AND attribute_exists({key2})"
or_result = self.result | self.other
assert or_result.render(
) == "{key} = {value__value} OR attribute_exists({key2})"
group_and_result = self.between & or_result
assert group_and_result.render() == (
"{key3} BETWEEN {value1__value} AND {value2__value}"
" AND {key} = {value__value} OR attribute_exists({key2})")
group_or_result = self.between | and_result
assert group_or_result.render() == (
"{key3} BETWEEN {value1__value} AND {value2__value}"
" OR {key} = {value__value} AND attribute_exists({key2})")
with pytest.raises(ExpressionError):
_ = self.result & Expression("{key}")
with pytest.raises(ExpressionError):
_ = self.result | Expression("{key}")
def test_query(self):
self.table_mock.query.return_value = {
"Items": [{
"pk": "my_pk",
"sk": "sk"
}, {
"pk": "my_pk2",
"sk": "sk2"
}]
}
filter_expression_mock = ConditionExpression("key")
assert list(
self.result.query(
partition_key="pk_value",
sort_key="sk_value",
filter_expression=filter_expression_mock,
data={"key": ["key_value"]},
limit=1,
)) == [{
"pk": "my_pk",
"sk": "sk"
}]
self.table_mock.query.assert_called_with(
KeyConditionExpression="#aab = :aaa AND #aac = :aab",
FilterExpression="#aaa = :aac___0",
ConsistentRead=False,
ScanIndexForward=True,
ExpressionAttributeNames={
"#aaa": "key",
"#aab": "pk",
"#aac": "sk"
},
ExpressionAttributeValues={
":aaa": "pk_value",
":aab": "sk_value",
":aac___0": "key_value",
},
Limit=1,
)
self.table_mock.reset_mock()
list(
self.result.query(partition_key="pk",
sort_key_prefix="sk_prefix",
limit=1))
self.table_mock.query.assert_called_with(
KeyConditionExpression="#aaa = :aaa AND begins_with(#aab, :aab)",
ConsistentRead=False,
ScanIndexForward=True,
ExpressionAttributeNames={
"#aaa": "pk",
"#aab": "sk"
},
ExpressionAttributeValues={
":aaa": "pk",
":aab": "sk_prefix"
},
Limit=1,
)
with pytest.raises(DynamoTableError):
list(self.result.query(partition_key=None, limit=1))
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)
def query(
self,
partition_key: Any,
index: Optional[DynamoTableIndex] = None,
sort_key: Optional[Any] = None,
sort_key_prefix: Optional[str] = None,
filter_expression: Optional[ConditionExpressionType] = None,
consistent_read: bool = False,
scan_index_forward: bool = True,
projection: Iterable[str] = tuple(),
data: Optional[Dict[str, Any]] = None,
limit: Optional[int] = None,
) -> Iterator[_RecordType]:
"""
Query table records by index.
Example:
```python
# UserTable is a subclass of a DynamoTable
user_table = UserTable()
user_records = user_table.query(
# query by our PK
partition_key="new_users",
# and SK starting with `email_`
sort_key_prefix="email_",
# get only users older than ...
# we will provide values in data
filter_expression=ConditionExpression("age", "<="),
# get only first 5 results
limit=5,
# get only name and email fields
projection=("name", "email"),
# ...older than 45 years
data= {"age": 45}
# start with the newest records
scan_index_forward=False,
)
for user_record in user_records:
print(user_record)
```
Arguments:
partition_key -- Partition key value.
index -- DynamoTableIndex instance, primary index is used if not provided.
sort_key -- Sort key value.
sort_key_prefix -- Sort key prefix value.
filter_expression -- Query filter expression.
consistent_read -- `ConsistentRead` boto3 parameter.
scan_index_forward -- Whether to scan index from the beginning.
projection -- Record fields to return, by default returns all fields.
limit -- Max number of results.
Yields:
Matching record.
"""
if not index:
index = self.primary_index
sort_key_operator: SortKeyOperatorTypeDef = "="
partition_key_operator: PartitionKeyOperatorTypeDef = "="
if sort_key_prefix is not None:
sort_key_operator = "begins_with"
sort_key = sort_key_prefix
if partition_key is None:
raise DynamoTableError("partition_key should be set.")
key_condition_expression: ConditionExpressionType = ConditionExpression(
index.partition_key_name, operator=partition_key_operator)
if sort_key is not None and index.sort_key_name is not None:
key_condition_expression = key_condition_expression & ConditionExpression(
index.sort_key_name,
operator=sort_key_operator,
)
query = self.dynamo_query_class.build_query(
index_name=index.name,
key_condition_expression=key_condition_expression,
filter_expression=filter_expression,
consistent_read=consistent_read,
scan_index_forward=scan_index_forward,
logger=self._logger,
)
if limit:
query.limit(limit)
if projection:
query.projection(*projection)
query_data = index.get_query_data(partition_key, sort_key)
if data:
query_data.update(data)
for record in self._yield_from_query(query,
data=query_data,
limit=limit):
yield self._convert_record(record)
def test_query() -> None:
table_resource_mock = MagicMock()
query = (DynamoQuery.build_query(
key_condition_expression=ConditionExpression("pk"),
index_name="my_index",
filter_expression=ConditionExpression("test"),
projection_expression=ProjectionExpression("test2"),
exclusive_start_key={
"pk": "my_pk",
"sk": "my_sk"
},
limit=100,
).table(table=table_resource_mock,
table_keys=("pk", "sk")).projection("test"))
result = query.execute_dict({"pk": "pk_value", "test": "data"})
table_resource_mock.query.assert_called_with(
ConsistentRead=False,
ExclusiveStartKey={
"pk": "my_pk",
"sk": "my_sk"
},
ExpressionAttributeNames={
"#aaa": "pk",
"#aab": "test"
},
ExpressionAttributeValues={
":aaa": "pk_value",
":aab": "data"
},
FilterExpression="#aab = :aab",
IndexName="my_index",
KeyConditionExpression="#aaa = :aaa",
Limit=100,
ProjectionExpression="#aab",
ScanIndexForward=True,
)
assert list(result.get_records()) == []
query.reset_start_key().execute(DataTable().add_record(
{
"pk": "pk_value",
"test": "data"
}, {
"pk": "pk_value2",
"test": "data"
}))
table_resource_mock.query.assert_called_with(
ConsistentRead=False,
ExclusiveStartKey=ANY,
ExpressionAttributeNames={
"#aaa": "pk",
"#aab": "test"
},
ExpressionAttributeValues={
":aaa": "pk_value2",
":aab": "data"
},
FilterExpression="#aab = :aab",
IndexName="my_index",
KeyConditionExpression="#aaa = :aaa",
Limit=100,
ProjectionExpression="#aab",
ScanIndexForward=True,
)
with pytest.raises(DynamoQueryError):
query = (DynamoQuery.build_query(
key_condition_expression=ConditionExpression("pk"),
index_name="my_index",
exclusive_start_key={
"pk": "my_pk"
},
limit=100,
).table(table=table_resource_mock,
table_keys=("pk",
"sk")).projection("test").execute_dict({}))