gragland · March 28, 2025 17:18 · Sep 24, 2023 · Sep 24, 2023 · Sep 24, 2023 · Sep 24, 2023
diff --git a/use-optimistic-mutation-example.ts b/use-optimistic-mutation-example.ts
@@ -46,7 +46,8 @@ function Component({items}: {items:Items}) {
             return currentData?.filter((item) => item.itemId !== variables.itemId)
           },
         },
-        // Update some local state and rollback on error
+        // 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),

diff --git a/use-optimistic-mutation-example.ts b/use-optimistic-mutation-example.ts
@@ -46,8 +46,7 @@ function Component({items}: {items:Items}) {
             return currentData?.filter((item) => item.itemId !== variables.itemId)
           },
         },
-        // Update some local state
-        // We can update non React Query data sources by defining `getData` and `setData`
+        // Update some local state and rollback on error
         {
           getData: () => history,
           setData: (data) => setHistory(data),

diff --git a/example.ts → use-optimistic-mutation-example.ts b/example.ts → use-optimistic-mutation-example.ts
diff --git a/example.ts b/example.ts
@@ -0,0 +1,69 @@
+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
+        // We can update non React Query data sources by defining `getData` and `setData`
+        {
+          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>
+  )
+}
diff --git a/useOptimisticMutation.ts b/useOptimisticMutation.ts
@@ -0,0 +1,114 @@
+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()
+          }
+        })
+      }
+    },
+  })
+}