fix(useQuery): make sure queryKeys are always an array at runtime when passed to the queryFn

this is a regression from #2074

the real culprit we tried to fix was typings being still just TQueryKey in the FetchContext, which is not what it is at runtime, because if your QueryKey is a string, it will be [string] there. See: https://github.com/tannerlinsley/react-query/pull/2074/files#r606362656

this fix re-adds the `ensureArray` call, but makes it narrow the type (compromised type assertion) so that the FetchContext will receive an `EnsuredQueryKey`, which is always an Array

Author: TkDodo

src/core/query.ts

@@ -5,13 +5,15 @@ import {
   noop,
   replaceEqualDeep,
   timeUntilStale,
+  ensureQueryKeyArray,
 } from './utils'
 import type {
   InitialDataFunction,
   QueryKey,
   QueryOptions,
   QueryStatus,
   QueryFunctionContext,
+  EnsuredQueryKey,
 } from './types'
 import type { QueryCache } from './queryCache'
 import type { QueryObserver } from './queryObserver'
@@ -58,8 +60,8 @@ export interface FetchContext<
 > {
   fetchFn: () => unknown | Promise<unknown>
   fetchOptions?: FetchOptions
-  options: QueryOptions<TQueryFnData, TError, TData, TQueryKey>
-  queryKey: TQueryKey
+  options: QueryOptions<TQueryFnData, TError, TData, any>
+  queryKey: EnsuredQueryKey<TQueryKey>
   state: QueryState<TData, TError>
 }
 
@@ -376,9 +378,11 @@ export class Query<
       }
     }
 
+    const queryKey = ensureQueryKeyArray(this.queryKey)
+
     // Create query function context
     const queryFnContext: QueryFunctionContext<TQueryKey> = {
-      queryKey: this.queryKey,
+      queryKey,
       pageParam: undefined,
     }
 
@@ -389,10 +393,10 @@ export class Query<
         : Promise.reject('Missing queryFn')
 
     // Trigger behavior hook
-    const context: FetchContext<TQueryFnData, TError, TData, any> = {
+    const context: FetchContext<TQueryFnData, TError, TData, TQueryKey> = {
       fetchOptions,
       options: this.options,
-      queryKey: this.queryKey,
+      queryKey: queryKey,
       state: this.state,
       fetchFn,
     }

src/core/types.ts

@@ -4,6 +4,9 @@ import type { RetryValue, RetryDelayValue } from './retryer'
 import type { QueryFilters } from './utils'
 
 export type QueryKey = string | readonly unknown[]
+export type EnsuredQueryKey<T extends QueryKey> = T extends string
+  ? [T]
+  : Exclude<T, string>
 
 export type QueryFunction<
   T = unknown,
@@ -14,14 +17,12 @@ export interface QueryFunctionContext<
   TQueryKey extends QueryKey = QueryKey,
   TPageParam = any
 > {
-  queryKey: TQueryKey
+  queryKey: EnsuredQueryKey<TQueryKey>
   pageParam?: TPageParam
 }
 
 export type InitialDataFunction<T> = () => T | undefined
 
-export type InitialStaleFunction = () => boolean
-
 export type PlaceholderDataFunction<TResult> = () => TResult | undefined
 
 export type QueryKeyHashFunction<TQueryKey extends QueryKey> = (
@@ -225,7 +226,13 @@ export interface FetchInfiniteQueryOptions<
   TError = unknown,
   TData = TQueryFnData,
   TQueryKey extends QueryKey = QueryKey
-> extends FetchQueryOptions<TQueryFnData, TError, InfiniteData<TData>, TQueryKey> {}
+>
+  extends FetchQueryOptions<
+    TQueryFnData,
+    TError,
+    InfiniteData<TData>,
+    TQueryKey
+  > {}
 
 export interface ResultOptions {
   throwOnError?: boolean
@@ -474,7 +481,9 @@ export interface MutationOptions<
   mutationFn?: MutationFunction<TData, TVariables>
   mutationKey?: MutationKey
   variables?: TVariables
-  onMutate?: (variables: TVariables) => Promise<TContext> | Promise<undefined> | TContext | undefined
+  onMutate?: (
+    variables: TVariables
+  ) => Promise<TContext> | Promise<undefined> | TContext | undefined
   onSuccess?: (
     data: TData,
     variables: TVariables,

src/core/utils.ts

@@ -1,5 +1,6 @@
 import type { Mutation } from './mutation'
 import type { Query } from './query'
+import { EnsuredQueryKey } from './types'
 import type {
   MutationFunction,
   MutationKey,
@@ -88,8 +89,12 @@ export function isValidTimeout(value: any): value is number {
   return typeof value === 'number' && value >= 0 && value !== Infinity
 }
 
-export function ensureArray<T>(value: T | T[]): T[] {
-  return Array.isArray(value) ? value : [value]
+export function ensureQueryKeyArray<T extends QueryKey>(
+  value: T
+): EnsuredQueryKey<T> {
+  return (Array.isArray(value)
+    ? value
+    : ([value] as unknown)) as EnsuredQueryKey<T>
 }
 
 export function difference<T>(array1: T[], array2: T[]): T[] {
@@ -256,7 +261,7 @@ export function hashQueryKeyByOptions<TQueryKey extends QueryKey = QueryKey>(
  * Default query keys hash function.
  */
 export function hashQueryKey(queryKey: QueryKey): string {
-  const asArray = Array.isArray(queryKey) ? queryKey : [queryKey]
+  const asArray = ensureQueryKeyArray(queryKey)
   return stableValueHash(asArray)
 }
 
@@ -280,7 +285,7 @@ export function stableValueHash(value: any): string {
  * Checks if key `b` partially matches with key `a`.
  */
 export function partialMatchKey(a: QueryKey, b: QueryKey): boolean {
-  return partialDeepEqual(ensureArray(a), ensureArray(b))
+  return partialDeepEqual(ensureQueryKeyArray(a), ensureQueryKeyArray(b))
 }
 
 /**

src/react/tests/useQuery.test.tsx

@@ -79,15 +79,25 @@ describe('useQuery', () => {
       type MyData = number
       type MyQueryKey = readonly ['my-data', number]
 
-      const getMyData: QueryFunction<MyData, MyQueryKey> = async ({
+      const getMyDataArrayKey: QueryFunction<MyData, MyQueryKey> = async ({
         queryKey: [, n],
       }) => {
         return n + 42
       }
 
       useQuery({
         queryKey: ['my-data', 100],
-        queryFn: getMyData,
+        queryFn: getMyDataArrayKey,
+      })
+
+      const getMyDataStringKey: QueryFunction<MyData, '1'> = async context => {
+        expectType<['1']>(context.queryKey)
+        return Number(context.queryKey[0]) + 42
+      }
+
+      useQuery({
+        queryKey: '1',
+        queryFn: getMyDataStringKey,
       })
     }
   })