Skip to content

Commit 44567cc

Browse files
ymc9ymc9
committed
feat(tanstack-query): generate fetch and prefetch functions
1 parent 1bc622a commit 44567cc

File tree

5 files changed

+473
-176
lines changed

5 files changed

+473
-176
lines changed

‎packages/plugins/tanstack-query/src/generator.ts

Lines changed: 171 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,8 @@ const supportedTargets = ['react', 'vue', 'svelte'];
2525
type TargetFramework = (typeof supportedTargets)[number];
2626
type TanStackVersion = 'v4' | 'v5';
2727

28+
// TODO: turn it into a class to simplify parameter passing
29+
2830
export async function generate(model: Model, options: PluginOptions, dmmf: DMMF.Document) {
2931
const project = createProject();
3032
const warnings: string[] = [];
@@ -40,6 +42,17 @@ export async function generate(model: Model, options: PluginOptions, dmmf: DMMF.
4042
throw new PluginError(name, `Unsupported version "${version}": use "v4" or "v5"`);
4143
}
4244

45+
if (options.generatePrefetch !== undefined && typeof options.generatePrefetch !== 'boolean') {
46+
throw new PluginError(
47+
name,
48+
`Invalid "generatePrefetch" option: expected boolean, got ${options.generatePrefetch}`
49+
);
50+
}
51+
52+
if (options.generatePrefetch === true && version === 'v4') {
53+
throw new PluginError(name, `"generatePrefetch" is not supported for version "v4"`);
54+
}
55+
4356
let outDir = requireOption<string>(options, 'output', name);
4457
outDir = resolvePath(outDir, options);
4558
ensureEmptyDir(outDir);
@@ -71,27 +84,20 @@ function generateQueryHook(
7184
model: string,
7285
operation: string,
7386
returnArray: boolean,
87+
returnNullable: boolean,
7488
optionalInput: boolean,
7589
overrideReturnType?: string,
7690
overrideInputType?: string,
7791
overrideTypeParameters?: string[],
7892
supportInfinite = false,
7993
supportOptimistic = false,
80-
supportPrefetching = false,
94+
generatePrefetch = false
8195
) {
8296
const generateModes: ('' | 'Infinite' | 'Suspense' | 'SuspenseInfinite' | 'Prefetch' | 'PrefetchInfinite')[] = [''];
8397
if (supportInfinite) {
8498
generateModes.push('Infinite');
8599
}
86100

87-
if (supportPrefetching) {
88-
generateModes.push('Prefetch');
89-
90-
if (supportInfinite) {
91-
generateModes.push('PrefetchInfinite');
92-
}
93-
}
94-
95101
if (target === 'react' && version === 'v5') {
96102
// react-query v5 supports suspense query
97103
generateModes.push('Suspense');
@@ -100,34 +106,45 @@ function generateQueryHook(
100106
}
101107
}
102108

103-
for (const generateMode of generateModes) {
104-
const capOperation = upperCaseFirst(operation);
105-
106-
const argsType = overrideInputType ?? `Prisma.${model}${capOperation}Args`;
107-
const inputType = makeQueryArgsType(target, argsType);
108-
109-
const infinite = generateMode.includes('Infinite');
110-
const suspense = generateMode.includes('Suspense');
111-
const prefetch = generateMode.includes('Prefetch');
112-
const prefetchInfinite = generateMode.includes('PrefetchInfinite');
109+
const getArgsType = () => {
110+
return overrideInputType ?? `Prisma.${model}${upperCaseFirst(operation)}Args`;
111+
};
113112

114-
const optimistic =
115-
supportOptimistic &&
116-
// infinite queries are not subject to optimistic updates
117-
!infinite;
113+
const getInputType = (prefetch: boolean) => {
114+
return makeQueryArgsType(target, getArgsType(), prefetch);
115+
};
118116

117+
const getReturnType = (optimistic: boolean) => {
119118
let defaultReturnType = `Prisma.${model}GetPayload<TArgs>`;
120119
if (optimistic) {
121120
defaultReturnType += '& { $optimistic?: boolean }';
122121
}
123122
if (returnArray) {
124123
defaultReturnType = `Array<${defaultReturnType}>`;
125124
}
126-
if (prefetch || prefetchInfinite) {
127-
defaultReturnType = `Promise<void>`;
125+
if (returnNullable) {
126+
defaultReturnType = `(${defaultReturnType}) | null`;
128127
}
129128

130129
const returnType = overrideReturnType ?? defaultReturnType;
130+
return returnType;
131+
};
132+
133+
const capOperation = upperCaseFirst(operation);
134+
135+
for (const generateMode of generateModes) {
136+
const argsType = getArgsType();
137+
const inputType = getInputType(false);
138+
139+
const infinite = generateMode.includes('Infinite');
140+
const suspense = generateMode.includes('Suspense');
141+
142+
const optimistic =
143+
supportOptimistic &&
144+
// infinite queries are not subject to optimistic updates
145+
!infinite;
146+
147+
const returnType = getReturnType(optimistic);
131148
const optionsType = makeQueryOptions(target, 'TQueryFnData', 'TData', infinite, suspense, version);
132149

133150
const func = sf.addFunction({
@@ -163,6 +180,58 @@ function generateQueryHook(
163180
)}/${operation}\`, args, options, fetch);`,
164181
]);
165182
}
183+
184+
if (generatePrefetch) {
185+
const argsType = getArgsType();
186+
const inputType = getInputType(true);
187+
const returnType = getReturnType(false);
188+
189+
const modes = [
190+
{ mode: 'prefetch', infinite: false },
191+
{ mode: 'fetch', infinite: false },
192+
];
193+
if (supportInfinite) {
194+
modes.push({ mode: 'prefetch', infinite: true }, { mode: 'fetch', infinite: true });
195+
}
196+
197+
for (const { mode, infinite } of modes) {
198+
const optionsType = makePrefetchQueryOptions(target, 'TQueryFnData', 'TData', infinite);
199+
200+
const func = sf.addFunction({
201+
name: `${mode}${infinite ? 'Infinite' : ''}${capOperation}${model}`,
202+
typeParameters: overrideTypeParameters ?? [
203+
`TArgs extends ${argsType}`,
204+
`TQueryFnData = ${returnType} `,
205+
'TData = TQueryFnData',
206+
'TError = DefaultError',
207+
],
208+
parameters: [
209+
{
210+
name: 'queryClient',
211+
type: 'QueryClient',
212+
},
213+
{
214+
name: optionalInput ? 'args?' : 'args',
215+
type: inputType,
216+
},
217+
{
218+
name: 'options?',
219+
type: optionsType,
220+
},
221+
],
222+
isExported: true,
223+
});
224+
225+
func.addStatements([
226+
makeGetContext(target),
227+
`return ${mode}${
228+
infinite ? 'Infinite' : ''
229+
}ModelQuery<TQueryFnData, TData, TError>(queryClient, '${model}', \`\${endpoint}/${lowerCaseFirst(
230+
model
231+
)}/${operation}\`, args, options, fetch);`,
232+
]);
233+
}
234+
}
166235
}
167236

168237
function generateMutationHook(
@@ -349,13 +418,15 @@ function generateModelHooks(
349418

350419
sf.addStatements('/* eslint-disable */');
351420

421+
const generatePrefetch = options.generatePrefetch === true;
422+
352423
const prismaImport = getPrismaClientImportSpec(outDir, options);
353424
sf.addImportDeclaration({
354425
namedImports: ['Prisma', model.name],
355426
isTypeOnly: true,
356427
moduleSpecifier: prismaImport,
357428
});
358-
sf.addStatements(makeBaseImports(target, version));
429+
sf.addStatements(makeBaseImports(target, version, generatePrefetch));
359430

360431
// Note: delegate models don't support create and upsert operations
361432

@@ -380,13 +451,14 @@ function generateModelHooks(
380451
model.name,
381452
'findMany',
382453
true,
454+
false,
383455
true,
384456
undefined,
385457
undefined,
386458
undefined,
387459
true,
388460
true,
389-
true
461+
generatePrefetch
390462
);
391463
}
392464

@@ -399,13 +471,14 @@ function generateModelHooks(
399471
model.name,
400472
'findUnique',
401473
false,
474+
true,
402475
false,
403476
undefined,
404477
undefined,
405478
undefined,
406479
false,
407480
true,
408-
true
481+
generatePrefetch
409482
);
410483
}
411484

@@ -419,12 +492,13 @@ function generateModelHooks(
419492
'findFirst',
420493
false,
421494
true,
495+
true,
422496
undefined,
423497
undefined,
424498
undefined,
425499
false,
426500
true,
427-
true
501+
generatePrefetch
428502
);
429503
}
430504

@@ -469,7 +543,13 @@ function generateModelHooks(
469543
'aggregate',
470544
false,
471545
false,
472-
`Prisma.Get${modelNameCap}AggregateType<TArgs>`
546+
false,
547+
`Prisma.Get${modelNameCap}AggregateType<TArgs>`,
548+
undefined,
549+
undefined,
550+
false,
551+
false,
552+
generatePrefetch
473553
);
474554
}
475555

@@ -553,9 +633,13 @@ function generateModelHooks(
553633
'groupBy',
554634
false,
555635
false,
636+
false,
556637
returnType,
557638
`Prisma.SubsetIntersection<TArgs, Prisma.${useName}GroupByArgs, OrderByArg> & InputErrors`,
558-
typeParameters
639+
typeParameters,
640+
false,
641+
false,
642+
generatePrefetch
559643
);
560644
}
561645

@@ -568,8 +652,14 @@ function generateModelHooks(
568652
model.name,
569653
'count',
570654
false,
655+
false,
571656
true,
572-
`TArgs extends { select: any; } ? TArgs['select'] extends true ? number : Prisma.GetScalarType<TArgs['select'], Prisma.${modelNameCap}CountAggregateOutputType> : number`
657+
`TArgs extends { select: any; } ? TArgs['select'] extends true ? number : Prisma.GetScalarType<TArgs['select'], Prisma.${modelNameCap}CountAggregateOutputType> : number`,
658+
undefined,
659+
undefined,
660+
false,
661+
false,
662+
generatePrefetch
573663
);
574664
}
575665

@@ -616,15 +706,25 @@ function makeGetContext(target: TargetFramework) {
616706
}
617707
}
618708

619-
function makeBaseImports(target: TargetFramework, version: TanStackVersion) {
709+
function makeBaseImports(target: TargetFramework, version: TanStackVersion, generatePrefetch: boolean) {
620710
const runtimeImportBase = makeRuntimeImportBase(version);
621711
const shared = [
622-
`import { useModelQuery, useInfiniteModelQuery, useModelMutation, usePrefetchModelQuery, usePrefetchInfiniteModelQuery } from '${runtimeImportBase}/${target}';`,
712+
`import { useModelQuery, useInfiniteModelQuery, useModelMutation } from '${runtimeImportBase}/${target}';`,
623713
`import type { PickEnumerable, CheckSelect, QueryError, ExtraQueryOptions, ExtraMutationOptions } from '${runtimeImportBase}';`,
624714
`import type { PolicyCrudKind } from '${RUNTIME_PACKAGE}'`,
625715
`import metadata from './__model_meta';`,
626716
`type DefaultError = QueryError;`,
627717
];
718+
719+
if (version === 'v5' && generatePrefetch) {
720+
shared.push(
721+
`import { fetchModelQuery, prefetchModelQuery, fetchInfiniteModelQuery, prefetchInfiniteModelQuery } from '${runtimeImportBase}/${target}';`
722+
);
723+
shared.push(
724+
`import type { QueryClient, FetchQueryOptions, FetchInfiniteQueryOptions } from '@tanstack/${target}-query';`
725+
);
726+
}
727+
628728
switch (target) {
629729
case 'react': {
630730
const suspense =
@@ -645,7 +745,8 @@ function makeBaseImports(target: TargetFramework, version: TanStackVersion) {
645745
return [
646746
`import type { UseMutationOptions, UseQueryOptions, UseInfiniteQueryOptions, InfiniteData } from '@tanstack/vue-query';`,
647747
`import { getHooksContext } from '${runtimeImportBase}/${target}';`,
648-
`import type { MaybeRefOrGetter, ComputedRef, UnwrapRef } from 'vue';`,
748+
`import type { MaybeRef, MaybeRefOrGetter, ComputedRef, UnwrapRef } from 'vue';`,
749+
...(generatePrefetch ? [`import { type MaybeRefDeep } from '${runtimeImportBase}/${target}';`] : []),
649750
...shared,
650751
];
651752
}
@@ -665,10 +766,14 @@ function makeBaseImports(target: TargetFramework, version: TanStackVersion) {
665766
}
666767
}
667768

668-
function makeQueryArgsType(target: string, argsType: string) {
769+
function makeQueryArgsType(target: string, argsType: string, prefetch: boolean) {
669770
const type = `Prisma.SelectSubset<TArgs, ${argsType}>`;
670771
if (target === 'vue') {
671-
return `MaybeRefOrGetter<${type}> | ComputedRef<${type}>`;
772+
if (prefetch) {
773+
return `MaybeRef<${type}>`;
774+
} else {
775+
return `MaybeRefOrGetter<${type}> | ComputedRef<${type}>`;
776+
}
672777
} else {
673778
return type;
674779
}
@@ -721,6 +826,35 @@ function makeQueryOptions(
721826
return result;
722827
}
723828

829+
function makePrefetchQueryOptions(target: string, returnType: string, dataType: string, infinite: boolean) {
830+
let result = match(target)
831+
.with('react', () =>
832+
infinite
833+
? `Omit<FetchInfiniteQueryOptions<${returnType}, TError, ${dataType}>, 'queryKey' | 'initialPageParam'>`
834+
: `Omit<FetchQueryOptions<${returnType}, TError, ${dataType}>, 'queryKey'>`
835+
)
836+
.with('vue', () =>
837+
infinite
838+
? `MaybeRefDeep<Omit<FetchInfiniteQueryOptions<${returnType}, TError, ${dataType}>, 'queryKey' | 'initialPageParam'>>`
839+
: `MaybeRefDeep<Omit<FetchQueryOptions<${returnType}, TError, ${dataType}>, 'queryKey'>>`
840+
)
841+
.with('svelte', () =>
842+
infinite
843+
? `Omit<FetchInfiniteQueryOptions<${returnType}, TError, ${dataType}>, 'queryKey' | 'initialPageParam'>`
844+
: `Omit<FetchQueryOptions<${returnType}, TError, ${dataType}>, 'queryKey'>`
845+
)
846+
.otherwise(() => {
847+
throw new PluginError(name, `Unsupported target: ${target}`);
848+
});
849+
850+
if (!infinite) {
851+
// non-infinite queries support extra options like optimistic updates
852+
result = `(${result} & ExtraQueryOptions)`;
853+
}
854+
855+
return result;
856+
}
857+
724858
function makeMutationOptions(target: string, returnType: string, argsType: string) {
725859
let result = match(target)
726860
.with('react', () => `UseMutationOptions<${returnType}, DefaultError, ${argsType}>`)

0 commit comments

Comments
 (0)