React Query CosmWasm native balance queries using Cosm.js

useNativeBalanceQueries.ts

import { useQueries, useQuery, UseQueryOptions } from '@tanstack/react-query'
import { CombinedQueries, useCombinedQueries } from 'api/queries/useCombinedQueries'
import { useNetworkCurrencies } from 'hooks/network/useNetworkCurrencies'
import { Coin } from '@cosmjs/stargate'
import { QueryOptions } from 'api/queries/query'
import { ReactQueryOptions } from 'api/queries/useReactQueryOptions'
import { CosmWasmClient } from '@cosmjs/cosmwasm-stargate'
import useReadonlyClients from 'contexts/ReadonlyClientContext'

interface UseNativeBalanceQueryOptionsParams<TData> extends ReactQueryOptions<Coin, TData> {
  client: CosmWasmClient | undefined
  account: string | undefined
  address: string
}

/**
 *
 * @param client
 * @param account
 * @param address address of the asset, like ujuno or ibc/aosetnuhaosentuhasoentuh
 * @param options
 * @return Coin in MICRO format
 */
const nativeBalanceQuery = <TData>({
  client,
  account,
  address,
  options,
}: UseNativeBalanceQueryOptionsParams<TData>): UseQueryOptions<Coin, Error, TData> => {
  return {
    ...options,
    queryKey: ['balance', account, address],
    queryFn: () => {
      if (!client || !account) return Promise.reject(new Error('Missing client or address'))
      return client.getBalance(account, address)
    },
    ...QueryOptions.DISABLE_WITHOUT(client && account, options?.enabled),
  }
}

interface BalancesQueries<TData> extends ReactQueryOptions<Coin, TData> {
  account?: string
}

/**
 * Hook to retrieve balances of multiple native assets for a given account.
 * @param address: user address
 * @param options
 */
export const useNativeCurrencyBalances = <TData = Coin>({
  account,
  options,
}: BalancesQueries<TData>): CombinedQueries<TData, string> => {
  const { map: coinMap } = useNetworkCurrencies()
  const { readOnlyClient } = useReadonlyClients()

  const denoms = []
  // add the chain native denoms
  for (const denom in coinMap) denoms.push(denom)

  const queries = useQueries({
    queries: denoms.map((denom: string) =>
      nativeBalanceQuery<TData>({ client: readOnlyClient, account, address: denom, options })
    ),
  })

  return useCombinedQueries<TData, string>({
    keys: denoms,
    // @ts-ignore TODO
    queries,
  })
}

interface NativeAssetBalancesQueries<TData> extends ReactQueryOptions<Coin, TData> {
  account?: string
  nativeAssets: readonly ChainAsset[]
}

/**
 * Hook to retrieve balances of multiple native assets for a given account.
 * @param address: user address
 * @param options
 * @param nativeAssets assets to query
 */
export const useNativeAssetBalancesQueries = <TData = Coin>({
  account,
  options,
  nativeAssets,
}: NativeAssetBalancesQueries<TData>): CombinedQueries<TData, string> => {
  const { readOnlyClient } = useReadonlyClients()

  const addresses = nativeAssets.map(({ address }) => address)

  const queries = useQueries({
    queries: addresses.map((address) =>
      nativeBalanceQuery<TData>({ client: readOnlyClient, account, address, options })
    ),
  })

  return useCombinedQueries<TData, string>({
    keys: addresses,
    // @ts-ignore TODO
    queries,
  })
}

export interface BalanceQuery<TData> extends ReactQueryOptions<Coin, TData> {
  account: string | undefined
  address: string
}

export const useNativeBalanceQuery = <TData>({
  account,
  options,
  address,
}: BalanceQuery<TData>) => {
  const { readOnlyClient } = useReadonlyClients()
  return useQuery(
    nativeBalanceQuery<TData>({
      client: readOnlyClient,
      account,
      address,
      options,
    })
  )
}