Add providesTags handling for upsertQueryEntries (#4872)

* Add providesTags handling for upsertQueryEntries

* Fix flakiness in infinite query polling test

Author: markerikson

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

@@ -1,4 +1,4 @@
-import type { Action, PayloadAction, UnknownAction } from '@reduxjs/toolkit'
+import type { PayloadAction } from '@reduxjs/toolkit'
 import {
   combineReducers,
   createAction,
@@ -23,7 +23,6 @@ import type {
   QueryCacheKey,
   SubscriptionState,
   ConfigState,
-  QueryKeys,
   InfiniteQuerySubState,
   InfiniteQueryDirection,
 } from './apiState'
@@ -36,18 +35,14 @@ import type {
   MutationThunk,
   QueryThunk,
   QueryThunkArg,
-  RejectedAction,
 } from './buildThunks'
 import { calculateProvidedByThunk } from './buildThunks'
 import {
   isInfiniteQueryDefinition,
   type AssertTagTypes,
-  type DefinitionType,
   type EndpointDefinitions,
   type FullTagDescription,
-  type QueryArgFrom,
   type QueryDefinition,
-  type ResultTypeFrom,
 } from '../endpointDefinitions'
 import type { Patch } from 'immer'
 import { isDraft } from 'immer'
@@ -61,6 +56,7 @@ import {
 import type { ApiContext } from '../apiTypes'
 import { isUpsertQuery } from './buildInitiate'
 import type { InternalSerializeQueryArgs } from '../defaultSerializeQueryArgs'
+import type { UnwrapPromise } from '../tsHelpers'
 
 /**
  * A typesafe single entry to be upserted into the cache
@@ -279,6 +275,7 @@ export function buildSlice({
       substate.fulfilledTimeStamp = meta.fulfilledTimeStamp
     })
   }
+
   const querySlice = createSlice({
     name: `${reducerPath}/queries`,
     initialState: initialState as QueryState<any>,
@@ -486,6 +483,11 @@ export function buildSlice({
     },
   })
 
+  type CalculateProvidedByAction = UnwrapPromise<
+    | ReturnType<ReturnType<QueryThunk>>
+    | ReturnType<ReturnType<InfiniteQueryThunk<any>>>
+  >
+
   const invalidationSlice = createSlice({
     name: `${reducerPath}/invalidation`,
     initialState: initialState as InvalidationState<string>,
@@ -562,26 +564,51 @@ export function buildSlice({
         .addMatcher(
           isAnyOf(isFulfilled(queryThunk), isRejectedWithValue(queryThunk)),
           (draft, action) => {
-            const providedTags = calculateProvidedByThunk(
-              action,
-              'providesTags',
-              definitions,
-              assertTagType,
-            )
-            const { queryCacheKey } = action.meta.arg
+            writeProvidedTagsForQuery(draft, action)
+          },
+        )
+        .addMatcher(
+          querySlice.actions.cacheEntriesUpserted.match,
+          (draft, action) => {
+            for (const { queryDescription: arg, value } of action.payload) {
+              const action: CalculateProvidedByAction = {
+                type: 'UNKNOWN',
+                payload: value,
+                meta: {
+                  requestStatus: 'fulfilled',
+                  requestId: 'UNKNOWN',
+                  arg,
+                },
+              }
 
-            invalidationSlice.caseReducers.updateProvidedBy(
-              draft,
-              invalidationSlice.actions.updateProvidedBy({
-                queryCacheKey,
-                providedTags,
-              }),
-            )
+              writeProvidedTagsForQuery(draft, action)
+            }
           },
         )
     },
   })
 
+  function writeProvidedTagsForQuery(
+    draft: InvalidationState<string>,
+    action: CalculateProvidedByAction,
+  ) {
+    const providedTags = calculateProvidedByThunk(
+      action,
+      'providesTags',
+      definitions,
+      assertTagType,
+    )
+    const { queryCacheKey } = action.meta.arg
+
+    invalidationSlice.caseReducers.updateProvidedBy(
+      draft,
+      invalidationSlice.actions.updateProvidedBy({
+        queryCacheKey,
+        providedTags,
+      }),
+    )
+  }
+
   // Dummy slice to generate actions
   const subscriptionSlice = createSlice({
     name: `${reducerPath}/subscriptions`,

packages/toolkit/src/query/tests/infiniteQueries.test.ts

@@ -567,10 +567,10 @@ describe('Infinite queries', () => {
     ])
 
     thirdPromise.updateSubscriptionOptions({
-      pollingInterval: 10,
+      pollingInterval: 50,
     })
 
-    await delay(5)
+    await delay(25)
 
     let entry = countersApi.endpoints.counters.select('item')(
       storeRef.store.getState(),
@@ -582,7 +582,7 @@ describe('Infinite queries', () => {
       { page: 5, hitCounter: 3 },
     ])
 
-    await delay(10)
+    await delay(50)
 
     entry = countersApi.endpoints.counters.select('item')(
       storeRef.store.getState(),

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

@@ -81,7 +81,12 @@ const api = createApi({
     }),
     postWithSideEffect: build.query<Post, string>({
       query: (id) => `post/${id}`,
-      providesTags: ['Post'],
+      providesTags: (result) => {
+        if (result) {
+          return [{ type: 'Post', id: result.id } as const]
+        }
+        return []
+      },
       async onCacheEntryAdded(arg, api) {
         // Verify that lifecycle promise resolution works
         const res = await api.cacheDataLoaded
@@ -419,15 +424,16 @@ describe('upsertQueryEntries', () => {
 
     expect(api.endpoints.getPosts.select()(state).data).toBe(posts)
 
-    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],
-    )
+    for (const post of posts) {
+      expect(api.endpoints.postWithSideEffect.select(post.id)(state).data).toBe(
+        post,
+      )
+
+      // Should have added tags
+      expect(state.api.provided.Post[post.id]).toEqual([
+        `postWithSideEffect("${post.id}")`,
+      ])
+    }
   })
 
   test('Triggers cache lifecycles and side effects', async () => {