Skip to content

Commit d9e481b

Browse files
authored
fix upsertQueryData race situations (#2646)
1 parent 50e1b8e commit d9e481b

File tree

4 files changed

+89
-22
lines changed

4 files changed

+89
-22
lines changed

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

Lines changed: 16 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@ import type {
66
ResultTypeFrom,
77
} from '../endpointDefinitions'
88
import { DefinitionType } from '../endpointDefinitions'
9-
import type { QueryThunk, MutationThunk } from './buildThunks'
9+
import type { QueryThunk, MutationThunk, QueryThunkArg } from './buildThunks'
1010
import type { AnyAction, ThunkAction, SerializedError } from '@reduxjs/toolkit'
1111
import type { SubscriptionOptions, RootState } from './apiState'
1212
import { QueryStatus } from './apiState'
@@ -35,6 +35,8 @@ declare module './module' {
3535
}
3636

3737
export const forceQueryFnSymbol = Symbol('forceQueryFn')
38+
export const isUpsertQuery = (arg: QueryThunkArg) =>
39+
typeof arg[forceQueryFnSymbol] === 'function'
3840

3941
export interface StartQueryActionCreatorOptions {
4042
subscribe?: boolean
@@ -301,13 +303,20 @@ Features like automatic cache collection, automatic refetching etc. will not be
301303
const skippedSynchronously = stateAfter.requestId !== requestId
302304

303305
const runningQuery = runningQueries[queryCacheKey]
306+
const selectFromState = () => selector(getState())
304307

305308
const statePromise: QueryActionCreatorResult<any> = Object.assign(
306-
skippedSynchronously && !runningQuery
307-
? Promise.resolve(stateAfter)
308-
: Promise.all([runningQuery, thunkResult]).then(() =>
309-
selector(getState())
310-
),
309+
forceQueryFn
310+
? // a query has been forced (upsertQueryData)
311+
// -> we want to resolve it once data has been written with the data that will be written
312+
thunkResult.then(selectFromState)
313+
: skippedSynchronously && !runningQuery
314+
? // a query has been skipped due to a condition and we do not have any currently running query
315+
// -> we want to resolve it immediately with the current data
316+
Promise.resolve(stateAfter)
317+
: // query just started or one is already in flight
318+
// -> wait for the running query, then resolve with data from after that
319+
Promise.all([runningQuery, thunkResult]).then(selectFromState),
311320
{
312321
arg,
313322
requestId,
@@ -350,7 +359,7 @@ Features like automatic cache collection, automatic refetching etc. will not be
350359
}
351360
)
352361

353-
if (!runningQuery && !skippedSynchronously) {
362+
if (!runningQuery && !skippedSynchronously && !forceQueryFn) {
354363
runningQueries[queryCacheKey] = statePromise
355364
statePromise.then(() => {
356365
delete runningQueries[queryCacheKey]

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

Lines changed: 13 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,7 @@ import {
3939
copyWithStructuralSharing,
4040
} from '../utils'
4141
import type { ApiContext } from '../apiTypes'
42+
import { isUpsertQuery } from './buildInitiate'
4243

4344
function updateQuerySubstateIfExists(
4445
state: QueryState<any>,
@@ -145,7 +146,13 @@ export function buildSlice({
145146

146147
updateQuerySubstateIfExists(draft, arg.queryCacheKey, (substate) => {
147148
substate.status = QueryStatus.pending
148-
substate.requestId = meta.requestId
149+
150+
substate.requestId =
151+
isUpsertQuery(arg) && substate.requestId
152+
? // for `upsertQuery` **updates**, keep the current `requestId`
153+
substate.requestId
154+
: // for normal queries or `upsertQuery` **inserts** always update the `requestId`
155+
meta.requestId
149156
if (arg.originalArgs !== undefined) {
150157
substate.originalArgs = arg.originalArgs
151158
}
@@ -157,14 +164,11 @@ export function buildSlice({
157164
draft,
158165
meta.arg.queryCacheKey,
159166
(substate) => {
160-
if (substate.requestId !== meta.requestId) {
161-
if (
162-
substate.fulfilledTimeStamp &&
163-
meta.fulfilledTimeStamp < substate.fulfilledTimeStamp
164-
) {
165-
return
166-
}
167-
}
167+
if (
168+
substate.requestId !== meta.requestId &&
169+
!isUpsertQuery(meta.arg)
170+
)
171+
return
168172
const { merge } = definitions[
169173
meta.arg.endpointName
170174
] as QueryDefinition<any, any, any, any>

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

Lines changed: 3 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -7,11 +7,11 @@ import type {
77
} from '../baseQueryTypes'
88
import type { RootState, QueryKeys, QuerySubstateIdentifier } from './apiState'
99
import { QueryStatus } from './apiState'
10-
import {
11-
forceQueryFnSymbol,
10+
import type {
1211
StartQueryActionCreatorOptions,
1312
QueryActionCreatorResult,
1413
} from './buildInitiate'
14+
import { forceQueryFnSymbol, isUpsertQuery } from './buildInitiate'
1515
import type {
1616
AssertTagTypes,
1717
EndpointDefinition,
@@ -482,9 +482,7 @@ In the case of an unhandled error, no tags will be "provided" or "invalidated".`
482482
// Order of these checks matters.
483483
// In order for `upsertQueryData` to successfully run while an existing request is in flight,
484484
/// we have to check for that first, otherwise `queryThunk` will bail out and not run at all.
485-
const isUpsertQuery =
486-
typeof arg[forceQueryFnSymbol] === 'function' && arg.forceRefetch
487-
if (isUpsertQuery) return true
485+
if (isUpsertQuery(arg)) return true
488486

489487
// Don't retry a request that's currently in-flight
490488
if (requestState?.status === 'pending') return false

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

Lines changed: 57 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -234,6 +234,62 @@ describe('upsertQueryData', () => {
234234
contents: 'I love cheese!',
235235
})
236236
})
237+
238+
test('upsert while a normal query is running (success)', async () => {
239+
const fetchedData = {
240+
id: '3',
241+
title: 'All about cheese.',
242+
contents: 'Yummy',
243+
}
244+
baseQuery.mockImplementation(() => delay(20).then(() => fetchedData))
245+
const upsertedData = {
246+
id: '3',
247+
title: 'Data from a SSR Render',
248+
contents: 'This is just some random data',
249+
}
250+
251+
const selector = api.endpoints.post.select('3')
252+
const fetchRes = storeRef.store.dispatch(api.endpoints.post.initiate('3'))
253+
const upsertRes = storeRef.store.dispatch(
254+
api.util.upsertQueryData('post', '3', upsertedData)
255+
)
256+
257+
await upsertRes
258+
let state = selector(storeRef.store.getState())
259+
expect(state.data).toEqual(upsertedData)
260+
261+
await fetchRes
262+
state = selector(storeRef.store.getState())
263+
expect(state.data).toEqual(fetchedData)
264+
})
265+
test('upsert while a normal query is running (rejected)', async () => {
266+
baseQuery.mockImplementation(async () => {
267+
await delay(20)
268+
// eslint-disable-next-line no-throw-literal
269+
throw 'Error!'
270+
})
271+
const upsertedData = {
272+
id: '3',
273+
title: 'Data from a SSR Render',
274+
contents: 'This is just some random data',
275+
}
276+
277+
const selector = api.endpoints.post.select('3')
278+
const fetchRes = storeRef.store.dispatch(api.endpoints.post.initiate('3'))
279+
const upsertRes = storeRef.store.dispatch(
280+
api.util.upsertQueryData('post', '3', upsertedData)
281+
)
282+
283+
await upsertRes
284+
let state = selector(storeRef.store.getState())
285+
expect(state.data).toEqual(upsertedData)
286+
expect(state.isSuccess).toBeTruthy()
287+
288+
await fetchRes
289+
state = selector(storeRef.store.getState())
290+
expect(state.data).toEqual(upsertedData)
291+
expect(state.isError).toBeTruthy()
292+
})
237293
})
238294

239295
describe('full integration', () => {
@@ -367,7 +423,7 @@ describe('full integration', () => {
367423
)
368424
})
369425

370-
test.only('Interop with in-flight requests', async () => {
426+
test('Interop with in-flight requests', async () => {
371427
await act(async () => {
372428
const fetchRes = storeRef.store.dispatch(
373429
api.endpoints.post2.initiate('3')

0 commit comments

Comments
 (0)