Add upsertQueryData functionality per issue #1720

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

@@ -1,4 +1,4 @@
-import type { AnyAction, PayloadAction } from '@reduxjs/toolkit'
+import type { PayloadAction } from '@reduxjs/toolkit'
 import {
   combineReducers,
   createAction,
@@ -22,7 +22,7 @@ import type {
   ConfigState,
 } from './apiState'
 import { QueryStatus } from './apiState'
-import type { MutationThunk, QueryThunk } from './buildThunks'
+import type { MutationThunk, QueryThunk, QueryThunkArg } from './buildThunks'
 import { calculateProvidedByThunk } from './buildThunks'
 import type {
   AssertTagTypes,
@@ -129,6 +129,31 @@ export function buildSlice({
           substate.data = applyPatches(substate.data as any, patches.concat())
         })
       },
+      insertQueryResult(
+        draft,
+        {
+          payload: {
+            data,
+            arg: { originalArgs, endpointName, queryCacheKey },
+            fulfilledTimeStamp,
+          },
+        }: PayloadAction<{
+          data: any
+          arg: QueryThunkArg
+          fulfilledTimeStamp: number
+        }>
+      ) {
+        draft[queryCacheKey] = {
+          status: QueryStatus.fulfilled,
+          endpointName,
+          requestId: '',
+          originalArgs,
+          startedTimeStamp: fulfilledTimeStamp,
+          fulfilledTimeStamp,
+          data,
+          error: undefined,
+        }
+      },
     },
     extraReducers(builder) {
       builder

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

@@ -163,6 +163,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]>,
+  data: ResultTypeFrom<Definitions[EndpointName]>
+) => ThunkAction<PatchCollection, PartialState, any, AnyAction>
+
 /**
  * An object returned from dispatching a `api.util.updateQueryData` call.
  */
@@ -200,16 +209,26 @@ export function buildThunks<
 }) {
   type State = RootState<any, string, ReducerPath>
 
+  const querySubstateIdentifier = <
+    EndpointName extends string & QueryKeys<Definitions>
+  >(
+    endpointName: EndpointName,
+    args: QueryArgFrom<Definitions[EndpointName]>
+  ): QuerySubstateIdentifier => {
+    const endpointDefinition = endpointDefinitions[endpointName]
+    const queryCacheKey = serializeQueryArgs({
+      queryArgs: args,
+      endpointDefinition,
+      endpointName,
+    })
+    return { queryCacheKey }
+  }
+
   const patchQueryData: PatchQueryDataThunk<EndpointDefinitions, State> =
     (endpointName, args, patches) => (dispatch) => {
-      const endpointDefinition = endpointDefinitions[endpointName]
       dispatch(
         api.internalActions.queryResultPatched({
-          queryCacheKey: serializeQueryArgs({
-            queryArgs: args,
-            endpointDefinition,
-            endpointName,
-          }),
+          ...querySubstateIdentifier(endpointName, args),
           patches,
         })
       )
@@ -255,6 +274,43 @@ export function buildThunks<
       return ret
     }
 
+  const upsertQueryData: UpsertQueryDataThunk<EndpointDefinitions, State> =
+    (endpointName, args, data) => (dispatch, getState) => {
+      const currentState = (
+        api.endpoints[endpointName] as ApiEndpointQuery<any, any>
+      ).select(args)(getState())
+      if ('data' in currentState) {
+        return dispatch(
+          api.util.updateQueryData(endpointName, args, () => data)
+        )
+      }
+      let ret: PatchCollection = {
+        patches: [{ op: 'replace', path: [], value: data }],
+        inversePatches: [{ op: 'replace', path: [], value: undefined }],
+        undo: () =>
+          dispatch(
+            api.internalActions.removeQueryResult(
+              querySubstateIdentifier(endpointName, args)
+            )
+          ),
+      }
+
+      dispatch(
+        api.internalActions.insertQueryResult({
+          arg: {
+            ...querySubstateIdentifier(endpointName, args),
+            endpointName,
+            originalArgs: args,
+            type: 'query',
+          },
+          data,
+          fulfilledTimeStamp: Date.now(),
+        })
+      )
+
+      return ret
+    }
+
   const executeEndpoint: AsyncThunkPayloadCreator<
     ThunkResult,
     QueryThunkArg | MutationThunkArg,
@@ -494,6 +550,7 @@ In the case of an unhandled error, no tags will be "provided" or "invalidated".`
     mutationThunk,
     prefetch,
     updateQueryData,
+    upsertQueryData,
     patchQueryData,
     buildMatchThunkActions,
   }

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,25 @@ declare module '../apiTypes' {
           Definitions,
           RootState<Definitions, string, ReducerPath>
         >
+        /**
+         * A Redux thunk that manually adds a 'fulfilled' result to the API cache state with the provided data.  Unlike `patchQueryData`, which can only update previously-fetched data, `upsertQueryData` can both update existing results and add completely new entries to the cache.
+         *
+         * The thunk action creator accepts three arguments: the name of the endpoint we are updating (such as `'getPost'`), any relevant query arguments, and the result data for this API call.
+         *
+         * Caution: This is an advanced function which should be avoided unless absolutely necessary.
+         *
+         *  @example
+         *
+         * ```ts
+         * dispatch(
+         *   api.util.updateQueryData('getPosts', '1', { id: 1, name: 'Teddy' })
+         * )
+         * ```
+         */
+        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.
          *
@@ -406,6 +429,7 @@ export const coreModule = (): Module<CoreModule> => ({
       mutationThunk,
       patchQueryData,
       updateQueryData,
+      upsertQueryData,
       prefetch,
       buildMatchThunkActions,
     } = buildThunks({
@@ -434,6 +458,7 @@ export const coreModule = (): Module<CoreModule> => ({
     safeAssign(api.util, {
       patchQueryData,
       updateQueryData,
+      upsertQueryData,
       prefetch,
       resetApiState: sliceActions.resetApiState,
     })

packages/toolkit/src/query/tests/optimisticUpdates.test.tsx

@@ -224,6 +224,85 @@ describe('updateQueryData', () => {
   })
 })
 
+describe('upsertQueryData', () => {
+  test('can add new cache entries', async () => {
+    const { result } = renderHook(
+      () => api.endpoints.post.useQuery('4', { skip: true }),
+      {
+        wrapper: storeRef.wrapper,
+      }
+    )
+    await hookWaitFor(() => expect(result.current.isUninitialized).toBeTruthy())
+
+    const dataBefore = result.current.data
+    expect(dataBefore).toBeUndefined()
+
+    let returnValue!: ReturnType<ReturnType<typeof api.util.updateQueryData>>
+    act(() => {
+      returnValue = storeRef.store.dispatch(
+        api.util.upsertQueryData('post', '4', {
+          id: '4',
+          title: 'All about cheese.',
+          contents: 'I love cheese!',
+        })
+      )
+    })
+
+    const selector = api.endpoints.post.select('4')
+
+    const queryStateAfter = selector(storeRef.store.getState())
+
+    expect(queryStateAfter.data).toEqual({
+      id: '4',
+      title: 'All about cheese.',
+      contents: 'I love cheese!',
+    })
+
+    // TODO: expect(returnValue).toEqual(???)
+  })
+
+  test('can update existing', async () => {
+    baseQuery
+      .mockImplementationOnce(async () => ({
+        id: '3',
+        title: 'All about cheese.',
+        contents: 'TODO',
+      }))
+      .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.updateQueryData>>
+    act(() => {
+      returnValue = storeRef.store.dispatch(
+        api.util.upsertQueryData('post', '3', {
+          id: '3',
+          title: 'All about cheese.',
+          contents: 'I love cheese!',
+        })
+      )
+    })
+
+    expect(result.current.data).toEqual({
+      id: '3',
+      title: 'All about cheese.',
+      contents: 'I love cheese!',
+    })
+
+    // TODO: expect(returnValue).toEqual(???)
+  })
+})
+
 describe('full integration', () => {
   test('success case', async () => {
     baseQuery