fix: tanstack-query, fix the incorrect query typing when user provide…

…s a custom selector
zenstackhq · Jan 29, 2024 · d78e80c · d78e80c
1 parent 8559a78
commit d78e80c
Show file tree

Hide file tree

Showing 7 changed files with 105 additions and 87 deletions.
diff --git a/packages/plugins/tanstack-query/src/generator.ts b/packages/plugins/tanstack-query/src/generator.ts
@@ -84,9 +84,9 @@ function generateQueryHook(
     const capOperation = upperCaseFirst(operation);
 
     const argsType = overrideInputType ?? `Prisma.${model}${capOperation}Args`;
-    const inputType = `Prisma.SelectSubset<T, ${argsType}>`;
+    const inputType = `Prisma.SelectSubset<TArgs, ${argsType}>`;
 
-    let defaultReturnType = `Prisma.${model}GetPayload<T>`;
+    let defaultReturnType = `Prisma.${model}GetPayload<TArgs>`;
     if (optimisticUpdate) {
         defaultReturnType += '& { $optimistic?: boolean }';
     }
@@ -95,11 +95,16 @@ function generateQueryHook(
     }
 
     const returnType = overrideReturnType ?? defaultReturnType;
-    const optionsType = makeQueryOptions(target, returnType, infinite, version);
+    const optionsType = makeQueryOptions(target, 'TQueryFnData', 'TData', infinite, version);
 
     const func = sf.addFunction({
         name: `use${infinite ? 'Infinite' : ''}${capOperation}${model}`,
-        typeParameters: overrideTypeParameters ?? [`T extends ${argsType}`],
+        typeParameters: overrideTypeParameters ?? [
+            `TArgs extends ${argsType}`,
+            `TQueryFnData = ${returnType} `,
+            'TData = TQueryFnData',
+            'TError = DefaultError',
+        ],
         parameters: [
             {
                 name: optionalInput ? 'args?' : 'args',
@@ -129,7 +134,9 @@ function generateQueryHook(
 
     func.addStatements([
         makeGetContext(target),
-        `return ${infinite ? 'useInfiniteModelQuery' : 'useModelQuery'}('${model}', \`\${endpoint}/${lowerCaseFirst(
+        `return ${
+            infinite ? 'useInfiniteModelQuery' : 'useModelQuery'
+        }<TQueryFnData, TData, TError>('${model}', \`\${endpoint}/${lowerCaseFirst(
             model
         )}/${operation}\`, args, options, fetch${optimisticUpdate ? ', optimisticUpdate' : ''});`,
     ]);
@@ -403,7 +410,7 @@ function generateModelHooks(
             'aggregate',
             false,
             false,
-            `Prisma.Get${modelNameCap}AggregateType<T>`
+            `Prisma.Get${modelNameCap}AggregateType<TArgs>`
         );
     }
 
@@ -415,16 +422,27 @@ function generateModelHooks(
             useName = model.name;
         }
 
+        const returnType = `{} extends InputErrors ? 
+        Array<PickEnumerable<Prisma.${modelNameCap}GroupByOutputType, TArgs['by']> &
+          {
+            [P in ((keyof TArgs) & (keyof Prisma.${modelNameCap}GroupByOutputType))]: P extends '_count'
+              ? TArgs[P] extends boolean
+                ? number
+                : Prisma.GetScalarType<TArgs[P], Prisma.${modelNameCap}GroupByOutputType[P]>
+              : Prisma.GetScalarType<TArgs[P], Prisma.${modelNameCap}GroupByOutputType[P]>
+          }
+        > : InputErrors`;
+
         const typeParameters = [
-            `T extends Prisma.${useName}GroupByArgs`,
-            `HasSelectOrTake extends Prisma.Or<Prisma.Extends<'skip', Prisma.Keys<T>>, Prisma.Extends<'take', Prisma.Keys<T>>>`,
+            `TArgs extends Prisma.${useName}GroupByArgs`,
+            `HasSelectOrTake extends Prisma.Or<Prisma.Extends<'skip', Prisma.Keys<TArgs>>, Prisma.Extends<'take', Prisma.Keys<TArgs>>>`,
             `OrderByArg extends Prisma.True extends HasSelectOrTake ? { orderBy: Prisma.${useName}GroupByArgs['orderBy'] }: { orderBy?: Prisma.${useName}GroupByArgs['orderBy'] },`,
-            `OrderFields extends Prisma.ExcludeUnderscoreKeys<Prisma.Keys<Prisma.MaybeTupleToUnion<T['orderBy']>>>`,
-            `ByFields extends Prisma.MaybeTupleToUnion<T['by']>`,
+            `OrderFields extends Prisma.ExcludeUnderscoreKeys<Prisma.Keys<Prisma.MaybeTupleToUnion<TArgs['orderBy']>>>`,
+            `ByFields extends Prisma.MaybeTupleToUnion<TArgs['by']>`,
             `ByValid extends Prisma.Has<ByFields, OrderFields>`,
-            `HavingFields extends Prisma.GetHavingFields<T['having']>`,
+            `HavingFields extends Prisma.GetHavingFields<TArgs['having']>`,
             `HavingValid extends Prisma.Has<ByFields, HavingFields>`,
-            `ByEmpty extends T['by'] extends never[] ? Prisma.True : Prisma.False`,
+            `ByEmpty extends TArgs['by'] extends never[] ? Prisma.True : Prisma.False`,
             `InputErrors extends ByEmpty extends Prisma.True
             ? \`Error: "by" must not be empty.\`
             : HavingValid extends Prisma.False
@@ -440,8 +458,8 @@ function generateModelHooks(
                     \` in "having" needs to be provided in "by"\`,
                     ]
             }[HavingFields]
-            : 'take' extends Prisma.Keys<T>
-            ? 'orderBy' extends Prisma.Keys<T>
+            : 'take' extends Prisma.Keys<TArgs>
+            ? 'orderBy' extends Prisma.Keys<TArgs>
             ? ByValid extends Prisma.True
                 ? {}
                 : {
@@ -450,8 +468,8 @@ function generateModelHooks(
                     : \`Error: Field "\${P}" in "orderBy" needs to be provided in "by"\`
                 }[OrderFields]
             : 'Error: If you provide "take", you also need to provide "orderBy"'
-            : 'skip' extends Prisma.Keys<T>
-            ? 'orderBy' extends Prisma.Keys<T>
+            : 'skip' extends Prisma.Keys<TArgs>
+            ? 'orderBy' extends Prisma.Keys<TArgs>
             ? ByValid extends Prisma.True
                 ? {}
                 : {
@@ -467,19 +485,11 @@ function generateModelHooks(
                 ? never
                 : \`Error: Field "\${P}" in "orderBy" needs to be provided in "by"\`
             }[OrderFields]`,
+            `TQueryFnData = ${returnType}`,
+            `TData = TQueryFnData`,
+            `TError = DefaultError`,
         ];
 
-        const returnType = `{} extends InputErrors ? 
-        Array<PickEnumerable<Prisma.${modelNameCap}GroupByOutputType, T['by']> &
-          {
-            [P in ((keyof T) & (keyof Prisma.${modelNameCap}GroupByOutputType))]: P extends '_count'
-              ? T[P] extends boolean
-                ? number
-                : Prisma.GetScalarType<T[P], Prisma.${modelNameCap}GroupByOutputType[P]>
-              : Prisma.GetScalarType<T[P], Prisma.${modelNameCap}GroupByOutputType[P]>
-          }
-        > : InputErrors`;
-
         generateQueryHook(
             target,
             version,
@@ -489,7 +499,7 @@ function generateModelHooks(
             false,
             false,
             returnType,
-            `Prisma.SubsetIntersection<T, Prisma.${useName}GroupByArgs, OrderByArg> & InputErrors`,
+            `Prisma.SubsetIntersection<TArgs, Prisma.${useName}GroupByArgs, OrderByArg> & InputErrors`,
             typeParameters
         );
     }
@@ -504,7 +514,7 @@ function generateModelHooks(
             'count',
             false,
             true,
-            `T extends { select: any; } ? T['select'] extends true ? number : Prisma.GetScalarType<T['select'], Prisma.${modelNameCap}CountAggregateOutputType> : number`
+            `TArgs extends { select: any; } ? TArgs['select'] extends true ? number : Prisma.GetScalarType<TArgs['select'], Prisma.${modelNameCap}CountAggregateOutputType> : number`
         );
     }
 }
@@ -552,12 +562,13 @@ function makeBaseImports(target: TargetFramework, version: TanStackVersion) {
         `import { useModelQuery, useInfiniteModelQuery, useModelMutation } from '${runtimeImportBase}/${target}';`,
         `import type { PickEnumerable, CheckSelect } from '${runtimeImportBase}';`,
         `import metadata from './__model_meta';`,
+        `type DefaultError = Error;`,
     ];
     switch (target) {
         case 'react':
             return [
                 `import type { UseMutationOptions, UseQueryOptions, UseInfiniteQueryOptions, InfiniteData } from '@tanstack/react-query';`,
-                `import { RequestHandlerContext, getHooksContext } from '${runtimeImportBase}/${target}';`,
+                `import { getHooksContext } from '${runtimeImportBase}/${target}';`,
                 ...shared,
             ];
         case 'vue':
@@ -573,32 +584,38 @@ function makeBaseImports(target: TargetFramework, version: TanStackVersion) {
                 ...(version === 'v5'
                     ? [`import type { InfiniteData, StoreOrVal } from '@tanstack/svelte-query';`]
                     : []),
-                `import { SvelteQueryContextKey, type RequestHandlerContext, getHooksContext } from '${runtimeImportBase}/${target}';`,
+                `import { getHooksContext } from '${runtimeImportBase}/${target}';`,
                 ...shared,
             ];
         default:
             throw new PluginError(name, `Unsupported target: ${target}`);
     }
 }
 
-function makeQueryOptions(target: string, returnType: string, infinite: boolean, version: TanStackVersion) {
+function makeQueryOptions(
+    target: string,
+    returnType: string,
+    dataType: string,
+    infinite: boolean,
+    version: TanStackVersion
+) {
     switch (target) {
         case 'react':
             return infinite
                 ? version === 'v4'
-                    ? `Omit<UseInfiniteQueryOptions<${returnType}>, 'queryKey'>`
-                    : `Omit<UseInfiniteQueryOptions<${returnType}, unknown, InfiniteData<${returnType}>>, 'queryKey'>`
-                : `Omit<UseQueryOptions<${returnType}>, 'queryKey'>`;
+                    ? `Omit<UseInfiniteQueryOptions<${returnType}, TError, ${dataType}>, 'queryKey'>`
+                    : `Omit<UseInfiniteQueryOptions<${returnType}, TError, InfiniteData<${dataType}>>, 'queryKey'>`
+                : `Omit<UseQueryOptions<${returnType}, TError, ${dataType}>, 'queryKey'>`;
         case 'vue':
-            return `Omit<Use${infinite ? 'Infinite' : ''}QueryOptions<${returnType}>, 'queryKey'>`;
+            return `Omit<Use${infinite ? 'Infinite' : ''}QueryOptions<${returnType}, TError, ${dataType}>, 'queryKey'>`;
         case 'svelte':
             return infinite
                 ? version === 'v4'
-                    ? `Omit<CreateInfiniteQueryOptions<${returnType}>, 'queryKey'>`
-                    : `StoreOrVal<Omit<CreateInfiniteQueryOptions<${returnType}, unknown, InfiniteData<${returnType}>>, 'queryKey'>>`
+                    ? `Omit<CreateInfiniteQueryOptions<${returnType}, TError, ${dataType}>, 'queryKey'>`
+                    : `StoreOrVal<Omit<CreateInfiniteQueryOptions<${returnType}, TError, InfiniteData<${dataType}>>, 'queryKey'>>`
                 : version === 'v4'
-                ? `Omit<CreateQueryOptions<${returnType}>, 'queryKey'>`
-                : `StoreOrVal<Omit<CreateQueryOptions<${returnType}>, 'queryKey'>>`;
+                ? `Omit<CreateQueryOptions<${returnType}, TError, ${dataType}>, 'queryKey'>`
+                : `StoreOrVal<Omit<CreateQueryOptions<${returnType}, TError, ${dataType}>, 'queryKey'>>`;
         default:
             throw new PluginError(name, `Unsupported target: ${target}`);
     }

diff --git a/packages/plugins/tanstack-query/src/runtime-v5/react.ts b/packages/plugins/tanstack-query/src/runtime-v5/react.ts
@@ -55,18 +55,18 @@ export const Provider = RequestHandlerContext.Provider;
  * @param optimisticUpdate Whether to enable automatic optimistic update
  * @returns useQuery hook
  */
-export function useModelQuery<R>(
+export function useModelQuery<TQueryFnData, TData, TError>(
     model: string,
     url: string,
     args?: unknown,
-    options?: Omit<UseQueryOptions<R>, 'queryKey'>,
+    options?: Omit<UseQueryOptions<TQueryFnData, TError, TData>, 'queryKey'>,
     fetch?: FetchFn,
     optimisticUpdate = false
 ) {
     const reqUrl = makeUrl(url, args);
     return useQuery({
         queryKey: getQueryKey(model, url, args, false, optimisticUpdate),
-        queryFn: () => fetcher<R, false>(reqUrl, undefined, fetch, false),
+        queryFn: () => fetcher<TQueryFnData, false>(reqUrl, undefined, fetch, false),
         ...options,
     });
 }
@@ -81,17 +81,17 @@ export function useModelQuery<R>(
  * @param fetch The fetch function to use for sending the HTTP request
  * @returns useInfiniteQuery hook
  */
-export function useInfiniteModelQuery<R>(
+export function useInfiniteModelQuery<TQueryFnData, TData, TError>(
     model: string,
     url: string,
     args: unknown,
-    options: Omit<UseInfiniteQueryOptions<R, unknown, InfiniteData<R>>, 'queryKey'>,
+    options: Omit<UseInfiniteQueryOptions<TQueryFnData, TError, InfiniteData<TData>>, 'queryKey'>,
     fetch?: FetchFn
 ) {
     return useInfiniteQuery({
         queryKey: getQueryKey(model, url, args, true),
         queryFn: ({ pageParam }) => {
-            return fetcher<R, false>(makeUrl(url, pageParam ?? args), undefined, fetch, false);
+            return fetcher<TQueryFnData, false>(makeUrl(url, pageParam ?? args), undefined, fetch, false);
         },
         ...options,
     });

diff --git a/packages/plugins/tanstack-query/src/runtime-v5/svelte.ts b/packages/plugins/tanstack-query/src/runtime-v5/svelte.ts
@@ -5,10 +5,10 @@ import {
     createQuery,
     useQueryClient,
     type CreateInfiniteQueryOptions,
+    type CreateQueryOptions,
     type InfiniteData,
     type MutationOptions,
     type StoreOrVal,
-    type CreateQueryOptions,
 } from '@tanstack/svelte-query-v5';
 import { ModelMeta } from '@zenstackhq/runtime/cross';
 import { getContext, setContext } from 'svelte';
@@ -58,17 +58,17 @@ export function getHooksContext() {
  * @param optimisticUpdate Whether to enable automatic optimistic update
  * @returns useQuery hook
  */
-export function useModelQuery<R>(
+export function useModelQuery<TQueryFnData, TData, TError>(
     model: string,
     url: string,
     args?: unknown,
-    options?: StoreOrVal<Omit<CreateQueryOptions<R>, 'queryKey'>>,
+    options?: StoreOrVal<Omit<CreateQueryOptions<TQueryFnData, TData, TError>, 'queryKey'>>,
     fetch?: FetchFn,
     optimisticUpdate = false
 ) {
     const reqUrl = makeUrl(url, args);
     const queryKey = getQueryKey(model, url, args, false, optimisticUpdate);
-    const queryFn = () => fetcher<R, false>(reqUrl, undefined, fetch, false);
+    const queryFn = () => fetcher<TQueryFnData, false>(reqUrl, undefined, fetch, false);
 
     let mergedOpt: any;
     if (isStore(options)) {
@@ -100,19 +100,19 @@ export function useModelQuery<R>(
  * @param options The svelte-query infinite query options object
  * @returns useQuery hook
  */
-export function useInfiniteModelQuery<R>(
+export function useInfiniteModelQuery<TQueryFnData, TData, TError>(
     model: string,
     url: string,
     args: unknown,
-    options: StoreOrVal<Omit<CreateInfiniteQueryOptions<R, unknown, InfiniteData<R>>, 'queryKey'>>,
+    options: StoreOrVal<Omit<CreateInfiniteQueryOptions<TQueryFnData, TError, InfiniteData<TData>>, 'queryKey'>>,
     fetch?: FetchFn
 ) {
     const queryKey = getQueryKey(model, url, args, true);
     const queryFn = ({ pageParam }: { pageParam: unknown }) =>
-        fetcher<R, false>(makeUrl(url, pageParam ?? args), undefined, fetch, false);
+        fetcher<TQueryFnData, false>(makeUrl(url, pageParam ?? args), undefined, fetch, false);
 
-    let mergedOpt: StoreOrVal<CreateInfiniteQueryOptions<R, unknown, InfiniteData<R>>>;
-    if (isStore<CreateInfiniteQueryOptions<R, unknown, InfiniteData<R>>>(options)) {
+    let mergedOpt: StoreOrVal<CreateInfiniteQueryOptions<TQueryFnData, TError, InfiniteData<TData>>>;
+    if (isStore<CreateInfiniteQueryOptions<TQueryFnData, TError, InfiniteData<TData>>>(options)) {
         // options is store
         mergedOpt = derived([options], ([$opt]) => {
             return {
@@ -129,7 +129,7 @@ export function useInfiniteModelQuery<R>(
             ...options,
         };
     }
-    return createInfiniteQuery<R, unknown, InfiniteData<R>>(mergedOpt);
+    return createInfiniteQuery<TQueryFnData, TError, InfiniteData<TData>>(mergedOpt);
 }
 
 function isStore<T>(opt: unknown): opt is Readable<T> {

diff --git a/packages/plugins/tanstack-query/src/runtime/react.ts b/packages/plugins/tanstack-query/src/runtime/react.ts
@@ -55,18 +55,18 @@ export function getHooksContext() {
  * @param optimisticUpdate Whether to enable automatic optimistic update
  * @returns useQuery hook
  */
-export function useModelQuery<R>(
+export function useModelQuery<TQueryFnData, TData, TError>(
     model: string,
     url: string,
     args?: unknown,
-    options?: Omit<UseQueryOptions<R>, 'queryKey'>,
+    options?: Omit<UseQueryOptions<TQueryFnData, TError, TData>, 'queryKey'>,
     fetch?: FetchFn,
     optimisticUpdate = false
 ) {
     const reqUrl = makeUrl(url, args);
-    return useQuery<R>({
+    return useQuery<TQueryFnData, TError, TData>({
         queryKey: getQueryKey(model, url, args, false, optimisticUpdate),
-        queryFn: () => fetcher<R, false>(reqUrl, undefined, fetch, false),
+        queryFn: () => fetcher<TQueryFnData, false>(reqUrl, undefined, fetch, false),
         ...options,
     });
 }
@@ -81,17 +81,17 @@ export function useModelQuery<R>(
  * @param fetch The fetch function to use for sending the HTTP request
  * @returns useInfiniteQuery hook
  */
-export function useInfiniteModelQuery<R>(
+export function useInfiniteModelQuery<TQueryFnData, TData, TError>(
     model: string,
     url: string,
     args?: unknown,
-    options?: Omit<UseInfiniteQueryOptions<R>, 'queryKey'>,
+    options?: Omit<UseInfiniteQueryOptions<TQueryFnData, TError, TData>, 'queryKey'>,
     fetch?: FetchFn
 ) {
-    return useInfiniteQuery<R>({
+    return useInfiniteQuery<TQueryFnData, TError, TData>({
         queryKey: getQueryKey(model, url, args, true),
         queryFn: ({ pageParam }) => {
-            return fetcher<R, false>(makeUrl(url, pageParam ?? args), undefined, fetch, false);
+            return fetcher<TQueryFnData, false>(makeUrl(url, pageParam ?? args), undefined, fetch, false);
         },
         ...options,
     });