feat: improved automatic query invalidation for tanstack-query

packages/plugins/tanstack-query/package.json

@@ -99,15 +99,20 @@
         "@tanstack/svelte-query": "^4.29.7",
         "@tanstack/svelte-query-v5": "npm:@tanstack/svelte-query@^5.0.0",
         "@tanstack/vue-query": "^4.37.0",
+        "@testing-library/react": "^14.0.0",
         "@types/jest": "^29.5.0",
+        "@types/nock": "^11.1.0",
         "@types/node": "^18.0.0",
         "@types/react": "18.2.0",
         "@types/semver": "^7.3.13",
         "@types/tmp": "^0.2.3",
         "@zenstackhq/testtools": "workspace:*",
         "copyfiles": "^2.4.1",
-        "jest": "^29.5.0",
+        "jest": "^29.7.0",
+        "jest-environment-jsdom": "^29.7.0",
+        "nock": "^13.3.6",
         "react": "18.2.0",
+        "react-test-renderer": "^18.2.0",
         "replace-in-file": "^7.0.1",
         "rimraf": "^3.0.2",
         "svelte": "^4.2.1",

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

@@ -2,6 +2,7 @@ import type { DMMF } from '@prisma/generator-helper';
 import {
     PluginError,
     PluginOptions,
+    generateModelMeta as _generateModelMeta,
     createProject,
     getDataModels,
     getPrismaClientImportSpec,
@@ -44,6 +45,8 @@ export async function generate(model: Model, options: PluginOptions, dmmf: DMMF.
         throw new PluginError(options.name, `Unsupported version "${version}": use "v4" or "v5"`);
     }
 
+    await generateModelMeta(project, outDir, models);
+
     generateIndex(project, outDir, models, target, version);
 
     models.forEach((dataModel) => {
@@ -158,11 +161,11 @@ function generateMutationHook(
             {
                 name: `_mutation`,
                 initializer: `
-                    ${httpVerb}Mutation<${argsType}, ${
+                    mutate<${argsType}, ${
                     overrideReturnType ?? model
-                }, ${checkReadBack}>('${model}', \`\${endpoint}/${lowerCaseFirst(
+                }, ${checkReadBack}>('${model}', '${httpVerb.toUpperCase()}', \`\${endpoint}/${lowerCaseFirst(
                     model
-                )}/${operation}\`, options, fetch, invalidateQueries, ${checkReadBack})
+                )}/${operation}\`, metadata, options, fetch, invalidateQueries, ${checkReadBack})
                 `,
             },
         ],
@@ -438,16 +441,21 @@ function generateModelHooks(
     }
 }
 
+async function generateModelMeta(project: Project, outDir: string, models: DataModel[]) {
+    await _generateModelMeta(project, models, path.join(outDir, '__model_meta.ts'), false, true);
+}
+
 function generateIndex(
     project: Project,
     outDir: string,
     models: DataModel[],
     target: string,
     version: TanStackVersion
 ) {
+    const runtimeImportBase = makeRuntimeImportBase(version);
     const sf = project.createSourceFile(path.join(outDir, 'index.ts'), undefined, { overwrite: true });
     sf.addStatements(models.map((d) => `export * from './${paramCase(d.name)}';`));
-    const runtimeImportBase = makeRuntimeImportBase(version);
+    sf.addStatements(`export { getQueryKey } from '${runtimeImportBase}';`);
     switch (target) {
         case 'react':
             sf.addStatements(`export { Provider } from '${runtimeImportBase}/react';`);
@@ -477,8 +485,9 @@ function makeGetContext(target: TargetFramework) {
 function makeBaseImports(target: TargetFramework, version: TanStackVersion) {
     const runtimeImportBase = makeRuntimeImportBase(version);
     const shared = [
-        `import { query, infiniteQuery, postMutation, putMutation, deleteMutation } from '${runtimeImportBase}/${target}';`,
+        `import { query, infiniteQuery, mutate } from '${runtimeImportBase}/${target}';`,
         `import type { PickEnumerable, CheckSelect } from '${runtimeImportBase}';`,
+        `import metadata from './__model_meta';`,
     ];
     switch (target) {
         case 'react':

packages/plugins/tanstack-query/src/runtime-v5/index.ts

@@ -1,2 +1,2 @@
 export * from '../runtime/prisma-types';
-export type { FetchFn } from '../runtime/common';
+export { type FetchFn, getQueryKey } from '../runtime/common';

packages/plugins/tanstack-query/src/runtime-v5/react.ts

@@ -5,20 +5,20 @@ import {
     useQuery,
     useQueryClient,
     type InfiniteData,
-    type MutateFunction,
-    type QueryClient,
     type UseInfiniteQueryOptions,
     type UseMutationOptions,
     type UseQueryOptions,
 } from '@tanstack/react-query-v5';
-import { createContext } from 'react';
+import type { ModelMeta } from '@zenstackhq/runtime/cross';
+import { createContext, useContext } from 'react';
 import {
     DEFAULT_QUERY_ENDPOINT,
     FetchFn,
-    QUERY_KEY_PREFIX,
     fetcher,
+    getQueryKey,
     makeUrl,
     marshal,
+    setupInvalidation,
     type APIContext,
 } from '../runtime/common';
 
@@ -53,7 +53,7 @@ export function query<R>(
 ) {
     const reqUrl = makeUrl(url, args);
     return useQuery({
-        queryKey: [QUERY_KEY_PREFIX + model, url, args],
+        queryKey: getQueryKey(model, url, args),
         queryFn: () => fetcher<R, false>(reqUrl, undefined, fetch, false),
         ...options,
     });
@@ -76,136 +76,54 @@ export function infiniteQuery<R>(
     fetch?: FetchFn
 ) {
     return useInfiniteQuery({
-        queryKey: [QUERY_KEY_PREFIX + model, url, args],
+        queryKey: getQueryKey(model, url, args),
         queryFn: ({ pageParam }) => {
             return fetcher<R, false>(makeUrl(url, pageParam ?? args), undefined, fetch, false);
         },
         ...options,
     });
 }
 
-/**
- * Creates a POST mutation with react-query.
- *
- * @param model The name of the model under mutation.
- * @param url The request URL.
- * @param options The react-query options.
- * @param invalidateQueries Whether to invalidate queries after mutation.
- * @returns useMutation hooks
- */
-export function postMutation<T, R = any, C extends boolean = boolean, Result = C extends true ? R | undefined : R>(
-    model: string,
-    url: string,
-    options?: Omit<UseMutationOptions<Result, unknown, T>, 'mutationFn'>,
-    fetch?: FetchFn,
-    invalidateQueries = true,
-    checkReadBack?: C
-) {
-    const queryClient = useQueryClient();
-    const mutationFn = (data: any) =>
-        fetcher<R, C>(
-            url,
-            {
-                method: 'POST',
-                headers: {
-                    'content-type': 'application/json',
-                },
-                body: marshal(data),
-            },
-            fetch,
-            checkReadBack
-        ) as Promise<Result>;
-
-    const finalOptions = mergeOptions(model, options, invalidateQueries, mutationFn, queryClient);
-    const mutation = useMutation(finalOptions);
-    return mutation;
-}
-
-/**
- * Creates a PUT mutation with react-query.
- *
- * @param model The name of the model under mutation.
- * @param url The request URL.
- * @param options The react-query options.
- * @param invalidateQueries Whether to invalidate queries after mutation.
- * @returns useMutation hooks
- */
-export function putMutation<T, R = any, C extends boolean = boolean, Result = C extends true ? R | undefined : R>(
+export function mutate<T, R = any, C extends boolean = boolean, Result = C extends true ? R | undefined : R>(
     model: string,
+    method: 'POST' | 'PUT' | 'DELETE',
     url: string,
+    modelMeta: ModelMeta,
     options?: Omit<UseMutationOptions<Result, unknown, T>, 'mutationFn'>,
     fetch?: FetchFn,
     invalidateQueries = true,
     checkReadBack?: C
 ) {
     const queryClient = useQueryClient();
-    const mutationFn = (data: any) =>
-        fetcher<R, C>(
-            url,
-            {
-                method: 'PUT',
+    const mutationFn = (data: any) => {
+        const reqUrl = method === 'DELETE' ? makeUrl(url, data) : url;
+        const fetchInit: RequestInit = {
+            method,
+            ...(method !== 'DELETE' && {
                 headers: {
                     'content-type': 'application/json',
                 },
                 body: marshal(data),
-            },
-            fetch,
-            checkReadBack
-        ) as Promise<Result>;
-
-    const finalOptions = mergeOptions(model, options, invalidateQueries, mutationFn, queryClient);
-    const mutation = useMutation(finalOptions);
-    return mutation;
-}
-
-/**
- * Creates a DELETE mutation with react-query.
- *
- * @param model The name of the model under mutation.
- * @param url The request URL.
- * @param options The react-query options.
- * @param invalidateQueries Whether to invalidate queries after mutation.
- * @returns useMutation hooks
- */
-export function deleteMutation<T, R = any, C extends boolean = boolean, Result = C extends true ? R | undefined : R>(
-    model: string,
-    url: string,
-    options?: Omit<UseMutationOptions<Result, unknown, T>, 'mutationFn'>,
-    fetch?: FetchFn,
-    invalidateQueries = true,
-    checkReadBack?: C
-) {
-    const queryClient = useQueryClient();
-    const mutationFn = (data: any) =>
-        fetcher<R, C>(
-            makeUrl(url, data),
-            {
-                method: 'DELETE',
-            },
-            fetch,
-            checkReadBack
-        ) as Promise<Result>;
-
-    const finalOptions = mergeOptions(model, options, invalidateQueries, mutationFn, queryClient);
-    const mutation = useMutation(finalOptions);
-    return mutation;
-}
-
-function mergeOptions<T, R = any>(
-    model: string,
-    options: Omit<UseMutationOptions<R, unknown, T, unknown>, 'mutationFn'> | undefined,
-    invalidateQueries: boolean,
-    mutationFn: MutateFunction<R, unknown, T>,
-    queryClient: QueryClient
-): UseMutationOptions<R, unknown, T, unknown> {
-    const result = { ...options, mutationFn };
-    if (options?.onSuccess || invalidateQueries) {
-        result.onSuccess = (...args) => {
-            if (invalidateQueries) {
-                queryClient.invalidateQueries({ queryKey: [QUERY_KEY_PREFIX + model] });
-            }
-            return options?.onSuccess?.(...args);
+            }),
         };
+        return fetcher<R, C>(reqUrl, fetchInit, fetch, checkReadBack) as Promise<Result>;
+    };
+
+    const finalOptions = { ...options, mutationFn };
+    if (invalidateQueries) {
+        const { logging } = useContext(RequestHandlerContext);
+        const operation = url.split('/').pop();
+        if (operation) {
+            setupInvalidation(
+                model,
+                operation,
+                modelMeta,
+                finalOptions,
+                (predicate) => queryClient.invalidateQueries({ predicate }),
+                logging
+            );
+        }
     }
-    return result;
+
+    return useMutation(finalOptions);
 }