useOptimisticMutation for React Query. Optimistically update data in multiple locations with rollback on error.

use-optimistic-mutation-example.ts

import axios from 'axios'
import { useOptimisticMutation } from "./useOptimisticMutation.ts"

type Response = boolean
type Error = unknown
type MutationVariables = {itemId: string}
type Items = {id: string; name: string}[]
type Likes = {itemId: string}[]
type History = {type: string}[]

function Component({items}: {items:Items}) {
  // Some local state for our example
  const [history, setHistory] = useState()

  // Mutation to delete an item and optimistically update data in three locations
  const {mutate: deleteItem} = useOptimisticMutation<
    Response,
    Error,
    MutationVariables,
    // Data types for our optimistic handlers
    [
      Items | undefined, 
      Likes | undefined, 
      History
    ]
  >({
    mutationFn: async (variables) => {
      return axios.post('/api/items/add', variables).then((res) => res.data)
    },
    // This is where the magic happens
    optimistic: (variables) => {
      return [
        // Remove from items
        {
          // The React Query key to find the cached data
          queryKey: ['items'],
          // Function to modify the cached data
          updater: (currentData) => {
            return currentData?.filter((item) => item.id !== variables.itemId)
          },
        },
        // Remove from likes
        {
          queryKey: ['likes'],
          updater: (currentData) => {
            return currentData?.filter((item) => item.itemId !== variables.itemId)
          },
        },
        // Update some local state by specifying `getData` and `setData`
        // Useful to handle with this hook so it gets rolled back on error
        {
          getData: () => history,
          setData: (data) => setHistory(data),
          updater: (currentData) => {
            return [...(currentData || []), {type: 'delete'}]
          },
        },
      ]
    },
  })
  
  return (
    <div>
      {items.map(item => (
        <Item item={item} onDelete={() => deleteItem({itemId: item.id}) } />
      )}
    </div>
  )
}

useOptimisticMutation.ts

import {useMutation, useQueryClient, QueryKey, MutationFunction} from '@tanstack/react-query'

type MutationContext = {
  results: {
    rollback: () => void
    invalidate?: () => void
    didCancelFetch?: boolean
  }[]
}

type HandlerReactQuery<TOptimisticData> = {
  queryKey: QueryKey
  updater: (data: TOptimisticData) => TOptimisticData | undefined
}

type Handler<TOptimisticData> = {
  getData: () => TOptimisticData
  setData: (data: TOptimisticData) => void
  updater: (data: TOptimisticData) => TOptimisticData | undefined
}

type OptimisticFunction<TOptimisticDataArray, TVariables> = (variables: TVariables) => {
  [K in keyof TOptimisticDataArray]:
    | HandlerReactQuery<TOptimisticDataArray[K]>
    | Handler<TOptimisticDataArray[K]>
}

type OptimisticMutationProps<TData, TVariables, TOptimisticDataArray> = {
  mutationFn: MutationFunction<TData, TVariables>
  optimistic: OptimisticFunction<TOptimisticDataArray, TVariables>
}

export function useOptimisticMutation<
  TData,
  TError,
  TVariables,
  TOptimisticDataArray extends unknown[]
>({
  mutationFn,
  optimistic,
}: OptimisticMutationProps<TData, TVariables, TOptimisticDataArray>) {
  const queryClient = useQueryClient()

  return useMutation<TData, TError, TVariables, MutationContext>({
    mutationFn,
    onMutate: async (variables) => {
      const results = []
      const handlers = optimistic(variables)

      for (const handler of handlers) {
        if ('queryKey' in handler) {
          const {queryKey, updater} = handler
          let didCancelFetch = false

          // If query is currently fetching, we cancel it to avoid overwriting our optimistic update.
          // This would happen if query responds with old data after our optimistic update is applied.
          const isFetching = queryClient.getQueryState(queryKey)?.fetchStatus === 'fetching'
          if (isFetching) {
            await queryClient.cancelQueries(queryKey)
            didCancelFetch = true
          }

          // Get previous data before optimistic update
          const previousData = queryClient.getQueryData(queryKey)
          // Rollback function we call if mutation fails
          const rollback = () => queryClient.setQueryData(queryKey, previousData)
          // Invalidate function to call after mutation is done if we cancelled a fetch.
          // This ensures that we get both the optimistic update and fresh data from the server.
          const invalidate = () => queryClient.invalidateQueries(queryKey)

          // Update data in React Query cache
          queryClient.setQueryData(queryKey, updater)

          // Add to results that we read in onError and onSettled
          results.push({
            rollback,
            invalidate,
            didCancelFetch,
          })
        } else {
          // If no query key then we're not operating on the React Query cache
          // We expect to have a `getData` and `setData` function
          const {getData, setData, updater} = handler
          const previousData = getData()
          const rollback = () => setData(previousData)
          setData(updater)
          results.push({
            rollback,
          })
        }
      }

      return {results}
    },
    // On error revert all queries to their previous data
    onError: (error, variables, context) => {
      if (context?.results) {
        context.results.forEach(({rollback}) => {
          rollback()
        })
      }
    },
    // When mutation is done invalidate cancelled queries so they get refetched
    onSettled: (data, error, variables, context) => {
      if (context?.results) {
        context.results.forEach(({didCancelFetch, invalidate}) => {
          if (didCancelFetch && invalidate) {
            invalidate()
          }
        })
      }
    },
  })
}

What's missing?

• partial query keys (use setQueriesData, support filter, exact, predicate)
• pass custom onError and onSettled
• what else?