From a36fe2a320b2a710173399e855789910cc8d5cc0 Mon Sep 17 00:00:00 2001
From: sam585456525 <sam585456525@gmail.com>
Date: Sun, 15 Oct 2023 16:08:06 +0800
Subject: [PATCH] fix: infiniteData type

---
 .../typescript/react-query/src/fetcher.ts     | 15 ++++++++-
 .../__snapshots__/react-query.spec.ts.snap    | 32 +++++++++++++++++--
 .../react-query/tests/react-query.spec.ts     |  3 +-
 3 files changed, 45 insertions(+), 5 deletions(-)

diff --git a/packages/plugins/typescript/react-query/src/fetcher.ts b/packages/plugins/typescript/react-query/src/fetcher.ts
index 451043871..9d7991ad7 100644
--- a/packages/plugins/typescript/react-query/src/fetcher.ts
+++ b/packages/plugins/typescript/react-query/src/fetcher.ts
@@ -21,6 +21,7 @@ type ReactQueryMethodMap = {
   [key: string]: {
     getHook: (operationName?: string) => string;
     getOptions: () => string;
+    getOtherTypes?: () => { [key: string]: string };
   };
 };
 
@@ -44,6 +45,7 @@ export abstract class FetcherRenderer {
       infiniteQuery: {
         getHook: (operationName = 'Query') => `use${suspenseText}Infinite${operationName}`,
         getOptions: () => `Use${suspenseText}InfiniteQueryOptions`,
+        getOtherTypes: () => ({ infiniteData: 'InfiniteData' }),
       },
       query: {
         getHook: (operationName = 'Query') => `use${suspenseText}${operationName}`,
@@ -63,8 +65,15 @@ export abstract class FetcherRenderer {
 
     const { infiniteQuery } = this.createQueryMethodMap(isSuspense);
 
+    const isNextVersion = this.visitor.config.reactQueryVersion >= 5;
+
     this.visitor.reactQueryHookIdentifiersInUse.add(infiniteQuery.getHook());
     this.visitor.reactQueryOptionsIdentifiersInUse.add(infiniteQuery.getOptions());
+    if (isNextVersion) {
+      this.visitor.reactQueryOptionsIdentifiersInUse.add(
+        infiniteQuery.getOtherTypes().infiniteData,
+      );
+    }
 
     const variables = this.generateInfiniteQueryVariablesSignature(config);
     const options = this.generateInfiniteQueryOptionsSignature(config, isSuspense);
@@ -80,7 +89,11 @@ export abstract class FetcherRenderer {
     `;
 
       return `export const ${infiniteQuery.getHook(operationName)} = <
-      TData = ${operationResultType},
+      TData = ${
+        isNextVersion
+          ? `${infiniteQuery.getOtherTypes().infiniteData}<${operationResultType}>`
+          : operationResultType
+      },
       TError = ${this.visitor.config.errorType}
     >(${argumentsResult}) => {
     ${implHookOuter}
diff --git a/packages/plugins/typescript/react-query/tests/__snapshots__/react-query.spec.ts.snap b/packages/plugins/typescript/react-query/tests/__snapshots__/react-query.spec.ts.snap
index 0eb94f400..ccfed1a09 100644
--- a/packages/plugins/typescript/react-query/tests/__snapshots__/react-query.spec.ts.snap
+++ b/packages/plugins/typescript/react-query/tests/__snapshots__/react-query.spec.ts.snap
@@ -1272,7 +1272,7 @@ function fetcher<TData, TVariables>(endpoint: string, requestInit: RequestInit,
 ]
 `;
 
-exports[`React-Query support v5 syntax 1`] = `
+exports[`React-Query support v5 syntax: content 1`] = `
 "
 
 export const TestDocument = \`
@@ -1326,7 +1326,7 @@ export const useSuspenseTestQuery = <
     )};
 
 export const useInfiniteTestQuery = <
-      TData = TestQuery,
+      TData = InfiniteData<TestQuery>,
       TError = unknown
     >(
       dataSource: { endpoint: string, fetchParams?: RequestInit },
@@ -1346,7 +1346,7 @@ export const useInfiniteTestQuery = <
     )};
 
 export const useSuspenseInfiniteTestQuery = <
-      TData = TestQuery,
+      TData = InfiniteData<TestQuery>,
       TError = unknown
     >(
       dataSource: { endpoint: string, fetchParams?: RequestInit },
@@ -1390,3 +1390,29 @@ export const useTestMutation = <
     )};
 "
 `;
+
+exports[`React-Query support v5 syntax: prepend 1`] = `
+[
+  "import { useQuery, useSuspenseQuery, useInfiniteQuery, useSuspenseInfiniteQuery, useMutation, UseQueryOptions, UseSuspenseQueryOptions, UseInfiniteQueryOptions, InfiniteData, UseSuspenseInfiniteQueryOptions, UseMutationOptions } from '@tanstack/react-query';",
+  "
+function fetcher<TData, TVariables>(endpoint: string, requestInit: RequestInit, query: string, variables?: TVariables) {
+  return async (): Promise<TData> => {
+    const res = await fetch(endpoint, {
+      method: 'POST',
+      ...requestInit,
+      body: JSON.stringify({ query, variables }),
+    });
+
+    const json = await res.json();
+
+    if (json.errors) {
+      const { message } = json.errors[0];
+
+      throw new Error(message);
+    }
+
+    return json.data;
+  }
+}",
+]
+`;
diff --git a/packages/plugins/typescript/react-query/tests/react-query.spec.ts b/packages/plugins/typescript/react-query/tests/react-query.spec.ts
index 83d123f4b..815432f1d 100644
--- a/packages/plugins/typescript/react-query/tests/react-query.spec.ts
+++ b/packages/plugins/typescript/react-query/tests/react-query.spec.ts
@@ -46,7 +46,8 @@ describe('React-Query', () => {
 
     const out = (await plugin(schema, docs, config)) as Types.ComplexPluginOutput;
 
-    expect(out.content).toMatchSnapshot();
+    expect(out.prepend).toMatchSnapshot('prepend');
+    expect(out.content).toMatchSnapshot('content');
     await validateTypeScript(mergeOutputs(out), schema, docs, config);
   });