Add `trpc.useContext().setQueryData()` callback

When updating values in the query cache which rely upon the old value, an updater callback comes handy, e.g.:

setQueryData(
  ["Page.findAllBySiteId", data.siteId],
  (prev) => [...prev, data],
);

Instead of:

const prevPagesBySiteId =
  getQueryData(["Page.findBySiteId", data.siteId]) ?? [];
setQueryData(
  ["Page.findAllBySiteId", data.siteId],
  [...prevPagesBySiteId, data],
);

This overload should be made available through trpc.useContext().setQueryData.

[kripod]

I’ve just updated the debouncer yet another time, as it turns out that the onError method of mutations discarded due to debouncing ran after the onMutate callback of newer mutations. The wrong scheduling behavior messed up optimistic updates, so I came up with this:

import type { TRPCClientError } from "@trpc/client";
import { createReactQueryHooks } from "@trpc/react";
import type { inferProcedureInput, inferProcedureOutput } from "@trpc/server";
import { createNanoEvents } from "nanoevents";
import * as React from "react";
import {
  CancelledError,
  useMutation,
  UseMutationOptions,
  UseMutationResult,
} from "react-query";
import useConstant from "use-constant";

import type { AppRouter } from "@/pages/api/trpc/[trpc]";
import { sleep } from "@/utils/sleep"; // https://github.com/trpc/trpc/issues/368#issuecomment-842275320

type TMutations = AppRouter["_def"]["mutations"];
type TError = TRPCClientError<AppRouter>;

function useDebouncedMutation<
  TPath extends keyof TMutations & string,
  TInput extends inferProcedureInput<TMutations[TPath]>,
  TOutput extends inferProcedureOutput<TMutations[TPath]>,
>(
  path: TPath,
  {
    delayMs,
    retry: retryBase,
    onMutate: onMutateBase,
    onSettled: onSettledBase,
    ...restOpts
  }: UseMutationOptions<TOutput, TError, TInput> & {
    delayMs: number;
  },
): UseMutationResult<TOutput, TError, TInput> {
  const abortControllerRef = React.useRef<AbortController | null>(null);
  const discardedMutationEmitter = useConstant(() =>
    createNanoEvents<{ settled: () => void }>(),
  );
  const { client } = trpc.useContext();

  const retry: UseMutationOptions<TOutput, TError, TInput>["retry"] = (
    failureCount,
    error,
  ) => {
    if (error instanceof CancelledError) return false;

    if (typeof retryBase === "function") return retryBase(failureCount, error);
    if (typeof retryBase === "number") return failureCount < retryBase;
    return retryBase ?? false;
  };

  return useMutation<TOutput, TError, TInput>(
    async (input) => {
      abortControllerRef.current = new AbortController();
      try {
        await sleep(delayMs, { signal: abortControllerRef.current.signal });
      } catch {
        // eslint-disable-next-line @typescript-eslint/no-throw-literal
        throw new CancelledError({ revert: true, silent: true });
      }
      abortControllerRef.current = null;

      // eslint-disable-next-line @typescript-eslint/no-explicit-any
      return (client.mutation as any)(path, input);
    },
    {
      retry,
      onMutate: async (...args) => {
        /* Wait for discarded mutation to roll back its optimistic changes */
        await new Promise<void>((resolve) => {
          if (abortControllerRef.current) {
            const unsubscribe = discardedMutationEmitter.on("settled", () => {
              resolve();
              unsubscribe();
            });
            abortControllerRef.current.abort();
          } else {
            resolve();
          }
        });

        return onMutateBase?.(...args);
      },
      onSettled: async (...args) => {
        const result = await onSettledBase?.(...args);

        const error = args[1];
        if (error instanceof CancelledError) {
          discardedMutationEmitter.emit("settled");
        }

        return result;
      },
      ...restOpts,
    },
  );
}

I’m not sure if React Query’s event handling can be invoked forcefully without relying upon a tiny event emitter from nanoevents. I just came across the notifyManager.flush method, but not sure if that’s applicable here.

[KATT]

Nice, you know some nice typescript tricks, would love to see you apply them on some contributions to tRPC 😻