From 75db6b42111c5b7df8a0cba32eb275fd97263a71 Mon Sep 17 00:00:00 2001
From: gwansikk <seorit98@gmail.com>
Date: Fri, 14 Jun 2024 21:27:53 +0900
Subject: [PATCH] refactor(react-query): improve type inference for
 useSuspenseQueries with skipToken

---
 .../__tests__/useSuspenseQueries.test-d.tsx   | 22 +++++++++++++++++--
 .../react-query/src/useSuspenseQueries.ts     | 14 +++++++-----
 2 files changed, 29 insertions(+), 7 deletions(-)

diff --git a/packages/react-query/src/__tests__/useSuspenseQueries.test-d.tsx b/packages/react-query/src/__tests__/useSuspenseQueries.test-d.tsx
index c844b8f119..3b3b48065b 100644
--- a/packages/react-query/src/__tests__/useSuspenseQueries.test-d.tsx
+++ b/packages/react-query/src/__tests__/useSuspenseQueries.test-d.tsx
@@ -1,8 +1,8 @@
 import { describe, expectTypeOf, it } from 'vitest'
-import { useSuspenseQueries } from '..'
+import { skipToken, useSuspenseQueries } from '..'
 import { queryOptions } from '../queryOptions'
 import type { OmitKeyof } from '..'
-import type { UseQueryOptions } from '../types'
+import type { UseQueryOptions, UseSuspenseQueryResult } from '../types'
 
 describe('UseSuspenseQueries config object overload', () => {
   it('TData should always be defined', () => {
@@ -113,4 +113,22 @@ describe('UseSuspenseQueries config object overload', () => {
       expectTypeOf(data).toEqualTypeOf<Data>()
     })
   })
+
+  it('TData should have correct type when conditional skipToken is passed', () => {
+    const queryResults = useSuspenseQueries({
+      queries: [
+        {
+          queryKey: ['withSkipToken'],
+          queryFn: Math.random() > 0.5 ? skipToken : () => Promise.resolve(5),
+        },
+      ],
+    })
+
+    const firstResult = queryResults[0]
+
+    expectTypeOf(firstResult).toEqualTypeOf<
+      UseSuspenseQueryResult<number, Error>
+    >()
+    expectTypeOf(firstResult.data).toEqualTypeOf<number>()
+  })
 })
diff --git a/packages/react-query/src/useSuspenseQueries.ts b/packages/react-query/src/useSuspenseQueries.ts
index 400944fb75..b316849be2 100644
--- a/packages/react-query/src/useSuspenseQueries.ts
+++ b/packages/react-query/src/useSuspenseQueries.ts
@@ -6,13 +6,15 @@ import type {
   DefaultError,
   QueryClient,
   QueryFunction,
-  SkipToken,
   ThrowOnError,
 } from '@tanstack/query-core'
 
 // Avoid TS depth-limit error in case of large array literal
 type MAXIMUM_DEPTH = 20
 
+// Widen the type of the symbol to enable type inference even if skipToken is not immutable.
+type SkipTokenForUseQueries = symbol
+
 type GetUseSuspenseQueryOptions<T> =
   // Part 1: responsible for applying explicit type parameter to function arguments, if object { queryFnData: TQueryFnData, error: TError, data: TData }
   T extends {
@@ -36,7 +38,7 @@ type GetUseSuspenseQueryOptions<T> =
                 T extends {
                     queryFn?:
                       | QueryFunction<infer TQueryFnData, infer TQueryKey>
-                      | SkipToken
+                      | SkipTokenForUseQueries
                     select?: (data: any) => infer TData
                     throwOnError?: ThrowOnError<any, infer TError, any, any>
                   }
@@ -49,7 +51,7 @@ type GetUseSuspenseQueryOptions<T> =
                 : T extends {
                       queryFn?:
                         | QueryFunction<infer TQueryFnData, infer TQueryKey>
-                        | SkipToken
+                        | SkipTokenForUseQueries
                       throwOnError?: ThrowOnError<any, infer TError, any, any>
                     }
                   ? UseSuspenseQueryOptions<
@@ -78,7 +80,9 @@ type GetUseSuspenseQueryResult<T> =
               ? UseSuspenseQueryResult<TQueryFnData>
               : // Part 3: responsible for mapping inferred type to results, if no explicit parameter was provided
                 T extends {
-                    queryFn?: QueryFunction<infer TQueryFnData, any> | SkipToken
+                    queryFn?:
+                      | QueryFunction<infer TQueryFnData, any>
+                      | SkipTokenForUseQueries
                     select?: (data: any) => infer TData
                     throwOnError?: ThrowOnError<any, infer TError, any, any>
                   }
@@ -89,7 +93,7 @@ type GetUseSuspenseQueryResult<T> =
                 : T extends {
                       queryFn?:
                         | QueryFunction<infer TQueryFnData, any>
-                        | SkipToken
+                        | SkipTokenForUseQueries
                       throwOnError?: ThrowOnError<any, infer TError, any, any>
                     }
                   ? UseSuspenseQueryResult<