Skip to content

Instantly share code, notes, and snippets.

@gragland
Last active July 1, 2024 07:39
Show Gist options
  • Save gragland/30a8282714bc6f4f0f6024fee7e9492f to your computer and use it in GitHub Desktop.
Save gragland/30a8282714bc6f4f0f6024fee7e9492f to your computer and use it in GitHub Desktop.
useOptimisticMutation for React Query. Optimistically update data in multiple locations with rollback on error.
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>
)
}
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()
}
})
}
},
})
}
@gragland
Copy link
Author

What's missing?

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

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment