react-query-optimistic-update.ts

import _ from "lodash";
import { QueryClient } from "react-query";

/**
 * optimisticUpdate
 *
 * A utility function to update the react-query cache with data during
 * a mutation.
 *
 * This clones the current data, uses the supplied function to update that
 * object, and then updates the cache with the new object.
 *
 * It returns a context that React Query stores for rolling back the changes
 * onError
 */
export async function optimisticUpdate<QueryData>(
  queryClient: QueryClient,
  queryKey: string[],
  fn: (data: QueryData) => QueryData
) {
  // Cancel any outgoing refetches (so they don't overwrite our optimistic update)
  await queryClient.cancelQueries(queryKey);

  // Snapshot the previous value
  const previous: QueryData | undefined = queryClient.getQueryData(queryKey);
  if (!previous) {
    // throw new Error("Query not found or was never loaded");
    return { previous: undefined };
  }

  const clone = _.cloneDeep(previous);

  const updated = fn(clone);

  // Optimistically update to the new value
  queryClient.setQueryData(queryKey, updated);

  // Return a context with the previous state for rollback on error
  return { previous };
}

/**
 * Reset react-query cache to the previous state before the mutation
 */
export function rollback<QueryData>(
  queryClient: QueryClient,
  queryKey: string[],
  context: { previous: QueryData | undefined } | undefined
) {
  if (context?.previous) {
    queryClient.setQueryData(queryKey, context.previous);
  }
}

usage.tsx

import _ from "lodash";
import { useMutation } from "react-query";

import queryClient from "services/queryClient";
import { optimisticUpdate, rollback } from "services/api/queryData";

import type {
  Model,
  PostModel,
} from "pages/api/model";

// Always write these 3 types:
type TVariables = {  // what the mutation is called with
  id: string;
} & Partial<PostModel>;
type TResponse = Model; // what the update response looks like
type TContext = { previous: Model[] | undefined }; // the context is where the query is stored in case of errors

/**
 * Example. Suppose that you have a list of "model" that were fetched with 
 * query key ['get-model']
 *
 * This mutation function will let you update any one of those, and will
 * optimistically update that model in the list.
 *
 * You can also do this for delete, push/create etc.
 */
export default function useUpdateModel(id: string) {
  const queryKey = ['get-model'];

  const mutation = useMutation<TResponse, Error, TVariables, TContext>(
    (variables) => {
      return postToServer(variables);
    },
    {
      onMutate: (updateData) => {
        return optimisticUpdate<Model[]>(
          queryClient,
          queryKey,
          (clone) => {
             // the clone is like immer's draft objects.
            const model = _.find(
              clone,
              (m) => m.id === id
            );
            if (model) {
              // just update it in place
              _.assign(model, updateData);
            }
            return clone;
          }
        );
      },
      onError: (err, updateAssessmentData, context) => {
        console.error(err);
        rollback(queryClient, queryKey, context);
      },
      // Always refetch after error or success:
      onSettled: () => {
        queryClient.refetchQueries(queryKey);
      },
    }
  );

  // this should probably be wrapped with useCallback ...
  return (
    id: string,
    updateData: Partial<Model>
  ) => {
    mutation.mutate({
      ...updateData,
      id,
    });
  };
}