Merge pull request #9928 from Chia-Network/get_keys_value

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

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] = []

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]

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