diff --git a/packages/toolkit/src/query/core/buildMiddleware/cacheCollection.ts b/packages/toolkit/src/query/core/buildMiddleware/cacheCollection.ts index 881fd9e11d..8b3ae6b198 100644 --- a/packages/toolkit/src/query/core/buildMiddleware/cacheCollection.ts +++ b/packages/toolkit/src/query/core/buildMiddleware/cacheCollection.ts @@ -44,12 +44,14 @@ export const buildCacheCollectionHandler: InternalHandlerBuilder = ({ context, internalState, }) => { - const { removeQueryResult, unsubscribeQueryResult } = api.internalActions + const { removeQueryResult, unsubscribeQueryResult, cacheEntriesUpserted } = + api.internalActions const canTriggerUnsubscribe = isAnyOf( unsubscribeQueryResult.match, queryThunk.fulfilled, queryThunk.rejected, + cacheEntriesUpserted.match, ) function anySubscriptionsRemainingForKey(queryCacheKey: string) { @@ -66,16 +68,27 @@ export const buildCacheCollectionHandler: InternalHandlerBuilder = ({ ) => { if (canTriggerUnsubscribe(action)) { const state = mwApi.getState()[reducerPath] - const { queryCacheKey } = unsubscribeQueryResult.match(action) - ? action.payload - : action.meta.arg - - handleUnsubscribe( - queryCacheKey, - state.queries[queryCacheKey]?.endpointName, - mwApi, - state.config, - ) + let queryCacheKeys: QueryCacheKey[] + + if (cacheEntriesUpserted.match(action)) { + queryCacheKeys = action.payload.map( + (entry) => entry.queryDescription.queryCacheKey, + ) + } else { + const { queryCacheKey } = unsubscribeQueryResult.match(action) + ? action.payload + : action.meta.arg + queryCacheKeys = [queryCacheKey] + } + + for (const queryCacheKey of queryCacheKeys) { + handleUnsubscribe( + queryCacheKey, + state.queries[queryCacheKey]?.endpointName, + mwApi, + state.config, + ) + } } if (api.util.resetApiState.match(action)) { @@ -132,6 +145,7 @@ export const buildCacheCollectionHandler: InternalHandlerBuilder = ({ if (currentTimeout) { clearTimeout(currentTimeout) } + currentRemovalTimeouts[queryCacheKey] = setTimeout(() => { if (!anySubscriptionsRemainingForKey(queryCacheKey)) { api.dispatch(removeQueryResult({ queryCacheKey })) diff --git a/packages/toolkit/src/query/core/buildMiddleware/cacheLifecycle.ts b/packages/toolkit/src/query/core/buildMiddleware/cacheLifecycle.ts index ebbbc727bb..72f84a0c54 100644 --- a/packages/toolkit/src/query/core/buildMiddleware/cacheLifecycle.ts +++ b/packages/toolkit/src/query/core/buildMiddleware/cacheLifecycle.ts @@ -183,6 +183,30 @@ export const buildCacheLifecycleHandler: InternalHandlerBuilder = ({ } const lifecycleMap: Record = {} + function resolveLifecycleEntry( + cacheKey: string, + data: unknown, + meta: unknown, + ) { + const lifecycle = lifecycleMap[cacheKey] + + if (lifecycle?.valueResolved) { + lifecycle.valueResolved({ + data, + meta, + }) + delete lifecycle.valueResolved + } + } + + function removeLifecycleEntry(cacheKey: string) { + const lifecycle = lifecycleMap[cacheKey] + if (lifecycle) { + delete lifecycleMap[cacheKey] + lifecycle.cacheEntryRemoved() + } + } + const handler: ApiMiddlewareInternalHandler = ( action, mwApi, @@ -190,17 +214,37 @@ export const buildCacheLifecycleHandler: InternalHandlerBuilder = ({ ) => { const cacheKey = getCacheKey(action) - if (queryThunk.pending.match(action)) { + function checkForNewCacheKey( + endpointName: string, + cacheKey: string, + requestId: string, + originalArgs: unknown, + ) { const oldState = stateBefore[reducerPath].queries[cacheKey] const state = mwApi.getState()[reducerPath].queries[cacheKey] if (!oldState && state) { - handleNewKey( - action.meta.arg.endpointName, - action.meta.arg.originalArgs, - cacheKey, - mwApi, + handleNewKey(endpointName, originalArgs, cacheKey, mwApi, requestId) + } + } + + if (queryThunk.pending.match(action)) { + checkForNewCacheKey( + action.meta.arg.endpointName, + cacheKey, + action.meta.requestId, + action.meta.arg.originalArgs, + ) + } else if (api.internalActions.cacheEntriesUpserted.match(action)) { + for (const { queryDescription, value } of action.payload) { + const { endpointName, originalArgs, queryCacheKey } = queryDescription + checkForNewCacheKey( + endpointName, + queryCacheKey, action.meta.requestId, + originalArgs, ) + + resolveLifecycleEntry(queryCacheKey, value, {}) } } else if (mutationThunk.pending.match(action)) { const state = mwApi.getState()[reducerPath].mutations[cacheKey] @@ -214,27 +258,15 @@ export const buildCacheLifecycleHandler: InternalHandlerBuilder = ({ ) } } else if (isFulfilledThunk(action)) { - const lifecycle = lifecycleMap[cacheKey] - if (lifecycle?.valueResolved) { - lifecycle.valueResolved({ - data: action.payload, - meta: action.meta.baseQueryMeta, - }) - delete lifecycle.valueResolved - } + resolveLifecycleEntry(cacheKey, action.payload, action.meta.baseQueryMeta) } else if ( api.internalActions.removeQueryResult.match(action) || api.internalActions.removeMutationResult.match(action) ) { - const lifecycle = lifecycleMap[cacheKey] - if (lifecycle) { - delete lifecycleMap[cacheKey] - lifecycle.cacheEntryRemoved() - } + removeLifecycleEntry(cacheKey) } else if (api.util.resetApiState.match(action)) { - for (const [cacheKey, lifecycle] of Object.entries(lifecycleMap)) { - delete lifecycleMap[cacheKey] - lifecycle.cacheEntryRemoved() + for (const cacheKey of Object.keys(lifecycleMap)) { + removeLifecycleEntry(cacheKey) } } } diff --git a/packages/toolkit/src/query/core/buildSlice.ts b/packages/toolkit/src/query/core/buildSlice.ts index 77daeb7729..c9652bf766 100644 --- a/packages/toolkit/src/query/core/buildSlice.ts +++ b/packages/toolkit/src/query/core/buildSlice.ts @@ -76,6 +76,11 @@ type NormalizedQueryUpsertEntryPayload = { value: any } +export type ProcessedQueryUpsertEntry = { + queryDescription: QueryThunkArg + value: unknown +} + /** * A typesafe representation of a util action creator that accepts cache entry descriptions to upsert */ @@ -90,7 +95,7 @@ export type UpsertEntries = < > }, ], -) => PayloadAction +) => PayloadAction function updateQuerySubstateIfExists( state: QueryState, @@ -276,7 +281,7 @@ export function buildSlice({ reducer( draft, action: PayloadAction< - NormalizedQueryUpsertEntryPayload[], + ProcessedQueryUpsertEntry[], string, { RTK_autoBatch: boolean @@ -286,19 +291,7 @@ export function buildSlice({ >, ) { for (const entry of action.payload) { - const { endpointName, args, value } = entry - const endpointDefinition = definitions[endpointName] - - const arg: QueryThunkArg = { - type: 'query', - endpointName: endpointName, - originalArgs: entry.args, - queryCacheKey: serializeQueryArgs({ - queryArgs: args, - endpointDefinition, - endpointName, - }), - } + const { queryDescription: arg, value } = entry writePendingCacheEntry(draft, arg, true, { arg, requestId: action.meta.requestId, @@ -313,13 +306,31 @@ export function buildSlice({ fulfilledTimeStamp: action.meta.timestamp, baseQueryMeta: {}, }, - entry.value, + value, ) } }, prepare: (payload: NormalizedQueryUpsertEntryPayload[]) => { + const queryDescriptions: ProcessedQueryUpsertEntry[] = payload.map( + (entry) => { + const { endpointName, args, value } = entry + const endpointDefinition = definitions[endpointName] + const queryDescription: QueryThunkArg = { + type: 'query', + endpointName: endpointName, + originalArgs: entry.args, + queryCacheKey: serializeQueryArgs({ + queryArgs: args, + endpointDefinition, + endpointName, + }), + } + return { queryDescription, value } + }, + ) + const result = { - payload, + payload: queryDescriptions, meta: { [SHOULD_AUTOBATCH]: true, requestId: nanoid(), diff --git a/packages/toolkit/src/query/tests/optimisticUpserts.test.tsx b/packages/toolkit/src/query/tests/optimisticUpserts.test.tsx index 673266ab16..5c0b244390 100644 --- a/packages/toolkit/src/query/tests/optimisticUpserts.test.tsx +++ b/packages/toolkit/src/query/tests/optimisticUpserts.test.tsx @@ -1,4 +1,5 @@ import { createApi } from '@reduxjs/toolkit/query/react' +import { createAction } from '@reduxjs/toolkit' import { actionsReducer, hookWaitFor, @@ -16,6 +17,8 @@ interface Post { const baseQuery = vi.fn() beforeEach(() => baseQuery.mockReset()) +const postAddedAction = createAction('postAdded') + const api = createApi({ baseQuery: (...args: any[]) => { const result = baseQuery(...args) @@ -65,6 +68,18 @@ const api = createApi({ } }, }), + postWithSideEffect: build.query({ + query: (id) => `post/${id}`, + providesTags: ['Post'], + async onCacheEntryAdded(arg, api) { + // Verify that lifecycle promise resolution works + const res = await api.cacheDataLoaded + + // and leave a side effect we can check in the test + api.dispatch(postAddedAction(res.data.id)) + }, + keepUnusedDataFor: 0.01, + }), }), }) @@ -331,47 +346,86 @@ describe('upsertQueryData', () => { }) describe('upsertEntries', () => { + const posts: Post[] = [ + { + id: '1', + contents: 'A', + title: 'A', + }, + { + id: '2', + contents: 'B', + title: 'B', + }, + { + id: '3', + contents: 'C', + title: 'C', + }, + ] + + const entriesAction = api.util.upsertEntries([ + { + endpointName: 'getPosts', + args: undefined, + value: posts, + }, + ...posts.map((post) => ({ + endpointName: 'postWithSideEffect' as const, + args: post.id, + value: post, + })), + ]) + test('Upserts many entries at once', async () => { - const posts: Post[] = [ - { - id: '1', - contents: 'A', - title: 'A', - }, - { - id: '2', - contents: 'B', - title: 'B', - }, - { - id: '3', - contents: 'C', - title: 'C', - }, - ] - - storeRef.store.dispatch( - api.util.upsertEntries([ - { - endpointName: 'getPosts', - args: undefined, - value: posts, - }, - ...posts.map((post) => ({ - endpointName: 'post' as const, - args: post.id, - value: post, - })), - ]), - ) + storeRef.store.dispatch(entriesAction) const state = storeRef.store.getState() expect(api.endpoints.getPosts.select()(state).data).toBe(posts) - expect(api.endpoints.post.select('1')(state).data).toBe(posts[0]) - expect(api.endpoints.post.select('2')(state).data).toBe(posts[1]) - expect(api.endpoints.post.select('3')(state).data).toBe(posts[2]) + expect(api.endpoints.postWithSideEffect.select('1')(state).data).toBe( + posts[0], + ) + expect(api.endpoints.postWithSideEffect.select('2')(state).data).toBe( + posts[1], + ) + expect(api.endpoints.postWithSideEffect.select('3')(state).data).toBe( + posts[2], + ) + }) + + test('Triggers cache lifecycles and side effects', async () => { + storeRef.store.dispatch(entriesAction) + + // Tricky timing. The cache data promises will be resolved + // in microtasks, so we just need any async delay here. + await delay(1) + + const state = storeRef.store.getState() + + // onCacheEntryAdded should have run for each post, + // including cache data being resolved + for (const post of posts) { + const matchingSideEffectAction = state.actions.find( + (action) => postAddedAction.match(action) && action.payload === post.id, + ) + expect(matchingSideEffectAction).toBeTruthy() + } + + expect(api.endpoints.postWithSideEffect.select('1')(state).data).toBe( + posts[0], + ) + + // The cache data should be removed after the keepUnusedDataFor time, + // so wait longer than that + await delay(20) + + const stateAfter = storeRef.store.getState() + + expect(api.endpoints.postWithSideEffect.select('1')(stateAfter).data).toBe( + undefined, + ) }) })