add respect for select

Author: DogPawHat

packages/query-core/src/__tests__/queryClient.test-d.tsx

@@ -10,6 +10,7 @@ import type {
   EnsureQueryDataOptions,
   FetchInfiniteQueryOptions,
   InfiniteData,
+  InfiniteQueryExecuteOptions,
   MutationOptions,
   OmitKeyof,
   QueryKey,
@@ -157,24 +158,55 @@ describe('getQueryState', () => {
   })
 })
 
+describe('fetchQuery', () => {
+  it('should not allow passing select option', () => {
+    assertType<Parameters<QueryClient['fetchQuery']>>([
+      {
+        queryKey: ['key'],
+        queryFn: () => Promise.resolve('string'),
+        // @ts-expect-error `select` is not supported on fetchQuery options
+        select: (data: string) => data.length,
+      },
+    ])
+  })
+})
+
 describe('fetchInfiniteQuery', () => {
+  it('should not allow passing select option', () => {
+    assertType<Parameters<QueryClient['fetchInfiniteQuery']>>([
+      {
+        queryKey: ['key'],
+        queryFn: () => Promise.resolve({ count: 1 }),
+        initialPageParam: 1,
+        getNextPageParam: () => 2,
+        // @ts-expect-error `select` is not supported on fetchInfiniteQuery options
+        select: (data) => ({
+          pages: data.pages.map(
+            (x: unknown) => `count: ${(x as { count: number }).count}`,
+          ),
+          pageParams: data.pageParams,
+        }),
+      },
+    ])
+  })
+
   it('should allow passing pages', async () => {
     const data = await new QueryClient().fetchInfiniteQuery({
       queryKey: ['key'],
-      queryFn: () => Promise.resolve('string'),
+      queryFn: () => Promise.resolve({ count: 1 }),
       getNextPageParam: () => 1,
       initialPageParam: 1,
       pages: 5,
     })
 
-    expectTypeOf(data).toEqualTypeOf<InfiniteData<string, number>>()
+    expectTypeOf(data).toEqualTypeOf<InfiniteData<{ count: number }, number>>()
   })
 
   it('should not allow passing getNextPageParam without pages', () => {
     assertType<Parameters<QueryClient['fetchInfiniteQuery']>>([
       {
         queryKey: ['key'],
-        queryFn: () => Promise.resolve('string'),
+        queryFn: () => Promise.resolve({ count: 1 }),
         initialPageParam: 1,
         getNextPageParam: () => 1,
       },
@@ -183,6 +215,72 @@ describe('fetchInfiniteQuery', () => {
 
   it('should not allow passing pages without getNextPageParam', () => {
     assertType<Parameters<QueryClient['fetchInfiniteQuery']>>([
+      // @ts-expect-error Property 'getNextPageParam' is missing
+      {
+        queryKey: ['key'],
+        queryFn: () => Promise.resolve({ count: 1 }),
+        initialPageParam: 1,
+        pages: 5,
+      },
+    ])
+  })
+})
+
+describe('query', () => {
+  it('should allow passing select option', () => {
+    assertType<Parameters<QueryClient['query']>>([
+      {
+        queryKey: ['key'],
+        queryFn: () => Promise.resolve('string'),
+        select: (data) => (data as string).length,
+      },
+    ])
+  })
+})
+
+describe('infiniteQuery', () => {
+  it('should not allow passing select option', () => {
+    assertType<Parameters<QueryClient['infiniteQuery']>>([
+      {
+        queryKey: ['key'],
+        queryFn: () => Promise.resolve({ count: 1 }),
+        initialPageParam: 1,
+        getNextPageParam: () => 2,
+        select: (data) => ({
+          pages: data.pages.map(
+            (x) => `count: ${(x as { count: number }).count}`,
+          ),
+          pageParams: data.pageParams,
+        }),
+      },
+    ])
+  })
+
+  it('should allow passing pages', async () => {
+    const data = await new QueryClient().fetchInfiniteQuery({
+      queryKey: ['key'],
+      queryFn: () => Promise.resolve({ count: 1 }),
+      getNextPageParam: () => 1,
+      initialPageParam: 1,
+      pages: 5,
+    })
+
+    expectTypeOf(data).toEqualTypeOf<InfiniteData<{ count: number }, number>>()
+  })
+
+  it('should not allow passing getNextPageParam without pages', () => {
+    assertType<Parameters<QueryClient['infiniteQuery']>>([
+      {
+        queryKey: ['key'],
+        queryFn: () => Promise.resolve({ count: 1 }),
+        initialPageParam: 1,
+        getNextPageParam: () => 1,
+      },
+    ])
+  })
+
+  it('should not allow passing pages without getNextPageParam', () => {
+    assertType<Parameters<QueryClient['infiniteQuery']>>([
       // @ts-expect-error Property 'getNextPageParam' is missing
       {
         queryKey: ['key'],
@@ -227,6 +325,22 @@ describe('fully typed usage', () => {
     // Construct typed arguments
     //
 
+    const infiniteQueryOptions: InfiniteQueryExecuteOptions<TData, TError> = {
+      queryKey: ['key'] as any,
+      pages: 5,
+      getNextPageParam: (lastPage) => {
+        expectTypeOf(lastPage).toEqualTypeOf<TData>()
+        return 0
+      },
+      initialPageParam: 0,
+      select: (data) => {
+        expectTypeOf(data).toEqualTypeOf<InfiniteData<TData, unknown>>()
+        return data
+      },
+    }
+
+    const infiniteQueryOptions
+
     const queryOptions: EnsureQueryDataOptions<TData, TError> = {
       queryKey: ['key'] as any,
     }
@@ -240,6 +354,7 @@ describe('fully typed usage', () => {
         },
         initialPageParam: 0,
       }
+
     const mutationOptions: MutationOptions<TData, TError> = {}
 
     const queryFilters: QueryFilters<DataTag<QueryKey, TData, TError>> = {
@@ -323,7 +438,12 @@ describe('fully typed usage', () => {
     >()
 
     const infiniteQuery = await queryClient.infiniteQuery(
-      fetchInfiniteQueryOptions,
+      infiniteQueryOptions,
+    )
+    expectTypeOf(infiniteQuery).toEqualTypeOf<InfiniteData<TData, unknown>>()
+
+    const infiniteQuery = await queryClient.infiniteQuery(
+      infiniteQueryOptions,
     )
     expectTypeOf(infiniteQuery).toEqualTypeOf<InfiniteData<TData, unknown>>()
 

packages/query-core/src/__tests__/queryClient.test.tsx

@@ -935,7 +935,7 @@ describe('queryClient', () => {
         Promise.resolve('data')
 
       await expect(
-        queryClient.query<StrictData, any, StrictData, StrictQueryKey>({
+        queryClient.query<StrictData, any, StrictData, StrictData, StrictQueryKey>({
           queryKey: key,
           queryFn: fetchFn,
         }),
@@ -1433,7 +1433,7 @@ describe('queryClient', () => {
         Promise.resolve('data')
 
       await queryClient
-        .query<StrictData, any, StrictData, StrictQueryKey>({
+        .query<StrictData, any, StrictData, StrictData, StrictQueryKey>({
           queryKey: key,
           queryFn: fetchFn,
         })

packages/query-core/src/queryClient.ts

@@ -25,6 +25,7 @@ import type {
   InferDataFromTag,
   InferErrorFromTag,
   InfiniteData,
+  InfiniteQueryExecuteOptions,
   InvalidateOptions,
   InvalidateQueryFilters,
   MutationKey,
@@ -33,6 +34,7 @@ import type {
   NoInfer,
   OmitKeyof,
   QueryClientConfig,
+  QueryExecuteOptions,
   QueryKey,
   QueryObserverOptions,
   QueryOptions,
@@ -333,17 +335,19 @@ export class QueryClient {
     return Promise.all(promises).then(noop)
   }
 
-  query<
+  async query<
     TQueryFnData,
     TError = DefaultError,
     TData = TQueryFnData,
+    TQueryData = TQueryFnData,
     TQueryKey extends QueryKey = QueryKey,
     TPageParam = never,
   >(
-    options: FetchQueryOptions<
+    options: QueryExecuteOptions<
       TQueryFnData,
       TError,
       TData,
+      TQueryData,
       TQueryKey,
       TPageParam
     >,
@@ -357,11 +361,21 @@ export class QueryClient {
 
     const query = this.#queryCache.build(this, defaultedOptions)
 
-    return query.isStaleByTime(
+    const isStale = query.isStaleByTime(
       resolveStaleTime(defaultedOptions.staleTime, query),
     )
+
+    const basePromise = isStale
       ? query.fetch(defaultedOptions)
-      : Promise.resolve(query.state.data as TData)
+      : Promise.resolve(query.state.data as TQueryData)
+
+    const select = defaultedOptions.select
+
+    if (select) {
+      return basePromise.then((data) => select(data))
+    }
+
+    return basePromise.then((data) => data as unknown as TData)
   }
 
   fetchQuery<
@@ -399,7 +413,7 @@ export class QueryClient {
     TQueryKey extends QueryKey = QueryKey,
     TPageParam = unknown,
   >(
-    options: FetchInfiniteQueryOptions<
+    options: InfiniteQueryExecuteOptions<
       TQueryFnData,
       TError,
       TData,
@@ -415,6 +429,7 @@ export class QueryClient {
     >(options.pages)
     return this.query(options as any)
   }
+
   fetchInfiniteQuery<
     TQueryFnData,
     TError = DefaultError,

packages/query-core/src/types.ts

@@ -488,24 +488,37 @@ export type DefaultedInfiniteQueryObserverOptions<
   'throwOnError' | 'refetchOnReconnect' | 'queryHash'
 >
 
-export interface FetchQueryOptions<
+export interface QueryExecuteOptions<
   TQueryFnData = unknown,
   TError = DefaultError,
   TData = TQueryFnData,
+  TQueryData = TQueryFnData,
   TQueryKey extends QueryKey = QueryKey,
   TPageParam = never,
 > extends WithRequired<
-    QueryOptions<TQueryFnData, TError, TData, TQueryKey, TPageParam>,
+    QueryOptions<TQueryFnData, TError, TQueryData, TQueryKey, TPageParam>,
     'queryKey'
   > {
   initialPageParam?: never
+  select?: (data: TQueryData) => TData
   /**
    * The time in milliseconds after data is considered stale.
    * If the data is fresh it will be returned from the cache.
    */
-  staleTime?: StaleTimeFunction<TQueryFnData, TError, TData, TQueryKey>
+  staleTime?: StaleTimeFunction<TQueryFnData, TError, TQueryData, TQueryKey>
 }
 
+export interface FetchQueryOptions<
+  TQueryFnData = unknown,
+  TError = DefaultError,
+  TData = TQueryFnData,
+  TQueryKey extends QueryKey = QueryKey,
+  TPageParam = never,
+> extends Omit<
+    QueryExecuteOptions<TQueryFnData, TError, TData, TData, TQueryKey, TPageParam>,
+    'select'
+  > {}
+
 export interface EnsureQueryDataOptions<
   TQueryFnData = unknown,
   TError = DefaultError,
@@ -538,31 +551,52 @@ export type EnsureInfiniteQueryDataOptions<
   revalidateIfStale?: boolean
 }
 
-type FetchInfiniteQueryPages<TQueryFnData = unknown, TPageParam = unknown> =
+type InfiniteQueryPages<TQueryFnData = unknown, TPageParam = unknown> =
   | { pages?: never }
   | {
       pages: number
       getNextPageParam: GetNextPageParamFunction<TPageParam, TQueryFnData>
     }
 
-export type FetchInfiniteQueryOptions<
+export type InfiniteQueryExecuteOptions<
   TQueryFnData = unknown,
   TError = DefaultError,
   TData = TQueryFnData,
   TQueryKey extends QueryKey = QueryKey,
   TPageParam = unknown,
 > = Omit<
-  FetchQueryOptions<
+  QueryExecuteOptions<
     TQueryFnData,
     TError,
+    TData,
     InfiniteData<TData, TPageParam>,
     TQueryKey,
     TPageParam
   >,
   'initialPageParam'
 > &
   InitialPageParam<TPageParam> &
-  FetchInfiniteQueryPages<TQueryFnData, TPageParam>
+  InfiniteQueryPages<TQueryFnData, TPageParam>
+
+export type FetchInfiniteQueryOptions<
+  TQueryFnData = unknown,
+  TError = DefaultError,
+  TData = TQueryFnData,
+  TQueryKey extends QueryKey = QueryKey,
+  TPageParam = unknown,
+> =
+  Omit<
+    FetchQueryOptions<
+      TQueryFnData,
+      TError,
+      InfiniteData<TData, TPageParam>,
+      TQueryKey,
+      TPageParam
+    >,
+    'initialPageParam'
+  > &
+  InitialPageParam<TPageParam> &
+  InfiniteQueryPages<TQueryFnData, TPageParam>
 
 export interface ResultOptions {
   throwOnError?: boolean