Feature/typed query key

packages/query-core/src/queryClient.ts

@@ -11,7 +11,6 @@ import { focusManager } from './focusManager'
 import { onlineManager } from './onlineManager'
 import { notifyManager } from './notifyManager'
 import { infiniteQueryBehavior } from './infiniteQueryBehavior'
-import type { DataTag, NoInfer } from './types'
 import type { QueryState } from './query'
 import type {
   CancelOptions,
@@ -34,6 +33,7 @@ import type {
   RefetchQueryFilters,
   ResetOptions,
   SetDataOptions,
+  TypedQueryKey,
 } from './types'
 import type { MutationFilters, QueryFilters, Updater } from './utils'
 
@@ -108,16 +108,9 @@ export class QueryClient {
     return this.#mutationCache.findAll({ ...filters, status: 'pending' }).length
   }
 
-  getQueryData<
-    TQueryFnData = unknown,
-    TaggedQueryKey extends QueryKey = QueryKey,
-    TInferredQueryFnData = TaggedQueryKey extends DataTag<
-      unknown,
-      infer TaggedValue
-    >
-      ? TaggedValue
-      : TQueryFnData,
-  >(queryKey: TaggedQueryKey): TInferredQueryFnData | undefined
+  getQueryData<TQueryFnData = unknown>(
+    queryKey: TypedQueryKey<TQueryFnData> | QueryKey,
+  ): TQueryFnData | undefined
   getQueryData(queryKey: QueryKey) {
     return this.#queryCache.find({ queryKey })?.state.data
   }
@@ -146,24 +139,12 @@ export class QueryClient {
       })
   }
 
-  setQueryData<
-    TQueryFnData = unknown,
-    TaggedQueryKey extends QueryKey = QueryKey,
-    TInferredQueryFnData = TaggedQueryKey extends DataTag<
-      unknown,
-      infer TaggedValue
-    >
-      ? TaggedValue
-      : TQueryFnData,
-  >(
-    queryKey: TaggedQueryKey,
-    updater: Updater<
-      NoInfer<TInferredQueryFnData> | undefined,
-      NoInfer<TInferredQueryFnData> | undefined
-    >,
+  setQueryData<TQueryFnData>(
+    queryKey: TypedQueryKey<TQueryFnData> | QueryKey,
+    updater: Updater<TQueryFnData | undefined, TQueryFnData | undefined>,
     options?: SetDataOptions,
-  ): TInferredQueryFnData | undefined {
-    const query = this.#queryCache.find<TInferredQueryFnData>({ queryKey })
+  ): TQueryFnData | undefined {
+    const query = this.#queryCache.find<TQueryFnData>({ queryKey })
     const prevData = query?.state.data
     const data = functionalUpdate(updater, prevData)
 

packages/query-core/src/tests/queryClient.types.test.tsx

@@ -2,12 +2,12 @@ import { describe, it } from 'vitest'
 import { QueryClient } from '../queryClient'
 import { doNotExecute } from './utils'
 import type { Equal, Expect } from './utils'
-import type { DataTag, InfiniteData } from '../types'
+import type { InfiniteData, TypedQueryKey } from '../types'
 
 describe('getQueryData', () => {
   it('should be typed if key is tagged', () => {
     doNotExecute(() => {
-      const queryKey = ['key'] as DataTag<Array<string>, number>
+      const queryKey = ['key'] as TypedQueryKey<number>
       const queryClient = new QueryClient()
       const data = queryClient.getQueryData(queryKey)
 
@@ -51,7 +51,7 @@ describe('getQueryData', () => {
 describe('setQueryData', () => {
   it('updater should be typed if key is tagged', () => {
     doNotExecute(() => {
-      const queryKey = ['key'] as DataTag<Array<string>, number>
+      const queryKey = ['key'] as TypedQueryKey<number>
       const queryClient = new QueryClient()
       const data = queryClient.setQueryData(queryKey, (prev) => {
         const result: Expect<Equal<typeof prev, number | undefined>> = true
@@ -65,7 +65,7 @@ describe('setQueryData', () => {
 
   it('value should be typed if key is tagged', () => {
     doNotExecute(() => {
-      const queryKey = ['key'] as DataTag<Array<string>, number>
+      const queryKey = ['key'] as TypedQueryKey<number>
       const queryClient = new QueryClient()
 
       // @ts-expect-error value should be a number
@@ -95,13 +95,13 @@ describe('setQueryData', () => {
     })
   })
 
-  it('should infer unknown for value if key is not tagged', () => {
+  it('should infer literal value if key is not tagged', () => {
     doNotExecute(() => {
       const queryKey = ['key'] as const
       const queryClient = new QueryClient()
       const data = queryClient.setQueryData(queryKey, 'foo')
 
-      const result: Expect<Equal<typeof data, unknown>> = true
+      const result: Expect<Equal<typeof data, 'foo' | undefined>> = true
       return result
     })
   })
@@ -172,3 +172,44 @@ describe('fetchInfiniteQuery', () => {
     })
   })
 })
+
+describe('fetchInfiniteQuery', () => {
+  it('should allow passing pages', () => {
+    doNotExecute(async () => {
+      const data = await new QueryClient().fetchInfiniteQuery({
+        queryKey: ['key'],
+        queryFn: () => Promise.resolve('string'),
+        getNextPageParam: () => 1,
+        initialPageParam: 1,
+        pages: 5,
+      })
+
+      const result: Expect<Equal<typeof data, InfiniteData<string, number>>> =
+        true
+      return result
+    })
+  })
+
+  it('should not allow passing getNextPageParam without pages', () => {
+    doNotExecute(async () => {
+      return new QueryClient().fetchInfiniteQuery({
+        queryKey: ['key'],
+        queryFn: () => Promise.resolve('string'),
+        initialPageParam: 1,
+        getNextPageParam: () => 1,
+      })
+    })
+  })
+
+  it('should not allow passing pages without getNextPageParam', () => {
+    doNotExecute(async () => {
+      // @ts-expect-error Property 'getNextPageParam' is missing
+      return new QueryClient().fetchInfiniteQuery({
+        queryKey: ['key'],
+        queryFn: () => Promise.resolve('string'),
+        initialPageParam: 1,
+        pages: 5,
+      })
+    })
+  })
+})

packages/query-core/src/types.ts

@@ -23,10 +23,7 @@ export type DefaultError = Register extends {
 
 export type QueryKey = ReadonlyArray<unknown>
 
-export declare const dataTagSymbol: unique symbol
-export type DataTag<Type, Value> = Type & {
-  [dataTagSymbol]: Value
-}
+export declare interface TypedQueryKey<Data> extends QueryKey {}
 
 export type QueryFunction<
   T = unknown,

packages/react-query/src/__tests__/infiniteQueryOptions.types.test.tsx

@@ -4,7 +4,7 @@ import { infiniteQueryOptions } from '../infiniteQueryOptions'
 import { useInfiniteQuery } from '../useInfiniteQuery'
 import { useSuspenseInfiniteQuery } from '../useSuspenseInfiniteQuery'
 import { doNotExecute } from './utils'
-import type { InfiniteData, dataTagSymbol } from '@tanstack/query-core'
+import type { InfiniteData, TypedQueryKey } from '@tanstack/query-core'
 import type { Equal, Expect } from './utils'
 
 describe('queryOptions', () => {
@@ -97,7 +97,7 @@ describe('queryOptions', () => {
       })
 
       const result: Expect<
-        Equal<(typeof queryKey)[typeof dataTagSymbol], InfiniteData<string>>
+        Equal<typeof queryKey, TypedQueryKey<InfiniteData<string>>>
       > = true
       return result
     })
@@ -112,7 +112,7 @@ describe('queryOptions', () => {
         })
 
         const result: Expect<
-          Equal<(typeof queryKey)[typeof dataTagSymbol], InfiniteData<string>>
+          Equal<typeof queryKey, TypedQueryKey<InfiniteData<string>>>
         > = true
         return result
       })

packages/react-query/src/__tests__/queryOptions.types.test.tsx

@@ -5,8 +5,8 @@ import { useQuery } from '../useQuery'
 import { useQueries } from '../useQueries'
 import { useSuspenseQuery } from '../useSuspenseQuery'
 import { doNotExecute } from './utils'
-import type { dataTagSymbol } from '@tanstack/query-core'
 import type { Equal, Expect } from './utils'
+import type { TypedQueryKey } from '@tanstack/query-core'
 
 describe('queryOptions', () => {
   it('should not allow excess properties', () => {
@@ -93,9 +93,7 @@ describe('queryOptions', () => {
         queryFn: () => Promise.resolve(5),
       })
 
-      const result: Expect<
-        Equal<(typeof queryKey)[typeof dataTagSymbol], number>
-      > = true
+      const result: Expect<Equal<typeof queryKey, TypedQueryKey<number>>> = true
       return result
     })
 
@@ -106,9 +104,8 @@ describe('queryOptions', () => {
           queryFn: () => 5,
         })
 
-        const result: Expect<
-          Equal<(typeof queryKey)[typeof dataTagSymbol], number>
-        > = true
+        const result: Expect<Equal<typeof queryKey, TypedQueryKey<number>>> =
+          true
         return result
       })
     })
@@ -119,9 +116,8 @@ describe('queryOptions', () => {
           queryKey: ['key'],
         })
 
-        const result: Expect<
-          Equal<(typeof queryKey)[typeof dataTagSymbol], unknown>
-        > = true
+        const result: Expect<Equal<typeof queryKey, TypedQueryKey<unknown>>> =
+          true
         return result
       })
     })
@@ -179,5 +175,26 @@ describe('queryOptions', () => {
         return result
       })
     })
+
+    it('should properly type value when passed to setQueryData', () => {
+      doNotExecute(() => {
+        const { queryKey } = queryOptions({
+          queryKey: ['key'],
+          queryFn: () => Promise.resolve(5),
+        })
+
+        const queryClient = new QueryClient()
+
+        // @ts-expect-error value should be a number
+        queryClient.setQueryData(queryKey, '5')
+        // @ts-expect-error value should be a number
+        queryClient.setQueryData(queryKey, () => '5')
+
+        const data = queryClient.setQueryData(queryKey, 5)
+
+        const result: Expect<Equal<typeof data, number | undefined>> = true
+        return result
+      })
+    })
   })
 })

packages/react-query/src/infiniteQueryOptions.ts

@@ -1,5 +1,4 @@
-import type { DataTag } from '@tanstack/query-core'
-import type { InfiniteData } from '@tanstack/query-core'
+import type { InfiniteData, TypedQueryKey } from '@tanstack/query-core'
 import type { UseInfiniteQueryOptions } from './types'
 import type { DefaultError, QueryKey } from '@tanstack/query-core'
 
@@ -41,52 +40,64 @@ export type DefinedInitialDataInfiniteOptions<
     | (() => NonUndefinedGuard<InfiniteData<TQueryFnData, TPageParam>>)
 }
 
+type ValidateInfiniteQueryOptions<T> = {
+  [K in keyof T]: K extends keyof UseInfiniteQueryOptions ? T[K] : never
+}
+
 export function infiniteQueryOptions<
   TQueryFnData,
   TError = DefaultError,
   TData = InfiniteData<TQueryFnData>,
-  TQueryKey extends QueryKey = QueryKey,
   TPageParam = unknown,
 >(
-  options: UndefinedInitialDataInfiniteOptions<
+  options: ValidateInfiniteQueryOptions<
+    UndefinedInitialDataInfiniteOptions<
+      TQueryFnData,
+      TError,
+      TData,
+      QueryKey,
+      TPageParam
+    >
+  >,
+): Omit<
+  UndefinedInitialDataInfiniteOptions<
     TQueryFnData,
     TError,
     TData,
-    TQueryKey,
+    QueryKey,
     TPageParam
   >,
-): UndefinedInitialDataInfiniteOptions<
-  TQueryFnData,
-  TError,
-  TData,
-  TQueryKey,
-  TPageParam
+  'queryKey'
 > & {
-  queryKey: DataTag<TQueryKey, TData>
+  queryKey: TypedQueryKey<InfiniteData<TQueryFnData> | undefined>
 }
 
 export function infiniteQueryOptions<
   TQueryFnData,
   TError = DefaultError,
   TData = InfiniteData<TQueryFnData>,
-  TQueryKey extends QueryKey = QueryKey,
   TPageParam = unknown,
 >(
-  options: DefinedInitialDataInfiniteOptions<
+  options: ValidateInfiniteQueryOptions<
+    DefinedInitialDataInfiniteOptions<
+      TQueryFnData,
+      TError,
+      TData,
+      QueryKey,
+      TPageParam
+    >
+  >,
+): Omit<
+  DefinedInitialDataInfiniteOptions<
     TQueryFnData,
     TError,
     TData,
-    TQueryKey,
+    QueryKey,
     TPageParam
   >,
-): DefinedInitialDataInfiniteOptions<
-  TQueryFnData,
-  TError,
-  TData,
-  TQueryKey,
-  TPageParam
+  'queryKey'
 > & {
-  queryKey: DataTag<TQueryKey, TData>
+  queryKey: TypedQueryKey<InfiniteData<TQueryFnData>>
 }
 
 export function infiniteQueryOptions(options: unknown) {