Enable upsertQueryData calls while a request is in flight

Author: markerikson

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

@@ -157,7 +157,14 @@ export function buildSlice({
             draft,
             meta.arg.queryCacheKey,
             (substate) => {
-              if (substate.requestId !== meta.requestId) return
+              if (substate.requestId !== meta.requestId) {
+                if (
+                  substate.fulfilledTimeStamp &&
+                  meta.fulfilledTimeStamp < substate.fulfilledTimeStamp
+                ) {
+                  return
+                }
+              }
               const { merge } = definitions[
                 meta.arg.endpointName
               ] as QueryDefinition<any, any, any, any>

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

@@ -10,6 +10,7 @@ import { QueryStatus } from './apiState'
 import {
   forceQueryFnSymbol,
   StartQueryActionCreatorOptions,
+  QueryActionCreatorResult,
 } from './buildInitiate'
 import type {
   AssertTagTypes,
@@ -176,7 +177,12 @@ export type UpsertQueryDataThunk<
   endpointName: EndpointName,
   args: QueryArgFrom<Definitions[EndpointName]>,
   value: ResultTypeFrom<Definitions[EndpointName]>
-) => ThunkAction<void, PartialState, any, AnyAction>
+) => ThunkAction<
+  QueryActionCreatorResult<Definitions[EndpointName]>,
+  PartialState,
+  any,
+  AnyAction
+>
 
 /**
  * An object returned from dispatching a `api.util.updateQueryData` call.
@@ -272,7 +278,7 @@ export function buildThunks<
 
   const upsertQueryData: UpsertQueryDataThunk<EndpointDefinitions, State> =
     (endpointName, args, value) => (dispatch) => {
-      dispatch(
+      return dispatch(
         (
           api.endpoints[endpointName] as ApiEndpointQuery<
             QueryDefinition<any, any, any, any, any>,
@@ -469,12 +475,17 @@ In the case of an unhandled error, no tags will be "provided" or "invalidated".`
       const requestState = state[reducerPath]?.queries?.[arg.queryCacheKey]
       const fulfilledVal = requestState?.fulfilledTimeStamp
 
-      // Don't retry a request that's currently in-flight
-      if (requestState?.status === 'pending') return false
+      // Order of these checks matters.
+      // In order for `upsertQueryData` to successfully run while an existing request is
+      /// in flight, we have to check `isForcedQuery` before `status === 'pending'`,
+      // otherwise `queryThunk` will bail out and not run at all.
 
       // if this is forced, continue
       if (isForcedQuery(arg, state)) return true
 
+      // Don't retry a request that's currently in-flight
+      if (requestState?.status === 'pending') return false
+
       // Pull from the cache unless we explicitly force refetch or qualify based on time
       if (fulfilledVal)
         // Value is cached and we didn't specify to refresh, skip it.

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

@@ -1,7 +1,7 @@
 import { createApi } from '@reduxjs/toolkit/query/react'
 import { actionsReducer, hookWaitFor, setupApiStore, waitMs } from './helpers'
 import { skipToken } from '../core/buildSelectors'
-import { renderHook, act } from '@testing-library/react'
+import { renderHook, act, waitFor } from '@testing-library/react'
 
 interface Post {
   id: string
@@ -12,6 +12,10 @@ interface Post {
 const baseQuery = jest.fn()
 beforeEach(() => baseQuery.mockReset())
 
+function delay(ms: number) {
+  return new Promise((resolve) => setTimeout(resolve, ms))
+}
+
 const api = createApi({
   baseQuery: (...args: any[]) => {
     const result = baseQuery(...args)
@@ -46,6 +50,18 @@ const api = createApi({
       },
       invalidatesTags: (result) => (result ? ['Post'] : []),
     }),
+    post2: build.query<Post, string>({
+      queryFn: async (id) => {
+        await delay(20)
+        return {
+          data: {
+            id,
+            title: 'All about cheese.',
+            contents: 'TODO',
+          },
+        }
+      },
+    }),
   }),
 })
 
@@ -195,10 +211,9 @@ describe('upsertQueryData', () => {
     )
     await hookWaitFor(() => expect(result.current.isError).toBeTruthy())
 
-    let returnValue!: ReturnType<ReturnType<typeof api.util.upsertQueryData>>
     // upsert the data
     act(() => {
-      returnValue = storeRef.store.dispatch(
+      storeRef.store.dispatch(
         api.util.upsertQueryData('post', '4', {
           id: '4',
           title: 'All about cheese',
@@ -351,4 +366,44 @@ describe('full integration', () => {
       50
     )
   })
+
+  test.only('Interop with in-flight requests', async () => {
+    await act(async () => {
+      const fetchRes = storeRef.store.dispatch(
+        api.endpoints.post2.initiate('3')
+      )
+
+      const upsertRes = storeRef.store.dispatch(
+        api.util.upsertQueryData('post2', '3', {
+          id: '3',
+          title: 'Upserted title',
+          contents: 'Upserted contents',
+        })
+      )
+
+      const selectEntry = api.endpoints.post2.select('3')
+      await waitFor(
+        () => {
+          const entry1 = selectEntry(storeRef.store.getState())
+          expect(entry1.data).toEqual({
+            id: '3',
+            title: 'Upserted title',
+            contents: 'Upserted contents',
+          })
+        },
+        { interval: 1, timeout: 15 }
+      )
+      await waitFor(
+        () => {
+          const entry2 = selectEntry(storeRef.store.getState())
+          expect(entry2.data).toEqual({
+            id: '3',
+            title: 'All about cheese.',
+            contents: 'TODO',
+          })
+        },
+        { interval: 1 }
+      )
+    })
+  })
 })