feat: batch query can split queries to avoid hitting query gas limits

securesecrets · Mar 28, 2024 · c4d8927 · c4d8927
1 parent 211e3a2
commit c4d8927
Show file tree

Hide file tree

Showing 6 changed files with 199 additions and 18 deletions.
diff --git a/.changeset/fifty-shirts-melt.md b/.changeset/fifty-shirts-melt.md
@@ -0,0 +1,5 @@
+---
+"@shadeprotocol/shadejs": patch
+---
+
+Batch query can split queries to avoid hitting query gas limits
diff --git a/src/contracts/services/batchQuery.test.ts b/src/contracts/services/batchQuery.test.ts
@@ -4,6 +4,7 @@ import {
   vi,
   beforeAll,
   afterAll,
+  afterEach,
 } from 'vitest';
 import { of } from 'rxjs';
 import batchPairConfigResponse from '~/test/mocks/batchQuery/batchPairConfigResponse.json';
@@ -12,10 +13,12 @@ import { BatchQueryParams } from '~/types/contracts/batchQuery/model';
 import { msgBatchQuery } from '~/contracts/definitions/batchQuery';
 import batchPricesWithError from '~/test/mocks/batchQuery/batchIndividualPricesWithErrorResponse.json';
 import { batchPricesWithErrorParsed } from '~/test/mocks/batchQuery/batchPricesWithErrorParsed';
+import { SecretNetworkClient } from 'secretjs';
 import {
   parseBatchQuery,
   batchQuery$,
   batchQuery,
+  batchQuerySingleBatch$,
 } from './batchQuery';
 
 const sendSecretClientContractQuery$ = vi.hoisted(() => vi.fn());
@@ -38,6 +41,10 @@ afterAll(() => {
   vi.clearAllMocks();
 });
 
+afterEach(() => {
+  vi.clearAllMocks();
+});
+
 test('it can parse the batch query success response', () => {
   expect(parseBatchQuery(
     batchPairConfigResponse,
@@ -50,13 +57,41 @@ test('it can parse the batch query mixed succes/error response', () => {
   )).toStrictEqual(batchPricesWithErrorParsed);
 });
 
-test('it can call the batch query service', async () => {
+test('it can call the single batch query service', async () => {
+  const input = {
+    contractAddress: 'CONTRACT_ADDRESS',
+    codeHash: 'CODE_HASH',
+    queries: ['BATCH_QUERY' as unknown as BatchQueryParams],
+    client: 'SECRET_CLIENT' as unknown as SecretNetworkClient,
+  };
+  // observables function
+  sendSecretClientContractQuery$.mockReturnValueOnce(of(batchPairConfigResponse));
+
+  let output;
+  batchQuerySingleBatch$(input).subscribe({
+    next: (response) => {
+      output = response;
+    },
+  });
+
+  expect(msgBatchQuery).toHaveBeenNthCalledWith(1, input.queries);
+  expect(output).toStrictEqual(batchPairConfigParsed);
+
+  // async/await function
+  sendSecretClientContractQuery$.mockReturnValueOnce(of(batchPairConfigResponse));
+  const response = await batchQuery(input);
+  expect(msgBatchQuery).toHaveBeenNthCalledWith(2, input.queries);
+  expect(response).toStrictEqual(batchPairConfigParsed);
+});
+
+test('it can call the multi-batch query service on a single batch', async () => {
   const input = {
     contractAddress: 'CONTRACT_ADDRESS',
     codeHash: 'CODE_HASH',
     lcdEndpoint: 'LCD_ENDPOINT',
     chainId: 'CHAIN_ID',
     queries: ['BATCH_QUERY' as unknown as BatchQueryParams],
+    // no batch param passed in, so it will process as a single batch
   };
   // observables function
   sendSecretClientContractQuery$.mockReturnValueOnce(of(batchPairConfigResponse));
@@ -77,3 +112,48 @@ test('it can call the batch query service', async () => {
   expect(msgBatchQuery).toHaveBeenNthCalledWith(2, input.queries);
   expect(response).toStrictEqual(batchPairConfigParsed);
 });
+
+test('it can call the multi-batch query service for multiple batches', async () => {
+  const input = {
+    contractAddress: 'CONTRACT_ADDRESS',
+    codeHash: 'CODE_HASH',
+    lcdEndpoint: 'LCD_ENDPOINT',
+    chainId: 'CHAIN_ID',
+    queries: [
+      'BATCH_QUERY_1' as unknown as BatchQueryParams,
+      'BATCH_QUERY_2' as unknown as BatchQueryParams,
+    ],
+    batchSize: 1,
+  };
+
+  // combined array of two of the outputs
+  const combinedOutput = batchPairConfigParsed.concat(batchPairConfigParsed);
+
+  // observables function
+  // provide two mocks of the same data
+  sendSecretClientContractQuery$
+    .mockReturnValueOnce(of(batchPairConfigResponse))
+    .mockReturnValueOnce(of(batchPairConfigResponse));
+
+  let output;
+  batchQuery$(input).subscribe({
+    next: (response) => {
+      output = response;
+    },
+  });
+
+  expect(msgBatchQuery).toHaveBeenNthCalledWith(1, [input.queries[0]]);
+  expect(msgBatchQuery).toHaveBeenNthCalledWith(2, [input.queries[1]]);
+
+  expect(output).toStrictEqual(combinedOutput);
+
+  // async/await function
+  sendSecretClientContractQuery$
+    .mockReturnValueOnce(of(batchPairConfigResponse))
+    .mockReturnValueOnce(of(batchPairConfigResponse));
+
+  const response = await batchQuery(input);
+  expect(msgBatchQuery).toHaveBeenNthCalledWith(3, [input.queries[0]]);
+  expect(msgBatchQuery).toHaveBeenNthCalledWith(4, [input.queries[1]]);
+  expect(response).toStrictEqual(combinedOutput);
+});
diff --git a/src/contracts/services/batchQuery.ts b/src/contracts/services/batchQuery.ts
@@ -3,6 +3,10 @@ import {
   first,
   map,
   lastValueFrom,
+  forkJoin,
+  concatAll,
+  reduce,
+  catchError,
 } from 'rxjs';
 import { sendSecretClientContractQuery$ } from '~/client/services/clientServices';
 import { getActiveQueryClient$ } from '~/client';
@@ -15,6 +19,7 @@ import {
 } from '~/types/contracts/batchQuery/model';
 import { BatchQueryResponse } from '~/types/contracts/batchQuery/response';
 import { decodeB64ToJson } from '~/lib/utils';
+import { SecretNetworkClient } from 'secretjs';
 
 /**
  * a parses the batch query response into a usable data model
@@ -42,59 +47,125 @@ function parseBatchQuery(response: BatchQueryResponse): BatchQueryParsedResponse
   });
 }
 
+// Function to divide an array of queries into batches of arrays
+function divideSingleBatchIntoArrayOfMultipleBatches(array: BatchQueryParams[], batchSize: number) {
+  const batches = [];
+  for (let i = 0; i < array.length; i += batchSize) {
+    batches.push(array.slice(i, i + batchSize));
+  }
+  return batches;
+}
+
 /**
  * batch query of multiple contracts/message at a time
  */
-const batchQuery$ = ({
+const batchQuerySingleBatch$ = ({
   contractAddress,
   codeHash,
-  lcdEndpoint,
-  chainId,
   queries,
+  client,
 }:{
   contractAddress: string,
   codeHash?: string,
   lcdEndpoint?: string,
   chainId?: string,
-  queries: BatchQueryParams[]
-}) => getActiveQueryClient$(lcdEndpoint, chainId).pipe(
-  switchMap(({ client }) => sendSecretClientContractQuery$({
-    queryMsg: msgBatchQuery(queries),
-    client,
-    contractAddress,
-    codeHash,
-  })),
+  queries: BatchQueryParams[],
+  client: SecretNetworkClient
+}) => sendSecretClientContractQuery$({
+  queryMsg: msgBatchQuery(queries),
+  client,
+  contractAddress,
+  codeHash,
+}).pipe(
   map((response) => parseBatchQuery(response as BatchQueryResponse)),
   first(),
+  catchError((err) => {
+    if (err.message.includes('{wasm contract}')) {
+      throw new Error('{wasm contract} error that typically occurs when batch size is too large and node gas query limits are exceeded. Consider reducing the batch size.');
+    } else {
+      throw new Error(err);
+    }
+  }),
 );
 
 /**
  * batch query of multiple contracts/message at a time
+ * @param batchSize defaults to processing all queries in a single batch
+ * when the batchSize is not passed in.
+ */
+const batchQuery$ = ({
+  contractAddress,
+  codeHash,
+  lcdEndpoint,
+  chainId,
+  queries,
+  batchSize,
+}:{
+  contractAddress: string,
+  codeHash?: string,
+  lcdEndpoint?: string,
+  chainId?: string,
+  queries: BatchQueryParams[],
+  batchSize?: number,
+}) => {
+  // if batch size is passed in, convert single batch into multiple batches,
+  // otherwise process all data in a single batch
+  const batches = batchSize
+    ? divideSingleBatchIntoArrayOfMultipleBatches(queries, batchSize)
+    : [queries]; // array of arrays required for the forkJoin
+
+  return getActiveQueryClient$(lcdEndpoint, chainId).pipe(
+    switchMap(({ client }) => forkJoin(
+      batches.map((batch) => batchQuerySingleBatch$({
+        contractAddress,
+        codeHash,
+        queries: batch,
+        client,
+      })),
+    ).pipe(
+      concatAll(),
+      reduce((
+        acc: BatchQueryParsedResponse,
+        curr: BatchQueryParsedResponse,
+      ) => acc.concat(curr), []), // Flatten nested arrays into a single array
+      first(),
+    )),
+  );
+};
+
+/**
+ * batch query of multiple contracts/message at a time
+ * @param batchSize defaults to processing all queries in a single batch
+ * when the batchSize is not passed in.
  */
 async function batchQuery({
   contractAddress,
   codeHash,
   lcdEndpoint,
   chainId,
   queries,
+  batchSize,
 }:{
   contractAddress: string,
   codeHash?: string,
   lcdEndpoint?: string,
   chainId?: string,
-  queries: BatchQueryParams[]
+  queries: BatchQueryParams[],
+  batchSize?: number,
 }) {
   return lastValueFrom(batchQuery$({
     contractAddress,
     codeHash,
     lcdEndpoint,
     chainId,
     queries,
+    batchSize,
   }));
 }
 
 export {
   parseBatchQuery,
   batchQuery$,
   batchQuery,
+  batchQuerySingleBatch$,
 };
diff --git a/src/contracts/services/config.ts b/src/contracts/services/config.ts
@@ -0,0 +1,8 @@
+const SERVICE_BATCH_SIZE = {
+  PAIR_INFO: 60,
+  PAIR_CONFIG: 60,
+};
+
+export {
+  SERVICE_BATCH_SIZE,
+};
diff --git a/src/contracts/services/swap.test.ts b/src/contracts/services/swap.test.ts
@@ -285,6 +285,7 @@ test('it can call the batch pairs info query service', async () => {
       },
       queryMsg: 'PAIR_INFO_MSG',
     }],
+    batchSize: 60,
   });
 
   expect(output).toStrictEqual(pairsInfoParsed);
@@ -305,6 +306,7 @@ test('it can call the batch pairs info query service', async () => {
       },
       queryMsg: 'PAIR_INFO_MSG',
     }],
+    batchSize: 60,
   });
   expect(response).toStrictEqual(pairsInfoParsed);
 });
@@ -340,7 +342,7 @@ test('it can call the batch staking info query service', async () => {
         address: input.stakingContracts[0].address,
         codeHash: input.stakingContracts[0].codeHash,
       },
-      queryMsg: 'PAIR_INFO_MSG',
+      queryMsg: 'STAKING_CONFIG_MSG',
     }],
   });
 
@@ -400,6 +402,7 @@ test('it can call the batch pair config query service', async () => {
       },
       queryMsg: 'PAIR_CONFIG_MSG',
     }],
+    batchSize: 60,
   });
 
   expect(output).toStrictEqual(batchPairsConfigParsed);
@@ -420,6 +423,7 @@ test('it can call the batch pair config query service', async () => {
       },
       queryMsg: 'PAIR_CONFIG_MSG',
     }],
+    batchSize: 60,
   });
   expect(response).toStrictEqual(batchPairsConfigParsed);
 });