add infinite query type support for selectCachedArgsForQuery (#4880)

Author: alexmotoc

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

@@ -6,7 +6,7 @@ import type {
   InfiniteQueryDefinition,
   MutationDefinition,
   PageParamFrom,
-  QueryArgFrom,
+  QueryArgFromAnyQuery,
   QueryDefinition,
   ResultTypeFrom,
 } from '../endpointDefinitions'
@@ -190,7 +190,7 @@ type BaseQuerySubState<
   /**
    * The argument originally passed into the hook or `initiate` action call
    */
-  originalArgs: QueryArgFrom<D>
+  originalArgs: QueryArgFromAnyQuery<D>
   /**
    * A unique ID associated with the request
    */

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

@@ -6,6 +6,7 @@ import type {
   InfiniteQueryDefinition,
   MutationDefinition,
   QueryArgFrom,
+  QueryArgFromAnyQuery,
   QueryDefinition,
   ReducerPathFrom,
   TagDescription,
@@ -29,7 +30,11 @@ import { QueryStatus, getRequestStatusFlags } from './apiState'
 import { getMutationCacheKey } from './buildSlice'
 import type { createSelector as _createSelector } from './rtkImports'
 import { createNextState } from './rtkImports'
-import { getNextPageParam, getPreviousPageParam } from './buildThunks'
+import {
+  type AllQueryKeys,
+  getNextPageParam,
+  getPreviousPageParam,
+} from './buildThunks'
 
 export type SkipToken = typeof skipToken
 /**
@@ -372,10 +377,12 @@ export function buildSelectors<
     )
   }
 
-  function selectCachedArgsForQuery<QueryName extends QueryKeys<Definitions>>(
+  function selectCachedArgsForQuery<
+    QueryName extends AllQueryKeys<Definitions>,
+  >(
     state: RootState,
     queryName: QueryName,
-  ): Array<QueryArgFrom<Definitions[QueryName]>> {
+  ): Array<QueryArgFromAnyQuery<Definitions[QueryName]>> {
     return Object.values(selectQueries(state) as QueryState<any>)
       .filter(
         (

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

@@ -19,6 +19,7 @@ import type {
   InfiniteQueryDefinition,
   MutationDefinition,
   QueryArgFrom,
+  QueryArgFromAnyQuery,
   QueryDefinition,
   TagDescription,
 } from '../endpointDefinitions'
@@ -389,10 +390,10 @@ export interface ApiModules<
        *
        * Can be used for mutations that want to do optimistic updates instead of invalidating a set of tags, but don't know exactly what they need to update.
        */
-      selectCachedArgsForQuery: <QueryName extends QueryKeys<Definitions>>(
+      selectCachedArgsForQuery: <QueryName extends AllQueryKeys<Definitions>>(
         state: RootState<Definitions, string, ReducerPath>,
         queryName: QueryName,
-      ) => Array<QueryArgFrom<Definitions[QueryName]>>
+      ) => Array<QueryArgFromAnyQuery<Definitions[QueryName]>>
     }
     /**
      * Endpoints based on the input endpoints provided to `createApi`, containing `select` and `action matchers`.

packages/toolkit/src/query/endpointDefinitions.ts

@@ -1059,6 +1059,15 @@ export type InfiniteQueryArgFrom<
   D extends BaseEndpointDefinition<any, any, any>,
 > = D extends InfiniteQueryDefinition<infer QA, any, any, any, any> ? QA : never
 
+export type QueryArgFromAnyQuery<
+  D extends BaseEndpointDefinition<any, any, any>,
+> =
+  D extends InfiniteQueryDefinition<any, any, any, any, any>
+    ? InfiniteQueryArgFrom<D>
+    : D extends QueryDefinition<any, any, any, any>
+      ? QueryArgFrom<D>
+      : never
+
 export type ResultTypeFrom<D extends BaseEndpointDefinition<any, any, any>> =
   D extends BaseEndpointDefinition<any, any, infer RT> ? RT : unknown
 

packages/toolkit/src/query/tests/buildSelector.test-d.ts

@@ -73,6 +73,21 @@ describe('type tests', () => {
         getTodos: build.query<Todos, string>({
           query: () => '/todos',
         }),
+        getInfiniteTodos: build.infiniteQuery<Todos, string, number>({
+          infiniteQueryOptions: {
+            initialPageParam: 0,
+            maxPages: 3,
+            getNextPageParam: (
+              lastPage,
+              allPages,
+              lastPageParam,
+              allPageParams,
+            ) => lastPageParam + 1,
+          },
+          query({ pageParam }) {
+            return `/todos?page=${pageParam}`
+          },
+        }),
       }),
     })
 
@@ -86,5 +101,11 @@ describe('type tests', () => {
     expectTypeOf(
       exampleApi.util.selectCachedArgsForQuery(store.getState(), 'getTodos'),
     ).toEqualTypeOf<string[]>()
+    expectTypeOf(
+      exampleApi.util.selectCachedArgsForQuery(
+        store.getState(),
+        'getInfiniteTodos',
+      ),
+    ).toEqualTypeOf<string[]>()
   })
 })