Fix updateQueryData for infinite queries (#4796)

* Fix updateQueryData for infinite queries

* Fix upsertQueryData for infinite queries

* Fix upsertQueryEntries for infinite queries

* Fix types in lifecycle middleware

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

@@ -154,6 +154,19 @@ export type QueryKeys<Definitions extends EndpointDefinitions> = {
     ? K
     : never
 }[keyof Definitions]
+
+export type InfiniteQueryKeys<Definitions extends EndpointDefinitions> = {
+  [K in keyof Definitions]: Definitions[K] extends InfiniteQueryDefinition<
+    any,
+    any,
+    any,
+    any,
+    any
+  >
+    ? K
+    : never
+}[keyof Definitions]
+
 export type MutationKeys<Definitions extends EndpointDefinitions> = {
   [K in keyof Definitions]: Definitions[K] extends MutationDefinition<
     any,

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

@@ -121,7 +121,7 @@ export type QueryActionCreatorResult<
 export type InfiniteQueryActionCreatorResult<
   D extends InfiniteQueryDefinition<any, any, any, any, any>,
 > = Promise<InfiniteQueryResultSelectorResult<D>> & {
-  arg: QueryArgFrom<D>
+  arg: InfiniteQueryArgFrom<D>
   requestId: string
   subscriptionOptions: SubscriptionOptions | undefined
   abort(): void

packages/toolkit/src/query/core/buildMiddleware/cacheLifecycle.ts

@@ -344,7 +344,7 @@ export const buildCacheLifecycleHandler: InternalHandlerBuilder = ({
             mwApi.dispatch(
               api.util.updateQueryData(
                 endpointName as never,
-                originalArgs,
+                originalArgs as never,
                 updateRecipe,
               ),
             )

packages/toolkit/src/query/core/buildMiddleware/queryLifecycle.ts

@@ -476,7 +476,7 @@ export const buildQueryLifecycleHandler: InternalHandlerBuilder = ({
                 mwApi.dispatch(
                   api.util.updateQueryData(
                     endpointName as never,
-                    originalArgs,
+                    originalArgs as never,
                     updateRecipe,
                   ),
                 )

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

@@ -29,6 +29,9 @@ import type {
 } from './apiState'
 import { QueryStatus } from './apiState'
 import type {
+  AllQueryKeys,
+  QueryArgFromAnyQueryDefinition,
+  DataFromAnyQueryDefinition,
   InfiniteQueryThunk,
   MutationThunk,
   QueryThunk,
@@ -64,11 +67,11 @@ import type { InternalSerializeQueryArgs } from '../defaultSerializeQueryArgs'
  */
 export type NormalizedQueryUpsertEntry<
   Definitions extends EndpointDefinitions,
-  EndpointName extends QueryKeys<Definitions>,
+  EndpointName extends AllQueryKeys<Definitions>,
 > = {
   endpointName: EndpointName
-  arg: QueryArgFrom<Definitions[EndpointName]>
-  value: ResultTypeFrom<Definitions[EndpointName]>
+  arg: QueryArgFromAnyQueryDefinition<Definitions, EndpointName>
+  value: DataFromAnyQueryDefinition<Definitions, EndpointName>
 }
 
 /**
@@ -89,7 +92,7 @@ export type ProcessedQueryUpsertEntry = {
  * 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>>,
+  EndpointNames extends Array<AllQueryKeys<Definitions>>,
 >(
   entries: [
     ...{

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

@@ -6,6 +6,7 @@ import type {
   ThunkDispatch,
   UnknownAction,
 } from '@reduxjs/toolkit'
+import util from 'util'
 import type { Patch } from 'immer'
 import { isDraftable, produceWithPatches } from 'immer'
 import type { Api, ApiContext } from '../apiTypes'
@@ -19,8 +20,10 @@ import type {
   AssertTagTypes,
   EndpointDefinition,
   EndpointDefinitions,
+  InfiniteQueryArgFrom,
   InfiniteQueryDefinition,
   MutationDefinition,
+  PageParamFrom,
   QueryArgFrom,
   QueryDefinition,
   ResultTypeFrom,
@@ -40,9 +43,11 @@ import type {
   InfiniteQueryConfigOptions,
   QueryCacheKey,
   InfiniteQueryDirection,
+  InfiniteQueryKeys,
 } from './apiState'
 import { QueryStatus } from './apiState'
 import type {
+  InfiniteQueryActionCreatorResult,
   QueryActionCreatorResult,
   StartInfiniteQueryActionCreatorOptions,
   StartQueryActionCreatorOptions,
@@ -194,29 +199,80 @@ export type PatchQueryDataThunk<
   updateProvided?: boolean,
 ) => ThunkAction<void, PartialState, any, UnknownAction>
 
+export type AllQueryKeys<Definitions extends EndpointDefinitions> =
+  | QueryKeys<Definitions>
+  | InfiniteQueryKeys<Definitions>
+
+export type QueryArgFromAnyQueryDefinition<
+  Definitions extends EndpointDefinitions,
+  EndpointName extends AllQueryKeys<Definitions>,
+> =
+  Definitions[EndpointName] extends InfiniteQueryDefinition<
+    any,
+    any,
+    any,
+    any,
+    any
+  >
+    ? InfiniteQueryArgFrom<Definitions[EndpointName]>
+    : Definitions[EndpointName] extends QueryDefinition<any, any, any, any>
+      ? QueryArgFrom<Definitions[EndpointName]>
+      : never
+
+export type DataFromAnyQueryDefinition<
+  Definitions extends EndpointDefinitions,
+  EndpointName extends AllQueryKeys<Definitions>,
+> =
+  Definitions[EndpointName] extends InfiniteQueryDefinition<
+    any,
+    any,
+    any,
+    any,
+    any
+  >
+    ? InfiniteData<
+        ResultTypeFrom<Definitions[EndpointName]>,
+        PageParamFrom<Definitions[EndpointName]>
+      >
+    : Definitions[EndpointName] extends QueryDefinition<any, any, any, any>
+      ? ResultTypeFrom<Definitions[EndpointName]>
+      : unknown
+
+export type UpsertThunkResult<
+  Definitions extends EndpointDefinitions,
+  EndpointName extends AllQueryKeys<Definitions>,
+> =
+  Definitions[EndpointName] extends InfiniteQueryDefinition<
+    any,
+    any,
+    any,
+    any,
+    any
+  >
+    ? InfiniteQueryActionCreatorResult<Definitions[EndpointName]>
+    : Definitions[EndpointName] extends QueryDefinition<any, any, any, any>
+      ? QueryActionCreatorResult<Definitions[EndpointName]>
+      : QueryActionCreatorResult<never>
+
 export type UpdateQueryDataThunk<
   Definitions extends EndpointDefinitions,
   PartialState,
-> = <EndpointName extends QueryKeys<Definitions>>(
+> = <EndpointName extends AllQueryKeys<Definitions>>(
   endpointName: EndpointName,
-  arg: QueryArgFrom<Definitions[EndpointName]>,
-  updateRecipe: Recipe<ResultTypeFrom<Definitions[EndpointName]>>,
+  arg: QueryArgFromAnyQueryDefinition<Definitions, EndpointName>,
+  updateRecipe: Recipe<DataFromAnyQueryDefinition<Definitions, EndpointName>>,
   updateProvided?: boolean,
 ) => ThunkAction<PatchCollection, PartialState, any, UnknownAction>
 
 export type UpsertQueryDataThunk<
   Definitions extends EndpointDefinitions,
   PartialState,
-> = <EndpointName extends QueryKeys<Definitions>>(
+> = <EndpointName extends AllQueryKeys<Definitions>>(
   endpointName: EndpointName,
-  arg: QueryArgFrom<Definitions[EndpointName]>,
-  value: ResultTypeFrom<Definitions[EndpointName]>,
+  arg: QueryArgFromAnyQueryDefinition<Definitions, EndpointName>,
+  value: DataFromAnyQueryDefinition<Definitions, EndpointName>,
 ) => ThunkAction<
-  QueryActionCreatorResult<
-    Definitions[EndpointName] extends QueryDefinition<any, any, any, any>
-      ? Definitions[EndpointName]
-      : never
-  >,
+  UpsertThunkResult<Definitions, EndpointName>,
   PartialState,
   any,
   UnknownAction
@@ -368,7 +424,8 @@ export function buildThunks<
 
   const upsertQueryData: UpsertQueryDataThunk<Definitions, State> =
     (endpointName, arg, value) => (dispatch) => {
-      return dispatch(
+      type EndpointName = typeof endpointName
+      const res = dispatch(
         (
           api.endpoints[endpointName] as ApiEndpointQuery<
             QueryDefinition<any, any, any, any, any>,
@@ -381,7 +438,9 @@ export function buildThunks<
             data: value,
           }),
         }),
-      )
+      ) as UpsertThunkResult<Definitions, EndpointName>
+
+      return res
     }
 
   // The generic async payload function for all of our thunks
@@ -582,6 +641,14 @@ export function buildThunks<
           // Fetch first page
           result = await fetchPage(existingData, firstPageParam, maxPages)
 
+          if (forceQueryFn) {
+            // HACK `upsertQueryData` expects the user to pass in the `{pages, pageParams}` structure,
+            // but `fetchPage` treats that as `pages[0]`. We have to manually un-nest it.
+            result = {
+              data: (result.data as InfiniteData<unknown, unknown>).pages[0],
+            } as QueryReturnValue
+          }
+
           // Fetch remaining pages
           for (let i = 1; i < totalPages; i++) {
             const param = getNextPageParam(

packages/toolkit/src/query/tests/infiniteQueries.test.ts

@@ -523,14 +523,14 @@ describe('Infinite queries', () => {
     )
 
     const res1 = storeRef.store.dispatch(
-      pokemonApi.endpoints.getInfinitePokemon.initiate('fire', {}),
+      pokemonApiWithRefetch.endpoints.getInfinitePokemon.initiate('fire', {}),
     )
 
     const entry1InitialLoad = await res1
     checkResultData(entry1InitialLoad, [[{ id: '0', name: 'Pokemon 0' }]])
 
     const res2 = storeRef.store.dispatch(
-      pokemonApi.endpoints.getInfinitePokemon.initiate('fire', {
+      pokemonApiWithRefetch.endpoints.getInfinitePokemon.initiate('fire', {
         direction: 'forward',
       }),
     )
@@ -541,4 +541,84 @@ describe('Infinite queries', () => {
       [{ id: '1', name: 'Pokemon 1' }],
     ])
   })
+
+  test('Works with cache manipulation utils', async () => {
+    const res1 = storeRef.store.dispatch(
+      pokemonApi.endpoints.getInfinitePokemon.initiate('fire', {}),
+    )
+
+    const entry1InitialLoad = await res1
+    checkResultData(entry1InitialLoad, [[{ id: '0', name: 'Pokemon 0' }]])
+
+    storeRef.store.dispatch(
+      pokemonApi.util.updateQueryData('getInfinitePokemon', 'fire', (draft) => {
+        draft.pages.push([{ id: '1', name: 'Pokemon 1' }])
+        draft.pageParams.push(1)
+      }),
+    )
+
+    const selectFire = pokemonApi.endpoints.getInfinitePokemon.select('fire')
+    const entry1Updated = selectFire(storeRef.store.getState())
+
+    expect(entry1Updated.data).toEqual({
+      pages: [
+        [{ id: '0', name: 'Pokemon 0' }],
+        [{ id: '1', name: 'Pokemon 1' }],
+      ],
+      pageParams: [0, 1],
+    })
+
+    const res2 = storeRef.store.dispatch(
+      pokemonApi.util.upsertQueryData('getInfinitePokemon', 'water', {
+        pages: [[{ id: '2', name: 'Pokemon 2' }]],
+        pageParams: [2],
+      }),
+    )
+
+    const entry2InitialLoad = await res2
+    const selectWater = pokemonApi.endpoints.getInfinitePokemon.select('water')
+    const entry2Updated = selectWater(storeRef.store.getState())
+
+    expect(entry2Updated.data).toEqual({
+      pages: [[{ id: '2', name: 'Pokemon 2' }]],
+      pageParams: [2],
+    })
+
+    storeRef.store.dispatch(
+      pokemonApi.util.upsertQueryEntries([
+        {
+          endpointName: 'getInfinitePokemon',
+          arg: 'air',
+          value: {
+            pages: [[{ id: '3', name: 'Pokemon 3' }]],
+            pageParams: [3],
+          },
+        },
+      ]),
+    )
+
+    const selectAir = pokemonApi.endpoints.getInfinitePokemon.select('air')
+    const entry3Initial = selectAir(storeRef.store.getState())
+
+    expect(entry3Initial.data).toEqual({
+      pages: [[{ id: '3', name: 'Pokemon 3' }]],
+      pageParams: [3],
+    })
+
+    await storeRef.store.dispatch(
+      pokemonApi.endpoints.getInfinitePokemon.initiate('air', {
+        direction: 'forward',
+      }),
+    )
+
+    const entry3Updated = selectAir(storeRef.store.getState())
+
+    expect(entry3Updated.data).toEqual({
+      pages: [
+        [{ id: '3', name: 'Pokemon 3' }],
+        [{ id: '4', name: 'Pokemon 4' }],
+      ],
+      pageParams: [3, 4],
+    })
+  })
 })