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