Custom useMutation implementation

custom-use-mutation.ts

import {
  type MutationOptions as RQ_MutationOptions,
  type InvalidateOptions,
  useMutation as RQ_useMutation,
} from "@tanstack/react-query"

type ValOrFunc<Value, Args extends any[]> = Value | ((...args: Args[]) => Value)
function extractFromValOrFunc<Value, Args extends any[]>(
  item: ValOrFunc<Value, Args>,
  ...args: Args
) {
  // @ts-expect-error: I know better
  if (typeof item == "function") return item(...args) as Value
  return item
}

export type MutationOptions<Params = void, Data = unknown> = RQ_MutationOptions<
  Data,
  Error,
  Params
> & {
  successMessage?: ValOrFunc<string, [Data]>
  errorMessage?: ValOrFunc<string, [Error]>
  invalidateQueries?: ValOrFunc<Array<readonly string[]>, [Data, Params]>
  invalidateOptions?: InvalidateOptions
}
export function useMutation<Params = void, Data = unknown>(
  opts: MutationOptions<Data, Params>,
) {
  const queryClient = useQueryClient()

  const {
    errorMessage,
    successMessage,
    invalidateQueries,
    invalidateOptions = {},
    ...rest
  } = opts

  return RQ_useMutation({
    ...rest,
    
    onSuccess: (data, vars, ...args) => {
      rest.onSuccess?.(data, vars, ...args)

      if (!successMessage) return
      const msg = extractFromValOrFunc(successMessage, data)
      toast.success(msg)

      if (invalidateQueries) {
        const queryKeys = extractFromValOrFunc(invalidateQueries, data, vars)
        for (const queryKey of queryKeys) {
          queryClient.invalidateQueries({ queryKey }, invalidateOptions)
        }
      }
    },
    
    onError: (error, ...args) => {
      rest.onError?.(error, ...args)

      if (!errorMessage) return
      const msg = extractFromValOrFunc(errorMessage, error as Error)
      toast.error(msg)
    },
  })
}