Update merge logic to make it optional and handle different use cases

- Make `merge` an optional callback
- Flip args order to put `currentCacheData` first
- Use Immer to handle the "mutate or return" behavior
- Only do structural sharing if there's no `merge`

Author: markerikson

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

@@ -6,6 +6,7 @@ import {
   isAnyOf,
   isFulfilled,
   isRejectedWithValue,
+  createNextState,
 } from '@reduxjs/toolkit'
 import type {
   CombinedState as CombinedQueryState,
@@ -157,16 +158,37 @@ export function buildSlice({
             meta.arg.queryCacheKey,
             (substate) => {
               if (substate.requestId !== meta.requestId) return
-              const { merge = (x: any) => x } = definitions[
+              const { merge } = definitions[
                 meta.arg.endpointName
               ] as QueryDefinition<any, any, any, any>
               substate.status = QueryStatus.fulfilled
-              let newData = merge(payload, substate.data)
 
-              substate.data =
-                definitions[meta.arg.endpointName].structuralSharing ?? true
-                  ? copyWithStructuralSharing(substate.data, newData)
-                  : newData
+              if (merge) {
+                if (substate.data !== undefined) {
+                  // There's existing cache data. Let the user merge it in themselves.
+                  // We're already inside an Immer-powered reducer, and the user could just mutate `substate.data`
+                  // themselves inside of `merge()`. But, they might also want to return a new value.
+                  // Try to let Immer figure that part out, save the result, and assign it to `substate.data`.
+                  let newData = createNextState(
+                    substate.data,
+                    (draftSubstateData) => {
+                      // As usual with Immer, you can mutate _or_ return inside here, but not both
+                      return merge(draftSubstateData, payload)
+                    }
+                  )
+                  substate.data = newData
+                } else {
+                  // Presumably a fresh request. Just cache the response data.
+                  substate.data = payload
+                }
+              } else {
+                // Assign or safely update the cache data.
+                substate.data =
+                  definitions[meta.arg.endpointName].structuralSharing ?? true
+                    ? copyWithStructuralSharing(substate.data, payload)
+                    : payload
+              }
+
               delete substate.error
               substate.fulfilledTimeStamp = meta.fulfilledTimeStamp
             }

packages/toolkit/src/query/endpointDefinitions.ts

@@ -278,15 +278,22 @@ export interface QueryExtraOptions<
 
   /**
    * Can be provided to merge the current cache value into the new cache value.
+   * If supplied, no automatic structural sharing will be applied - it's up to
+   * you to update the cache appropriately.
+   *
+   * Since this is wrapped with Immer, you , you may either mutate the `currentCacheValue` directly,
+   * or return a new value, but _not_ both at once.   *
+   *
+   * Will only be called if the existing `currentCacheValue` is not `undefined`.
    *
    * Useful if you don't want a new request to completely override the current cache value,
    * maybe because you have manually updated it from another source and don't want those
    * updates to get lost.
    */
   merge?(
-    newCacheValue: ResultType,
-    currentCacheValue: ResultType | undefined
-  ): ResultType
+    currentCacheData: ResultType,
+    responseData: ResultType
+  ): ResultType | void
 }
 
 export type QueryDefinition<

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

@@ -29,6 +29,10 @@ const authSlice = createSlice({
 
 const storeRef = setupApiStore(api, { auth: authSlice.reducer })
 
+function delay(ms: number) {
+  return new Promise((resolve) => setTimeout(resolve, ms))
+}
+
 it('only resets the api state when resetApiState is dispatched', async () => {
   storeRef.store.dispatch({ type: 'unrelated' }) // trigger "registered middleware" into place
   const initialState = storeRef.store.getState()
@@ -77,3 +81,100 @@ it('only resets the api state when resetApiState is dispatched', async () => {
 
   expect(storeRef.store.getState()).toEqual(initialState)
 })
+
+describe.only('`merge` callback', () => {
+  const baseQuery = (args?: any) => ({ data: args })
+
+  interface Todo {
+    id: string
+    text: string
+  }
+
+  it('Calls `merge` once there is existing data, and allows mutations of cache state', async () => {
+    let mergeCalled = false
+    let queryFnCalls = 0
+    const todoTexts = ['A', 'B', 'C', 'D']
+
+    const api = createApi({
+      baseQuery,
+      endpoints: (build) => ({
+        getTodos: build.query<Todo[], void>({
+          async queryFn() {
+            const text = todoTexts[queryFnCalls]
+            return { data: [{ id: `${queryFnCalls++}`, text }] }
+          },
+          merge(currentCacheValue, responseData) {
+            mergeCalled = true
+            currentCacheValue.push(...responseData)
+          },
+        }),
+      }),
+    })
+
+    const storeRef = setupApiStore(api, undefined, {
+      withoutTestLifecycles: true,
+    })
+
+    const selectTodoEntry = api.endpoints.getTodos.select()
+
+    const res = storeRef.store.dispatch(api.endpoints.getTodos.initiate())
+    await res
+    expect(mergeCalled).toBe(false)
+    const todoEntry1 = selectTodoEntry(storeRef.store.getState())
+    expect(todoEntry1.data).toEqual([{ id: '0', text: 'A' }])
+
+    res.refetch()
+
+    await delay(10)
+
+    expect(mergeCalled).toBe(true)
+    const todoEntry2 = selectTodoEntry(storeRef.store.getState())
+
+    expect(todoEntry2.data).toEqual([
+      { id: '0', text: 'A' },
+      { id: '1', text: 'B' },
+    ])
+  })
+
+  it('Allows returning a different value from `merge`', async () => {
+    let firstQueryFnCall = true
+
+    const api = createApi({
+      baseQuery,
+      endpoints: (build) => ({
+        getTodos: build.query<Todo[], void>({
+          async queryFn() {
+            const item = firstQueryFnCall
+              ? { id: '0', text: 'A' }
+              : { id: '1', text: 'B' }
+            firstQueryFnCall = false
+            return { data: [item] }
+          },
+          merge(currentCacheValue, responseData) {
+            return responseData
+          },
+        }),
+      }),
+    })
+
+    const storeRef = setupApiStore(api, undefined, {
+      withoutTestLifecycles: true,
+    })
+
+    const selectTodoEntry = api.endpoints.getTodos.select()
+
+    const res = storeRef.store.dispatch(api.endpoints.getTodos.initiate())
+    await res
+
+    const todoEntry1 = selectTodoEntry(storeRef.store.getState())
+    expect(todoEntry1.data).toEqual([{ id: '0', text: 'A' }])
+
+    res.refetch()
+
+    await delay(10)
+
+    const todoEntry2 = selectTodoEntry(storeRef.store.getState())
+
+    expect(todoEntry2.data).toEqual([{ id: '1', text: 'B' }])
+  })
+})