add reset method to useMutation hook (#1476)

Co-authored-by: Shrugsy <joshfraser91@gmail.com>

Author: phryneas

docs/rtk-query/api/created-api/hooks.mdx

@@ -316,9 +316,9 @@ type UseMutationStateOptions = {
   selectFromResult?: (result: UseMutationStateDefaultResult) => any
 }
 
-type UseMutationTrigger<T> = (
-  arg: any
-) => Promise<{ data: T } | { error: BaseQueryError | SerializedError }> & {
+type UseMutationTrigger<T> = (arg: any) => Promise<
+  { data: T } | { error: BaseQueryError | SerializedError }
+> & {
   requestId: string // A string generated by RTK Query
   abort: () => void // A method to cancel the mutation promise
   unwrap: () => Promise<T> // A method to unwrap the mutation call and provide the raw response/error
@@ -360,7 +360,10 @@ selectFromResult: () => ({})
 
 - **Returns**: A tuple containing:
   - `trigger`: A function that triggers an update to the data based on the provided argument. The trigger function returns a promise with the properties shown above that may be used to handle the behavior of the promise
-  - `mutationState`: A query status object containing the current loading state and metadata about the request, or the values returned by the `selectFromResult` option where applicable
+  - `mutationState`: A query status object containing the current loading state and metadata about the request, or the values returned by the `selectFromResult` option where applicable.
+    Additionally, this object will contain
+    - a `reset` method to reset the hook back to it's original state and remove the current result from the cache
+    - an `originalArgs` property that contains the argument passed to the last call of the `trigger` function.
 
 #### Description
 
@@ -459,9 +462,8 @@ type UseQuerySubscriptionResult = {
 ## `useLazyQuery`
 
 ```ts title="Accessing a useLazyQuery hook" no-transpile
-const [trigger, result, lastPromiseInfo] = api.endpoints.getPosts.useLazyQuery(
-  options
-)
+const [trigger, result, lastPromiseInfo] =
+  api.endpoints.getPosts.useLazyQuery(options)
 // or
 const [trigger, result, lastPromiseInfo] = api.useLazyGetPostsQuery(options)
 ```
@@ -520,9 +522,8 @@ type UseLazyQueryLastPromiseInfo = {
 ## `useLazyQuerySubscription`
 
 ```ts title="Accessing a useLazyQuerySubscription hook" no-transpile
-const [trigger, lastArg] = api.endpoints.getPosts.useLazyQuerySubscription(
-  options
-)
+const [trigger, lastArg] =
+  api.endpoints.getPosts.useLazyQuerySubscription(options)
 ```
 
 #### Signature

packages/toolkit/src/query/core/buildInitiate.ts

@@ -177,6 +177,8 @@ export type MutationActionCreatorResult<
    * A method to manually unsubscribe from the mutation call, meaning it will be removed from cache after the usual caching grace period.
    The value returned by the hook will reset to `isUninitialized` afterwards.
    */
+  reset(): void
+  /** @deprecated has been renamed to `reset` */
   unsubscribe(): void
 }
 
@@ -193,7 +195,7 @@ export function buildInitiate({
 }) {
   const {
     unsubscribeQueryResult,
-    unsubscribeMutationResult,
+    removeMutationResult,
     updateSubscriptionOptions,
   } = api.internalActions
   return { buildInitiateQuery, buildInitiateMutation }
@@ -299,14 +301,18 @@ Features like automatic cache collection, automatic refetching etc. will not be
           .unwrap()
           .then((data) => ({ data }))
           .catch((error) => ({ error }))
+
+        const reset = () => {
+          if (track) dispatch(removeMutationResult({ requestId }))
+        }
+
         return Object.assign(returnValuePromise, {
           arg: thunkResult.arg,
           requestId,
           abort,
           unwrap: thunkResult.unwrap,
-          unsubscribe() {
-            if (track) dispatch(unsubscribeMutationResult({ requestId }))
-          },
+          unsubscribe: reset,
+          reset,
         })
       }
   }

packages/toolkit/src/query/core/buildMiddleware/cacheLifecycle.ts

@@ -235,7 +235,7 @@ export const build: SubMiddlewareBuilder = ({
           }
         } else if (
           api.internalActions.removeQueryResult.match(action) ||
-          api.internalActions.unsubscribeMutationResult.match(action)
+          api.internalActions.removeMutationResult.match(action)
         ) {
           const lifecycle = lifecycleMap[cacheKey]
           if (lifecycle) {
@@ -257,7 +257,7 @@ export const build: SubMiddlewareBuilder = ({
       if (isMutationThunk(action)) return action.meta.requestId
       if (api.internalActions.removeQueryResult.match(action))
         return action.payload.queryCacheKey
-      if (api.internalActions.unsubscribeMutationResult.match(action))
+      if (api.internalActions.removeMutationResult.match(action))
         return action.payload.requestId
       return ''
     }

packages/toolkit/src/query/core/buildSlice.ts

@@ -172,7 +172,7 @@ export function buildSlice({
     name: `${reducerPath}/mutations`,
     initialState: initialState as MutationState<any>,
     reducers: {
-      unsubscribeMutationResult(
+      removeMutationResult(
         draft,
         action: PayloadAction<MutationSubstateIdentifier>
       ) {
@@ -379,6 +379,8 @@ export function buildSlice({
     ...querySlice.actions,
     ...subscriptionSlice.actions,
     ...mutationSlice.actions,
+    /** @deprecated has been renamed to `removeMutationResult` */
+    unsubscribeMutationResult: mutationSlice.actions.removeMutationResult,
     resetApiState,
   }
 

packages/toolkit/src/query/react/buildHooks.ts

@@ -408,6 +408,11 @@ export type UseMutationStateResult<
   R
 > = NoInfer<R> & {
   originalArgs?: QueryArgFrom<D>
+  /**
+   * Resets the hook state to it's initial `uninitialized` state.
+   * This will also remove the last result from the cache.
+   */
+  reset: () => void
 }
 
 /**
@@ -424,10 +429,28 @@ export type UseMutation<D extends MutationDefinition<any, any, any, any>> = <
   R extends Record<string, any> = MutationResultSelectorResult<D>
 >(
   options?: UseMutationStateOptions<D, R>
-) => [
-  (arg: QueryArgFrom<D>) => MutationActionCreatorResult<D>,
-  UseMutationStateResult<D, R>
-]
+) => [MutationTrigger<D>, UseMutationStateResult<D, R>]
+
+export type MutationTrigger<D extends MutationDefinition<any, any, any, any>> =
+  {
+    /**
+     * Triggers the mutation and returns a Promise.
+     * @remarks
+     * If you need to access the error or success payload immediately after a mutation, you can chain .unwrap().
+     *
+     * @example
+     * ```ts
+     * // codeblock-meta title="Using .unwrap with async await"
+     * try {
+     *   const payload = await addPost({ id: 1, name: 'Example' }).unwrap();
+     *   console.log('fulfilled', payload)
+     * } catch (error) {
+     *   console.error('rejected', error);
+     * }
+     * ```
+     */
+    (arg: QueryArgFrom<D>): MutationActionCreatorResult<D>
+  }
 
 const defaultQueryStateSelector: QueryStateSelector<any, any> = (x) => x
 const defaultMutationStateSelector: MutationStateSelector<any, any> = (x) => x
@@ -764,47 +787,32 @@ export function buildHooks<Definitions extends EndpointDefinitions>({
         Definitions
       >
       const dispatch = useDispatch<ThunkDispatch<any, any, AnyAction>>()
-      const [requestId, setRequestId] = useState<string>()
+      const [promise, setPromise] = useState<MutationActionCreatorResult<any>>()
 
-      const promiseRef = useRef<MutationActionCreatorResult<any>>()
-
-      useEffect(() => {
-        return () => {
-          promiseRef.current?.unsubscribe()
-          promiseRef.current = undefined
-        }
-      }, [])
+      useEffect(() => () => promise?.reset(), [promise])
 
       const triggerMutation = useCallback(
         function (arg) {
-          let promise: MutationActionCreatorResult<any>
-          batch(() => {
-            promiseRef?.current?.unsubscribe()
-            promise = dispatch(initiate(arg))
-            promiseRef.current = promise
-            setRequestId(promise.requestId)
-          })
-          return promise!
+          const promise = dispatch(initiate(arg))
+          setPromise(promise)
+          return promise
         },
         [dispatch, initiate]
       )
 
+      const { requestId } = promise || {}
       const mutationSelector = useMemo(
         () =>
-          createSelector([select(requestId || skipToken)], (subState) =>
-            selectFromResult(subState)
-          ),
+          createSelector([select(requestId || skipToken)], selectFromResult),
         [select, requestId, selectFromResult]
       )
 
       const currentState = useSelector(mutationSelector, shallowEqual)
-      const originalArgs = promiseRef.current?.arg.originalArgs
+      const originalArgs = promise?.arg.originalArgs
+      const reset = useCallback(() => setPromise(undefined), [])
       const finalState = useMemo(
-        () => ({
-          ...currentState,
-          originalArgs,
-        }),
-        [currentState, originalArgs]
+        () => ({ ...currentState, originalArgs, reset }),
+        [currentState, originalArgs, reset]
       )
 
       return useMemo(

packages/toolkit/src/query/tests/buildHooks.test.tsx

@@ -906,6 +906,63 @@ describe('hooks tests', () => {
       expect(screen.queryByText(/Successfully updated user/i)).toBeNull()
       screen.getByText('Request was aborted')
     })
+
+    test('useMutation return value contains originalArgs', async () => {
+      const { result } = renderHook(api.endpoints.updateUser.useMutation, {
+        wrapper: storeRef.wrapper,
+      })
+      const arg = { name: 'Foo' }
+
+      const firstRenderResult = result.current
+      expect(firstRenderResult[1].originalArgs).toBe(undefined)
+      act(() => void firstRenderResult[0](arg))
+      const secondRenderResult = result.current
+      expect(firstRenderResult[1].originalArgs).toBe(undefined)
+      expect(secondRenderResult[1].originalArgs).toBe(arg)
+    })
+
+    test('`reset` sets state back to original state', async () => {
+      function User() {
+        const [updateUser, result] = api.endpoints.updateUser.useMutation()
+        return (
+          <>
+            <span>
+              {result.isUninitialized
+                ? 'isUninitialized'
+                : result.isSuccess
+                ? 'isSuccess'
+                : 'other'}
+            </span>
+            <span>{result.originalArgs?.name}</span>
+            <button onClick={() => updateUser({ name: 'Yay' })}>trigger</button>
+            <button onClick={result.reset}>reset</button>
+          </>
+        )
+      }
+      render(<User />, { wrapper: storeRef.wrapper })
+
+      await screen.findByText(/isUninitialized/i)
+      expect(screen.queryByText('Yay')).toBeNull()
+      expect(Object.keys(storeRef.store.getState().api.mutations).length).toBe(
+        0
+      )
+
+      userEvent.click(screen.getByRole('button', { name: 'trigger' }))
+
+      await screen.findByText(/isSuccess/i)
+      expect(screen.queryByText('Yay')).not.toBeNull()
+      expect(Object.keys(storeRef.store.getState().api.mutations).length).toBe(
+        1
+      )
+
+      userEvent.click(screen.getByRole('button', { name: 'reset' }))
+
+      await screen.findByText(/isUninitialized/i)
+      expect(screen.queryByText('Yay')).toBeNull()
+      expect(Object.keys(storeRef.store.getState().api.mutations).length).toBe(
+        0
+      )
+    })
   })
 
   describe('usePrefetch', () => {
@@ -1880,8 +1937,8 @@ describe('hooks with createApi defaults set', () => {
         api.internalActions.middlewareRegistered.match,
         increment.matchPending,
         increment.matchFulfilled,
-        api.internalActions.unsubscribeMutationResult.match,
         increment.matchPending,
+        api.internalActions.unsubscribeMutationResult.match,
         increment.matchFulfilled
       )
     })
@@ -1966,19 +2023,6 @@ describe('hooks with createApi defaults set', () => {
       expect(getRenderCount()).toBe(5)
     })
 
-    test('useMutation return value contains originalArgs', async () => {
-      const { result } = renderHook(api.endpoints.increment.useMutation, {
-        wrapper: storeRef.wrapper,
-      })
-
-      const firstRenderResult = result.current
-      expect(firstRenderResult[1].originalArgs).toBe(undefined)
-      firstRenderResult[0](5)
-      const secondRenderResult = result.current
-      expect(firstRenderResult[1].originalArgs).toBe(undefined)
-      expect(secondRenderResult[1].originalArgs).toBe(5)
-    })
-
     it('useMutation with selectFromResult option has a type error if the result is not an object', async () => {
       function Counter() {
         const [increment] = api.endpoints.increment.useMutation({

packages/toolkit/src/query/tests/unionTypes.test.ts

@@ -497,6 +497,7 @@ describe.skip('TS only tests', () => {
       isLoading: true,
       isSuccess: false,
       isError: false,
+      reset: () => {},
     })(result)
   })