Fix updateQueryData for infinite queries

Author: markerikson

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/buildThunks.ts

@@ -19,8 +19,10 @@ import type {
   AssertTagTypes,
   EndpointDefinition,
   EndpointDefinitions,
+  InfiniteQueryArgFrom,
   InfiniteQueryDefinition,
   MutationDefinition,
+  PageParamFrom,
   QueryArgFrom,
   QueryDefinition,
   ResultTypeFrom,
@@ -40,6 +42,7 @@ import type {
   InfiniteQueryConfigOptions,
   QueryCacheKey,
   InfiniteQueryDirection,
+  InfiniteQueryKeys,
 } from './apiState'
 import { QueryStatus } from './apiState'
 import type {
@@ -194,13 +197,39 @@ export type PatchQueryDataThunk<
   updateProvided?: boolean,
 ) => ThunkAction<void, PartialState, any, UnknownAction>
 
+type AllQueryKeys<Definitions extends EndpointDefinitions> =
+  | QueryKeys<Definitions>
+  | InfiniteQueryKeys<Definitions>
+
 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]>>,
+  // TODO Find a way to deduplicate this ugly conditional type
+  arg: Definitions[EndpointName] extends InfiniteQueryDefinition<
+    any,
+    any,
+    any,
+    any,
+    any
+  >
+    ? InfiniteQueryArgFrom<Definitions[EndpointName]>
+    : QueryArgFrom<Definitions[EndpointName]>,
+  updateRecipe: Recipe<
+    Definitions[EndpointName] extends InfiniteQueryDefinition<
+      any,
+      any,
+      any,
+      any,
+      any
+    >
+      ? InfiniteData<
+          ResultTypeFrom<Definitions[EndpointName]>,
+          PageParamFrom<Definitions[EndpointName]>
+        >
+      : ResultTypeFrom<Definitions[EndpointName]>
+  >,
   updateProvided?: boolean,
 ) => ThunkAction<PatchCollection, PartialState, any, UnknownAction>
 

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,32 @@ 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 entry1Updated = pokemonApi.endpoints.getInfinitePokemon.select(
+      'fire',
+    )(storeRef.store.getState())
+
+    expect(entry1Updated.data).toEqual({
+      pages: [
+        [{ id: '0', name: 'Pokemon 0' }],
+        [{ id: '1', name: 'Pokemon 1' }],
+      ],
+      pageParams: [0, 1],
+    })
+  })
 })