fix(query): type return of useMutation query

orval-labs · Feb 11, 2024 · 1fa21d6 · 1fa21d6
1 parent 288f488
commit 1fa21d6
Show file tree

Hide file tree

Showing 2 changed files with 67 additions and 19 deletions.
diff --git a/packages/query/src/index.ts b/packages/query/src/index.ts
@@ -186,6 +186,7 @@ const REACT_QUERY_DEPENDENCIES: GeneratorDependency[] = [
       { name: 'QueryKey' },
       { name: 'QueryClient' },
       { name: 'InfiniteData' },
+      { name: 'UseMutationResult' },
     ],
     dependency: '@tanstack/react-query',
   },
@@ -728,6 +729,37 @@ const generateQueryReturnType = ({
   }
 };
 
+const generateMutatorReturnType = ({
+  outputClient,
+  dataType,
+  variableType,
+  packageJson,
+}: {
+  outputClient: OutputClient | OutputClientFunc;
+  dataType: unknown;
+  variableType: unknown;
+  packageJson?: PackageJson;
+}) => {
+  if (!packageJson || outputClient !== OutputClient.REACT_QUERY) {
+    return '';
+  }
+
+  const hasReactQueryV4 =
+    packageJson?.dependencies?.['@tanstack/react-query'] ??
+    packageJson?.devDependencies?.['@tanstack/react-query'];
+
+  if (!hasReactQueryV4) {
+    return '';
+  }
+
+  return `: UseMutationResult<
+        Awaited<ReturnType<${dataType}>>,
+        TError,
+        ${variableType},
+        TContext
+      >`;
+};
+
 const getQueryOptions = ({
   isRequestOptions,
   mutator,
@@ -1373,17 +1405,17 @@ const generateQueryHook = async (
 
     const mutationOptionsFn = `export const ${mutationOptionsFnName} = <TError = ${errorType},
     TContext = unknown>(${mutationArguments}): ${mutationOptionFnReturnType} => {
- ${
-   isRequestOptions
-     ? `const {mutation: mutationOptions${
-         !mutator
-           ? `, axios: axiosOptions`
-           : mutator?.hasSecondArg
-             ? ', request: requestOptions'
-             : ''
-       }} = options ?? {};`
-     : ''
- }
+${
+  isRequestOptions
+    ? `const {mutation: mutationOptions${
+        !mutator
+          ? `, axios: axiosOptions`
+          : mutator?.hasSecondArg
+            ? ', request: requestOptions'
+            : ''
+      }} = options ?? {};`
+    : ''
+}
 
       ${
         mutator?.isHook
@@ -1421,11 +1453,11 @@ const generateQueryHook = async (
         }
 
 
-   return  ${
-     !mutationOptionsMutator
-       ? '{ mutationFn, ...mutationOptions }'
-       : 'customOptions'
-   }}`;
+  return  ${
+    !mutationOptionsMutator
+      ? '{ mutationFn, ...mutationOptions }'
+      : 'customOptions'
+  }}`;
 
     const operationPrefix = hasSvelteQueryV4 ? 'create' : 'use';
 
@@ -1449,7 +1481,12 @@ ${mutationOptionsFn}
     ${doc}export const ${camel(
       `${operationPrefix}-${operationName}`,
     )} = <TError = ${errorType},
-    TContext = unknown>(${mutationArguments}) => {
+    TContext = unknown>(${mutationArguments})${generateMutatorReturnType({
+      outputClient,
+      dataType,
+      variableType: definitions ? `{${definitions}}` : 'void',
+      packageJson: context.output.packageJson,
+    })} => {
 
       const ${mutationOptionsVarName} = ${mutationOptionsFnName}(${
         isRequestOptions ? 'options' : 'mutationOptions'

diff --git a/samples/react-query/basic/src/api/endpoints/petstoreFromFileSpecWithTransformer.ts b/samples/react-query/basic/src/api/endpoints/petstoreFromFileSpecWithTransformer.ts
@@ -19,6 +19,7 @@ import type {
   UseInfiniteQueryOptions,
   UseInfiniteQueryResult,
   UseMutationOptions,
+  UseMutationResult,
   UseQueryOptions,
   UseQueryResult,
   UseSuspenseInfiniteQueryOptions,
@@ -476,7 +477,12 @@ export const useCreatePets = <
     { data: CreatePetsBody; version?: number },
     TContext
   >;
-}) => {
+}): UseMutationResult<
+  Awaited<ReturnType<typeof createPets>>,
+  TError,
+  { data: CreatePetsBody; version?: number },
+  TContext
+> => {
   const mutationOptions = getCreatePetsMutationOptions(options);
 
   return useMutation(mutationOptions);
@@ -543,7 +549,12 @@ export const useUpdatePets = <
     { data: NonReadonly<Pet>; version?: number },
     TContext
   >;
-}) => {
+}): UseMutationResult<
+  Awaited<ReturnType<typeof updatePets>>,
+  TError,
+  { data: NonReadonly<Pet>; version?: number },
+  TContext
+> => {
   const mutationOptions = getUpdatePetsMutationOptions(options);
 
   return useMutation(mutationOptions);