Merge pull request #211 from opentensor/chore/thewhaleking/more-efficient-query-map

more efficient query map

Author: thewhaleking

async_substrate_interface/async_substrate.py

@@ -55,6 +55,7 @@
     RuntimeCache,
     SubstrateMixin,
     Preprocessed,
+    RequestResults,
 )
 from async_substrate_interface.utils import (
     hex_to_bytes,
@@ -2391,7 +2392,7 @@ async def _make_rpc_request(
         attempt: int = 1,
         runtime: Optional[Runtime] = None,
         force_legacy_decode: bool = False,
-    ) -> RequestManager.RequestResults:
+    ) -> RequestResults:
         request_manager = RequestManager(payloads)
 
         subscription_added = False
@@ -3669,34 +3670,40 @@ async def query_map(
                         self.decode_ss58,
                     )
             else:
-                all_responses = []
                 page_batches = [
                     result_keys[i : i + page_size]
                     for i in range(0, len(result_keys), page_size)
                 ]
                 changes = []
-                for batch_group in [
-                    # run five concurrent batch pulls; could go higher, but it's good to be a good citizens
-                    # of the ecosystem
-                    page_batches[i : i + 5]
-                    for i in range(0, len(page_batches), 5)
-                ]:
-                    all_responses.extend(
-                        await asyncio.gather(
-                            *[
-                                self.rpc_request(
-                                    method="state_queryStorageAt",
-                                    params=[batch_keys, block_hash],
-                                    runtime=runtime,
-                                )
-                                for batch_keys in batch_group
-                            ]
+                payloads = []
+                for idx, page_batch in enumerate(page_batches):
+                    payloads.append(
+                        self.make_payload(
+                            str(idx), "state_queryStorageAt", [page_batch, block_hash]
                         )
                     )
-                for response in all_responses:
-                    for result_group in response["result"]:
-                        changes.extend(result_group["changes"])
-
+                results: RequestResults = await self._make_rpc_request(
+                    payloads, runtime=runtime
+                )
+                for result in results.values():
+                    res = result[0]
+                    if "error" in res:
+                        err_msg = res["error"]["message"]
+                        if (
+                            "Client error: Api called for an unknown Block: State already discarded"
+                            in err_msg
+                        ):
+                            bh = err_msg.split("State already discarded for ")[
+                                1
+                            ].strip()
+                            raise StateDiscardedError(bh)
+                        else:
+                            raise SubstrateRequestException(err_msg)
+                    elif "result" not in res:
+                        raise SubstrateRequestException(res)
+                    else:
+                        for result_group in res["result"]:
+                            changes.extend(result_group["changes"])
                 result = decode_query_map(
                     changes,
                     prefix,

async_substrate_interface/sync_substrate.py

@@ -34,6 +34,7 @@
     RequestManager,
     Preprocessed,
     ScaleObj,
+    RequestResults,
 )
 from async_substrate_interface.utils import (
     hex_to_bytes,
@@ -1892,7 +1893,7 @@ def _make_rpc_request(
         result_handler: Optional[ResultHandler] = None,
         attempt: int = 1,
         force_legacy_decode: bool = False,
-    ) -> RequestManager.RequestResults:
+    ) -> RequestResults:
         request_manager = RequestManager(payloads)
         _received = {}
         subscription_added = False

async_substrate_interface/types.py

@@ -369,9 +369,10 @@ def resolve_type_definition(type_id_):
         self.type_id_to_name = type_id_to_name
 
 
-class RequestManager:
-    RequestResults = dict[Union[str, int], list[Union[ScaleType, dict]]]
+RequestResults = dict[Union[str, int], list[Union[ScaleType, dict]]]
+
 
+class RequestManager:
     def __init__(self, payloads):
         self.response_map = {}
         self.responses = defaultdict(

async_substrate_interface/utils/decoding.py

@@ -113,7 +113,9 @@ def concat_hash_len(key_hasher: str) -> int:
 
     for item in result_group_changes:
         pre_decoded_keys.append(bytes.fromhex(item[0][len(prefix) :]))
-        pre_decoded_values.append(hex_to_bytes_(item[1]))
+        pre_decoded_values.append(
+            hex_to_bytes_(item[1]) if item[1] is not None else b""
+        )
     all_decoded = _decode_scale_list_with_runtime(
         pre_decoded_key_types + pre_decoded_value_types,
         pre_decoded_keys + pre_decoded_values,