From 93d442cdbce4542638ce21103a365bc26200c61c Mon Sep 17 00:00:00 2001
From: Dominik Dorfmeister <office@dorfmeister.cc>
Date: Sat, 2 Mar 2024 21:14:09 +0100
Subject: [PATCH] fix(types): make getQueryState types work with queryOptions
 (#7010)

---
 packages/query-core/src/queryClient.ts        | 19 +++++++++++++++----
 .../src/__tests__/queryOptions.test-d.tsx     | 10 ++++++++++
 2 files changed, 25 insertions(+), 4 deletions(-)

diff --git a/packages/query-core/src/queryClient.ts b/packages/query-core/src/queryClient.ts
index a09d838d4d..74f8b23352 100644
--- a/packages/query-core/src/queryClient.ts
+++ b/packages/query-core/src/queryClient.ts
@@ -204,11 +204,22 @@ export class QueryClient {
     )
   }
 
-  getQueryState<TQueryFnData = unknown, TError = DefaultError>(
-    queryKey: QueryKey,
-  ): QueryState<TQueryFnData, TError> | undefined {
+  getQueryState<
+    TQueryFnData = unknown,
+    TError = DefaultError,
+    TTaggedQueryKey extends QueryKey = QueryKey,
+    TInferredQueryFnData = TTaggedQueryKey extends DataTag<
+      unknown,
+      infer TaggedValue
+    >
+      ? TaggedValue
+      : TQueryFnData,
+  >(
+    queryKey: TTaggedQueryKey,
+  ): QueryState<TInferredQueryFnData, TError> | undefined {
     const options = this.defaultQueryOptions({ queryKey })
-    return this.#queryCache.get<TQueryFnData, TError>(options.queryHash)?.state
+    return this.#queryCache.get<TInferredQueryFnData, TError>(options.queryHash)
+      ?.state
   }
 
   removeQueries(filters?: QueryFilters): void {
diff --git a/packages/react-query/src/__tests__/queryOptions.test-d.tsx b/packages/react-query/src/__tests__/queryOptions.test-d.tsx
index 23259bc169..449e4923e0 100644
--- a/packages/react-query/src/__tests__/queryOptions.test-d.tsx
+++ b/packages/react-query/src/__tests__/queryOptions.test-d.tsx
@@ -108,6 +108,16 @@ describe('queryOptions', () => {
     const data = queryClient.getQueryData(queryKey)
     expectTypeOf(data).toEqualTypeOf<number | undefined>()
   })
+  it('should return the proper type when passed to getQueryState', () => {
+    const { queryKey } = queryOptions({
+      queryKey: ['key'],
+      queryFn: () => Promise.resolve(5),
+    })
+
+    const queryClient = new QueryClient()
+    const state = queryClient.getQueryState(queryKey)
+    expectTypeOf(state?.data).toEqualTypeOf<number | undefined>()
+  })
   it('should properly type updaterFn when passed to setQueryData', () => {
     const { queryKey } = queryOptions({
       queryKey: ['key'],