Merge pull request #4689 from alexmotoc/lazy-query-reset

Add ability to reset lazy query hooks

Author: markerikson

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

@@ -1398,10 +1398,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
 
@@ -1490,6 +1488,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)
  })