From 5e9c21440b4cd92613a9cf0dd55b5388e502228d Mon Sep 17 00:00:00 2001
From: Anshudhar Kumar Singh <anshudhar2001@gmail.com>
Date: Sat, 12 Oct 2024 01:14:57 +0530
Subject: [PATCH] fix: query-raw pagination

---
 packages/andrjs/src/AndromedaClient.ts     | 26 ++++++++-------
 packages/andrjs/src/clients/ChainClient.ts |  1 +
 packages/andrjs/src/clients/CosmClient.ts  |  1 +
 packages/cli/src/handlers/wasm.ts          | 38 +++++++++++++++++-----
 4 files changed, 46 insertions(+), 20 deletions(-)

diff --git a/packages/andrjs/src/AndromedaClient.ts b/packages/andrjs/src/AndromedaClient.ts
index 3d41a18..8b9eed5 100644
--- a/packages/andrjs/src/AndromedaClient.ts
+++ b/packages/andrjs/src/AndromedaClient.ts
@@ -6,6 +6,7 @@ import {
   SigningStargateClientOptions,
   StdFee,
   calculateFee,
+  createProtobufRpcClient,
 } from "@cosmjs/stargate";
 import { ADOAPI } from "./api";
 
@@ -18,6 +19,7 @@ import type { Fee, Msg } from "./types";
 import ADOSchemaAPI from "api/ADOSchemaAPI";
 import { RpcClient } from "@cosmjs/tendermint-rpc";
 import { PageRequest } from "cosmjs-types/cosmos/base/query/v1beta1/pagination";
+import { QueryClientImpl as WasmQueryClientImpl } from 'cosmjs-types/cosmwasm/wasm/v1/query'
 
 /**
  * A helper class for interacting with the Andromeda ecosystem
@@ -221,18 +223,20 @@ export default class AndromedaClient {
  */
   async queryContractRawAll(address: string, pagination: Partial<PageRequest>) {
     this.preMessage();
-    const result = await this.chainClient!.rawQueryClient!.wasm.getAllContractState(
-      address,
-      await this.encodePagination(pagination)
-    );
-    return result.models.map(model => ({
-      key: Buffer.from(model.key).toString('utf8'),
-      value: Buffer.from(model.value).toString('utf8')
-    }))
-  }
+    const rpcClient = createProtobufRpcClient(this.chainClient!.rawQueryClient!);
+    const wasmQueryClient = new WasmQueryClientImpl(rpcClient);
 
-  async encodePagination(pagination: Partial<PageRequest>) {
-    return PageRequest.encode(PageRequest.fromPartial(pagination)).finish();
+    const result = await wasmQueryClient.AllContractState({
+      address,
+      pagination: PageRequest.fromPartial(pagination)
+    });
+    return {
+      states: result.models.map(model => ({
+        key: Buffer.from(model.key).toString('utf8'),
+        value: Buffer.from(model.value).toString('utf8')
+      })),
+      pagination: result.pagination
+    }
   }
 
   /**
diff --git a/packages/andrjs/src/clients/ChainClient.ts b/packages/andrjs/src/clients/ChainClient.ts
index 34b668a..361d373 100644
--- a/packages/andrjs/src/clients/ChainClient.ts
+++ b/packages/andrjs/src/clients/ChainClient.ts
@@ -32,6 +32,7 @@ export default interface ChainClient {
   // The client used to query the chain
   queryClient?: CosmWasmClient;
 
+
   rawQueryClient?: QueryClient & TxExtension & WasmExtension
 
   // The current signer address
diff --git a/packages/andrjs/src/clients/CosmClient.ts b/packages/andrjs/src/clients/CosmClient.ts
index 9c68cf6..0e7a833 100644
--- a/packages/andrjs/src/clients/CosmClient.ts
+++ b/packages/andrjs/src/clients/CosmClient.ts
@@ -72,6 +72,7 @@ export default class CosmClient extends BaseChainClient implements ChainClient {
 
   /** Client for querying blockchain data */
   public queryClient?: ChainClient["queryClient"];
+
   /** Client for querying data with raw implementation */
   public rawQueryClient?: ChainClient["rawQueryClient"];
   /** Client for interacting with the Comet BFT consensus engine */
diff --git a/packages/cli/src/handlers/wasm.ts b/packages/cli/src/handlers/wasm.ts
index 9f55769..ca34b58 100644
--- a/packages/cli/src/handlers/wasm.ts
+++ b/packages/cli/src/handlers/wasm.ts
@@ -68,13 +68,21 @@ export const commands: Commands = {
     ],
     flags: {
       limit: {
-        description: "Paginate response, 0 mean no limit",
+        description: "Limit number of results",
         usage: "--limit 10",
       },
       offset: {
-        description: "Paginate response",
+        description: "Offset number of results (This might not be supported by all chains)",
         usage: "--offset 0",
       },
+      ['next-key']: {
+        description: "Next key to paginate from",
+        usage: "--next-key <key>",
+      },
+      ['next-key-bytes']: {
+        description: "Next key to paginate from (bytes)",
+        usage: "--next-key-bytes <key>",
+      }
     }
   },
   execute: {
@@ -258,22 +266,34 @@ async function queryRawHandler(input: string[], flags: Flags) {
   const [contractAddr] = input;
   const limit = BigInt(flags.limit ?? '10');
   const offset = BigInt(flags.offset ?? '0');
+  const nextKey = flags['next-key'] as string | undefined;
+  const nextKeyBytes = flags['next-key-bytes'] as string | undefined;
+
+  if (nextKeyBytes && nextKey) {
+    throw new Error("Cannot provide both next-key and next-key-bytes");
+  }
+
+  const key = nextKeyBytes ? Uint8Array.from(Buffer.from(nextKeyBytes, 'hex')) : nextKey ? Uint8Array.from(Buffer.from(nextKey, 'utf8')) : undefined;
 
   const states = await displaySpinnerAsync(
     "Querying contrac key...",
-    async () => await State.client.queryContractRawAll(contractAddr, { limit, offset })
+    async () => await State.client.queryContractRawAll(contractAddr, { limit, offset, key })
   );
   console.log();
-  const infoTable = new Table(logTableConfig);
 
 
-  infoTable.push([pc.bold("Key"), ":", pc.bold("Value")]);
-  for (const state of states) {
-    infoTable.push([pc.blue(state.key), ":", state.value]);
+  for (const state of states.states) {
+    console.log(pc.blue(state.key));
+    console.log(state.value);
+    console.log();
   }
-  console.log(infoTable.toString());
   console.log();
-  console.log(pc.gray(`Limit - ${limit}, Offset - ${offset}`))
+  console.log(pc.gray(`Limit - ${limit}, Offset - ${offset}, key - ${key ? Buffer.from(key).toString('utf8') : 'None'}`))
+  if (states.pagination) {
+    console.log(pc.gray(`Next Key - ${Buffer.from(states.pagination.nextKey).toString('utf8')}`));
+    console.log(pc.gray(`Next Key Bytes - ${Buffer.from(states.pagination.nextKey).toString('hex')}`));
+    console.log(pc.gray(`Total - ${states.pagination.total.toString()}`));
+  }
   console.log()