From eaf7d5edc18d89bd31b02a83f5bb85478be9f8db Mon Sep 17 00:00:00 2001
From: Mark Erikson <mark@isquaredsoftware.com>
Date: Tue, 29 Nov 2022 23:44:38 -0500
Subject: [PATCH] Pass additional metadata to `merge`

---
 packages/toolkit/src/query/core/buildSlice.ts |  9 ++-
 .../toolkit/src/query/endpointDefinitions.ts  |  8 ++-
 .../toolkit/src/query/tests/createApi.test.ts | 55 +++++++++++++++++--
 3 files changed, 65 insertions(+), 7 deletions(-)

diff --git a/packages/toolkit/src/query/core/buildSlice.ts b/packages/toolkit/src/query/core/buildSlice.ts
index ef4fc26f1..3343b6dc3 100644
--- a/packages/toolkit/src/query/core/buildSlice.ts
+++ b/packages/toolkit/src/query/core/buildSlice.ts
@@ -181,6 +181,8 @@ export function buildSlice({
 
               if (merge) {
                 if (substate.data !== undefined) {
+                  const { fulfilledTimeStamp, arg, baseQueryMeta, requestId } =
+                    meta
                   // 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.
@@ -189,7 +191,12 @@ export function buildSlice({
                     substate.data,
                     (draftSubstateData) => {
                       // As usual with Immer, you can mutate _or_ return inside here, but not both
-                      return merge(draftSubstateData, payload)
+                      return merge(draftSubstateData, payload, {
+                        arg: arg.originalArgs,
+                        baseQueryMeta,
+                        fulfilledTimeStamp,
+                        requestId,
+                      })
                     }
                   )
                   substate.data = newData
diff --git a/packages/toolkit/src/query/endpointDefinitions.ts b/packages/toolkit/src/query/endpointDefinitions.ts
index f5243f299..5ace421d4 100644
--- a/packages/toolkit/src/query/endpointDefinitions.ts
+++ b/packages/toolkit/src/query/endpointDefinitions.ts
@@ -445,7 +445,13 @@ export interface QueryExtraOptions<
    */
   merge?(
     currentCacheData: ResultType,
-    responseData: ResultType
+    responseData: ResultType,
+    otherArgs: {
+      arg: QueryArg
+      baseQueryMeta: BaseQueryMeta<BaseQuery>
+      requestId: string
+      fulfilledTimeStamp: number
+    }
   ): ResultType | void
 
   /**
diff --git a/packages/toolkit/src/query/tests/createApi.test.ts b/packages/toolkit/src/query/tests/createApi.test.ts
index 37c1f52ca..75d7690ad 100644
--- a/packages/toolkit/src/query/tests/createApi.test.ts
+++ b/packages/toolkit/src/query/tests/createApi.test.ts
@@ -35,6 +35,11 @@ afterAll(() => {
   spy.mockRestore()
 })
 
+function paginate<T>(array: T[], page_size: number, page_number: number) {
+  // human-readable page numbers usually start with 1, so we reduce 1 in the first argument
+  return array.slice((page_number - 1) * page_size, page_number * page_size)
+}
+
 test('sensible defaults', () => {
   const api = createApi({
     baseQuery: fetchBaseQuery(),
@@ -923,6 +928,22 @@ describe('custom serializeQueryArgs per endpoint', () => {
           return currentArg !== previousArg
         },
       }),
+      listItems2: build.query<{ items: string[]; meta?: any }, number>({
+        query: (pageNumber) => `/listItems2?page=${pageNumber}`,
+        serializeQueryArgs: ({ endpointName }) => {
+          return endpointName
+        },
+        transformResponse(items: string[]) {
+          return { items }
+        },
+        merge: (currentCache, newData, meta) => {
+          currentCache.items.push(...newData.items)
+          currentCache.meta = meta
+        },
+        forceRefetch({ currentArg, previousArg }) {
+          return currentArg !== previousArg
+        },
+      }),
     }),
   })
 
@@ -1010,11 +1031,6 @@ describe('custom serializeQueryArgs per endpoint', () => {
     const allItems = ['a', 'b', 'c', 'd', 'e', 'f', 'g', 'i']
     const PAGE_SIZE = 3
 
-    function paginate<T>(array: T[], page_size: number, page_number: number) {
-      // human-readable page numbers usually start with 1, so we reduce 1 in the first argument
-      return array.slice((page_number - 1) * page_size, page_number * page_size)
-    }
-
     server.use(
       rest.get('https://example.com/listItems', (req, res, ctx) => {
         const pageString = req.url.searchParams.get('page')
@@ -1038,4 +1054,33 @@ describe('custom serializeQueryArgs per endpoint', () => {
     const updatedEntry = selectListItems(storeRef.store.getState())
     expect(updatedEntry.data).toEqual(['a', 'b', 'c', 'd', 'e', 'f'])
   })
+
+  test('merge receives a meta object as an argument', async () => {
+    const allItems = ['a', 'b', 'c', 'd', 'e', 'f', 'g', 'i']
+    const PAGE_SIZE = 3
+
+    server.use(
+      rest.get('https://example.com/listItems2', (req, res, ctx) => {
+        const pageString = req.url.searchParams.get('page')
+        const pageNum = parseInt(pageString || '0')
+
+        const results = paginate(allItems, PAGE_SIZE, pageNum)
+        return res(ctx.json(results))
+      })
+    )
+
+    const selectListItems = api.endpoints.listItems2.select(0)
+
+    await storeRef.store.dispatch(api.endpoints.listItems2.initiate(1))
+    await storeRef.store.dispatch(api.endpoints.listItems2.initiate(2))
+    const cacheEntry = selectListItems(storeRef.store.getState())
+
+    // Should have passed along the third arg from `merge` containing these fields
+    expect(cacheEntry.data?.meta).toEqual({
+      requestId: expect.any(String),
+      fulfilledTimeStamp: expect.any(Number),
+      arg: 2,
+      baseQueryMeta: expect.any(Object),
+    })
+  })
 })