First draft implementation of upsertCacheEntries

Author: markerikson

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

@@ -8,6 +8,8 @@ import {
   isRejectedWithValue,
   createNextState,
   prepareAutoBatched,
+  SHOULD_AUTOBATCH,
+  nanoid,
 } from './rtkImports'
 import type {
   QuerySubstateIdentifier,
@@ -21,15 +23,24 @@ import type {
   QueryCacheKey,
   SubscriptionState,
   ConfigState,
+  QueryKeys,
 } from './apiState'
 import { QueryStatus } from './apiState'
-import type { MutationThunk, QueryThunk, RejectedAction } from './buildThunks'
+import type {
+  MutationThunk,
+  QueryThunk,
+  QueryThunkArg,
+  RejectedAction,
+} from './buildThunks'
 import { calculateProvidedByThunk } from './buildThunks'
 import type {
   AssertTagTypes,
+  DefinitionType,
   EndpointDefinitions,
   FullTagDescription,
+  QueryArgFrom,
   QueryDefinition,
+  ResultTypeFrom,
 } from '../endpointDefinitions'
 import type { Patch } from 'immer'
 import { isDraft } from 'immer'
@@ -42,6 +53,44 @@ import {
 } from '../utils'
 import type { ApiContext } from '../apiTypes'
 import { isUpsertQuery } from './buildInitiate'
+import type { InternalSerializeQueryArgs } from '../defaultSerializeQueryArgs'
+
+/**
+ * A typesafe single entry to be upserted into the cache
+ */
+export type NormalizedQueryUpsertEntry<
+  Definitions extends EndpointDefinitions,
+  EndpointName extends QueryKeys<Definitions>,
+> = {
+  endpointName: EndpointName
+  args: QueryArgFrom<Definitions[EndpointName]>
+  value: ResultTypeFrom<Definitions[EndpointName]>
+}
+
+/**
+ * The internal version that is not typesafe since we can't carry the generics through `createSlice`
+ */
+type NormalizedQueryUpsertEntryPayload = {
+  endpointName: string
+  args: any
+  value: any
+}
+
+/**
+ * A typesafe representation of a util action creator that accepts cache entry descriptions to upsert
+ */
+export type UpsertEntries<Definitions extends EndpointDefinitions> = <
+  EndpointNames extends Array<QueryKeys<Definitions>>,
+>(
+  entries: [
+    ...{
+      [I in keyof EndpointNames]: NormalizedQueryUpsertEntry<
+        Definitions,
+        EndpointNames[I]
+      >
+    },
+  ],
+) => PayloadAction<NormalizedQueryUpsertEntryPayload>
 
 function updateQuerySubstateIfExists(
   state: QueryState<any>,
@@ -92,6 +141,7 @@ export function buildSlice({
   reducerPath,
   queryThunk,
   mutationThunk,
+  serializeQueryArgs,
   context: {
     endpointDefinitions: definitions,
     apiUid,
@@ -104,6 +154,7 @@ export function buildSlice({
   reducerPath: string
   queryThunk: QueryThunk
   mutationThunk: MutationThunk
+  serializeQueryArgs: InternalSerializeQueryArgs
   context: ApiContext<EndpointDefinitions>
   assertTagType: AssertTagTypes
   config: Omit<
@@ -221,6 +272,63 @@ export function buildSlice({
         },
         prepare: prepareAutoBatched<QuerySubstateIdentifier>(),
       },
+      cacheEntriesUpserted: {
+        reducer(
+          draft,
+          action: PayloadAction<
+            NormalizedQueryUpsertEntryPayload[],
+            string,
+            {
+              RTK_autoBatch: boolean
+              requestId: string
+              timestamp: number
+            }
+          >,
+        ) {
+          for (const entry of action.payload) {
+            const { endpointName, args, value } = entry
+            const endpointDefinition = definitions[endpointName]
+
+            const arg: QueryThunkArg = {
+              type: 'query',
+              endpointName: endpointName,
+              originalArgs: entry.args,
+              queryCacheKey: serializeQueryArgs({
+                queryArgs: args,
+                endpointDefinition,
+                endpointName,
+              }),
+            }
+            writePendingCacheEntry(draft, arg, true, {
+              arg,
+              requestId: action.meta.requestId,
+              startedTimeStamp: action.meta.timestamp,
+            })
+
+            writeFulfilledCacheEntry(
+              draft,
+              {
+                arg,
+                requestId: action.meta.requestId,
+                fulfilledTimeStamp: action.meta.timestamp,
+                baseQueryMeta: {},
+              },
+              entry.value,
+            )
+          }
+        },
+        prepare: (payload: NormalizedQueryUpsertEntryPayload[]) => {
+          const result = {
+            payload,
+            meta: {
+              [SHOULD_AUTOBATCH]: true,
+              requestId: nanoid(),
+              timestamp: Date.now(),
+            },
+          }
+          return result
+        },
+      },
       queryResultPatched: {
         reducer(
           draft,

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

@@ -47,7 +47,7 @@ import type {
   BuildSelectorsApiEndpointQuery,
 } from './buildSelectors'
 import { buildSelectors } from './buildSelectors'
-import type { SliceActions } from './buildSlice'
+import type { SliceActions, UpsertEntries } from './buildSlice'
 import { buildSlice } from './buildSlice'
 import type {
   BuildThunksApiEndpointMutation,
@@ -320,6 +320,9 @@ export interface ApiModules<
        * ```
        */
       resetApiState: SliceActions['resetApiState']
+
+      upsertEntries: UpsertEntries<Definitions>
+
       /**
        * A Redux action creator that can be used to manually invalidate cache tags for [automated re-fetching](../../usage/automated-refetching.mdx).
        *
@@ -527,6 +530,7 @@ export const coreModule = ({
       context,
       queryThunk,
       mutationThunk,
+      serializeQueryArgs,
       reducerPath,
       assertTagType,
       config: {
@@ -545,6 +549,7 @@ export const coreModule = ({
       upsertQueryData,
       prefetch,
       resetApiState: sliceActions.resetApiState,
+      upsertEntries: sliceActions.cacheEntriesUpserted as any,
     })
     safeAssign(api.internalActions, sliceActions)
 

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

@@ -27,6 +27,9 @@ const api = createApi({
   },
   tagTypes: ['Post'],
   endpoints: (build) => ({
+    getPosts: build.query<Post[], void>({
+      query: () => '/posts',
+    }),
     post: build.query<Post, string>({
       query: (id) => `post/${id}`,
       providesTags: ['Post'],
@@ -327,6 +330,51 @@ describe('upsertQueryData', () => {
   })
 })
 
+describe('upsertEntries', () => {
+  test('Upserts many entries at once', async () => {
+    const posts: Post[] = [
+      {
+        id: '1',
+        contents: 'A',
+        title: 'A',
+      },
+      {
+        id: '2',
+        contents: 'B',
+        title: 'B',
+      },
+      {
+        id: '3',
+        contents: 'C',
+        title: 'C',
+      },
+    ]
+
+    storeRef.store.dispatch(
+      api.util.upsertEntries([
+        {
+          endpointName: 'getPosts',
+          args: undefined,
+          value: posts,
+        },
+        ...posts.map((post) => ({
+          endpointName: 'post' as const,
+          args: post.id,
+          value: post,
+        })),
+      ]),
+    )
+
+    const state = storeRef.store.getState()
+
+    expect(api.endpoints.getPosts.select()(state).data).toBe(posts)
+
+    expect(api.endpoints.post.select('1')(state).data).toBe(posts[0])
+    expect(api.endpoints.post.select('2')(state).data).toBe(posts[1])
+    expect(api.endpoints.post.select('3')(state).data).toBe(posts[2])
+  })
+})
+
 describe('full integration', () => {
   test('success case', async () => {
     baseQuery