Add unwrap to QueryActionCreatorResult and update LazyQueryTrigger (#1701)

* Add unwrap to QueryActionCreatorResult and update LazyQueryTrigger
* Add test for the return value of lazy query's trigger
* Fix types for lazy query :fingers_crossed:
* minor refactor, add test for positive case
* Use QueryResultSelectorResult instead of QuerySubState
* Update buildHooks test and type expectations for useLazyQuery unwrapping
* Update UseLazyQueryTrigger signature
* Remove extra unwrap and add queryString types
* remove import, fix test

Co-authored-by: Lenz Weber <[email protected]>

Author: msutkowski

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

@@ -490,7 +490,18 @@ type UseLazyQueryOptions = {
   selectFromResult?: (result: UseQueryStateDefaultResult) => any
 }
 
-type UseLazyQueryTrigger = (arg: any) => void
+type UseLazyQueryTrigger<T> = (arg: any) => Promise<
+  QueryResultSelectorResult
+> & {
+  arg: unknown // Whatever argument was provided to the query
+  requestId: string // A string generated by RTK Query
+  subscriptionOptions: SubscriptionOptions // The values used for the query subscription
+  abort: () => void // A method to cancel the query promise
+  unwrap: () => Promise<T> // A method to unwrap the query call and provide the raw response/error
+  unsubscribe: () => void // A method used to manually unsubscribe from the query results
+  refetch: () => void // A method used to re-run the query. In most cases when using a lazy query, you will never use this and should prefer to call the trigger again.
+  updateSubscriptionOptions: (options: SubscriptionOptions) () => void // A method used to update the subscription options (eg. pollingInterval)
+}
 
 type UseQueryStateResult<T> = {
   // Base query state

packages/toolkit/package.json

@@ -37,6 +37,7 @@
     "@types/json-stringify-safe": "^5.0.0",
     "@types/nanoid": "^2.1.0",
     "@types/node": "^10.14.4",
+    "@types/query-string": "^6.3.0",
     "@types/react": "^17.0.3",
     "@types/react-dom": "^17.0.3",
     "@types/react-redux": "^7.1.16",

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

@@ -8,11 +8,12 @@ import type {
 import { DefinitionType } from '../endpointDefinitions'
 import type { QueryThunk, MutationThunk } from './buildThunks'
 import type { AnyAction, ThunkAction, SerializedError } from '@reduxjs/toolkit'
-import type { QuerySubState, SubscriptionOptions, RootState } from './apiState'
+import type { SubscriptionOptions, RootState } from './apiState'
 import type { InternalSerializeQueryArgs } from '../defaultSerializeQueryArgs'
 import type { Api, ApiContext } from '../apiTypes'
 import type { ApiEndpointQuery } from './module'
 import type { BaseQueryError } from '../baseQueryTypes'
+import type { QueryResultSelectorResult } from './buildSelectors'
 
 declare module './module' {
   export interface ApiEndpointQuery<
@@ -47,11 +48,12 @@ type StartQueryActionCreator<
 
 export type QueryActionCreatorResult<
   D extends QueryDefinition<any, any, any, any>
-> = Promise<QuerySubState<D>> & {
+> = Promise<QueryResultSelectorResult<D>> & {
   arg: QueryArgFrom<D>
   requestId: string
   subscriptionOptions: SubscriptionOptions | undefined
   abort(): void
+  unwrap(): Promise<ResultTypeFrom<D>>
   unsubscribe(): void
   refetch(): void
   updateSubscriptionOptions(options: SubscriptionOptions): void
@@ -273,7 +275,7 @@ Features like automatic cache collection, automatic refetching etc. will not be
         })
         const thunkResult = dispatch(thunk)
         middlewareWarning(getState)
-        const { requestId, abort } = thunkResult
+        const { requestId, abort, unwrap } = thunkResult
         const statePromise: QueryActionCreatorResult<any> = Object.assign(
           Promise.all([runningQueries[queryCacheKey], thunkResult]).then(() =>
             (api.endpoints[endpointName] as ApiEndpointQuery<any, any>).select(
@@ -285,6 +287,7 @@ Features like automatic cache collection, automatic refetching etc. will not be
             requestId,
             subscriptionOptions,
             abort,
+            unwrap,
             refetch() {
               dispatch(
                 queryAction(arg, { subscribe: false, forceRefetch: true })
@@ -339,7 +342,7 @@ Features like automatic cache collection, automatic refetching etc. will not be
         })
         const thunkResult = dispatch(thunk)
         middlewareWarning(getState)
-        const { requestId, abort } = thunkResult
+        const { requestId, abort, unwrap } = thunkResult
         const returnValuePromise = thunkResult
           .unwrap()
           .then((data) => ({ data }))
@@ -353,7 +356,7 @@ Features like automatic cache collection, automatic refetching etc. will not be
           arg: thunkResult.arg,
           requestId,
           abort,
-          unwrap: thunkResult.unwrap,
+          unwrap,
           unsubscribe: reset,
           reset,
         })

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

@@ -202,8 +202,25 @@ export type LazyQueryTrigger<D extends QueryDefinition<any, any, any, any>> = {
    *
    * By default, this will start a new request even if there is already a value in the cache.
    * If you want to use the cache value and only start a request if there is no cache value, set the second argument to `true`.
+   *
+   * @remarks
+   * If you need to access the error or success payload immediately after a lazy query, you can chain .unwrap().
+   *
+   * @example
+   * ```ts
+   * // codeblock-meta title="Using .unwrap with async await"
+   * try {
+   *   const payload = await getUserById(1).unwrap();
+   *   console.log('fulfilled', payload)
+   * } catch (error) {
+   *   console.error('rejected', error);
+   * }
+   * ```
    */
-  (arg: QueryArgFrom<D>, preferCacheValue?: boolean): void
+  (
+    arg: QueryArgFrom<D>,
+    preferCacheValue?: boolean
+  ): QueryActionCreatorResult<D>
 }
 
 /**
@@ -221,10 +238,7 @@ export type UseLazyQuerySubscription<
   D extends QueryDefinition<any, any, any, any>
 > = (
   options?: SubscriptionOptions
-) => readonly [
-  (arg: QueryArgFrom<D>) => void,
-  QueryArgFrom<D> | UninitializedValue
-]
+) => readonly [LazyQueryTrigger<D>, QueryArgFrom<D> | UninitializedValue]
 
 export type QueryStateSelector<
   R extends Record<string, any>,
@@ -681,17 +695,22 @@ export function buildHooks<Definitions extends EndpointDefinitions>({
 
       const trigger = useCallback(
         function (arg: any, preferCacheValue = false) {
+          let promise: QueryActionCreatorResult<any>
+
           batch(() => {
             promiseRef.current?.unsubscribe()
 
-            promiseRef.current = dispatch(
+            promiseRef.current = promise = dispatch(
               initiate(arg, {
                 subscriptionOptions: subscriptionOptionsRef.current,
                 forceRefetch: !preferCacheValue,
               })
             )
+
             setArg(arg)
           })
+
+          return promise!
         },
         [dispatch, initiate]
       )

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

@@ -10,6 +10,7 @@ import userEvent from '@testing-library/user-event'
 import { rest } from 'msw'
 import {
   actionsReducer,
+  ANY,
   expectExactType,
   expectType,
   setupApiStore,
@@ -48,8 +49,10 @@ const api = createApi({
     }
   },
   endpoints: (build) => ({
-    getUser: build.query<any, number>({
-      query: (arg) => arg,
+    getUser: build.query<{ name: string }, number>({
+      query: () => ({
+        body: { name: 'Timmy' },
+      }),
     }),
     getIncrementedAmount: build.query<any, void>({
       query: () => ({
@@ -774,6 +777,98 @@ describe('hooks tests', () => {
           .actions.filter(api.internalActions.unsubscribeQueryResult.match)
       ).toHaveLength(4)
     })
+
+    test('useLazyQuery hook callback returns various properties to handle the result', async () => {
+      function User() {
+        const [getUser] = api.endpoints.getUser.useLazyQuery()
+        const [{ successMsg, errMsg, isAborted }, setValues] = React.useState({
+          successMsg: '',
+          errMsg: '',
+          isAborted: false,
+        })
+
+        const handleClick = (abort: boolean) => async () => {
+          const res = getUser(1)
+
+          // no-op simply for clearer type assertions
+          res.then((result) => {
+            if (result.isSuccess) {
+              expectType<{
+                data: {
+                  name: string
+                }
+              }>(result)
+            }
+            if (result.isError) {
+              expectType<{
+                error: { status: number; data: unknown } | SerializedError
+              }>(result)
+            }
+          })
+
+          expectType<number>(res.arg)
+          expectType<string>(res.requestId)
+          expectType<() => void>(res.abort)
+          expectType<() => Promise<{ name: string }>>(res.unwrap)
+          expectType<() => void>(res.unsubscribe)
+          expectType<(options: SubscriptionOptions) => void>(
+            res.updateSubscriptionOptions
+          )
+          expectType<() => void>(res.refetch)
+
+          // abort the query immediately to force an error
+          if (abort) res.abort()
+          res
+            .unwrap()
+            .then((result) => {
+              expectType<{ name: string }>(result)
+              setValues({
+                successMsg: `Successfully fetched user ${result.name}`,
+                errMsg: '',
+                isAborted: false,
+              })
+            })
+            .catch((err) => {
+              setValues({
+                successMsg: '',
+                errMsg: `An error has occurred fetching userId: ${res.arg}`,
+                isAborted: err.name === 'AbortError',
+              })
+            })
+        }
+
+        return (
+          <div>
+            <button onClick={handleClick(false)}>
+              Fetch User successfully
+            </button>
+            <button onClick={handleClick(true)}>Fetch User and abort</button>
+            <div>{successMsg}</div>
+            <div>{errMsg}</div>
+            <div>{isAborted ? 'Request was aborted' : ''}</div>
+          </div>
+        )
+      }
+
+      render(<User />, { wrapper: storeRef.wrapper })
+      expect(screen.queryByText(/An error has occurred/i)).toBeNull()
+      expect(screen.queryByText(/Successfully fetched user/i)).toBeNull()
+      expect(screen.queryByText('Request was aborted')).toBeNull()
+
+      fireEvent.click(
+        screen.getByRole('button', { name: 'Fetch User and abort' })
+      )
+      await screen.findByText('An error has occurred fetching userId: 1')
+      expect(screen.queryByText(/Successfully fetched user/i)).toBeNull()
+      screen.getByText('Request was aborted')
+
+      fireEvent.click(
+        screen.getByRole('button', { name: 'Fetch User successfully' })
+      )
+      await screen.findByText('Successfully fetched user Timmy')
+      expect(screen.queryByText(/An error has occurred/i)).toBeNull()
+      expect(screen.queryByText('Request was aborted')).toBeNull()
+    })
   })
 
   describe('useMutation', () => {
@@ -844,7 +939,7 @@ describe('hooks tests', () => {
 
           // no-op simply for clearer type assertions
           res.then((result) => {
-            expectExactType<
+            expectType<
               | {
                   error: { status: number; data: unknown } | SerializedError
                 }
@@ -998,7 +1093,7 @@ describe('hooks tests', () => {
       expect(
         api.endpoints.getUser.select(USER_ID)(storeRef.store.getState() as any)
       ).toEqual({
-        data: {},
+        data: { name: 'Timmy' },
         endpointName: 'getUser',
         error: undefined,
         fulfilledTimeStamp: expect.any(Number),
@@ -1019,7 +1114,7 @@ describe('hooks tests', () => {
       expect(
         api.endpoints.getUser.select(USER_ID)(storeRef.store.getState() as any)
       ).toEqual({
-        data: {},
+        data: { name: 'Timmy' },
         endpointName: 'getUser',
         fulfilledTimeStamp: expect.any(Number),
         isError: false,
@@ -1067,7 +1162,7 @@ describe('hooks tests', () => {
       expect(
         api.endpoints.getUser.select(USER_ID)(storeRef.store.getState() as any)
       ).toEqual({
-        data: {},
+        data: { name: 'Timmy' },
         endpointName: 'getUser',
         fulfilledTimeStamp: expect.any(Number),
         isError: false,
@@ -1085,7 +1180,7 @@ describe('hooks tests', () => {
       expect(
         api.endpoints.getUser.select(USER_ID)(storeRef.store.getState() as any)
       ).toEqual({
-        data: {},
+        data: { name: 'Timmy' },
         endpointName: 'getUser',
         fulfilledTimeStamp: expect.any(Number),
         isError: false,
@@ -1135,7 +1230,7 @@ describe('hooks tests', () => {
       expect(
         api.endpoints.getUser.select(USER_ID)(storeRef.store.getState() as any)
       ).toEqual({
-        data: {},
+        data: { name: 'Timmy' },
         endpointName: 'getUser',
         fulfilledTimeStamp: expect.any(Number),
         isError: false,
@@ -1155,7 +1250,7 @@ describe('hooks tests', () => {
       expect(
         api.endpoints.getUser.select(USER_ID)(storeRef.store.getState() as any)
       ).toEqual({
-        data: {},
+        data: { name: 'Timmy' },
         endpointName: 'getUser',
         fulfilledTimeStamp: expect.any(Number),
         isError: false,

yarn.lock

@@ -5058,6 +5058,7 @@ __metadata:
     "@types/json-stringify-safe": ^5.0.0
     "@types/nanoid": ^2.1.0
     "@types/node": ^10.14.4
+    "@types/query-string": ^6.3.0
     "@types/react": ^17.0.3
     "@types/react-dom": ^17.0.3
     "@types/react-redux": ^7.1.16
@@ -6133,6 +6134,15 @@ __metadata:
   languageName: node
   linkType: hard
 
+"@types/query-string@npm:^6.3.0":
+  version: 6.3.0
+  resolution: "@types/query-string@npm:6.3.0"
+  dependencies:
+    query-string: "*"
+  checksum: 7d507aea24e650548bc8a164ae695deb1eaf7a566fdaeb53ec3f89605dbd2c5201daf9b76c31843bc8accc01e866df8b2a06557514ad8014c3af99c664cd0b5c
+  languageName: node
+  linkType: hard
+
 "@types/react-dom@npm:17.0.0":
   version: 17.0.0
   resolution: "@types/react-dom@npm:17.0.0"
@@ -20272,17 +20282,7 @@ fsevents@^1.2.7:
   languageName: node
   linkType: hard
 
-"query-string@npm:^4.1.0":
-  version: 4.3.4
-  resolution: "query-string@npm:4.3.4"
-  dependencies:
-    object-assign: ^4.1.0
-    strict-uri-encode: ^1.0.0
-  checksum: 3b2bae6a8454cf0edf11cf1aa4d1f920398bbdabc1c39222b9bb92147e746fcd97faf00e56f494728fb66b2961b495ba0fde699d5d3bd06b11472d664b36c6cf
-  languageName: node
-  linkType: hard
-
-"query-string@npm:^7.0.1":
+"query-string@npm:*, query-string@npm:^7.0.1":
   version: 7.0.1
   resolution: "query-string@npm:7.0.1"
   dependencies:
@@ -20294,6 +20294,16 @@ fsevents@^1.2.7:
   languageName: node
   linkType: hard
 
+"query-string@npm:^4.1.0":
+  version: 4.3.4
+  resolution: "query-string@npm:4.3.4"
+  dependencies:
+    object-assign: ^4.1.0
+    strict-uri-encode: ^1.0.0
+  checksum: 3b2bae6a8454cf0edf11cf1aa4d1f920398bbdabc1c39222b9bb92147e746fcd97faf00e56f494728fb66b2961b495ba0fde699d5d3bd06b11472d664b36c6cf
+  languageName: node
+  linkType: hard
+
 "querystring-es3@npm:^0.2.0":
   version: 0.2.1
   resolution: "querystring-es3@npm:0.2.1"