Fix upsertQueryData for infinite queries

Author: markerikson

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

@@ -46,6 +46,7 @@ import type {
 } from './apiState'
 import { QueryStatus } from './apiState'
 import type {
+  InfiniteQueryActionCreatorResult,
   QueryActionCreatorResult,
   StartInfiniteQueryActionCreatorOptions,
   StartQueryActionCreatorOptions,
@@ -201,51 +202,76 @@ type AllQueryKeys<Definitions extends EndpointDefinitions> =
   | QueryKeys<Definitions>
   | InfiniteQueryKeys<Definitions>
 
-export type UpdateQueryDataThunk<
+type QueryArgFromAnyQueryDefinition<
   Definitions extends EndpointDefinitions,
-  PartialState,
-> = <EndpointName extends AllQueryKeys<Definitions>>(
-  endpointName: EndpointName,
-  // TODO Find a way to deduplicate this ugly conditional type
-  arg: Definitions[EndpointName] extends InfiniteQueryDefinition<
+  EndpointName extends AllQueryKeys<Definitions>,
+> =
+  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]>
-  >,
+    : Definitions[EndpointName] extends QueryDefinition<any, any, any, any>
+      ? QueryArgFrom<Definitions[EndpointName]>
+      : never
+
+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
+
+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 AllQueryKeys<Definitions>>(
+  endpointName: 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
@@ -397,7 +423,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>,
@@ -410,7 +437,9 @@ export function buildThunks<
             data: value,
           }),
         }),
-      )
+      ) as UpsertThunkResult<Definitions, EndpointName>
+
+      return res
     }
 
   // The generic async payload function for all of our thunks
@@ -611,6 +640,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

@@ -568,5 +568,27 @@ describe('Infinite queries', () => {
       ],
       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 entry2Updated = pokemonApi.endpoints.getInfinitePokemon.select(
+      'water',
+    )(storeRef.store.getState())
+
+    console.log(
+      'entry2Updated',
+      util.inspect(entry2Updated, { depth: Infinity }),
+    )
+
+    expect(entry2Updated.data).toEqual({
+      pages: [[{ id: '2', name: 'Pokemon 2' }]],
+      pageParams: [2],
+    })
   })
 })