Implement a util function to batch-upsert cache entries

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 }))

packages/toolkit/src/query/core/buildMiddleware/cacheLifecycle.ts

@@ -183,24 +183,68 @@ export const buildCacheLifecycleHandler: InternalHandlerBuilder = ({
   }
   const lifecycleMap: Record<string, CacheLifecycle> = {}
 
+  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,
     stateBefore,
   ) => {
     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)
       }
     }
   }