feat: add insertMany to repository API (Batch accounts 2/5)

packages/snap/CHANGELOG.md

@@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Added
+
+- Add `insertMany` to repository API ([#609](https://github.com/MetaMask/snap-bitcoin-wallet/pull/609))
+
 ### Changed
 
 - Show a confirmation dialog before signing a PSBT from KeyringHandler and sending a transfer ([#591](https://github.com/MetaMask/snap-bitcoin-wallet/pull/591))

packages/snap/src/entities/account.ts

@@ -123,6 +123,13 @@ export type BitcoinAccount = {
    */
   takeStaged(): ChangeSet | undefined;
 
+  /**
+   * Check whether a change set exists without making it unavailable for extraction.
+   *
+   * @returns true if there is a change set
+   */
+  hasStaged(): boolean;
+
   /**
    * Returns a Transaction Builder.
    *
@@ -262,6 +269,16 @@ export type BitcoinAccountRepository = {
    */
   getByDerivationPath(derivationPath: string[]): Promise<BitcoinAccount | null>;
 
+  /**
+   * Get accounts by derivation path.
+   *
+   * @param derivationPaths - derivation paths.
+   * @returns the accounts or null if they do not exist, in input order
+   */
+  getByDerivationPaths(
+    derivationPaths: string[][],
+  ): Promise<(BitcoinAccount | null)[]>;
+
   /**
    * Create a new account, without persisting it.
    *
@@ -283,6 +300,13 @@ export type BitcoinAccountRepository = {
    */
   insert(account: BitcoinAccount): Promise<BitcoinAccount>;
 
+  /**
+   * Insert accounts.
+   *
+   * @param accounts - Bitcoin accounts.
+   */
+  insertMany(accounts: BitcoinAccount[]): Promise<BitcoinAccount[]>;
+
   /**
    * Update an account.
    *

packages/snap/src/entities/snap.ts

@@ -13,7 +13,7 @@ import type { Inscription } from './meta-protocols';
 
 export type SnapState = {
   // accountId -> account state. This is the main state of the snap.
-  accounts: Record<string, AccountState>;
+  accounts: Record<string, AccountState | null>;
   // derivationPath -> accountId. Only needed for fast lookup.
   derivationPaths: Record<string, string>;
 };

packages/snap/src/infra/BdkAccountAdapter.ts

@@ -41,6 +41,8 @@ export class BdkAccountAdapter implements BitcoinAccount {
 
   readonly #wallet: Wallet;
 
+  #staged?: ChangeSet;
+
   readonly #capabilities: AccountCapability[];
 
   constructor(id: string, derivationPath: string[], wallet: Wallet) {
@@ -157,7 +159,13 @@ export class BdkAccountAdapter implements BitcoinAccount {
   }
 
   takeStaged(): ChangeSet | undefined {
-    return this.#wallet.take_staged();
+    const staged = this.#collectStaged();
+    this.#staged = undefined;
+    return staged;
+  }
+
+  hasStaged(): boolean {
+    return Boolean(this.#collectStaged());
   }
 
   buildTx(): TransactionBuilder {
@@ -243,4 +251,18 @@ export class BdkAccountAdapter implements BitcoinAccount {
       new UnconfirmedTx(tx, BigInt(lastSeen)),
     ]);
   }
+
+  #collectStaged(): ChangeSet | undefined {
+    const staged = this.#wallet.take_staged();
+    if (!staged) {
+      return this.#staged;
+    }
+
+    if (this.#staged) {
+      staged.merge(this.#staged);
+    }
+
+    this.#staged = staged;
+    return this.#staged;
+  }
 }

packages/snap/src/store/BdkAccountRepository.test.ts

@@ -77,6 +77,7 @@ describe('BdkAccountRepository', () => {
     (mockAccount.takeStaged as jest.Mock) = jest
       .fn()
       .mockReturnValue(mockChangeSet);
+    (mockAccount.hasStaged as jest.Mock) = jest.fn().mockReturnValue(true);
     (mockChangeSet.to_json as jest.Mock) = jest
       .fn()
       .mockReturnValue(mockWalletData);
@@ -185,6 +186,94 @@ describe('BdkAccountRepository', () => {
     });
   });
 
+  describe('getByDerivationPaths', () => {
+    const derivationPath1 = ['m', "84'", "0'", "1'"];
+    const derivationPath2 = ['m', "84'", "0'", "2'"];
+    const accountState1 = {
+      ...mockAccountState,
+      derivationPath: derivationPath1,
+    };
+    const accountState2 = {
+      ...mockAccountState,
+      derivationPath: derivationPath2,
+    };
+    const mockAccount1 = mock<BitcoinAccount>({
+      ...mockAccount,
+      id: 'some-id-1',
+      derivationPath: derivationPath1,
+    });
+    const mockAccount2 = mock<BitcoinAccount>({
+      ...mockAccount,
+      id: 'some-id-2',
+      derivationPath: derivationPath2,
+    });
+
+    it('returns accounts in derivation path order with one state read per namespace', async () => {
+      mockSnapClient.getState
+        .mockResolvedValueOnce({
+          "m/84'/0'/1'": 'some-id-1',
+          "m/84'/0'/2'": 'some-id-2',
+        })
+        .mockResolvedValueOnce({
+          'some-id-1': accountState1,
+          'some-id-2': accountState2,
+        });
+      (BdkAccountAdapter.load as jest.Mock)
+        .mockReturnValueOnce(mockAccount2)
+        .mockReturnValueOnce(mockAccount1);
+
+      const result = await repo.getByDerivationPaths([
+        derivationPath2,
+        derivationPath1,
+      ]);
+
+      expect(mockSnapClient.getState).toHaveBeenCalledWith('derivationPaths');
+      expect(mockSnapClient.getState).toHaveBeenCalledWith('accounts');
+      expect(mockSnapClient.getState).toHaveBeenCalledTimes(2);
+      expect(result).toStrictEqual([mockAccount2, mockAccount1]);
+      expect(mockSnapClient.setState).not.toHaveBeenCalled();
+    });
+
+    it('repairs missing derivation path indexes from account state', async () => {
+      mockSnapClient.getState
+        .mockResolvedValueOnce({
+          "m/84'/0'/1'": 'some-id-1',
+        })
+        .mockResolvedValueOnce({
+          'some-id-1': accountState1,
+          'some-id-2': accountState2,
+        });
+      (BdkAccountAdapter.load as jest.Mock)
+        .mockReturnValueOnce(mockAccount1)
+        .mockReturnValueOnce(mockAccount2);
+
+      const result = await repo.getByDerivationPaths([
+        derivationPath1,
+        derivationPath2,
+      ]);
+
+      expect(result).toStrictEqual([mockAccount1, mockAccount2]);
+      expect(mockSnapClient.setState).toHaveBeenCalledWith('derivationPaths', {
+        "m/84'/0'/1'": 'some-id-1',
+        "m/84'/0'/2'": 'some-id-2',
+      });
+    });
+
+    it('repairs a missing derivation path index for a single lookup', async () => {
+      mockSnapClient.getState.mockResolvedValueOnce({}).mockResolvedValueOnce({
+        'some-id-1': accountState1,
+      });
+      (BdkAccountAdapter.load as jest.Mock).mockReturnValueOnce(mockAccount1);
+
+      const result = await repo.getByDerivationPaths([derivationPath1]);
+
+      expect(result).toStrictEqual([mockAccount1]);
+      expect(mockSnapClient.setState).toHaveBeenCalledWith('derivationPaths', {
+        "m/84'/0'/1'": 'some-id-1',
+      });
+    });
+  });
+
   describe('getWithSigner', () => {
     it('returns null if account not found', async () => {
       mockSnapClient.getState.mockResolvedValue(null);
@@ -258,6 +347,105 @@ describe('BdkAccountRepository', () => {
     });
   });
 
+  describe('insertMany', () => {
+    it('returns an empty array when there are no accounts to insert', async () => {
+      const result = await repo.insertMany([]);
+
+      expect(result).toStrictEqual([]);
+      expect(mockSnapClient.getState).not.toHaveBeenCalled();
+      expect(mockSnapClient.setState).not.toHaveBeenCalled();
+    });
+
+    it('throws an error without consuming staged data if any account has no wallet data', async () => {
+      const accountWithWalletData = mock<BitcoinAccount>({
+        id: 'some-id-1',
+        derivationPath: ['m', "84'", "0'", "1'"],
+      });
+      (accountWithWalletData.hasStaged as jest.Mock) = jest
+        .fn()
+        .mockReturnValue(true);
+      (accountWithWalletData.takeStaged as jest.Mock) = jest
+        .fn()
+        .mockReturnValue(mockChangeSet);
+      const missingWalletAccount = mock<BitcoinAccount>({
+        id: 'missing-wallet',
+        derivationPath: ['m', "84'", "0'", "2'"],
+      });
+      (missingWalletAccount.hasStaged as jest.Mock) = jest
+        .fn()
+        .mockReturnValue(false);
+      (missingWalletAccount.takeStaged as jest.Mock) = jest.fn();
+
+      await expect(
+        repo.insertMany([accountWithWalletData, missingWalletAccount]),
+      ).rejects.toThrow(
+        'Missing changeset data for account "missing-wallet" for insertion.',
+      );
+
+      expect(accountWithWalletData.hasStaged).toHaveBeenCalled();
+      expect(missingWalletAccount.hasStaged).toHaveBeenCalled();
+      expect(accountWithWalletData.takeStaged).not.toHaveBeenCalled();
+      expect(missingWalletAccount.takeStaged).not.toHaveBeenCalled();
+      expect(mockSnapClient.getState).not.toHaveBeenCalled();
+      expect(mockSnapClient.setState).not.toHaveBeenCalled();
+    });
+
+    it('inserts multiple accounts with one accounts write and one derivation path write', async () => {
+      const existingAccountState: AccountState = {
+        wallet: mockWalletData,
+        inscriptions: [],
+        derivationPath: mockDerivationPath,
+      };
+      const account1 = mock<BitcoinAccount>();
+      account1.id = 'some-id-1';
+      account1.derivationPath = ['m', "84'", "0'", "1'"];
+      const account2 = mock<BitcoinAccount>();
+      account2.id = 'some-id-2';
+      account2.derivationPath = ['m', "84'", "0'", "2'"];
+      (account1.takeStaged as jest.Mock) = jest
+        .fn()
+        .mockReturnValue(mockChangeSet);
+      (account1.hasStaged as jest.Mock) = jest.fn().mockReturnValue(true);
+      (account2.takeStaged as jest.Mock) = jest
+        .fn()
+        .mockReturnValue(mockChangeSet);
+      (account2.hasStaged as jest.Mock) = jest.fn().mockReturnValue(true);
+      mockSnapClient.getState
+        .mockResolvedValueOnce({
+          'existing-id': existingAccountState,
+        })
+        .mockResolvedValueOnce({
+          "m/84'/0'/0'": 'existing-id',
+        });
+
+      const result = await repo.insertMany([account1, account2]);
+
+      expect(result).toStrictEqual([account1, account2]);
+      expect(mockSnapClient.setState).toHaveBeenNthCalledWith(1, 'accounts', {
+        'existing-id': existingAccountState,
+        'some-id-1': {
+          wallet: mockWalletData,
+          inscriptions: [],
+          derivationPath: ['m', "84'", "0'", "1'"],
+        },
+        'some-id-2': {
+          wallet: mockWalletData,
+          inscriptions: [],
+          derivationPath: ['m', "84'", "0'", "2'"],
+        },
+      });
+      expect(mockSnapClient.setState).toHaveBeenNthCalledWith(
+        2,
+        'derivationPaths',
+        {
+          "m/84'/0'/0'": 'existing-id',
+          "m/84'/0'/1'": 'some-id-1',
+          "m/84'/0'/2'": 'some-id-2',
+        },
+      );
+    });
+  });
+
   describe('update', () => {
     it('does nothing if no wallet data', async () => {
       await repo.update({