Implement functions to return associated context when fetching accounts. (#2237)

ts/packages/anchor/src/program/namespace/account.ts

@@ -8,6 +8,8 @@ import {
   Commitment,
   GetProgramAccountsFilter,
   AccountInfo,
+  RpcResponseAndContext,
+  Context,
 } from "@solana/web3.js";
 import Provider, { getProvider } from "../../provider.js";
 import { Idl, IdlAccountDef } from "../../idl.js";
@@ -133,14 +135,30 @@ export class AccountClient<
     address: Address,
     commitment?: Commitment
   ): Promise<T | null> {
-    const accountInfo = await this.getAccountInfo(address, commitment);
-    if (accountInfo === null) {
-      return null;
-    }
-    return this._coder.accounts.decode<T>(
-      this._idlAccount.name,
-      accountInfo.data
+    const { data } = await this.fetchNullableAndContext(address, commitment);
+    return data;
+  }
+
+  /**
+   * Returns a deserialized account along with the associated rpc response context, returning null if it doesn't exist.
+   *
+   * @param address The address of the account to fetch.
+   */
+  async fetchNullableAndContext(
+    address: Address,
+    commitment?: Commitment
+  ): Promise<{ data: T | null; context: Context }> {
+    const accountInfo = await this.getAccountInfoAndContext(
+      address,
+      commitment
     );
+    const { value, context } = accountInfo;
+    return {
+      data: value
+        ? this._coder.accounts.decode<T>(this._idlAccount.name, value.data)
+        : null,
+      context,
+    };
   }
 
   /**
@@ -149,13 +167,32 @@ export class AccountClient<
    * @param address The address of the account to fetch.
    */
   async fetch(address: Address, commitment?: Commitment): Promise<T> {
-    const data = await this.fetchNullable(address, commitment);
+    const { data } = await this.fetchNullableAndContext(address, commitment);
     if (data === null) {
       throw new Error(`Account does not exist ${address.toString()}`);
     }
     return data;
   }
 
+  /**
+   * Returns a deserialized account along with the associated rpc response context.
+   *
+   * @param address The address of the account to fetch.
+   */
+  async fetchAndContext(
+    address: Address,
+    commitment?: Commitment
+  ): Promise<{ data: T | null; context: Context }> {
+    const { data, context } = await this.fetchNullableAndContext(
+      address,
+      commitment
+    );
+    if (data === null) {
+      throw new Error(`Account does not exist ${address.toString()}`);
+    }
+    return { data, context };
+  }
+
   /**
    * Returns multiple deserialized accounts.
    * Accounts not found or with wrong discriminator are returned as null.
@@ -166,21 +203,36 @@ export class AccountClient<
     addresses: Address[],
     commitment?: Commitment
   ): Promise<(Object | null)[]> {
-    const accounts = await rpcUtil.getMultipleAccounts(
+    const accounts = await this.fetchMultipleAndContext(addresses, commitment);
+    return accounts.map((account) => (account ? account.data : null));
+  }
+
+  /**
+   * Returns multiple deserialized accounts.
+   * Accounts not found or with wrong discriminator are returned as null.
+   *
+   * @param addresses The addresses of the accounts to fetch.
+   */
+  async fetchMultipleAndContext(
+    addresses: Address[],
+    commitment?: Commitment
+  ): Promise<({ data: Object; context: Context } | null)[]> {
+    const accounts = await rpcUtil.getMultipleAccountsAndContext(
       this._provider.connection,
       addresses.map((address) => translateAddress(address)),
       commitment
     );
 
     // Decode accounts where discriminator is correct, null otherwise
-    return accounts.map((account) => {
-      if (account == null) {
+    return accounts.map((result) => {
+      if (result == null) {
         return null;
       }
-      return this._coder.accounts.decode(
-        this._idlAccount.name,
-        account?.account.data
-      );
+      const { account, context } = result;
+      return {
+        data: this._coder.accounts.decode(this._idlAccount.name, account.data),
+        context,
+      };
     });
   }
 
@@ -346,6 +398,16 @@ export class AccountClient<
       commitment
     );
   }
+
+  async getAccountInfoAndContext(
+    address: Address,
+    commitment?: Commitment
+  ): Promise<RpcResponseAndContext<AccountInfo<Buffer> | null>> {
+    return await this._provider.connection.getAccountInfoAndContext(
+      translateAddress(address),
+      commitment
+    );
+  }
 }
 
 /**

ts/packages/anchor/src/utils/rpc.ts

@@ -13,6 +13,7 @@ import {
   RpcResponseAndContext,
   SimulatedTransactionResponse,
   SendTransactionError,
+  Context,
 } from "@solana/web3.js";
 import { chunks } from "../utils/common.js";
 import { Address, translateAddress } from "../program/common.js";
@@ -78,43 +79,79 @@ export async function getMultipleAccounts(
   commitment?: Commitment
 ): Promise<
   Array<null | { publicKey: PublicKey; account: AccountInfo<Buffer> }>
+> {
+  const results = await getMultipleAccountsAndContext(
+    connection,
+    publicKeys,
+    commitment
+  );
+  return results.map((result) => {
+    return result
+      ? { publicKey: result.publicKey, account: result.account }
+      : null;
+  });
+}
+
+export async function getMultipleAccountsAndContext(
+  connection: Connection,
+  publicKeys: PublicKey[],
+  commitment?: Commitment
+): Promise<
+  Array<null | {
+    context: Context;
+    publicKey: PublicKey;
+    account: AccountInfo<Buffer>;
+  }>
 > {
   if (publicKeys.length <= GET_MULTIPLE_ACCOUNTS_LIMIT) {
-    return await getMultipleAccountsCore(connection, publicKeys, commitment);
+    return await getMultipleAccountsAndContextCore(
+      connection,
+      publicKeys,
+      commitment
+    );
   } else {
     const batches = chunks(publicKeys, GET_MULTIPLE_ACCOUNTS_LIMIT);
     const results = await Promise.all<
-      Array<null | { publicKey: PublicKey; account: AccountInfo<Buffer> }>
+      Array<null | {
+        publicKey: PublicKey;
+        account: AccountInfo<Buffer>;
+        context: Context;
+      }>
     >(
       batches.map((batch) =>
-        getMultipleAccountsCore(connection, batch, commitment)
+        getMultipleAccountsAndContextCore(connection, batch, commitment)
       )
     );
     return results.flat();
   }
 }
 
-async function getMultipleAccountsCore(
+async function getMultipleAccountsAndContextCore(
   connection: Connection,
   publicKeys: PublicKey[],
   commitmentOverride?: Commitment
 ): Promise<
-  Array<null | { publicKey: PublicKey; account: AccountInfo<Buffer> }>
+  Array<null | {
+    publicKey: PublicKey;
+    account: AccountInfo<Buffer>;
+    context: Context;
+  }>
 > {
   const commitment = commitmentOverride ?? connection.commitment;
-  const accounts = await connection.getMultipleAccountsInfo(
-    publicKeys,
-    commitment
-  );
-  return accounts.map((account, idx) => {
+  const { value: accountInfos, context } =
+    await connection.getMultipleAccountsInfoAndContext(publicKeys, commitment);
+  const accounts = accountInfos.map((account, idx) => {
     if (account === null) {
       return null;
     }
     return {
       publicKey: publicKeys[idx],
       account,
+      context,
     };
   });
+
+  return accounts;
 }
 
 // copy from @solana/web3.js that has a commitment param