build a patchQueriesData thunk to batch patching of endpoints

reduxjs · Jan 22, 2024 · fc4e644 · fc4e644
1 parent c3cf5a7
commit fc4e644
Show file tree

Hide file tree

Showing 4 changed files with 143 additions and 72 deletions.
diff --git a/packages/toolkit/package.json b/packages/toolkit/package.json
@@ -101,6 +101,7 @@
     "format:check": "prettier --list-different \"(src|examples)/**/*.{ts,tsx}\" \"docs/*/**.md\"",
     "lint": "eslint src examples",
     "test": "vitest --run",
+    "test:watch": "vitest --watch",
     "type-tests": "yarn tsc -p src/tests/tsconfig.typetests.json && yarn tsc -p src/query/tests/tsconfig.typetests.json",
     "prepack": "yarn build"
   },

diff --git a/packages/toolkit/src/query/core/buildSlice.ts b/packages/toolkit/src/query/core/buildSlice.ts
@@ -125,22 +125,28 @@ export function buildSlice({
         },
         prepare: prepareAutoBatched<QuerySubstateIdentifier>(),
       },
-      queryResultPatched: {
+      queryResultsPatched: {
         reducer(
           draft,
           {
-            payload: { queryCacheKey, patches },
+            payload,
           }: PayloadAction<
-            QuerySubstateIdentifier & { patches: readonly Patch[] }
+            Array<QuerySubstateIdentifier & { patches: readonly Patch[] }>
           >
         ) {
-          updateQuerySubstateIfExists(draft, queryCacheKey, (substate) => {
-            substate.data = applyPatches(substate.data as any, patches.concat())
-          })
+          for (const { queryCacheKey, patches } of payload) {
+            updateQuerySubstateIfExists(draft, queryCacheKey, (substate) => {
+              substate.data = applyPatches(
+                substate.data as any,
+                patches.concat()
+              )
+            })
+          }
         },
-        prepare: prepareAutoBatched<
-          QuerySubstateIdentifier & { patches: readonly Patch[] }
-        >(),
+        prepare:
+          prepareAutoBatched<
+            Array<QuerySubstateIdentifier & { patches: readonly Patch[] }>
+          >(),
       },
     },
     extraReducers(builder) {
@@ -328,39 +334,46 @@ export function buildSlice({
     name: `${reducerPath}/invalidation`,
     initialState: initialState as InvalidationState<string>,
     reducers: {
-      updateProvidedBy: {
+      updateProvidedBys: {
         reducer(
           draft,
-          action: PayloadAction<{
-            queryCacheKey: QueryCacheKey
-            providedTags: readonly FullTagDescription<string>[]
-          }>
+          action: PayloadAction<
+            Array<{
+              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 { queryCacheKey, providedTags } of 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)
+            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>[]
-        }>(),
+        prepare: prepareAutoBatched<
+          Array<{
+            queryCacheKey: QueryCacheKey
+            providedTags: readonly FullTagDescription<string>[]
+          }>
+        >(),
       },
     },
     extraReducers(builder) {
@@ -408,12 +421,14 @@ export function buildSlice({
             )
             const { queryCacheKey } = action.meta.arg
 
-            invalidationSlice.caseReducers.updateProvidedBy(
+            invalidationSlice.caseReducers.updateProvidedBys(
               draft,
-              invalidationSlice.actions.updateProvidedBy({
-                queryCacheKey,
-                providedTags,
-              })
+              invalidationSlice.actions.updateProvidedBys([
+                {
+                  queryCacheKey,
+                  providedTags,
+                },
+              ])
             )
           }
         )

diff --git a/packages/toolkit/src/query/core/buildThunks.ts b/packages/toolkit/src/query/core/buildThunks.ts
@@ -169,6 +169,20 @@ export type PatchQueryDataThunk<
   updateProvided?: boolean
 ) => ThunkAction<void, PartialState, any, UnknownAction>
 
+export type PatchQueriesDataThunk<
+  Definitions extends EndpointDefinitions,
+  PartialState
+> = (
+  patchesByEndpointName: {
+    [EndpointName in QueryKeys<Definitions>]?: Array<{
+      args: QueryArgFrom<Definitions[EndpointName]>
+      patches: readonly Patch[]
+      updateProvided?: boolean
+    }>
+  },
+  defaultUpdateProvided?: boolean
+) => ThunkAction<void, PartialState, any, UnknownAction>
+
 export type UpdateQueryDataThunk<
   Definitions extends EndpointDefinitions,
   PartialState
@@ -236,43 +250,88 @@ export function buildThunks<
 }) {
   type State = RootState<any, string, ReducerPath>
 
-  const patchQueryData: PatchQueryDataThunk<EndpointDefinitions, State> =
-    (endpointName, args, patches, updateProvided) => (dispatch, getState) => {
-      const endpointDefinition = endpointDefinitions[endpointName]
-
-      const queryCacheKey = serializeQueryArgs({
-        queryArgs: args,
-        endpointDefinition,
-        endpointName,
-      })
+  const patchQueriesData: PatchQueriesDataThunk<Definitions, State> =
+    (patchesByEndpointName, defaultUpdateProvided) => (dispatch, getState) => {
+      const queryResultPatches: Parameters<
+        typeof api.internalActions.queryResultsPatched
+      >[0] = []
+
+      const arrayified = Object.entries<
+        | {
+            args: any
+            patches: readonly Patch[]
+            updateProvided?: boolean | undefined
+          }[]
+        | undefined
+      >(patchesByEndpointName)
+      for (const [endpointName, patches] of arrayified) {
+        if (!patches) continue
+        for (const { args, patches: endpointPatches } of patches) {
+          const endpointDefinition = endpointDefinitions[endpointName]
+
+          const queryCacheKey = serializeQueryArgs({
+            queryArgs: args,
+            endpointDefinition,
+            endpointName,
+          })
 
-      dispatch(
-        api.internalActions.queryResultPatched({ queryCacheKey, patches })
-      )
+          queryResultPatches.push({ queryCacheKey, patches: endpointPatches })
+        }
+      }
 
-      if (!updateProvided) {
-        return
+      if (queryResultPatches.length) {
+        dispatch(api.internalActions.queryResultsPatched(queryResultPatches))
       }
 
-      const newValue = api.endpoints[endpointName].select(args)(
-        // Work around TS 4.1 mismatch
-        getState() as RootState<any, any, any>
-      )
+      // now that the state is updated, we can update the tags
 
-      const providedTags = calculateProvidedBy(
-        endpointDefinition.providesTags,
-        newValue.data,
-        undefined,
-        args,
-        {},
-        assertTagType
-      )
+      const providedPatches: Parameters<
+        typeof api.internalActions.updateProvidedBys
+      >[0] = []
 
-      dispatch(
-        api.internalActions.updateProvidedBy({ queryCacheKey, providedTags })
-      )
+      for (const [endpointName, patches] of arrayified) {
+        if (!patches) continue
+        for (const { args, updateProvided } of patches) {
+          if (!(updateProvided ?? defaultUpdateProvided)) {
+            continue
+          }
+          const endpointDefinition = endpointDefinitions[endpointName]
+
+          const queryCacheKey = serializeQueryArgs({
+            queryArgs: args,
+            endpointDefinition,
+            endpointName,
+          })
+
+          const newValue = api.endpoints[endpointName].select(args)(
+            // Work around TS 4.1 mismatch
+            getState() as RootState<any, any, any>
+          )
+
+          const providedTags = calculateProvidedBy(
+            endpointDefinition.providesTags,
+            newValue.data,
+            undefined,
+            args,
+            {},
+            assertTagType
+          )
+
+          providedPatches.push({ queryCacheKey, providedTags })
+        }
+      }
+      if (providedPatches.length) {
+        dispatch(api.internalActions.updateProvidedBys(providedPatches))
+      }
     }
 
+  const patchQueryData: PatchQueryDataThunk<EndpointDefinitions, State> = (
+    endpointName,
+    args,
+    patches,
+    updateProvided
+  ) => patchQueriesData({ [endpointName]: [{ args, patches, updateProvided }] })
+
   const updateQueryData: UpdateQueryDataThunk<EndpointDefinitions, State> =
     (endpointName, args, updateRecipe, updateProvided = true) =>
     (dispatch, getState) => {

diff --git a/packages/toolkit/src/query/defaultSerializeQueryArgs.ts b/packages/toolkit/src/query/defaultSerializeQueryArgs.ts
@@ -42,8 +42,4 @@ export type SerializeQueryArgs<QueryArgs, ReturnType = string> = (_: {
   endpointName: string
 }) => ReturnType
 
-export type InternalSerializeQueryArgs = (_: {
-  queryArgs: any
-  endpointDefinition: EndpointDefinition<any, any, any, any>
-  endpointName: string
-}) => QueryCacheKey
+export type InternalSerializeQueryArgs = SerializeQueryArgs<any, QueryCacheKey>