Merge pull request #3255 from dutzi/pr/updateQueryData-provide

markerikson · web-flow · commit 767fedeea55d · 2023-09-23T20:38:55.000-04:00
diff --git a/docs/rtk-query/api/created-api/api-slice-utils.mdx b/docs/rtk-query/api/created-api/api-slice-utils.mdx
@@ -29,7 +29,8 @@ Some of the TS types on this page are pseudocode to illustrate intent, as the ac
 const updateQueryData = (
   endpointName: string,
   args: any,
-  updateRecipe: (draft: Draft<CachedState>) => void
+  updateRecipe: (draft: Draft<CachedState>) => void,
+  updateProvided?: boolean
 ) => ThunkAction<PatchCollection, PartialState, any, AnyAction>;
 
 interface PatchCollection {
@@ -43,6 +44,7 @@ interface PatchCollection {
   - `endpointName`: a string matching an existing endpoint name
   - `args`: an argument matching that used for a previous query call, used to determine which cached dataset needs to be updated
   - `updateRecipe`: an Immer `produce` callback that can apply changes to the cached state
+  - `updateProvided`: a boolean indicating whether the endpoint's provided tags should be re-calculated based on the updated cache. Defaults to `false`.
 
 #### Description
 
@@ -155,14 +157,16 @@ await dispatch(
 const patchQueryData = (
   endpointName: string,
   args: any
-  patches: Patch[]
+  patches: Patch[],
+  updateProvided?: boolean
 ) => ThunkAction<void, PartialState, any, AnyAction>;
 ```
 
 - **Parameters**
   - `endpointName`: a string matching an existing endpoint name
   - `args`: a cache key, used to determine which cached dataset needs to be updated
   - `patches`: an array of patches (or inverse patches) to apply to cached state. These would typically be obtained from the result of dispatching [`updateQueryData`](#updatequerydata)
+  - `updateProvided`: a boolean indicating whether the endpoint's provided tags should be re-calculated based on the updated cache. Defaults to `false`.
 
 #### Description
 
@@ -229,42 +233,42 @@ dispatch(api.util.prefetch('getPosts', undefined, { force: true }))
 ```
 
 ### `selectInvalidatedBy`
- 
+
 #### Signature
- 
+
 ```ts no-transpile
- function selectInvalidatedBy(
-   state: RootState,
-   tags: ReadonlyArray<TagDescription<string>>
- ): Array<{
-   endpointName: string
-   originalArgs: any
-   queryCacheKey: QueryCacheKey
- }>
+function selectInvalidatedBy(
+  state: RootState,
+  tags: ReadonlyArray<TagDescription<string>>
+): Array<{
+  endpointName: string
+  originalArgs: any
+  queryCacheKey: QueryCacheKey
+}>
 ```
- 
+
 - **Parameters**
   - `state`: the root state
   - `tags`: a readonly array of invalidated tags, where the provided `TagDescription` is one of the strings provided to the [`tagTypes`](../createApi.mdx#tagtypes) property of the api. e.g.
     - `[TagType]`
     - `[{ type: TagType }]`
     - `[{ type: TagType, id: number | string }]`
- 
+
 #### Description
- 
+
 A function that can select query parameters to be invalidated.
- 
+
 The function accepts two arguments
-  - the root state and
-  - the cache tags to be invalidated.
- 
+- the root state and
+- the cache tags to be invalidated.
+
 It returns an array that contains
-  - the endpoint name,
-  - the original args and
-  - the queryCacheKey.
- 
+- the endpoint name,
+- the original args and
+- the queryCacheKey.
+
 #### Example
- 
+
 ```ts no-transpile
 dispatch(api.util.selectInvalidatedBy(state, ['Post']))
 dispatch(api.util.selectInvalidatedBy(state, [{ type: 'Post', id: 1 }]))
diff --git a/packages/toolkit/src/query/core/buildSlice.ts b/packages/toolkit/src/query/core/buildSlice.ts
@@ -29,6 +29,7 @@ import { calculateProvidedByThunk } from './buildThunks'
 import type {
   AssertTagTypes,
   EndpointDefinitions,
+  FullTagDescription,
   QueryDefinition,
 } from '../endpointDefinitions'
 import type { Patch } from 'immer'
@@ -125,17 +126,22 @@ export function buildSlice({
         },
         prepare: prepareAutoBatched<QuerySubstateIdentifier>(),
       },
-      queryResultPatched(
-        draft,
-        {
-          payload: { queryCacheKey, patches },
-        }: PayloadAction<
+      queryResultPatched: {
+        reducer(
+          draft,
+          {
+            payload: { queryCacheKey, patches },
+          }: PayloadAction<
+            QuerySubstateIdentifier & { patches: readonly Patch[] }
+          >
+        ) {
+          updateQuerySubstateIfExists(draft, queryCacheKey, (substate) => {
+            substate.data = applyPatches(substate.data as any, patches.concat())
+          })
+        },
+        prepare: prepareAutoBatched<
           QuerySubstateIdentifier & { patches: readonly Patch[] }
-        >
-      ) {
-        updateQuerySubstateIfExists(draft, queryCacheKey, (substate) => {
-          substate.data = applyPatches(substate.data as any, patches.concat())
-        })
+        >(),
       },
     },
     extraReducers(builder) {
@@ -325,7 +331,42 @@ export function buildSlice({
   const invalidationSlice = createSlice({
     name: `${reducerPath}/invalidation`,
     initialState: initialState as InvalidationState<string>,
-    reducers: {},
+    reducers: {
+      updateProvidedBy: {
+        reducer(
+          draft,
+          action: PayloadAction<{
+            queryCacheKey: QueryCacheKey
+            providedTags: readonly FullTagDescription<string>[]
+          }>
+        ) {
+          const { queryCacheKey, providedTags } = action.payload
+
+          for (const tagTypeSubscriptions of Object.values(draft)) {
+            for (const idSubscriptions of Object.values(tagTypeSubscriptions)) {
+              const foundAt = idSubscriptions.indexOf(queryCacheKey)
+              if (foundAt !== -1) {
+                idSubscriptions.splice(foundAt, 1)
+              }
+            }
+          }
+
+          for (const { type, id } of providedTags) {
+            const subscribedQueries = ((draft[type] ??= {})[
+              id || '__internal_without_id'
+            ] ??= [])
+            const alreadySubscribed = subscribedQueries.includes(queryCacheKey)
+            if (!alreadySubscribed) {
+              subscribedQueries.push(queryCacheKey)
+            }
+          }
+        },
+        prepare: prepareAutoBatched<{
+          queryCacheKey: QueryCacheKey
+          providedTags: readonly FullTagDescription<string>[]
+        }>(),
+      },
+    },
     extraReducers(builder) {
       builder
         .addCase(
@@ -371,27 +412,13 @@ export function buildSlice({
             )
             const { queryCacheKey } = action.meta.arg
 
-            for (const tagTypeSubscriptions of Object.values(draft)) {
-              for (const idSubscriptions of Object.values(
-                tagTypeSubscriptions
-              )) {
-                const foundAt = idSubscriptions.indexOf(queryCacheKey)
-                if (foundAt !== -1) {
-                  idSubscriptions.splice(foundAt, 1)
-                }
-              }
-            }
-
-            for (const { type, id } of providedTags) {
-              const subscribedQueries = ((draft[type] ??= {})[
-                id || '__internal_without_id'
-              ] ??= [])
-              const alreadySubscribed =
-                subscribedQueries.includes(queryCacheKey)
-              if (!alreadySubscribed) {
-                subscribedQueries.push(queryCacheKey)
-              }
-            }
+            invalidationSlice.caseReducers.updateProvidedBy(
+              draft,
+              invalidationSlice.actions.updateProvidedBy({
+                queryCacheKey,
+                providedTags,
+              })
+            )
           }
         )
     },
@@ -497,6 +524,7 @@ export function buildSlice({
     ...subscriptionSlice.actions,
     ...internalSubscriptionsSlice.actions,
     ...mutationSlice.actions,
+    ...invalidationSlice.actions,
     /** @deprecated has been renamed to `removeMutationResult` */
     unsubscribeMutationResult: mutationSlice.actions.removeMutationResult,
     resetApiState,
diff --git a/packages/toolkit/src/query/core/buildThunks.ts b/packages/toolkit/src/query/core/buildThunks.ts
@@ -20,6 +20,7 @@ import type {
   QueryArgFrom,
   QueryDefinition,
   ResultTypeFrom,
+  FullTagDescription,
 } from '../endpointDefinitions'
 import { isQueryDefinition } from '../endpointDefinitions'
 import { calculateProvidedBy } from '../endpointDefinitions'
@@ -164,7 +165,8 @@ export type PatchQueryDataThunk<
 > = <EndpointName extends QueryKeys<Definitions>>(
   endpointName: EndpointName,
   args: QueryArgFrom<Definitions[EndpointName]>,
-  patches: readonly Patch[]
+  patches: readonly Patch[],
+  updateProvided?: boolean
 ) => ThunkAction<void, PartialState, any, AnyAction>
 
 export type UpdateQueryDataThunk<
@@ -173,7 +175,8 @@ export type UpdateQueryDataThunk<
 > = <EndpointName extends QueryKeys<Definitions>>(
   endpointName: EndpointName,
   args: QueryArgFrom<Definitions[EndpointName]>,
-  updateRecipe: Recipe<ResultTypeFrom<Definitions[EndpointName]>>
+  updateRecipe: Recipe<ResultTypeFrom<Definitions[EndpointName]>>,
+  updateProvided?: boolean
 ) => ThunkAction<PatchCollection, PartialState, any, AnyAction>
 
 export type UpsertQueryDataThunk<
@@ -222,57 +225,87 @@ export function buildThunks<
   context: { endpointDefinitions },
   serializeQueryArgs,
   api,
+  assertTagType,
 }: {
   baseQuery: BaseQuery
   reducerPath: ReducerPath
   context: ApiContext<Definitions>
   serializeQueryArgs: InternalSerializeQueryArgs
   api: Api<BaseQuery, Definitions, ReducerPath, any>
+  assertTagType: AssertTagTypes
 }) {
   type State = RootState<any, string, ReducerPath>
 
   const patchQueryData: PatchQueryDataThunk<EndpointDefinitions, State> =
-    (endpointName, args, patches) => (dispatch) => {
+    (endpointName, args, patches, updateProvided) => (dispatch, getState) => {
       const endpointDefinition = endpointDefinitions[endpointName]
+
+      const queryCacheKey = serializeQueryArgs({
+        queryArgs: args,
+        endpointDefinition,
+        endpointName,
+      })
+
       dispatch(
-        api.internalActions.queryResultPatched({
-          queryCacheKey: serializeQueryArgs({
-            queryArgs: args,
-            endpointDefinition,
-            endpointName,
-          }),
-          patches,
-        })
+        api.internalActions.queryResultPatched({ queryCacheKey, patches })
+      )
+
+      if (!updateProvided) {
+        return
+      }
+
+      const newValue = api.endpoints[endpointName].select(args)(getState())
+
+      const providedTags = calculateProvidedBy(
+        endpointDefinition.providesTags,
+        newValue.data,
+        undefined,
+        args,
+        {},
+        assertTagType
+      )
+
+      dispatch(
+        api.internalActions.updateProvidedBy({ queryCacheKey, providedTags })
       )
     }
 
   const updateQueryData: UpdateQueryDataThunk<EndpointDefinitions, State> =
-    (endpointName, args, updateRecipe) => (dispatch, getState) => {
-      const currentState = (
-        api.endpoints[endpointName] as ApiEndpointQuery<any, any>
-      ).select(args)(getState())
+    (endpointName, args, updateRecipe, updateProvided = true) =>
+    (dispatch, getState) => {
+      const endpointDefinition = api.endpoints[endpointName]
+
+      const currentState = endpointDefinition.select(args)(getState())
+
       let ret: PatchCollection = {
         patches: [],
         inversePatches: [],
         undo: () =>
           dispatch(
-            api.util.patchQueryData(endpointName, args, ret.inversePatches)
+            api.util.patchQueryData(
+              endpointName,
+              args,
+              ret.inversePatches,
+              updateProvided
+            )
           ),
       }
       if (currentState.status === QueryStatus.uninitialized) {
         return ret
       }
+      let newValue
       if ('data' in currentState) {
         if (isDraftable(currentState.data)) {
-          const [, patches, inversePatches] = produceWithPatches(
+          const [value, patches, inversePatches] = produceWithPatches(
             currentState.data,
             updateRecipe
           )
           ret.patches.push(...patches)
           ret.inversePatches.push(...inversePatches)
+          newValue = value
         } else {
-          const value = updateRecipe(currentState.data)
-          ret.patches.push({ op: 'replace', path: [], value })
+          newValue = updateRecipe(currentState.data)
+          ret.patches.push({ op: 'replace', path: [], value: newValue })
           ret.inversePatches.push({
             op: 'replace',
             path: [],
@@ -281,7 +314,9 @@ export function buildThunks<
         }
       }
 
-      dispatch(api.util.patchQueryData(endpointName, args, ret.patches))
+      dispatch(
+        api.util.patchQueryData(endpointName, args, ret.patches, updateProvided)
+      )
 
       return ret
     }
diff --git a/packages/toolkit/src/query/core/module.ts b/packages/toolkit/src/query/core/module.ts
@@ -518,6 +518,7 @@ export const coreModule = (): Module<CoreModule> => ({
       context,
       api,
       serializeQueryArgs,
+      assertTagType,
     })
 
     const { reducer, actions: sliceActions } = buildSlice({
diff --git a/packages/toolkit/src/query/tests/optimisticUpdates.test.tsx b/packages/toolkit/src/query/tests/optimisticUpdates.test.tsx
