From e502bc8ebbe083ccd9210dd08be7b702d3001ab3 Mon Sep 17 00:00:00 2001
From: Barnabas Jovanovics <barnabas.jovanovics@gmail.com>
Date: Wed, 20 Apr 2022 15:10:51 +0200
Subject: [PATCH 1/9] add upsertQueryDataThunk

---
 .../toolkit/src/query/core/buildInitiate.ts   |  9 ++-
 .../toolkit/src/query/core/buildThunks.ts     | 63 ++++++++++++++++++-
 2 files changed, 68 insertions(+), 4 deletions(-)

diff --git a/packages/toolkit/src/query/core/buildInitiate.ts b/packages/toolkit/src/query/core/buildInitiate.ts
index 6014feea47..9690744de8 100644
--- a/packages/toolkit/src/query/core/buildInitiate.ts
+++ b/packages/toolkit/src/query/core/buildInitiate.ts
@@ -13,7 +13,7 @@ import { QueryStatus } from './apiState'
 import type { InternalSerializeQueryArgs } from '../defaultSerializeQueryArgs'
 import type { Api, ApiContext } from '../apiTypes'
 import type { ApiEndpointQuery } from './module'
-import type { BaseQueryError } from '../baseQueryTypes'
+import type { BaseQueryError, QueryReturnValue } from '../baseQueryTypes'
 import type { QueryResultSelectorResult } from './buildSelectors'
 
 declare module './module' {
@@ -34,10 +34,11 @@ declare module './module' {
   }
 }
 
-export interface StartQueryActionCreatorOptions {
+export interface StartQueryActionCreatorOptions  {
   subscribe?: boolean
   forceRefetch?: boolean | number
   subscriptionOptions?: SubscriptionOptions
+  forceQueryFn?: () => QueryReturnValue
 }
 
 type StartQueryActionCreator<
@@ -179,6 +180,7 @@ export type MutationActionCreatorResult<
   unsubscribe(): void
 }
 
+
 export function buildInitiate({
   serializeQueryArgs,
   queryThunk,
@@ -259,7 +261,7 @@ Features like automatic cache collection, automatic refetching etc. will not be
     endpointDefinition: QueryDefinition<any, any, any, any>
   ) {
     const queryAction: StartQueryActionCreator<any> =
-      (arg, { subscribe = true, forceRefetch, subscriptionOptions } = {}) =>
+      (arg, { subscribe = true, forceRefetch, subscriptionOptions, forceQueryFn } = {}) =>
       (dispatch, getState) => {
         const queryCacheKey = serializeQueryArgs({
           queryArgs: arg,
@@ -270,6 +272,7 @@ Features like automatic cache collection, automatic refetching etc. will not be
           type: 'query',
           subscribe,
           forceRefetch,
+          forceQueryFn,
           subscriptionOptions,
           endpointName,
           originalArgs: arg,
diff --git a/packages/toolkit/src/query/core/buildThunks.ts b/packages/toolkit/src/query/core/buildThunks.ts
index e873bde369..f1c6c7364b 100644
--- a/packages/toolkit/src/query/core/buildThunks.ts
+++ b/packages/toolkit/src/query/core/buildThunks.ts
@@ -144,6 +144,7 @@ function defaultTransformResponse(baseQueryReturnValue: unknown) {
 
 export type MaybeDrafted<T> = T | Draft<T>
 export type Recipe<T> = (data: MaybeDrafted<T>) => void | MaybeDrafted<T>
+export type UpsertRecipe<T> = (data: MaybeDrafted<T> | undefined) => void | MaybeDrafted<T>
 
 export type PatchQueryDataThunk<
   Definitions extends EndpointDefinitions,
@@ -163,6 +164,15 @@ export type UpdateQueryDataThunk<
   updateRecipe: Recipe<ResultTypeFrom<Definitions[EndpointName]>>
 ) => ThunkAction<PatchCollection, PartialState, any, AnyAction>
 
+export type UpsertQueryDataThunk<
+  Definitions extends EndpointDefinitions,
+  PartialState
+> = <EndpointName extends QueryKeys<Definitions>>(
+  endpointName: EndpointName,
+  args: QueryArgFrom<Definitions[EndpointName]>,
+  upsertRecipe: UpsertRecipe<ResultTypeFrom<Definitions[EndpointName]>>
+) => ThunkAction<PatchCollection, PartialState, any, AnyAction>
+
 /**
  * An object returned from dispatching a `api.util.updateQueryData` call.
  */
@@ -252,6 +262,53 @@ export function buildThunks<
 
       dispatch(api.util.patchQueryData(endpointName, args, ret.patches))
 
+      return ret
+    }
+
+  const upsertQueryData: UpsertQueryDataThunk<EndpointDefinitions, State> =
+    (endpointName, args, upsertRecipe) => (dispatch, getState) => {
+      const currentState = (
+        api.endpoints[endpointName] as ApiEndpointQuery<any, any>
+      ).select(args)(getState())
+      let ret: PatchCollection = {
+        patches: [],
+        inversePatches: [],
+        undo: () =>
+          dispatch(
+            api.util.patchQueryData(endpointName, args, ret.inversePatches)
+          ),
+      }
+      if (currentState.status === QueryStatus.uninitialized) {
+        return ret
+      }
+      if ('data' in currentState) {
+        if (isDraftable(currentState.data)) {
+          const [, patches, inversePatches] = produceWithPatches(
+            currentState.data,
+            upsertRecipe
+          )
+          ret.patches.push(...patches)
+          ret.inversePatches.push(...inversePatches)
+        } else {
+          const value = upsertRecipe(currentState.data)
+          ret.patches.push({ op: 'replace', path: [], value })
+          ret.inversePatches.push({
+            op: 'replace',
+            path: [],
+            value: currentState.data,
+          })
+        }
+        dispatch(api.util.patchQueryData(endpointName, args, ret.patches))
+      } else {
+        ret.inversePatches.push({
+          op: 'replace',
+          path: [],
+          value: undefined,
+        })
+        dispatch(api.endpoints[endpointName].initiate(args, {subscribe: false, forceRefetch: true, }))
+      }
+
+
       return ret
     }
 
@@ -291,7 +348,10 @@ export function buildThunks<
         forced:
           arg.type === 'query' ? isForcedQuery(arg, getState()) : undefined,
       }
-      if (endpointDefinition.query) {
+      if ('forceQueryFn' in arg && arg.forceQueryFn) {
+        result = arg.forceQueryFn()
+        
+      }else if (endpointDefinition.query) {
         result = await baseQuery(
           endpointDefinition.query(arg.originalArgs),
           baseQueryApi,
@@ -527,6 +587,7 @@ In the case of an unhandled error, no tags will be "provided" or "invalidated".`
     mutationThunk,
     prefetch,
     updateQueryData,
+    upsertQueryData,
     patchQueryData,
     buildMatchThunkActions,
   }

From 8c3f5dd6a5eb968fed648716e221a5711a167566 Mon Sep 17 00:00:00 2001
From: Barnabas Jovanovics <barnabas.jovanovics@gmail.com>
Date: Wed, 20 Apr 2022 16:43:51 +0200
Subject: [PATCH 2/9] add tests

---
 .../toolkit/src/query/core/buildThunks.ts     |  19 +-
 packages/toolkit/src/query/core/module.ts     |  12 +-
 .../query/tests/optimisticUpserts.test.tsx    | 385 ++++++++++++++++++
 3 files changed, 410 insertions(+), 6 deletions(-)
 create mode 100644 packages/toolkit/src/query/tests/optimisticUpserts.test.tsx

diff --git a/packages/toolkit/src/query/core/buildThunks.ts b/packages/toolkit/src/query/core/buildThunks.ts
index f1c6c7364b..cf68d87476 100644
--- a/packages/toolkit/src/query/core/buildThunks.ts
+++ b/packages/toolkit/src/query/core/buildThunks.ts
@@ -278,9 +278,6 @@ export function buildThunks<
             api.util.patchQueryData(endpointName, args, ret.inversePatches)
           ),
       }
-      if (currentState.status === QueryStatus.uninitialized) {
-        return ret
-      }
       if ('data' in currentState) {
         if (isDraftable(currentState.data)) {
           const [, patches, inversePatches] = produceWithPatches(
@@ -305,10 +302,22 @@ export function buildThunks<
           path: [],
           value: undefined,
         })
-        dispatch(api.endpoints[endpointName].initiate(args, {subscribe: false, forceRefetch: true, }))
+        dispatch(
+          (
+            api.endpoints[endpointName] as ApiEndpointQuery<
+              QueryDefinition<any, any, any, any, any>,
+              Definitions
+            >
+          ).initiate(args, {
+            subscribe: false,
+            forceRefetch: true,
+            forceQueryFn: () => ({
+              data: upsertRecipe(undefined),
+            }),
+          })
+        )
       }
 
-
       return ret
     }
 
diff --git a/packages/toolkit/src/query/core/module.ts b/packages/toolkit/src/query/core/module.ts
index d419d8689f..c0c58582f2 100644
--- a/packages/toolkit/src/query/core/module.ts
+++ b/packages/toolkit/src/query/core/module.ts
@@ -1,7 +1,11 @@
 /**
  * Note: this file should import all other files for type discovery and declaration merging
  */
-import type { PatchQueryDataThunk, UpdateQueryDataThunk } from './buildThunks'
+import type {
+  PatchQueryDataThunk,
+  UpdateQueryDataThunk,
+  UpsertQueryDataThunk,
+} from './buildThunks'
 import { buildThunks } from './buildThunks'
 import type {
   ActionCreatorWithPayload,
@@ -210,6 +214,10 @@ declare module '../apiTypes' {
           Definitions,
           RootState<Definitions, string, ReducerPath>
         >
+        upsertQueryData: UpsertQueryDataThunk<
+          Definitions,
+          RootState<Definitions, string, ReducerPath>
+        >
         /**
          * A Redux thunk that applies a JSON diff/patch array to the cached data for a given query result. This immediately updates the Redux state with those changes.
          *
@@ -416,6 +424,7 @@ export const coreModule = (): Module<CoreModule> => ({
       mutationThunk,
       patchQueryData,
       updateQueryData,
+      upsertQueryData,
       prefetch,
       buildMatchThunkActions,
     } = buildThunks({
@@ -444,6 +453,7 @@ export const coreModule = (): Module<CoreModule> => ({
     safeAssign(api.util, {
       patchQueryData,
       updateQueryData,
+      upsertQueryData,
       prefetch,
       resetApiState: sliceActions.resetApiState,
     })
diff --git a/packages/toolkit/src/query/tests/optimisticUpserts.test.tsx b/packages/toolkit/src/query/tests/optimisticUpserts.test.tsx
new file mode 100644
index 0000000000..9461e3483a
--- /dev/null
+++ b/packages/toolkit/src/query/tests/optimisticUpserts.test.tsx
@@ -0,0 +1,385 @@
+import { createApi } from '@reduxjs/toolkit/query/react'
+import { actionsReducer, hookWaitFor, setupApiStore, waitMs } from './helpers'
+import { skipToken } from '../core/buildSelectors'
+import { renderHook, act } from '@testing-library/react-hooks'
+
+interface Post {
+  id: string
+  title: string
+  contents: string
+}
+
+const baseQuery = jest.fn()
+beforeEach(() => baseQuery.mockReset())
+
+const api = createApi({
+  baseQuery: (...args: any[]) => {
+    const result = baseQuery(...args)
+    if (typeof result === 'object' && 'then' in result)
+      return result
+        .then((data: any) => ({ data, meta: 'meta' }))
+        .catch((e: any) => ({ error: e }))
+    return { data: result, meta: 'meta' }
+  },
+  tagTypes: ['Post'],
+  endpoints: (build) => ({
+    post: build.query<Post, string>({
+      query: (id) => `post/${id}`,
+      providesTags: ['Post'],
+    }),
+    updatePost: build.mutation<void, Pick<Post, 'id'> & Partial<Post>>({
+      query: ({ id, ...patch }) => ({
+        url: `post/${id}`,
+        method: 'PATCH',
+        body: patch,
+      }),
+      async onQueryStarted({ id, ...patch }, { dispatch, queryFulfilled }) {
+        const { undo } = dispatch(
+          api.util.upsertQueryData('post', id, (draft) => {
+            Object.assign(draft, patch)
+          })
+        )
+        queryFulfilled.catch(undo)
+      },
+      invalidatesTags: (result) => (result ? ['Post'] : []),
+    }),
+  }),
+})
+
+const storeRef = setupApiStore(api, {
+  ...actionsReducer,
+})
+
+describe('basic lifecycle', () => {
+  let onStart = jest.fn(),
+    onError = jest.fn(),
+    onSuccess = jest.fn()
+
+  const extendedApi = api.injectEndpoints({
+    endpoints: (build) => ({
+      test: build.mutation({
+        query: (x) => x,
+        async onQueryStarted(arg, api) {
+          onStart(arg)
+          try {
+            const result = await api.queryFulfilled
+            onSuccess(result)
+          } catch (e) {
+            onError(e)
+          }
+        },
+      }),
+    }),
+    overrideExisting: true,
+  })
+
+  beforeEach(() => {
+    onStart.mockReset()
+    onError.mockReset()
+    onSuccess.mockReset()
+  })
+
+  test('success', async () => {
+    const { result } = renderHook(
+      () => extendedApi.endpoints.test.useMutation(),
+      {
+        wrapper: storeRef.wrapper,
+      }
+    )
+
+    baseQuery.mockResolvedValue('success')
+
+    expect(onStart).not.toHaveBeenCalled()
+    expect(baseQuery).not.toHaveBeenCalled()
+    act(() => void result.current[0]('arg'))
+    expect(onStart).toHaveBeenCalledWith('arg')
+    expect(baseQuery).toHaveBeenCalledWith('arg', expect.any(Object), undefined)
+
+    expect(onError).not.toHaveBeenCalled()
+    expect(onSuccess).not.toHaveBeenCalled()
+    await act(() => waitMs(5))
+    expect(onError).not.toHaveBeenCalled()
+    expect(onSuccess).toHaveBeenCalledWith({ data: 'success', meta: 'meta' })
+  })
+
+  test('error', async () => {
+    const { result } = renderHook(
+      () => extendedApi.endpoints.test.useMutation(),
+      {
+        wrapper: storeRef.wrapper,
+      }
+    )
+
+    baseQuery.mockRejectedValue('error')
+
+    expect(onStart).not.toHaveBeenCalled()
+    expect(baseQuery).not.toHaveBeenCalled()
+    act(() => void result.current[0]('arg'))
+    expect(onStart).toHaveBeenCalledWith('arg')
+    expect(baseQuery).toHaveBeenCalledWith('arg', expect.any(Object), undefined)
+
+    expect(onError).not.toHaveBeenCalled()
+    expect(onSuccess).not.toHaveBeenCalled()
+    await act(() => waitMs(5))
+    expect(onError).toHaveBeenCalledWith({
+      error: 'error',
+      isUnhandledError: false,
+      meta: undefined,
+    })
+    expect(onSuccess).not.toHaveBeenCalled()
+  })
+})
+
+describe('upsertQueryData', () => {
+  test('updates cache values, can apply inverse patch', async () => {
+    baseQuery
+      .mockResolvedValueOnce({
+        id: '3',
+        title: 'All about cheese.',
+        contents: 'TODO',
+      })
+      // TODO I have no idea why the query is getting called multiple times,
+      // but passing an additional mocked value (_any_ value)
+      // seems to silence some annoying "got an undefined result" logging
+      .mockResolvedValueOnce(42)
+    const { result } = renderHook(() => api.endpoints.post.useQuery('3'), {
+      wrapper: storeRef.wrapper,
+    })
+    await hookWaitFor(() => expect(result.current.isSuccess).toBeTruthy())
+
+    const dataBefore = result.current.data
+    expect(dataBefore).toEqual({
+      id: '3',
+      title: 'All about cheese.',
+      contents: 'TODO',
+    })
+
+    let returnValue!: ReturnType<ReturnType<typeof api.util.upsertQueryData>>
+    act(() => {
+      returnValue = storeRef.store.dispatch(
+        api.util.upsertQueryData('post', '3', (draft) => {
+          if (draft) {
+            draft.contents = 'I love cheese!'
+          }
+        })
+      )
+    })
+
+    expect(result.current.data).not.toBe(dataBefore)
+    expect(result.current.data).toEqual({
+      id: '3',
+      title: 'All about cheese.',
+      contents: 'I love cheese!',
+    })
+
+    expect(returnValue).toEqual({
+      inversePatches: [{ op: 'replace', path: ['contents'], value: 'TODO' }],
+      patches: [{ op: 'replace', path: ['contents'], value: 'I love cheese!' }],
+      undo: expect.any(Function),
+    })
+
+    act(() => {
+      storeRef.store.dispatch(
+        api.util.patchQueryData('post', '3', returnValue.inversePatches)
+      )
+    })
+
+    expect(result.current.data).toEqual(dataBefore)
+  })
+
+  test('does update non-existing values', async () => {
+    baseQuery
+      // throw an error to make sure there is no cached data
+      .mockImplementationOnce(async () => {
+        throw new Error('failed to load')
+      })
+      .mockResolvedValueOnce(42)
+
+    // a subscriber is needed to have the data stay in the cache
+    // Not sure if this is the wanted behaviour, I would have liked
+    // it to stay in the cache for the x amount of time the cache
+    // is preserved normally after the last subscriber was unmounted
+    const { result, rerender } = renderHook(
+      () => api.endpoints.post.useQuery('4'),
+      {
+        wrapper: storeRef.wrapper,
+      }
+    )
+    await hookWaitFor(() => expect(result.current.isError).toBeTruthy())
+
+    let returnValue!: ReturnType<ReturnType<typeof api.util.upsertQueryData>>
+    // upsert the data
+    act(() => {
+      returnValue = storeRef.store.dispatch(
+        api.util.upsertQueryData('post', '4', (draft) => {
+          if (draft) {
+            draft.contents = 'I love cheese!'
+          } else {
+            return {
+              id: '4',
+              title: 'All about cheese',
+              contents: 'I love cheese!',
+            }
+          }
+        })
+      )
+    })
+
+    // the patch would remove the result again
+    // maybe this needs to be implemented differently in order
+    // for the whole thing to work correctly, I suppose that
+    // this would only revert the data but would not invalidate it
+    expect(returnValue).toEqual({
+      inversePatches: [{ op: 'replace', path: [], value: undefined }],
+      patches: [],
+      undo: expect.any(Function),
+    })
+
+    // rerender the hook
+    rerender()
+    // wait until everything has settled
+    await hookWaitFor(() => expect(result.current.isSuccess).toBeTruthy())
+
+    // the cached data is returned as the result
+    expect(result.current.data).toStrictEqual({
+      id: '4',
+      title: 'All about cheese',
+      contents: 'I love cheese!',
+    })
+  })
+})
+
+describe('full integration', () => {
+  test('success case', async () => {
+    baseQuery
+      .mockResolvedValueOnce({
+        id: '3',
+        title: 'All about cheese.',
+        contents: 'TODO',
+      })
+      .mockResolvedValueOnce({
+        id: '3',
+        title: 'Meanwhile, this changed server-side.',
+        contents: 'Delicious cheese!',
+      })
+      .mockResolvedValueOnce({
+        id: '3',
+        title: 'Meanwhile, this changed server-side.',
+        contents: 'Delicious cheese!',
+      })
+      .mockResolvedValueOnce(42)
+    const { result } = renderHook(
+      () => ({
+        query: api.endpoints.post.useQuery('3'),
+        mutation: api.endpoints.updatePost.useMutation(),
+      }),
+      {
+        wrapper: storeRef.wrapper,
+      }
+    )
+    await hookWaitFor(() => expect(result.current.query.isSuccess).toBeTruthy())
+
+    expect(result.current.query.data).toEqual({
+      id: '3',
+      title: 'All about cheese.',
+      contents: 'TODO',
+    })
+
+    act(() => {
+      result.current.mutation[0]({ id: '3', contents: 'Delicious cheese!' })
+    })
+
+    expect(result.current.query.data).toEqual({
+      id: '3',
+      title: 'All about cheese.',
+      contents: 'Delicious cheese!',
+    })
+
+    await hookWaitFor(() =>
+      expect(result.current.query.data).toEqual({
+        id: '3',
+        title: 'Meanwhile, this changed server-side.',
+        contents: 'Delicious cheese!',
+      })
+    )
+  })
+
+  test('error case', async () => {
+    baseQuery
+      .mockResolvedValueOnce({
+        id: '3',
+        title: 'All about cheese.',
+        contents: 'TODO',
+      })
+      .mockRejectedValueOnce('some error!')
+      .mockResolvedValueOnce({
+        id: '3',
+        title: 'Meanwhile, this changed server-side.',
+        contents: 'TODO',
+      })
+      .mockResolvedValueOnce(42)
+
+    const { result } = renderHook(
+      () => ({
+        query: api.endpoints.post.useQuery('3'),
+        mutation: api.endpoints.updatePost.useMutation(),
+      }),
+      {
+        wrapper: storeRef.wrapper,
+      }
+    )
+    await hookWaitFor(() => expect(result.current.query.isSuccess).toBeTruthy())
+
+    expect(result.current.query.data).toEqual({
+      id: '3',
+      title: 'All about cheese.',
+      contents: 'TODO',
+    })
+
+    act(() => {
+      result.current.mutation[0]({ id: '3', contents: 'Delicious cheese!' })
+    })
+
+    // optimistic update
+    expect(result.current.query.data).toEqual({
+      id: '3',
+      title: 'All about cheese.',
+      contents: 'Delicious cheese!',
+    })
+
+    // rollback
+    await hookWaitFor(() =>
+      expect(result.current.query.data).toEqual({
+        id: '3',
+        title: 'All about cheese.',
+        contents: 'TODO',
+      })
+    )
+
+    // mutation failed - will not invalidate query and not refetch data from the server
+    await expect(() =>
+      hookWaitFor(
+        () =>
+          expect(result.current.query.data).toEqual({
+            id: '3',
+            title: 'Meanwhile, this changed server-side.',
+            contents: 'TODO',
+          }),
+        50
+      )
+    ).rejects.toBeTruthy()
+
+    act(() => void result.current.query.refetch())
+
+    // manually refetching gives up-to-date data
+    await hookWaitFor(
+      () =>
+        expect(result.current.query.data).toEqual({
+          id: '3',
+          title: 'Meanwhile, this changed server-side.',
+          contents: 'TODO',
+        }),
+      50
+    )
+  })
+})

From ef461bce749c9a68438d97b89261ee092aeedf90 Mon Sep 17 00:00:00 2001
From: Barnabas Jovanovics <barnabas.jovanovics@gmail.com>
Date: Thu, 21 Apr 2022 08:50:03 +0200
Subject: [PATCH 3/9] use symbol as key for forceQueryFn

---
 .../toolkit/src/query/core/buildInitiate.ts   | 19 ++++++++++++-----
 .../toolkit/src/query/core/buildThunks.ts     | 21 ++++++++++++-------
 2 files changed, 28 insertions(+), 12 deletions(-)

diff --git a/packages/toolkit/src/query/core/buildInitiate.ts b/packages/toolkit/src/query/core/buildInitiate.ts
index 9690744de8..0f80123205 100644
--- a/packages/toolkit/src/query/core/buildInitiate.ts
+++ b/packages/toolkit/src/query/core/buildInitiate.ts
@@ -34,11 +34,13 @@ declare module './module' {
   }
 }
 
-export interface StartQueryActionCreatorOptions  {
+export const forceQueryFnSymbol = Symbol('forceQueryFn')
+
+export interface StartQueryActionCreatorOptions {
   subscribe?: boolean
   forceRefetch?: boolean | number
   subscriptionOptions?: SubscriptionOptions
-  forceQueryFn?: () => QueryReturnValue
+  [forceQueryFnSymbol]?: () => QueryReturnValue
 }
 
 type StartQueryActionCreator<
@@ -180,7 +182,6 @@ export type MutationActionCreatorResult<
   unsubscribe(): void
 }
 
-
 export function buildInitiate({
   serializeQueryArgs,
   queryThunk,
@@ -261,7 +262,15 @@ Features like automatic cache collection, automatic refetching etc. will not be
     endpointDefinition: QueryDefinition<any, any, any, any>
   ) {
     const queryAction: StartQueryActionCreator<any> =
-      (arg, { subscribe = true, forceRefetch, subscriptionOptions, forceQueryFn } = {}) =>
+      (
+        arg,
+        {
+          subscribe = true,
+          forceRefetch,
+          subscriptionOptions,
+          [forceQueryFnSymbol]: forceQueryFn,
+        } = {}
+      ) =>
       (dispatch, getState) => {
         const queryCacheKey = serializeQueryArgs({
           queryArgs: arg,
@@ -272,11 +281,11 @@ Features like automatic cache collection, automatic refetching etc. will not be
           type: 'query',
           subscribe,
           forceRefetch,
-          forceQueryFn,
           subscriptionOptions,
           endpointName,
           originalArgs: arg,
           queryCacheKey,
+          [forceQueryFnSymbol]: forceQueryFn,
         })
         const selector = (
           api.endpoints[endpointName] as ApiEndpointQuery<any, any>
diff --git a/packages/toolkit/src/query/core/buildThunks.ts b/packages/toolkit/src/query/core/buildThunks.ts
index cf68d87476..5d2b2ab580 100644
--- a/packages/toolkit/src/query/core/buildThunks.ts
+++ b/packages/toolkit/src/query/core/buildThunks.ts
@@ -7,7 +7,10 @@ import type {
 } from '../baseQueryTypes'
 import type { RootState, QueryKeys, QuerySubstateIdentifier } from './apiState'
 import { QueryStatus } from './apiState'
-import type { StartQueryActionCreatorOptions } from './buildInitiate'
+import {
+  forceQueryFnSymbol,
+  StartQueryActionCreatorOptions,
+} from './buildInitiate'
 import type {
   AssertTagTypes,
   EndpointDefinition,
@@ -144,7 +147,9 @@ function defaultTransformResponse(baseQueryReturnValue: unknown) {
 
 export type MaybeDrafted<T> = T | Draft<T>
 export type Recipe<T> = (data: MaybeDrafted<T>) => void | MaybeDrafted<T>
-export type UpsertRecipe<T> = (data: MaybeDrafted<T> | undefined) => void | MaybeDrafted<T>
+export type UpsertRecipe<T> = (
+  data: MaybeDrafted<T> | undefined
+) => void | MaybeDrafted<T>
 
 export type PatchQueryDataThunk<
   Definitions extends EndpointDefinitions,
@@ -311,7 +316,7 @@ export function buildThunks<
           ).initiate(args, {
             subscribe: false,
             forceRefetch: true,
-            forceQueryFn: () => ({
+            [forceQueryFnSymbol]: () => ({
               data: upsertRecipe(undefined),
             }),
           })
@@ -357,10 +362,12 @@ export function buildThunks<
         forced:
           arg.type === 'query' ? isForcedQuery(arg, getState()) : undefined,
       }
-      if ('forceQueryFn' in arg && arg.forceQueryFn) {
-        result = arg.forceQueryFn()
-        
-      }else if (endpointDefinition.query) {
+
+      const forceQueryFn =
+        arg.type === 'query' ? arg[forceQueryFnSymbol] : undefined
+      if (forceQueryFn) {
+        result = forceQueryFn()
+      } else if (endpointDefinition.query) {
         result = await baseQuery(
           endpointDefinition.query(arg.originalArgs),
           baseQueryApi,

From 525b52a4bb13472092baa2042c6ba82ac5b914a9 Mon Sep 17 00:00:00 2001
From: Barnabas Jovanovics <barnabas.jovanovics@gmail.com>
Date: Sun, 26 Jun 2022 13:51:20 +0200
Subject: [PATCH 4/9] dumb down upsertQueryData and add emptyCache boolean to
 patch result

---
 .../toolkit/src/query/core/buildThunks.ts     | 78 ++++++-------------
 1 file changed, 24 insertions(+), 54 deletions(-)

diff --git a/packages/toolkit/src/query/core/buildThunks.ts b/packages/toolkit/src/query/core/buildThunks.ts
index 5d2b2ab580..e3b1c0a547 100644
--- a/packages/toolkit/src/query/core/buildThunks.ts
+++ b/packages/toolkit/src/query/core/buildThunks.ts
@@ -175,13 +175,19 @@ export type UpsertQueryDataThunk<
 > = <EndpointName extends QueryKeys<Definitions>>(
   endpointName: EndpointName,
   args: QueryArgFrom<Definitions[EndpointName]>,
-  upsertRecipe: UpsertRecipe<ResultTypeFrom<Definitions[EndpointName]>>
-) => ThunkAction<PatchCollection, PartialState, any, AnyAction>
+  value: ResultTypeFrom<Definitions[EndpointName]>
+) => ThunkAction<void, PartialState, any, AnyAction>
 
 /**
  * An object returned from dispatching a `api.util.updateQueryData` call.
  */
 export type PatchCollection = {
+  /**
+   * A boolean stating if there was already data in the query cache
+   *
+   * If there was no data in the cache no update operation is performed
+   */
+  emptyCache: boolean
   /**
    * An `immer` Patch describing the cache update.
    */
@@ -236,6 +242,7 @@ export function buildThunks<
         api.endpoints[endpointName] as ApiEndpointQuery<any, any>
       ).select(args)(getState())
       let ret: PatchCollection = {
+        emptyCache: true,
         patches: [],
         inversePatches: [],
         undo: () =>
@@ -247,6 +254,7 @@ export function buildThunks<
         return ret
       }
       if ('data' in currentState) {
+        ret.emptyCache = false
         if (isDraftable(currentState.data)) {
           const [, patches, inversePatches] = produceWithPatches(
             currentState.data,
@@ -271,59 +279,21 @@ export function buildThunks<
     }
 
   const upsertQueryData: UpsertQueryDataThunk<EndpointDefinitions, State> =
-    (endpointName, args, upsertRecipe) => (dispatch, getState) => {
-      const currentState = (
-        api.endpoints[endpointName] as ApiEndpointQuery<any, any>
-      ).select(args)(getState())
-      let ret: PatchCollection = {
-        patches: [],
-        inversePatches: [],
-        undo: () =>
-          dispatch(
-            api.util.patchQueryData(endpointName, args, ret.inversePatches)
-          ),
-      }
-      if ('data' in currentState) {
-        if (isDraftable(currentState.data)) {
-          const [, patches, inversePatches] = produceWithPatches(
-            currentState.data,
-            upsertRecipe
-          )
-          ret.patches.push(...patches)
-          ret.inversePatches.push(...inversePatches)
-        } else {
-          const value = upsertRecipe(currentState.data)
-          ret.patches.push({ op: 'replace', path: [], value })
-          ret.inversePatches.push({
-            op: 'replace',
-            path: [],
-            value: currentState.data,
-          })
-        }
-        dispatch(api.util.patchQueryData(endpointName, args, ret.patches))
-      } else {
-        ret.inversePatches.push({
-          op: 'replace',
-          path: [],
-          value: undefined,
+    (endpointName, args, value) => (dispatch) => {
+      dispatch(
+        (
+          api.endpoints[endpointName] as ApiEndpointQuery<
+            QueryDefinition<any, any, any, any, any>,
+            Definitions
+          >
+        ).initiate(args, {
+          subscribe: false,
+          forceRefetch: true,
+          [forceQueryFnSymbol]: () => ({
+            data: value,
+          }),
         })
-        dispatch(
-          (
-            api.endpoints[endpointName] as ApiEndpointQuery<
-              QueryDefinition<any, any, any, any, any>,
-              Definitions
-            >
-          ).initiate(args, {
-            subscribe: false,
-            forceRefetch: true,
-            [forceQueryFnSymbol]: () => ({
-              data: upsertRecipe(undefined),
-            }),
-          })
-        )
-      }
-
-      return ret
+      )
     }
 
   const executeEndpoint: AsyncThunkPayloadCreator<

From 36feac51f0cf685ff7f694d010c3032757a046ff Mon Sep 17 00:00:00 2001
From: Barnabas Jovanovics <barnabas.jovanovics@gmail.com>
Date: Mon, 27 Jun 2022 17:16:18 +0200
Subject: [PATCH 5/9] use a better key name

---
 packages/toolkit/src/query/core/buildThunks.ts | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/packages/toolkit/src/query/core/buildThunks.ts b/packages/toolkit/src/query/core/buildThunks.ts
index e3b1c0a547..fd0cfc95e9 100644
--- a/packages/toolkit/src/query/core/buildThunks.ts
+++ b/packages/toolkit/src/query/core/buildThunks.ts
@@ -187,7 +187,7 @@ export type PatchCollection = {
    *
    * If there was no data in the cache no update operation is performed
    */
-  emptyCache: boolean
+  cacheEntryFound: boolean
   /**
    * An `immer` Patch describing the cache update.
    */
@@ -242,7 +242,7 @@ export function buildThunks<
         api.endpoints[endpointName] as ApiEndpointQuery<any, any>
       ).select(args)(getState())
       let ret: PatchCollection = {
-        emptyCache: true,
+        cacheEntryFound: true,
         patches: [],
         inversePatches: [],
         undo: () =>
@@ -254,7 +254,7 @@ export function buildThunks<
         return ret
       }
       if ('data' in currentState) {
-        ret.emptyCache = false
+        ret.cacheEntryFound = false
         if (isDraftable(currentState.data)) {
           const [, patches, inversePatches] = produceWithPatches(
             currentState.data,

From bac1e033f8c955869682399dcecc5c86631b5bf4 Mon Sep 17 00:00:00 2001
From: Barnabas Jovanovics <barnabas.jovanovics@gmail.com>
Date: Mon, 27 Jun 2022 17:18:14 +0200
Subject: [PATCH 6/9] fix logic

---
 packages/toolkit/src/query/core/buildThunks.ts | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/packages/toolkit/src/query/core/buildThunks.ts b/packages/toolkit/src/query/core/buildThunks.ts
index fd0cfc95e9..1ab4e8920c 100644
--- a/packages/toolkit/src/query/core/buildThunks.ts
+++ b/packages/toolkit/src/query/core/buildThunks.ts
@@ -242,7 +242,7 @@ export function buildThunks<
         api.endpoints[endpointName] as ApiEndpointQuery<any, any>
       ).select(args)(getState())
       let ret: PatchCollection = {
-        cacheEntryFound: true,
+        cacheEntryFound: false,
         patches: [],
         inversePatches: [],
         undo: () =>
@@ -254,7 +254,7 @@ export function buildThunks<
         return ret
       }
       if ('data' in currentState) {
-        ret.cacheEntryFound = false
+        ret.cacheEntryFound = true
         if (isDraftable(currentState.data)) {
           const [, patches, inversePatches] = produceWithPatches(
             currentState.data,

From 936493e8756c0044377b2be7711ceef3dd68d4de Mon Sep 17 00:00:00 2001
From: Mark Erikson <mark@isquaredsoftware.com>
Date: Sat, 27 Aug 2022 15:37:03 -0400
Subject: [PATCH 7/9] Update `upsertQueryData` tests to match current behavior

---
 .../toolkit/src/query/core/buildThunks.ts     |  8 --
 .../query/tests/optimisticUpserts.test.tsx    | 95 +++++++------------
 2 files changed, 32 insertions(+), 71 deletions(-)

diff --git a/packages/toolkit/src/query/core/buildThunks.ts b/packages/toolkit/src/query/core/buildThunks.ts
index 1ab4e8920c..4535892b9c 100644
--- a/packages/toolkit/src/query/core/buildThunks.ts
+++ b/packages/toolkit/src/query/core/buildThunks.ts
@@ -182,12 +182,6 @@ export type UpsertQueryDataThunk<
  * An object returned from dispatching a `api.util.updateQueryData` call.
  */
 export type PatchCollection = {
-  /**
-   * A boolean stating if there was already data in the query cache
-   *
-   * If there was no data in the cache no update operation is performed
-   */
-  cacheEntryFound: boolean
   /**
    * An `immer` Patch describing the cache update.
    */
@@ -242,7 +236,6 @@ export function buildThunks<
         api.endpoints[endpointName] as ApiEndpointQuery<any, any>
       ).select(args)(getState())
       let ret: PatchCollection = {
-        cacheEntryFound: false,
         patches: [],
         inversePatches: [],
         undo: () =>
@@ -254,7 +247,6 @@ export function buildThunks<
         return ret
       }
       if ('data' in currentState) {
-        ret.cacheEntryFound = true
         if (isDraftable(currentState.data)) {
           const [, patches, inversePatches] = produceWithPatches(
             currentState.data,
diff --git a/packages/toolkit/src/query/tests/optimisticUpserts.test.tsx b/packages/toolkit/src/query/tests/optimisticUpserts.test.tsx
index 9461e3483a..0884e63507 100644
--- a/packages/toolkit/src/query/tests/optimisticUpserts.test.tsx
+++ b/packages/toolkit/src/query/tests/optimisticUpserts.test.tsx
@@ -1,7 +1,7 @@
 import { createApi } from '@reduxjs/toolkit/query/react'
 import { actionsReducer, hookWaitFor, setupApiStore, waitMs } from './helpers'
 import { skipToken } from '../core/buildSelectors'
-import { renderHook, act } from '@testing-library/react-hooks'
+import { renderHook, act } from '@testing-library/react'
 
 interface Post {
   id: string
@@ -33,13 +33,16 @@ const api = createApi({
         method: 'PATCH',
         body: patch,
       }),
-      async onQueryStarted({ id, ...patch }, { dispatch, queryFulfilled }) {
-        const { undo } = dispatch(
-          api.util.upsertQueryData('post', id, (draft) => {
-            Object.assign(draft, patch)
-          })
-        )
-        queryFulfilled.catch(undo)
+      async onQueryStarted(arg, { dispatch, queryFulfilled, getState }) {
+        const currentItem = api.endpoints.post.select(arg.id)(getState())
+        if (currentItem?.data) {
+          dispatch(
+            api.util.upsertQueryData('post', arg.id, {
+              ...currentItem.data,
+              ...arg,
+            })
+          )
+        }
       },
       invalidatesTags: (result) => (result ? ['Post'] : []),
     }),
@@ -154,13 +157,12 @@ describe('upsertQueryData', () => {
       contents: 'TODO',
     })
 
-    let returnValue!: ReturnType<ReturnType<typeof api.util.upsertQueryData>>
-    act(() => {
-      returnValue = storeRef.store.dispatch(
-        api.util.upsertQueryData('post', '3', (draft) => {
-          if (draft) {
-            draft.contents = 'I love cheese!'
-          }
+    await act(async () => {
+      storeRef.store.dispatch(
+        api.util.upsertQueryData('post', '3', {
+          id: '3',
+          title: 'All about cheese.',
+          contents: 'I love cheese!',
         })
       )
     })
@@ -171,20 +173,6 @@ describe('upsertQueryData', () => {
       title: 'All about cheese.',
       contents: 'I love cheese!',
     })
-
-    expect(returnValue).toEqual({
-      inversePatches: [{ op: 'replace', path: ['contents'], value: 'TODO' }],
-      patches: [{ op: 'replace', path: ['contents'], value: 'I love cheese!' }],
-      undo: expect.any(Function),
-    })
-
-    act(() => {
-      storeRef.store.dispatch(
-        api.util.patchQueryData('post', '3', returnValue.inversePatches)
-      )
-    })
-
-    expect(result.current.data).toEqual(dataBefore)
   })
 
   test('does update non-existing values', async () => {
@@ -211,30 +199,14 @@ describe('upsertQueryData', () => {
     // upsert the data
     act(() => {
       returnValue = storeRef.store.dispatch(
-        api.util.upsertQueryData('post', '4', (draft) => {
-          if (draft) {
-            draft.contents = 'I love cheese!'
-          } else {
-            return {
-              id: '4',
-              title: 'All about cheese',
-              contents: 'I love cheese!',
-            }
-          }
+        api.util.upsertQueryData('post', '4', {
+          id: '4',
+          title: 'All about cheese',
+          contents: 'I love cheese!',
         })
       )
     })
 
-    // the patch would remove the result again
-    // maybe this needs to be implemented differently in order
-    // for the whole thing to work correctly, I suppose that
-    // this would only revert the data but would not invalidate it
-    expect(returnValue).toEqual({
-      inversePatches: [{ op: 'replace', path: [], value: undefined }],
-      patches: [],
-      undo: expect.any(Function),
-    })
-
     // rerender the hook
     rerender()
     // wait until everything has settled
@@ -285,13 +257,16 @@ describe('full integration', () => {
       contents: 'TODO',
     })
 
-    act(() => {
-      result.current.mutation[0]({ id: '3', contents: 'Delicious cheese!' })
+    await act(async () => {
+      await result.current.mutation[0]({
+        id: '3',
+        contents: 'Delicious cheese!',
+      })
     })
 
     expect(result.current.query.data).toEqual({
       id: '3',
-      title: 'All about cheese.',
+      title: 'Meanwhile, this changed server-side.',
       contents: 'Delicious cheese!',
     })
 
@@ -336,8 +311,11 @@ describe('full integration', () => {
       contents: 'TODO',
     })
 
-    act(() => {
-      result.current.mutation[0]({ id: '3', contents: 'Delicious cheese!' })
+    await act(async () => {
+      await result.current.mutation[0]({
+        id: '3',
+        contents: 'Delicious cheese!',
+      })
     })
 
     // optimistic update
@@ -347,15 +325,6 @@ describe('full integration', () => {
       contents: 'Delicious cheese!',
     })
 
-    // rollback
-    await hookWaitFor(() =>
-      expect(result.current.query.data).toEqual({
-        id: '3',
-        title: 'All about cheese.',
-        contents: 'TODO',
-      })
-    )
-
     // mutation failed - will not invalidate query and not refetch data from the server
     await expect(() =>
       hookWaitFor(

From 9f12efd55e673a37591ba00b035f1db117e22d45 Mon Sep 17 00:00:00 2001
From: Mark Erikson <mark@isquaredsoftware.com>
Date: Sat, 27 Aug 2022 19:32:14 -0400
Subject: [PATCH 8/9] Enable `upsertQueryData` calls while a request is in
 flight

---
 packages/toolkit/src/query/core/buildSlice.ts |  9 ++-
 .../toolkit/src/query/core/buildThunks.ts     | 19 ++++--
 .../query/tests/optimisticUpserts.test.tsx    | 61 ++++++++++++++++++-
 3 files changed, 81 insertions(+), 8 deletions(-)

diff --git a/packages/toolkit/src/query/core/buildSlice.ts b/packages/toolkit/src/query/core/buildSlice.ts
index cdfa6625f4..c9db2acbc5 100644
--- a/packages/toolkit/src/query/core/buildSlice.ts
+++ b/packages/toolkit/src/query/core/buildSlice.ts
@@ -157,7 +157,14 @@ export function buildSlice({
             draft,
             meta.arg.queryCacheKey,
             (substate) => {
-              if (substate.requestId !== meta.requestId) return
+              if (substate.requestId !== meta.requestId) {
+                if (
+                  substate.fulfilledTimeStamp &&
+                  meta.fulfilledTimeStamp < substate.fulfilledTimeStamp
+                ) {
+                  return
+                }
+              }
               const { merge } = definitions[
                 meta.arg.endpointName
               ] as QueryDefinition<any, any, any, any>
diff --git a/packages/toolkit/src/query/core/buildThunks.ts b/packages/toolkit/src/query/core/buildThunks.ts
index 4535892b9c..3248bcfb7c 100644
--- a/packages/toolkit/src/query/core/buildThunks.ts
+++ b/packages/toolkit/src/query/core/buildThunks.ts
@@ -10,6 +10,7 @@ import { QueryStatus } from './apiState'
 import {
   forceQueryFnSymbol,
   StartQueryActionCreatorOptions,
+  QueryActionCreatorResult,
 } from './buildInitiate'
 import type {
   AssertTagTypes,
@@ -176,7 +177,12 @@ export type UpsertQueryDataThunk<
   endpointName: EndpointName,
   args: QueryArgFrom<Definitions[EndpointName]>,
   value: ResultTypeFrom<Definitions[EndpointName]>
-) => ThunkAction<void, PartialState, any, AnyAction>
+) => ThunkAction<
+  QueryActionCreatorResult<Definitions[EndpointName]>,
+  PartialState,
+  any,
+  AnyAction
+>
 
 /**
  * An object returned from dispatching a `api.util.updateQueryData` call.
@@ -272,7 +278,7 @@ export function buildThunks<
 
   const upsertQueryData: UpsertQueryDataThunk<EndpointDefinitions, State> =
     (endpointName, args, value) => (dispatch) => {
-      dispatch(
+      return dispatch(
         (
           api.endpoints[endpointName] as ApiEndpointQuery<
             QueryDefinition<any, any, any, any, any>,
@@ -469,12 +475,17 @@ In the case of an unhandled error, no tags will be "provided" or "invalidated".`
       const requestState = state[reducerPath]?.queries?.[arg.queryCacheKey]
       const fulfilledVal = requestState?.fulfilledTimeStamp
 
-      // Don't retry a request that's currently in-flight
-      if (requestState?.status === 'pending') return false
+      // Order of these checks matters.
+      // In order for `upsertQueryData` to successfully run while an existing request is
+      /// in flight, we have to check `isForcedQuery` before `status === 'pending'`,
+      // otherwise `queryThunk` will bail out and not run at all.
 
       // if this is forced, continue
       if (isForcedQuery(arg, state)) return true
 
+      // Don't retry a request that's currently in-flight
+      if (requestState?.status === 'pending') return false
+
       // Pull from the cache unless we explicitly force refetch or qualify based on time
       if (fulfilledVal)
         // Value is cached and we didn't specify to refresh, skip it.
diff --git a/packages/toolkit/src/query/tests/optimisticUpserts.test.tsx b/packages/toolkit/src/query/tests/optimisticUpserts.test.tsx
index 0884e63507..56693a1e1b 100644
--- a/packages/toolkit/src/query/tests/optimisticUpserts.test.tsx
+++ b/packages/toolkit/src/query/tests/optimisticUpserts.test.tsx
@@ -1,7 +1,7 @@
 import { createApi } from '@reduxjs/toolkit/query/react'
 import { actionsReducer, hookWaitFor, setupApiStore, waitMs } from './helpers'
 import { skipToken } from '../core/buildSelectors'
-import { renderHook, act } from '@testing-library/react'
+import { renderHook, act, waitFor } from '@testing-library/react'
 
 interface Post {
   id: string
@@ -12,6 +12,10 @@ interface Post {
 const baseQuery = jest.fn()
 beforeEach(() => baseQuery.mockReset())
 
+function delay(ms: number) {
+  return new Promise((resolve) => setTimeout(resolve, ms))
+}
+
 const api = createApi({
   baseQuery: (...args: any[]) => {
     const result = baseQuery(...args)
@@ -46,6 +50,18 @@ const api = createApi({
       },
       invalidatesTags: (result) => (result ? ['Post'] : []),
     }),
+    post2: build.query<Post, string>({
+      queryFn: async (id) => {
+        await delay(20)
+        return {
+          data: {
+            id,
+            title: 'All about cheese.',
+            contents: 'TODO',
+          },
+        }
+      },
+    }),
   }),
 })
 
@@ -195,10 +211,9 @@ describe('upsertQueryData', () => {
     )
     await hookWaitFor(() => expect(result.current.isError).toBeTruthy())
 
-    let returnValue!: ReturnType<ReturnType<typeof api.util.upsertQueryData>>
     // upsert the data
     act(() => {
-      returnValue = storeRef.store.dispatch(
+      storeRef.store.dispatch(
         api.util.upsertQueryData('post', '4', {
           id: '4',
           title: 'All about cheese',
@@ -351,4 +366,44 @@ describe('full integration', () => {
       50
     )
   })
+
+  test.only('Interop with in-flight requests', async () => {
+    await act(async () => {
+      const fetchRes = storeRef.store.dispatch(
+        api.endpoints.post2.initiate('3')
+      )
+
+      const upsertRes = storeRef.store.dispatch(
+        api.util.upsertQueryData('post2', '3', {
+          id: '3',
+          title: 'Upserted title',
+          contents: 'Upserted contents',
+        })
+      )
+
+      const selectEntry = api.endpoints.post2.select('3')
+      await waitFor(
+        () => {
+          const entry1 = selectEntry(storeRef.store.getState())
+          expect(entry1.data).toEqual({
+            id: '3',
+            title: 'Upserted title',
+            contents: 'Upserted contents',
+          })
+        },
+        { interval: 1, timeout: 15 }
+      )
+      await waitFor(
+        () => {
+          const entry2 = selectEntry(storeRef.store.getState())
+          expect(entry2.data).toEqual({
+            id: '3',
+            title: 'All about cheese.',
+            contents: 'TODO',
+          })
+        },
+        { interval: 1 }
+      )
+    })
+  })
 })

From 06ee327e7f777be2320b77bffdd2e86614238420 Mon Sep 17 00:00:00 2001
From: Mark Erikson <mark@isquaredsoftware.com>
Date: Sat, 27 Aug 2022 19:49:52 -0400
Subject: [PATCH 9/9] Fix `upsertQueryData` thunk types

---
 packages/toolkit/src/query/core/buildThunks.ts | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/packages/toolkit/src/query/core/buildThunks.ts b/packages/toolkit/src/query/core/buildThunks.ts
index 3248bcfb7c..e912922d1f 100644
--- a/packages/toolkit/src/query/core/buildThunks.ts
+++ b/packages/toolkit/src/query/core/buildThunks.ts
@@ -178,7 +178,11 @@ export type UpsertQueryDataThunk<
   args: QueryArgFrom<Definitions[EndpointName]>,
   value: ResultTypeFrom<Definitions[EndpointName]>
 ) => ThunkAction<
-  QueryActionCreatorResult<Definitions[EndpointName]>,
+  QueryActionCreatorResult<
+    Definitions[EndpointName] extends QueryDefinition<any, any, any, any>
+      ? Definitions[EndpointName]
+      : never
+  >,
   PartialState,
   any,
   AnyAction
@@ -276,7 +280,7 @@ export function buildThunks<
       return ret
     }
 
-  const upsertQueryData: UpsertQueryDataThunk<EndpointDefinitions, State> =
+  const upsertQueryData: UpsertQueryDataThunk<Definitions, State> =
     (endpointName, args, value) => (dispatch) => {
       return dispatch(
         (