From 0c2f82144d6411c1875b17c1f261cd989b48aab6 Mon Sep 17 00:00:00 2001
From: woody <125113430+woodenfurniture@users.noreply.github.com>
Date: Tue, 7 May 2024 11:42:03 +1000
Subject: [PATCH] fix: state corruption of import accounts drawer (#6817)

* chore: migrate account import queries to useInfiniteQuery

* fix: ensure import account component state is not shared across chains

* fix: disable buttons when auto-fetching
---
 .../components/ImportAccounts.tsx             | 133 +++++++++---------
 .../ManageAccountsDrawer/helpers.ts           |  28 +++-
 .../queries/accountManagement.ts              |  84 +----------
 3 files changed, 95 insertions(+), 150 deletions(-)

diff --git a/src/components/ManageAccountsDrawer/components/ImportAccounts.tsx b/src/components/ManageAccountsDrawer/components/ImportAccounts.tsx
index 2ca36368503..d08754641c3 100644
--- a/src/components/ManageAccountsDrawer/components/ImportAccounts.tsx
+++ b/src/components/ManageAccountsDrawer/components/ImportAccounts.tsx
@@ -12,11 +12,10 @@ import {
 } from '@chakra-ui/react'
 import type { ChainId } from '@shapeshiftoss/caip'
 import { type AccountId, fromAccountId } from '@shapeshiftoss/caip'
-import type { AccountMetadata, Asset } from '@shapeshiftoss/types'
-import { useIsFetching, useQuery, useQueryClient } from '@tanstack/react-query'
+import type { Asset } from '@shapeshiftoss/types'
+import { useInfiniteQuery, useQuery } from '@tanstack/react-query'
 import { useCallback, useEffect, useMemo, useState } from 'react'
 import { useTranslate } from 'react-polyglot'
-import { reactQueries } from 'react-queries'
 import { accountManagement } from 'react-queries/queries/accountManagement'
 import { Amount } from 'components/Amount/Amount'
 import { MiddleEllipsis } from 'components/MiddleEllipsis/MiddleEllipsis'
@@ -34,6 +33,7 @@ import {
 } from 'state/slices/selectors'
 import { store, useAppDispatch, useAppSelector } from 'state/store'
 
+import { getAccountIdsWithActivityAndMetadata } from '../helpers'
 import { DrawerContentWrapper } from './DrawerContent'
 
 // The number of additional empty accounts to include in the initial fetch
@@ -160,55 +160,58 @@ export const ImportAccounts = ({ chainId, onClose }: ImportAccountsProps) => {
     selectHighestAccountNumberForChainId(state, highestAccountNumberForChainIdFilter),
   )
   const chainNamespaceDisplayName = asset?.networkName ?? ''
-  const [accounts, setAccounts] = useState<
-    { accountId: AccountId; accountMetadata: AccountMetadata; hasActivity: boolean }[][]
-  >([])
-  const queryClient = useQueryClient()
-  const isLoading =
-    useIsFetching({
-      predicate: query => {
-        return (
-          query.queryKey[0] === 'accountManagement' &&
-          ['accountIdWithActivityAndMetadata', 'firstAccountIdsWithActivityAndMetadata'].some(
-            str => str === query.queryKey[1],
-          )
-        )
-      },
-    }) > 0
+  const [autoFetching, setAutoFetching] = useState(true)
   const [toggledActiveAccountIds, setToggledActiveAccountIds] = useState<Set<AccountId>>(new Set())
 
+  // reset component state when chainId changes
+  useEffect(() => {
+    setAutoFetching(true)
+    setToggledActiveAccountIds(new Set())
+  }, [chainId])
+
   // initial fetch to detect the number of accounts based on the "first empty account" heuristic
-  const { data: allAccountIdsWithActivity } = useQuery(
-    accountManagement.firstAccountIdsWithActivityAndMetadata(
-      chainId,
-      wallet,
-      walletDeviceId,
-      // Account numbers are 0-indexed, so we need to add 1 to the highest account number.
-      // Add additional empty accounts to show more accounts without having to load more.
-      highestAccountNumber + 1 + NUM_ADDITIONAL_EMPTY_ACCOUNTS,
-    ),
-  )
+  const {
+    data: accounts,
+    fetchNextPage,
+    isLoading,
+  } = useInfiniteQuery({
+    queryKey: ['accountIdWithActivityAndMetadata', chainId, walletDeviceId, wallet !== null],
+    queryFn: async ({ pageParam: accountNumber }) => {
+      return {
+        accountNumber,
+        accountIdWithActivityAndMetadata: await getAccountIdsWithActivityAndMetadata(
+          accountNumber,
+          chainId,
+          wallet,
+        ),
+      }
+    },
+    initialPageParam: 0,
+    getNextPageParam: lastPage => {
+      return lastPage.accountNumber + 1
+    },
+  })
 
+  // Handle initial automatic loading
   useEffect(() => {
-    setAccounts(allAccountIdsWithActivity ?? [])
-  }, [allAccountIdsWithActivity])
-
-  const handleLoadMore = useCallback(async () => {
-    if (!wallet) return
-    const accountNumber = accounts.length
-    const accountResults = await queryClient.fetchQuery(
-      reactQueries.accountManagement.accountIdWithActivityAndMetadata(
-        accountNumber,
-        chainId,
-        wallet,
-        walletDeviceId,
-      ),
-    )
-    if (!accountResults.length) return
-    setAccounts(previousAccounts => {
-      return [...previousAccounts, accountResults]
-    })
-  }, [accounts, chainId, queryClient, wallet, walletDeviceId])
+    if (!autoFetching || !accounts) return
+
+    // Account numbers are 0-indexed, so we need to add 1 to the highest account number.
+    // Add additional empty accounts to show more accounts without having to load more.
+    const numAccountsToLoad = highestAccountNumber + 1 + NUM_ADDITIONAL_EMPTY_ACCOUNTS
+
+    if (accounts.pages.length < numAccountsToLoad) {
+      fetchNextPage()
+    } else {
+      // Stop auto-fetching and switch to manual mode
+      setAutoFetching(false)
+    }
+  }, [accounts, highestAccountNumber, fetchNextPage, autoFetching])
+
+  const handleLoadMore = useCallback(() => {
+    if (isLoading || autoFetching) return
+    fetchNextPage()
+  }, [autoFetching, isLoading, fetchNextPage])
 
   const handleToggleAccountIds = useCallback((accountIds: AccountId[]) => {
     setToggledActiveAccountIds(previousState => {
@@ -234,10 +237,15 @@ export const ImportAccounts = ({ chainId, onClose }: ImportAccountsProps) => {
       await dispatch(portfolioApi.endpoints.getAccount.initiate({ accountId, upsertOnFetch: true }))
     }
 
-    const accountMetadataByAccountId = accounts.reduce((accumulator, accounts) => {
-      const obj = accounts.reduce((innerAccumulator, { accountId, accountMetadata }) => {
-        return { ...innerAccumulator, [accountId]: accountMetadata }
-      }, {})
+    if (!accounts) return
+
+    const accountMetadataByAccountId = accounts.pages.reduce((accumulator, accounts) => {
+      const obj = accounts.accountIdWithActivityAndMetadata.reduce(
+        (innerAccumulator, { accountId, accountMetadata }) => {
+          return { ...innerAccumulator, [accountId]: accountMetadata }
+        },
+        {},
+      )
       return { ...accumulator, ...obj }
     }, {})
 
@@ -256,12 +264,13 @@ export const ImportAccounts = ({ chainId, onClose }: ImportAccountsProps) => {
   }, [toggledActiveAccountIds, accounts, dispatch, onClose, walletDeviceId])
 
   const accountRows = useMemo(() => {
-    if (!asset) return null
-    return accounts.map((accountsForAccountNumber, accountNumber) => {
-      const accountIds = accountsForAccountNumber.map(({ accountId }) => accountId)
+    if (!asset || !accounts) return null
+    return accounts.pages.map(({ accountIdWithActivityAndMetadata }, accountNumber) => {
+      const accountIds = accountIdWithActivityAndMetadata.map(({ accountId }) => accountId)
+      const key = accountIds.join('-')
       return (
         <TableRow
-          key={accountNumber}
+          key={key}
           accountNumber={accountNumber}
           accountIds={accountIds}
           asset={asset}
@@ -282,19 +291,13 @@ export const ImportAccounts = ({ chainId, onClose }: ImportAccountsProps) => {
       description={translate('accountManagement.importAccounts.description')}
       footer={
         <>
-          <Button
-            colorScheme='gray'
-            mr={3}
-            onClick={onClose}
-            isDisabled={isLoading}
-            _disabled={disabledProps}
-          >
+          <Button colorScheme='gray' mr={3} onClick={onClose}>
             {translate('common.cancel')}
           </Button>
           <Button
             colorScheme='blue'
             onClick={handleDone}
-            isDisabled={isLoading}
+            isDisabled={isLoading || autoFetching}
             _disabled={disabledProps}
           >
             {translate('common.done')}
@@ -307,14 +310,14 @@ export const ImportAccounts = ({ chainId, onClose }: ImportAccountsProps) => {
             <Table variant='simple'>
               <Tbody>
                 {accountRows}
-                {isLoading && <LoadingRow />}
+                {(isLoading || autoFetching) && <LoadingRow />}
               </Tbody>
             </Table>
           </TableContainer>
           <Button
             colorScheme='gray'
             onClick={handleLoadMore}
-            isDisabled={isLoading}
+            isDisabled={isLoading || autoFetching}
             _disabled={disabledProps}
           >
             {translate('common.loadMore')}
diff --git a/src/components/ManageAccountsDrawer/helpers.ts b/src/components/ManageAccountsDrawer/helpers.ts
index b1e149466cb..19724675aa4 100644
--- a/src/components/ManageAccountsDrawer/helpers.ts
+++ b/src/components/ManageAccountsDrawer/helpers.ts
@@ -1,7 +1,10 @@
-import type { ChainId } from '@shapeshiftoss/caip'
+import { type ChainId, fromAccountId } from '@shapeshiftoss/caip'
+import type { HDWallet } from '@shapeshiftoss/hdwallet-core'
 import { matchSorter } from 'match-sorter'
 import { getChainAdapterManager } from 'context/PluginProvider/chainAdapterSingleton'
-import { isSome } from 'lib/utils'
+import { deriveAccountIdsAndMetadata } from 'lib/account/account'
+import { assertGetChainAdapter, isSome } from 'lib/utils'
+import { checkAccountHasActivity } from 'state/slices/portfolioSlice/utils'
 
 export const filterChainIdsBySearchTerm = (search: string, chainIds: ChainId[]) => {
   if (!chainIds.length) return []
@@ -25,3 +28,24 @@ export const filterChainIdsBySearchTerm = (search: string, chainIds: ChainId[])
     threshold: matchSorter.rankings.CONTAINS,
   }).map(({ chainId }) => chainId)
 }
+
+export const getAccountIdsWithActivityAndMetadata = async (
+  accountNumber: number,
+  chainId: ChainId,
+  wallet: HDWallet | null,
+) => {
+  if (!wallet) return []
+  const input = { accountNumber, chainIds: [chainId], wallet }
+  const accountIdsAndMetadata = await deriveAccountIdsAndMetadata(input)
+
+  return Promise.all(
+    Object.entries(accountIdsAndMetadata).map(async ([accountId, accountMetadata]) => {
+      const { account: pubkey } = fromAccountId(accountId)
+      const adapter = assertGetChainAdapter(chainId)
+      const account = await adapter.getAccount(pubkey)
+      const hasActivity = checkAccountHasActivity(account)
+
+      return { accountId, accountMetadata, hasActivity }
+    }),
+  )
+}
diff --git a/src/react-queries/queries/accountManagement.ts b/src/react-queries/queries/accountManagement.ts
index 13ec8f57b4d..995770bcd28 100644
--- a/src/react-queries/queries/accountManagement.ts
+++ b/src/react-queries/queries/accountManagement.ts
@@ -1,31 +1,7 @@
 import { createQueryKeys } from '@lukemorales/query-key-factory'
 import type { AccountId } from '@shapeshiftoss/caip'
-import { type ChainId, fromAccountId } from '@shapeshiftoss/caip'
-import type { HDWallet } from '@shapeshiftoss/hdwallet-core'
-import type { AccountMetadata } from '@shapeshiftoss/types'
-import { deriveAccountIdsAndMetadata } from 'lib/account/account'
+import { fromAccountId } from '@shapeshiftoss/caip'
 import { assertGetChainAdapter } from 'lib/utils'
-import { checkAccountHasActivity } from 'state/slices/portfolioSlice/utils'
-
-const getAccountIdsWithActivityAndMetadata = async (
-  accountNumber: number,
-  chainId: ChainId,
-  wallet: HDWallet,
-) => {
-  const input = { accountNumber, chainIds: [chainId], wallet }
-  const accountIdsAndMetadata = await deriveAccountIdsAndMetadata(input)
-
-  return Promise.all(
-    Object.entries(accountIdsAndMetadata).map(async ([accountId, accountMetadata]) => {
-      const { account: pubkey } = fromAccountId(accountId)
-      const adapter = assertGetChainAdapter(chainId)
-      const account = await adapter.getAccount(pubkey)
-      const hasActivity = checkAccountHasActivity(account)
-
-      return { accountId, accountMetadata, hasActivity }
-    }),
-  )
-}
 
 export const accountManagement = createQueryKeys('accountManagement', {
   getAccount: (accountId: AccountId) => ({
@@ -37,62 +13,4 @@ export const accountManagement = createQueryKeys('accountManagement', {
       return account
     },
   }),
-  accountIdWithActivityAndMetadata: (
-    accountNumber: number,
-    chainId: ChainId,
-    wallet: HDWallet,
-    walletDeviceId: string,
-  ) => ({
-    queryKey: ['accountIdWithActivityAndMetadata', chainId, walletDeviceId, accountNumber],
-    queryFn: () => {
-      return getAccountIdsWithActivityAndMetadata(accountNumber, chainId, wallet)
-    },
-  }),
-  firstAccountIdsWithActivityAndMetadata: (
-    chainId: ChainId,
-    wallet: HDWallet | null,
-    walletDeviceId: string,
-    accountNumberLimit: number,
-  ) => ({
-    queryKey: [
-      'firstAccountIdsWithActivityAndMetadata',
-      chainId,
-      walletDeviceId,
-      accountNumberLimit,
-    ],
-    queryFn: async () => {
-      let accountNumber = 0
-
-      const accounts: {
-        accountId: AccountId
-        accountMetadata: AccountMetadata
-        hasActivity: boolean
-      }[][] = []
-
-      if (!wallet) return []
-
-      while (true) {
-        try {
-          if (accountNumber >= accountNumberLimit) {
-            break
-          }
-
-          const accountResults = await getAccountIdsWithActivityAndMetadata(
-            accountNumber,
-            chainId,
-            wallet,
-          )
-
-          accounts.push(accountResults)
-        } catch (error) {
-          console.error(error)
-          break
-        }
-
-        accountNumber++
-      }
-
-      return accounts
-    },
-  }),
 })