Merge pull request #1059 from reduxjs/feature/queryDefinition-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,
@@ -27,6 +28,7 @@ import { calculateProvidedByThunk } from './buildThunks'
 import type {
   AssertTagTypes,
   EndpointDefinitions,
+  QueryDefinition,
 } from '../endpointDefinitions'
 import type { Patch } from 'immer'
 import { applyPatches } from 'immer'
@@ -156,11 +158,37 @@ export function buildSlice({
             meta.arg.queryCacheKey,
             (substate) => {
               if (substate.requestId !== meta.requestId) return
+              const { merge } = definitions[
+                meta.arg.endpointName
+              ] as QueryDefinition<any, any, any, any>
               substate.status = QueryStatus.fulfilled
-              substate.data =
-                definitions[meta.arg.endpointName].structuralSharing ?? true
-                  ? copyWithStructuralSharing(substate.data, payload)
-                  : payload
+
+              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

@@ -275,6 +275,25 @@ export interface QueryExtraOptions<
   invalidatesTags?: never
 
   serializeQueryArgs?: SerializeQueryArgs<any>
+
+  /**
+   * 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?(
+    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' }])
+  })
+})