def test_filter_records_not_equals() -> None:
data_table = DataTable({"a": [1, 2, 1], "b": [3, 4, 5]})
assert data_table.filter_records({"a": 1},
operand=Filter.NOT_EQUALS) == {
"a": [2],
"b": [4],
}
assert data_table.filter_records({
"a": 2,
"b": 4
},
operand=Filter.NOT_EQUALS) == {
"a": [1, 1],
"b": [3, 5],
}
assert data_table.filter_records({
"a": 1,
"b": 4
},
operand=Filter.NOT_EQUALS) == {
"a": [1, 2, 1],
"b": [3, 4, 5],
}
with pytest.raises(DataTableError):
DataTable({
"a": [1, 2, 1],
"b": [3, 4]
}).filter_records({"a": 1}, operand=Filter.NOT_EQUALS)
def test_cached_batch_get(self):
self.client_mock.batch_get_item.return_value = {
"Responses": {
"my_table_name": [{
"pk": "my_pk",
"sk": "my_sk",
"data": "value"
}]
}
}
data_table = DataTable().add_record({"pk": "my_pk", "sk": "my_sk"})
result = list(self.result.cached_batch_get(data_table).get_records())
assert result == [{"pk": "my_pk", "sk": "my_sk", "data": "value"}]
self.client_mock.batch_get_item.assert_called_with(
RequestItems={
"my_table_name": {
"Keys": [{
"pk": "my_pk",
"sk": "my_sk"
}]
}
},
ReturnConsumedCapacity="NONE",
)
self.client_mock.batch_get_item.reset_mock()
result = list(self.result.cached_batch_get(data_table).get_records())
self.client_mock.batch_get_item.assert_not_called()
assert list(self.result.cached_batch_get(
DataTable()).get_records()) == []
def test_drop_duplicates(self):
class MyRecord(DictClass):
first_name: str
last_name: str
sex: str
data_table = DataTable(
record_class=MyRecord,
base_dict={
"first_name": ["ABC", "ABC", "DEF", "DEF", "DEF"],
"last_name": ["XYZ", "XYZ", "MNO", "MNO", "MNO"],
"sex": ["M", "M", "F", "M", "F"],
},
)
assert isinstance(data_table.get_record(0), MyRecord)
deduplicated_data_table = data_table.drop_duplicates()
assert isinstance(deduplicated_data_table.get_record(0), MyRecord)
assert {
"first_name": ["ABC", "DEF", "DEF"],
"last_name": ["XYZ", "MNO", "MNO"],
"sex": ["M", "F", "M"],
} == deduplicated_data_table
deduplicated_data_table = data_table.drop_duplicates(
subset=("first_name", ))
assert isinstance(deduplicated_data_table.get_record(0), MyRecord)
assert {
"first_name": ["ABC", "DEF"],
"last_name": ["XYZ", "MNO"],
"sex": ["M", "F"],
} == deduplicated_data_table
deduplicated_data_table = data_table.drop_duplicates(
subset=("last_name", "sex"))
assert isinstance(deduplicated_data_table.get_record(0), MyRecord)
assert {
"first_name": ["ABC", "DEF", "DEF"],
"last_name": ["XYZ", "MNO", "MNO"],
"sex": ["M", "F", "M"],
} == deduplicated_data_table
# it should throw an error as the column name provied is invalid
with pytest.raises(DataTableError):
_ = data_table.drop_duplicates(subset=("invalid_column", ))
# not normalized table
data_table = DataTable(
record_class=MyRecord,
base_dict={
"first_name": ["ABC", "ABC", "DEF", "DEF", "DEF"],
"last_name": ["XYZ", "XYZ", "MNO", "MNO"],
"sex": ["M", "M", "F", "M", "F"],
},
)
# It should throw an error as table is not normalized
with pytest.raises(DataTableError):
_ = data_table.drop_duplicates()
def test_batch_delete(self):
self.client_mock.batch_write_item.return_value = {
"Responses": {
"my_table_name": [{
"pk": "my_pk",
"sk": "my_sk",
"data": "value"
}]
}
}
data_table = DataTable().add_record({"pk": "my_pk", "sk": "my_sk"})
assert list(self.result.batch_delete(data_table).get_records()) == [{
"pk":
"my_pk",
"sk":
"my_sk"
}]
self.client_mock.batch_write_item.assert_called_with(
RequestItems={
"my_table_name": [{
"DeleteRequest": {
"Key": {
"pk": "my_pk",
"sk": "my_sk"
}
}
}]
},
ReturnConsumedCapacity="NONE",
ReturnItemCollectionMetrics="NONE",
)
assert list(self.result.batch_delete(DataTable()).get_records()) == []
def test_get_column() -> None:
data_table = DataTable({"a": [1, 2], "b": [3, DataTable.NOT_SET]})
assert data_table.get_column("a") == [1, 2]
assert data_table.get_column("b") == [3, None]
assert data_table.get_column("c") == [None, None]
with pytest.raises(DataTableError):
DataTable({"a": [1], "b": []}).get_column("a")
def test_is_normalized() -> None:
assert DataTable({"a": [1, 2], "b": [3, 4]}).is_normalized()
assert not DataTable({"a": [1, 2], "b": [3]}).is_normalized()
assert not DataTable({
"a": [1, 2],
"b": [3, 4],
"c": []
}).is_normalized()
def test_get_lenghts() -> None:
data_table = DataTable({"a": [1, 2, 3], "b": [1, 2]})
assert data_table["a"] == [1, 2, 3]
assert data_table["b"] == [1, 2]
assert data_table.max_length == 3
assert data_table.min_length == 2
assert data_table.get_lengths() == [3, 2]
assert DataTable().get_lengths() == []
assert DataTable({"a": [None]}).get_lengths() == [1]
def batch_get(self,
data_table: DataTable[_RecordType],
consistent_read: bool = False) -> DataTable[_RecordType]:
"""
Get multuple records as a DataTable from DB.
`data_table` must have all columns to calculate table keys.
Example:
```python
# UserTable is a subclass of a DynamoTable
user_table = UserTable()
# we should provide table keys or fields to calculate them
# in our case, PK is calculated from `email` field.
users_table = DataTable[UserRecord]().add_record(
{
"email": "*****@*****.**",
},
{
"email": "*****@*****.**",
},
)
user_records = user_table.batch_get(users_table)
for user_record in user_records:
# print found records
# if record was not found - it will still be returned
# but only with the data you provided
print(user_record)
```
Arguments:
data_table -- Request data table.
consistent_read -- `ConsistentRead` boto3 parameter.
Returns:
DataTable with existing records.
"""
if not data_table:
return data_table.copy()
get_data_table = DataTable()
for record in data_table.get_records():
record = self._convert_record(record)
record = self.normalize_record(record)
record.update(self._get_record_keys(record))
get_data_table.add_record(record)
results: DataTable[Any] = (
self.dynamo_query_class.build_batch_get_item(
consistent_read=consistent_read, logger=self._logger).table(
table_keys=self.table_keys,
table=self.table).execute(data_table=get_data_table))
return DataTable(record_class=self.record_class).add_table(results)
def test_get_record() -> None:
data_table = DataTable({"a": [1, 2], "b": [3, 4]})
assert data_table.get_record(0) == {"a": 1, "b": 3}
assert data_table.get_record(1) == {"a": 2, "b": 4}
with pytest.raises(DataTableError):
data_table.get_record(2)
with pytest.raises(DataTableError):
DataTable({"a": [1, 2], "b": [3]}).get_record(2)
def test_set() -> None:
data_table = DataTable({"a": [1, 2], "b": [DataTable.NOT_SET]})
result = data_table.set("a", 1, "value_a").set("b", 0, "value_b")
assert result is data_table
assert data_table == DataTable({"a": [1, "value_a"], "b": ["value_b"]})
with pytest.raises(DataTableError):
data_table.set("b", 1, "value_b")
with pytest.raises(DataTableError):
data_table.set("c", 0, "value_c")
def test_add_record() -> None:
data_table = DataTable({"a": [1], "b": [3]})
result = data_table.add_record({"a": 5, "c": 4}, {"c": 5})
assert result is data_table
assert data_table == {
"a": [1, 5, data_table.NOT_SET],
"b": [3, data_table.NOT_SET, data_table.NOT_SET],
"c": [data_table.NOT_SET, 4, 5],
}
with pytest.raises(DataTableError):
DataTable({"a": [1], "b": []}).add_record({"a": 1})
def test_append() -> None:
data_table = DataTable({"a": [1], "b": [], "c": [2, 3]})
assert data_table.append("a", [2, 3]) is data_table
assert data_table == {"a": [1, 2, 3], "b": [], "c": [2, 3]}
data_table = DataTable({"a": [1]})
data_table.append("b", [2, 3])
assert data_table == {"a": [1], "b": [2, 3]}
data_table = DataTable({"a": [1]})
with pytest.raises(DataTableError):
data_table.append("b", None)
def test_normalize() -> None:
data_table = DataTable({"a": [1, 2, 3], "b": [3, 4], "c": []})
data_table.normalize()
assert data_table == {
"a": [1, 2, 3],
"b": [3, 4, data_table.NOT_SET],
"c": [data_table.NOT_SET, data_table.NOT_SET, data_table.NOT_SET],
}
assert data_table.get_record(0) == {"a": 1, "b": 3, "c": None}
data_table = DataTable({"a": [1, 2], "b": [3, 4]})
data_table.normalize()
assert data_table == {"a": [1, 2], "b": [3, 4]}
def test_batch_get_records(self):
self.client_mock.batch_get_item.return_value = {
"Responses": {
"my_table_name": [{
"pk": "my_pk",
"sk": "my_sk",
"data": "value"
}]
}
}
assert list(
self.result.batch_get_records([{
"pk": "my_pk",
"sk": "my_sk"
}])) == [{
"pk": "my_pk",
"sk": "my_sk",
"data": "value"
}]
self.client_mock.batch_get_item.assert_called_with(
RequestItems={
"my_table_name": {
"Keys": [{
"pk": "my_pk",
"sk": "my_sk"
}]
}
},
ReturnConsumedCapacity="NONE",
)
assert list(self.result.batch_get(DataTable()).get_records()) == []
def test_get_records() -> None:
data_table = DataTable({"a": [1, 2], "b": [3, 4]})
records = data_table.get_records()
assert next(records) == {"a": 1, "b": 3}
assert next(records) == {"a": 2, "b": 4}
with pytest.raises(StopIteration):
next(records)
def test_has_column() -> None:
data_table = DataTable({"a": [1, 2], "b": [3, 4]})
assert data_table.has_column()
assert data_table.has_column("a")
assert data_table.has_column("b")
assert data_table.has_column("b", "a")
assert not data_table.has_column("c")
assert not data_table.has_column("a", "c")
def test_extend() -> None:
data_table = DataTable({"a": [1], "b": [], "c": [2, 3]})
assert data_table.extend({
"a": [2],
"b": [1]
}, {"d": [1]}) is data_table
assert data_table == {"a": [1, 2], "b": [1], "c": [2, 3], "d": [1]}
data_table = DataTable({})
data_table.extend({"a": [2], "b": [1]}, {"d": [1]}, {})
assert data_table == {"a": [2], "b": [1], "d": [1]}
data_table = DataTable({})
with pytest.raises(DataTableError):
data_table.extend({"d": None})
with pytest.raises(DataTableError):
data_table.extend([])
def test_has_set_column() -> None:
data_table = DataTable({"a": [1, 2], "b": [DataTable.NOT_SET, 3]})
assert data_table.has_set_column()
assert data_table.has_set_column("a")
assert not data_table.has_set_column("b")
assert not data_table.has_set_column("c")
assert not data_table.has_set_column("b", "a")
assert not data_table.has_set_column("c")
assert not data_table.has_set_column("a", "c")
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 cached_batch_get(
self,
data_table: DataTable[_RecordType]) -> DataTable[_RecordType]:
"""
Get multuple records as a DataTable from DB with caching.
`data_table` must have all columns to calculate table keys.
Can be used instead of `batch_get` method.
Arguments:
data_table -- Request data table.
Returns:
DataTable with existing records.
"""
if not data_table:
return data_table.copy()
result = DataTable(record_class=self.record_class)
non_cached_data_table = DataTable(record_class=self.record_class)
cached_data_table = DataTable(record_class=self.record_class)
for record in data_table.get_records():
record = self._convert_record(record)
record = self.normalize_record(record)
record_keys = self._get_record_keys(record)
cached_record = self._get_cached_record(record_keys)
if cached_record is self.NO_RECORD:
non_cached_data_table.add_record(record)
continue
if cached_record and not isinstance(cached_record, SentinelValue):
cached_data_table.add_record(cached_record)
non_cached_results = self.batch_get(non_cached_data_table)
for record in non_cached_results:
record_keys = self._get_record_keys(record)
self._cache_record(record_keys, record)
result.add_table(cached_data_table)
result.add_table(non_cached_results)
return result
def test_add_table() -> None:
data_table = DataTable({"a": [5], "b": [6]})
assert DataTable({
"a": [1, 2],
"b": [3, 4]
}).add_table(data_table) == {
"a": [1, 2, 5],
"b": [3, 4, 6],
}
assert DataTable({
"a": [1, 2],
"b": [3, 4]
}).add_table(data_table, data_table) == {
"a": [1, 2, 5, 5],
"b": [3, 4, 6, 6],
}
with pytest.raises(DataTableError):
DataTable({
"a": [1, 2],
"b": [3, 4]
}).add_table(DataTable({
"a": [5],
"b": []
}))
with pytest.raises(DataTableError):
DataTable({"a": [5], "b": []}).add_table(data_table)
def test_custom_record(self):
data_table = DataTable(record_class=UserRecord)
data_table.add_record({"name": "Jon"})
data_table.add_record(UserRecord(name="test", age=12))
assert isinstance(data_table.get_record(0), UserRecord)
assert list(data_table.get_records()) == [
UserRecord(name="Jon"),
UserRecord(name="test", age=12),
]
with pytest.raises(ValueError):
data_table.add_record({"unknown": "Jon"})
with pytest.raises(ValueError):
data_table.add_record({})
def test_filter_records_equals() -> None:
data_table = DataTable({"a": [1, 2, 1], "b": [3, 4, 5]})
assert data_table.filter_records({"a": 1}) == {
"a": [1, 1],
"b": [3, 5]
}
assert data_table.filter_records({
"a": 2,
"b": 4
}) == {
"a": [2],
"b": [4]
}
assert data_table.filter_records({
"a": 1,
"b": 4
}) == {
"a": [],
"b": []
}
with pytest.raises(DataTableError):
DataTable({"a": [1, 2, 1], "b": [3, 4]}).filter_records({"a": 1})
def test_builtin_copy() -> None:
base_dict = {"a": [[1, 2, 3]]}
data_table = DataTable(base_dict)
data_table_copy = copy(data_table)
assert isinstance(data_table_copy, DataTable)
assert data_table_copy is not data_table
assert data_table_copy["a"] is not data_table["a"]
assert data_table_copy["a"][0] is base_dict["a"][0]
data_table_deepcopy = deepcopy(data_table)
assert isinstance(data_table_deepcopy, DataTable)
assert data_table_deepcopy is not data_table
assert data_table_deepcopy["a"] is not data_table["a"]
assert data_table_deepcopy["a"][0] is not base_dict["a"][0]
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 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 main() -> None:
users_table = DataTable(record_class=UserRecord)
users_table.add_record(
{
"email": "*****@*****.**",
"name": "John",
"age": 34
},
{
"email": "*****@*****.**",
"company": "CiscoSystems",
"name": "Mary",
"age": 34
},
)
print("Get John's record:")
print(users_table.get_record(0))
print("All records as a list:")
print(list(users_table.get_records()))
print("Find record with name=Mary:")
print(users_table.filter_records({"name": "Mary"}).get_record(0))
def test_init() -> None:
base_dict = {"a": [1, 2, 3], "b": [1, 2]}
data_table = DataTable(base_dict)
DataTable.create(base_dict)
assert list(data_table.copy().normalize()) == [
{
"a": 1,
"b": 1
},
{
"a": 2,
"b": 2
},
{
"a": 3,
"b": None
},
]
assert list(data_table.keys()) == ["a", "b"]
assert list(data_table.values()) == [[1, 2, 3], [1, 2]]
assert list(data_table.items()) == [("a", [1, 2, 3]), ("b", [1, 2])]
assert data_table["a"] is not base_dict["a"]
assert data_table["a"] == [1, 2, 3]
assert data_table["b"] == [1, 2]
with pytest.raises(KeyError):
_ = data_table["c"]
assert data_table == {"a": [1, 2, 3], "b": [1, 2]}
assert not DataTable()
assert not DataTable({"a": []})
assert DataTable({"a": [1]})
with pytest.raises(DataTableError):
DataTable({"a": "b"})
with pytest.raises(DataTableError):
DataTable([1, 2, 3])
def test_batch_upsert(self, _patch_datetime):
self.client_mock.batch_write_item.return_value = {
"Responses": {
"my_table_name": [{
"pk": "my_pk",
"sk": "my_sk",
"data": "value1",
"preserve": "p1"
}]
}
}
self.client_mock.batch_get_item.return_value = {
"Responses": {
"my_table_name": [{
"pk": "my_pk",
"sk": "my_sk",
"data": "value1",
"preserve": "p1"
}]
}
}
data_table = DataTable().add_record({
"pk": "my_pk",
"sk": "my_sk",
"data": "value2",
"preserve": "p2",
"preserve2": "p3",
"gsi_pk": "gsi_pk",
"gsi_sk": "gsi_sk",
"lsi_pk": "lsi_pk",
"dt_created": None,
})
assert list(
self.result.batch_upsert(data_table,
set_if_not_exists_keys=[
"preserve", "preserve2"
]).get_records()) == [{
"pk":
"my_pk",
"sk":
"my_sk",
"data":
"value2",
"preserve":
"p1",
"preserve2":
"p3",
"gsi_pk":
"gsi_pk",
"gsi_sk":
"gsi_sk",
"lsi_pk":
"lsi_pk",
"dt_created":
"utcnow",
"dt_modified":
"utcnow",
}]
self.client_mock.batch_write_item.assert_called_with(
RequestItems={
"my_table_name": [{
"PutRequest": {
"Item": {
"pk": "my_pk",
"sk": "my_sk",
"data": "value2",
"preserve": "p1",
"preserve2": "p3",
"gsi_pk": "gsi_pk",
"gsi_sk": "gsi_sk",
"lsi_pk": "lsi_pk",
"dt_created": "utcnow",
"dt_modified": "utcnow",
}
}
}]
},
ReturnConsumedCapacity="NONE",
ReturnItemCollectionMetrics="NONE",
)
assert list(self.result.batch_upsert(DataTable()).get_records()) == []
with pytest.raises(DynamoTableError):
self.result.batch_upsert(DataTable().add_record({
"pk": "my_pk",
"sk": "my_sk",
"data": "value2",
"preserve": "p2",
"preserve2": "p3",
"gsi_pk": "gsi_pk",
"gsi_sk": "gsi_sk",
}))
with pytest.raises(DynamoTableError):
self.result.batch_upsert(DataTable().add_record({
"pk": "my_pk",
"sk": "my_sk",
"data": "value2",
"preserve": "p2",
"preserve2": "p3",
"gsi_pk": "gsi_pk",
"gsi_sk": "gsi_sk",
"lsi_pk": 12,
}))
def execute(
self,
data_table: Union[DataTable[_RecordType], RecordsType],
table: Optional[Table] = None,
table_keys: Optional[TableKeys] = TABLE_KEYS,
) -> DataTable[_RecordType]:
"""
Execute a query and get results. To get raw AWS responses, use
`query.get_raw_responses()` after this method. To get `LastEvaluatedKey`, use
`query.get_last_evaluated_key()` after this method.
If `table_keys` were not provided, method gets them from table schema. It is
slow, so it is better to pass them explicitly.
```python
input_data_table = DataTable()
input_data_table.add_record({
'name': 'John',
}, {
'name': 'Mike',
})
results = DynamoQuery.build_scan(
filter_expression=ConditionExpression('name'),
).execute(
table_keys=['pk'],
table=table_resource,
data_table=input_data_table,
)
```
Arguments:
table -- `boto3_resource.Table('my_table')`.
data_table -- `DataTable` with input data.
table_keys -- Primary and sort keys for table.
Returns:
A `DataTable` with query results.
"""
if not isinstance(data_table, DataTable):
data_table = DataTable().add_record(*data_table)
if not data_table.is_normalized():
raise DynamoQueryError("Input DataTable is not normalized.")
self.table(
table=table or self.table_resource,
table_keys=table_keys or self._table_keys,
)
self._logger.debug(
f"Execute {self._query_type.value} on {self.table_resource.name}")
if self._table_keys is None:
self._logger.warning(
"Table keys were not set, use `table_keys` argument, getting from schema."
)
self._table_keys = self.get_table_keys(self.table_resource)
self._logger.debug(f"Got table keys {set(self._table_keys)}")
self._raw_responses = []
method_map: Dict[QueryType,
Callable[[DataTable[_RecordType]],
DataTable[_RecordType]]] = {
QueryType.QUERY:
self._execute_method_query,
QueryType.SCAN:
self._execute_method_scan,
QueryType.GET_ITEM:
self._execute_method_get_item,
QueryType.UPDATE_ITEM:
self._execute_method_update_item,
QueryType.DELETE_ITEM:
self._execute_method_delete_item,
QueryType.BATCH_GET_ITEM:
self._execute_method_batch_get_item,
QueryType.BATCH_UPDATE_ITEM:
self._execute_method_batch_update_item,
QueryType.BATCH_DELETE_ITEM:
self._execute_method_batch_delete_item,
}
return method_map[self._query_type](data_table) # type: ignore