diff --git a/chia/data_layer/data_layer.py b/chia/data_layer/data_layer.py index 43dcd79d3e..870a459606 100644 --- a/chia/data_layer/data_layer.py +++ b/chia/data_layer/data_layer.py @@ -119,8 +119,8 @@ class DataLayer: return None return res.value - async def get_keys_values(self, store_id: bytes32) -> List[TerminalNode]: - res = await self.data_store.get_keys_values(store_id) + async def get_keys_values(self, store_id: bytes32, root_hash: Optional[bytes32]) -> List[TerminalNode]: + res = await self.data_store.get_keys_values(store_id, root_hash) if res is None: self.log.error("Failed to fetch keys values") return res diff --git a/chia/data_layer/data_store.py b/chia/data_layer/data_store.py index 3b0c691dc1..6438b0bb5f 100644 --- a/chia/data_layer/data_store.py +++ b/chia/data_layer/data_store.py @@ -374,13 +374,15 @@ class DataStore: return ancestors - async def get_keys_values(self, tree_id: bytes32, *, lock: bool = True) -> List[TerminalNode]: + async def get_keys_values( + self, tree_id: bytes32, root_hash: Optional[bytes32] = None, *, lock: bool = True + ) -> List[TerminalNode]: async with self.db_wrapper.locked_transaction(lock=lock): - root = await self.get_tree_root(tree_id=tree_id, lock=False) - - if root.node_hash is None: - return [] - + if root_hash is None: + root = await self.get_tree_root(tree_id=tree_id, lock=False) + if root.node_hash is None: + raise Exception(f"Root hash is unspecified for tree ID: {tree_id.hex()}") + root_hash = root.node_hash cursor = await self.db.execute( """ WITH RECURSIVE @@ -402,7 +404,7 @@ class DataStore: WHERE node_type == :node_type ORDER BY depth ASC, rights ASC """, - {"root_hash": root.node_hash.hex(), "node_type": NodeType.TERMINAL}, + {"root_hash": root_hash.hex(), "node_type": NodeType.TERMINAL}, ) terminal_nodes: List[TerminalNode] = [] diff --git a/chia/rpc/data_layer_rpc_api.py b/chia/rpc/data_layer_rpc_api.py index c7952213e2..7c4d89adc4 100644 --- a/chia/rpc/data_layer_rpc_api.py +++ b/chia/rpc/data_layer_rpc_api.py @@ -75,9 +75,12 @@ class DataLayerRpcApi: async def get_keys_values(self, request: Dict[str, Any]) -> Dict[str, Any]: store_id = bytes32(hexstr_to_bytes(request["id"])) + root_hash = request.get("root_hash") + if root_hash is not None: + root_hash = bytes32.from_hexstr(root_hash) if self.service is None: raise Exception("Data layer not created") - res = await self.service.get_keys_values(store_id) + res = await self.service.get_keys_values(store_id, root_hash) json_nodes = [] for node in res: json = recurse_jsonify(dataclasses.asdict(node)) # type: ignore[no-untyped-call] diff --git a/tests/core/data_layer/test_data_rpc.py b/tests/core/data_layer/test_data_rpc.py index 1e77ff6949..e742072004 100644 --- a/tests/core/data_layer/test_data_rpc.py +++ b/tests/core/data_layer/test_data_rpc.py @@ -257,6 +257,29 @@ async def test_keys_values_ancestors(one_wallet_node_and_rpc: nodes) -> None: val = await data_rpc_api.get_ancestors({"id": store_id.hex(), "hash": val["keys_values"][4]["hash"]}) # todo better assertions for get_ancestors result assert len(val["ancestors"]) == 3 + res_before = await data_rpc_api.get_roots({"ids": [store_id.hex()]}) + key6 = b"tasdfsd" + value6 = b"\x08\x02" + changelist = [{"action": "insert", "key": key6.hex(), "value": value6.hex()}] + key7 = b"basdff" + value7 = b"\x09\x02" + changelist.append({"action": "insert", "key": key7.hex(), "value": value7.hex()}) + res = await data_rpc_api.batch_update({"id": store_id.hex(), "changelist": changelist}) + update_tx_rec0 = res["tx_id"] + await asyncio.sleep(1) + for i in range(0, num_blocks): + await full_node_api.farm_new_transaction_block(FarmNewBlockProtocol(ph)) + await asyncio.sleep(0.2) + await time_out_assert(15, is_transaction_confirmed, True, "this is unused", wallet_rpc_api, update_tx_rec0) + res_after = await data_rpc_api.get_roots({"ids": [store_id.hex()]}) + pairs_before = await data_rpc_api.get_keys_values( + {"id": store_id.hex(), "root_hash": res_before["hashes"][0].hex()} + ) + pairs_after = await data_rpc_api.get_keys_values( + {"id": store_id.hex(), "root_hash": res_after["hashes"][0].hex()} + ) + assert len(pairs_before["keys_values"]) == 5 + assert len(pairs_after["keys_values"]) == 7 @pytest.mark.asyncio