add ability to reset lazy query hooks

Author: alexmotoc

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

@@ -18,6 +18,7 @@ import type {
   PrefetchOptions,
   QueryActionCreatorResult,
   QueryArgFrom,
+  QueryCacheKey,
   QueryDefinition,
   QueryKeys,
   QueryResultSelectorResult,
@@ -255,7 +256,7 @@ export type UseLazyQuery<D extends QueryDefinition<any, any, any, any>> = <
   options?: SubscriptionOptions & Omit<UseQueryStateOptions<D, R>, 'skip'>,
 ) => [
   LazyQueryTrigger<D>,
-  UseQueryStateResult<D, R>,
+  UseLazyQueryStateResult<D, R>,
   UseLazyQueryLastPromiseInfo<D>,
 ]
 
@@ -267,6 +268,33 @@ export type TypedUseLazyQuery<
   QueryDefinition<QueryArg, BaseQuery, string, ResultType, string>
 >
 
+export type UseLazyQueryStateResult<
+  D extends QueryDefinition<any, any, any, any>,
+  R = UseQueryStateDefaultResult<D>,
+> = UseQueryStateResult<D, R> & {
+  /**
+   * Resets the hook state to its initial `uninitialized` state.
+   * This will also remove the last result from the cache.
+   */
+  reset: () => void
+}
+
+/**
+ * Helper type to manually type the result
+ * of the `useLazyQuery` hook in userland code.
+ */
+export type TypedUseLazyQueryStateResult<
+  ResultType,
+  QueryArg,
+  BaseQuery extends BaseQueryFn,
+  R = UseQueryStateDefaultResult<
+    QueryDefinition<QueryArg, BaseQuery, string, ResultType, string>
+  >,
+> = UseLazyQueryStateResult<
+  QueryDefinition<QueryArg, BaseQuery, string, ResultType, string>,
+  R
+>
+
 export type LazyQueryTrigger<D extends QueryDefinition<any, any, any, any>> = {
   /**
    * Triggers a lazy query.
@@ -317,7 +345,11 @@ export type UseLazyQuerySubscription<
   D extends QueryDefinition<any, any, any, any>,
 > = (
   options?: SubscriptionOptions,
-) => readonly [LazyQueryTrigger<D>, QueryArgFrom<D> | UninitializedValue]
+) => readonly [
+  LazyQueryTrigger<D>,
+  QueryArgFrom<D> | UninitializedValue,
+  { reset: () => void },
+]
 
 export type TypedUseLazyQuerySubscription<
   ResultType,
@@ -1162,6 +1194,16 @@ export function buildHooks<Definitions extends EndpointDefinitions>({
         [dispatch, initiate],
       )
 
+      const reset = useCallback(() => {
+        if (promiseRef.current?.queryCacheKey) {
+          dispatch(
+            api.internalActions.removeQueryResult({
+              queryCacheKey: promiseRef.current?.queryCacheKey as QueryCacheKey,
+            }),
+          )
+        }
+      }, [dispatch])
+
       /* cleanup on unmount */
       useEffect(() => {
         return () => {
@@ -1176,7 +1218,10 @@ export function buildHooks<Definitions extends EndpointDefinitions>({
         }
       }, [arg, trigger])
 
-      return useMemo(() => [trigger, arg] as const, [trigger, arg])
+      return useMemo(
+        () => [trigger, arg, { reset }] as const,
+        [trigger, arg, reset],
+      )
     }
 
     const useQueryState: UseQueryState<any> = (
@@ -1249,16 +1294,16 @@ export function buildHooks<Definitions extends EndpointDefinitions>({
       useQuerySubscription,
       useLazyQuerySubscription,
       useLazyQuery(options) {
-        const [trigger, arg] = useLazyQuerySubscription(options)
+        const [trigger, arg, { reset }] = useLazyQuerySubscription(options)
         const queryStateResults = useQueryState(arg, {
           ...options,
           skip: arg === UNINITIALIZED_VALUE,
         })
 
         const info = useMemo(() => ({ lastArg: arg }), [arg])
         return useMemo(
-          () => [trigger, queryStateResults, info],
-          [trigger, queryStateResults, info],
+          () => [trigger, { ...queryStateResults, reset }, info],
+          [trigger, queryStateResults, reset, info],
         )
       },
       useQuery(arg, options) {

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

@@ -28,6 +28,7 @@ export type {
   TypedUseQuerySubscription,
   TypedUseLazyQuerySubscription,
   TypedUseQueryStateOptions,
+  TypedUseLazyQueryStateResult,
 } from './buildHooks'
 export { UNINITIALIZED_VALUE } from './constants'
 export { createApi, reactHooksModule, reactHooksModuleName }

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

@@ -1386,10 +1386,8 @@ describe('hooks tests', () => {
 
     test('useLazyQuery trigger promise returns the correctly updated data', async () => {
       const LazyUnwrapUseEffect = () => {
-        const [
-          triggerGetIncrementedAmount,
-          { isFetching, isSuccess, isError, error, data },
-        ] = api.endpoints.getIncrementedAmount.useLazyQuery()
+        const [triggerGetIncrementedAmount, { isFetching, isSuccess, data }] =
+          api.endpoints.getIncrementedAmount.useLazyQuery()
 
         type AmountData = { amount: number } | undefined
 
@@ -1478,6 +1476,50 @@ describe('hooks tests', () => {
         expect(screen.getByText('Unwrap data: 2')).toBeTruthy()
       })
     })
+
+    test('`reset` sets state back to original state', async () => {
+      function User() {
+        const [getUser, { isSuccess, isUninitialized, reset }, _lastInfo] =
+          api.endpoints.getUser.useLazyQuery()
+
+        const handleFetchClick = async () => {
+          await getUser(1).unwrap()
+        }
+
+        return (
+          <div>
+            <span>
+              {isUninitialized
+                ? 'isUninitialized'
+                : isSuccess
+                  ? 'isSuccess'
+                  : 'other'}
+            </span>
+            <button onClick={handleFetchClick}>Fetch User</button>
+            <button onClick={reset}>Reset</button>
+          </div>
+        )
+      }
+
+      render(<User />, { wrapper: storeRef.wrapper })
+
+      await screen.findByText(/isUninitialized/i)
+      expect(countObjectKeys(storeRef.store.getState().api.queries)).toBe(0)
+
+      userEvent.click(screen.getByRole('button', { name: 'Fetch User' }))
+
+      await screen.findByText(/isSuccess/i)
+      expect(countObjectKeys(storeRef.store.getState().api.queries)).toBe(1)
+
+      userEvent.click(
+        screen.getByRole('button', {
+          name: 'Reset',
+        }),
+      )
+
+      await screen.findByText(/isUninitialized/i)
+      expect(countObjectKeys(storeRef.store.getState().api.queries)).toBe(0)
+    })
   })
 
   describe('useMutation', () => {

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

@@ -9,6 +9,7 @@ import type {
   TypedUseQueryStateResult,
   TypedUseQuerySubscriptionResult,
   TypedLazyQueryTrigger,
+  TypedUseLazyQueryStateResult,
   TypedUseLazyQuery,
   TypedUseLazyQuerySubscription,
   TypedUseMutation,
@@ -820,7 +821,7 @@ describe('"Typed" helper types', () => {
     >().toMatchTypeOf(trigger)
 
     expectTypeOf<
-      TypedUseQueryHookResult<string, void, typeof baseQuery>
+      TypedUseLazyQueryStateResult<string, void, typeof baseQuery>
     >().toMatchTypeOf(result)
   })
 
@@ -834,7 +835,12 @@ describe('"Typed" helper types', () => {
     >().toMatchTypeOf(trigger)
 
     expectTypeOf<
-      TypedUseQueryHookResult<string, void, typeof baseQuery, { x: boolean }>
+      TypedUseLazyQueryStateResult<
+        string,
+        void,
+        typeof baseQuery,
+        { x: boolean }
+      >
     >().toMatchTypeOf(result)
   })