feat: improved automatic query invalidation for tanstack-query

packages/plugins/openapi/src/rest-generator.ts

@@ -19,6 +19,7 @@ import * as path from 'path';
 import pluralize from 'pluralize';
 import invariant from 'tiny-invariant';
 import YAML from 'yaml';
+import { name } from '.';
 import { OpenAPIGeneratorBase } from './generator-base';
 import { getModelResourceMeta } from './meta';
 
@@ -31,7 +32,7 @@ export class RESTfulOpenAPIGenerator extends OpenAPIGeneratorBase {
     private warnings: string[] = [];
 
     generate() {
-        let output = requireOption<string>(this.options, 'output');
+        let output = requireOption<string>(this.options, 'output', name);
         output = resolvePath(output, this.options);
 
         const components = this.generateComponents();

packages/plugins/openapi/src/rpc-generator.ts

@@ -18,6 +18,7 @@ import * as path from 'path';
 import invariant from 'tiny-invariant';
 import { upperCaseFirst } from 'upper-case-first';
 import YAML from 'yaml';
+import { name } from '.';
 import { OpenAPIGeneratorBase } from './generator-base';
 import { getModelResourceMeta } from './meta';
 
@@ -32,7 +33,7 @@ export class RPCOpenAPIGenerator extends OpenAPIGeneratorBase {
     private warnings: string[] = [];
 
     generate() {
-        let output = requireOption<string>(this.options, 'output');
+        let output = requireOption<string>(this.options, 'output', name);
         output = resolvePath(output, this.options);
 
         // input types

packages/plugins/swr/package.json

@@ -29,6 +29,7 @@
         "@zenstackhq/runtime": "workspace:*",
         "@zenstackhq/sdk": "workspace:*",
         "change-case": "^4.1.2",
+        "cross-fetch": "^4.0.0",
         "decimal.js": "^10.4.2",
         "lower-case-first": "^2.0.2",
         "semver": "^7.3.8",
@@ -37,6 +38,7 @@
     },
     "devDependencies": {
         "@tanstack/react-query": "^4.28.0",
+        "@testing-library/react": "^14.0.0",
         "@types/jest": "^29.5.0",
         "@types/node": "^18.0.0",
         "@types/react": "18.2.0",
@@ -45,6 +47,7 @@
         "@zenstackhq/testtools": "workspace:*",
         "copyfiles": "^2.4.1",
         "jest": "^29.5.0",
+        "nock": "^13.3.6",
         "react": "18.2.0",
         "rimraf": "^3.0.2",
         "swr": "^2.0.3",

packages/plugins/swr/src/generator.ts

@@ -2,6 +2,7 @@ import type { DMMF } from '@prisma/generator-helper';
 import {
     PluginOptions,
     createProject,
+    generateModelMeta,
     getDataModels,
     getPrismaClientImportSpec,
     getPrismaVersion,
@@ -16,9 +17,10 @@ import path from 'path';
 import semver from 'semver';
 import { FunctionDeclaration, OptionalKind, ParameterDeclarationStructure, Project, SourceFile } from 'ts-morph';
 import { upperCaseFirst } from 'upper-case-first';
+import { name } from '.';
 
 export async function generate(model: Model, options: PluginOptions, dmmf: DMMF.Document) {
-    let outDir = requireOption<string>(options, 'output');
+    let outDir = requireOption<string>(options, 'output', name);
     outDir = resolvePath(outDir, options);
 
     const project = createProject();
@@ -32,6 +34,8 @@ export async function generate(model: Model, options: PluginOptions, dmmf: DMMF.
 
     const models = getDataModels(model);
 
+    await generateModelMeta(project, models, path.join(outDir, '__model_meta.ts'), false, true);
+
     generateIndex(project, outDir, models);
 
     models.forEach((dataModel) => {
@@ -60,24 +64,20 @@ function generateModelHooks(project: Project, outDir: string, model: DataModel,
         moduleSpecifier: prismaImport,
     });
     sf.addStatements([
-        `import { useContext } from 'react';`,
-        `import { RequestHandlerContext, type GetNextArgs, type RequestOptions, type InfiniteRequestOptions, type PickEnumerable, type CheckSelect } from '@zenstackhq/swr/runtime';`,
+        `import { RequestHandlerContext, type GetNextArgs, type RequestOptions, type InfiniteRequestOptions, type PickEnumerable, type CheckSelect, useHooksContext } from '@zenstackhq/swr/runtime';`,
+        `import metadata from './__model_meta';`,
         `import * as request from '@zenstackhq/swr/runtime';`,
     ]);
 
     const modelNameCap = upperCaseFirst(model.name);
     const prismaVersion = getPrismaVersion();
 
-    const prefixesToMutate = ['find', 'aggregate', 'count', 'groupBy'];
     const useMutation = sf.addFunction({
         name: `useMutate${model.name}`,
         isExported: true,
         statements: [
-            'const { endpoint, fetch } = useContext(RequestHandlerContext);',
-            `const prefixesToMutate = [${prefixesToMutate
-                .map((prefix) => '`${endpoint}/' + lowerCaseFirst(model.name) + '/' + prefix + '`')
-                .join(', ')}];`,
-            'const mutate = request.getMutate(prefixesToMutate);',
+            'const { endpoint, fetch, logging } = useHooksContext();',
+            `const mutate = request.useMutate('${model.name}', metadata, logging);`,
         ],
     });
     const mutationFuncs: string[] = [];
@@ -297,8 +297,6 @@ function generateQueryHook(
     typeParameters?: string[],
     infinite = false
 ) {
-    const modelRouteName = lowerCaseFirst(model.name);
-
     const typeParams = typeParameters ? [...typeParameters] : [`T extends ${argsType}`];
     if (infinite) {
         typeParams.push(`R extends ${returnType}`);
@@ -329,10 +327,10 @@ function generateQueryHook(
     })
         .addBody()
         .addStatements([
-            'const { endpoint, fetch } = useContext(RequestHandlerContext);',
+            'const { endpoint, fetch } = useHooksContext();',
             !infinite
-                ? `return request.get<${returnType}>(\`\${endpoint}/${modelRouteName}/${operation}\`, args, options, fetch);`
-                : `return request.infiniteGet<${inputType} | undefined, ${returnType}>(\`\${endpoint}/${modelRouteName}/${operation}\`, getNextArgs, options, fetch);`,
+                ? `return request.useGet<${returnType}>('${model.name}', '${operation}', endpoint, args, options, fetch);`
+                : `return request.useInfiniteGet<${inputType} | undefined, ${returnType}>('${model.name}', '${operation}', endpoint, getNextArgs, options, fetch);`,
         ]);
 }
 

packages/plugins/swr/src/runtime/index.ts

@@ -1,6 +1,9 @@
 /* eslint-disable @typescript-eslint/no-explicit-any */
 import { deserialize, serialize } from '@zenstackhq/runtime/browser';
-import { createContext } from 'react';
+import { getMutatedModels, getReadModels, type ModelMeta, type PrismaWriteActionType } from '@zenstackhq/runtime/cross';
+import * as crossFetch from 'cross-fetch';
+import { lowerCaseFirst } from 'lower-case-first';
+import { createContext, useContext } from 'react';
 import type { Fetcher, MutatorCallback, MutatorOptions, SWRConfiguration, SWRResponse } from 'swr';
 import useSWR, { useSWRConfig } from 'swr';
 import useSWRInfinite, { SWRInfiniteConfiguration, SWRInfiniteFetcher, SWRInfiniteResponse } from 'swr/infinite';
@@ -18,19 +21,26 @@ export type RequestHandlerContext = {
     /**
      * The endpoint to use for the queries.
      */
-    endpoint: string;
+    endpoint?: string;
 
     /**
      * A custom fetch function for sending the HTTP requests.
      */
     fetch?: FetchFn;
+
+    /**
+     * If logging is enabled.
+     */
+    logging?: boolean;
 };
 
+const DEFAULT_QUERY_ENDPOINT = '/api/model';
+
 /**
  * Context for configuring react hooks.
  */
 export const RequestHandlerContext = createContext<RequestHandlerContext>({
-    endpoint: '/api/model',
+    endpoint: DEFAULT_QUERY_ENDPOINT,
     fetch: undefined,
 });
 
@@ -39,6 +49,14 @@ export const RequestHandlerContext = createContext<RequestHandlerContext>({
  */
 export const Provider = RequestHandlerContext.Provider;
 
+/**
+ * Hooks context.
+ */
+export function useHooksContext() {
+    const { endpoint, ...rest } = useContext(RequestHandlerContext);
+    return { endpoint: endpoint ?? DEFAULT_QUERY_ENDPOINT, ...rest };
+}
+
 /**
  * Client request options for regular query.
  */
@@ -69,6 +87,29 @@ export type InfiniteRequestOptions<Result, Error = any> = {
     initialData?: Result[];
 } & SWRInfiniteConfiguration<Result, Error, SWRInfiniteFetcher<Result>>;
 
+export const QUERY_KEY_PREFIX = 'zenstack';
+
+type QueryKey = { prefix: typeof QUERY_KEY_PREFIX; model: string; operation: string; args: unknown };
+
+export function getQueryKey(model: string, operation: string, args?: unknown) {
+    return JSON.stringify({ prefix: QUERY_KEY_PREFIX, model, operation, args });
+}
+
+export function parseQueryKey(key: unknown) {
+    if (typeof key !== 'string') {
+        return undefined;
+    }
+    try {
+        const parsed = JSON.parse(key);
+        if (!parsed || parsed.prefix !== QUERY_KEY_PREFIX) {
+            return undefined;
+        }
+        return parsed as QueryKey;
+    } catch {
+        return undefined;
+    }
+}
+
 /**
  * Makes a GET request with SWR.
  *
@@ -79,14 +120,17 @@ export type InfiniteRequestOptions<Result, Error = any> = {
  * @returns SWR response
  */
 // eslint-disable-next-line @typescript-eslint/no-explicit-any
-export function get<Result, Error = any>(
-    url: string | null,
+export function useGet<Result, Error = any>(
+    model: string,
+    operation: string,
+    endpoint: string,
     args?: unknown,
     options?: RequestOptions<Result, Error>,
     fetch?: FetchFn
 ): SWRResponse<Result, Error> {
-    const reqUrl = options?.disabled ? null : url ? makeUrl(url, args) : null;
-    return useSWR<Result, Error>(reqUrl, (url) => fetcher<Result, false>(url, undefined, fetch, false), {
+    const key = options?.disabled ? null : getQueryKey(model, operation, args);
+    const url = makeUrl(`${endpoint}/${lowerCaseFirst(model)}/${operation}`, args);
+    return useSWR<Result, Error>(key, () => fetcher<Result, false>(url, undefined, fetch, false), {
         ...options,
         fallbackData: options?.initialData ?? options?.fallbackData,
     });
@@ -107,26 +151,40 @@ export type GetNextArgs<Args, Result> = (pageIndex: number, previousPageData: Re
  * @returns SWR infinite query response
  */
 // eslint-disable-next-line @typescript-eslint/no-explicit-any
-export function infiniteGet<Args, Result, Error = any>(
-    url: string | null,
+export function useInfiniteGet<Args, Result, Error = any>(
+    model: string,
+    operation: string,
+    endpoint: string,
     getNextArgs: GetNextArgs<Args, any>,
     options?: InfiniteRequestOptions<Result, Error>,
     fetch?: FetchFn
 ): SWRInfiniteResponse<Result, Error> {
     const getKey = (pageIndex: number, previousPageData: Result | null) => {
-        if (options?.disabled || !url) {
+        if (options?.disabled) {
             return null;
         }
         const nextArgs = getNextArgs(pageIndex, previousPageData);
         return nextArgs !== null // null means reached the end
-            ? makeUrl(url, nextArgs)
+            ? getQueryKey(model, operation, nextArgs)
             : null;
     };
 
-    return useSWRInfinite<Result, Error>(getKey, (url) => fetcher<Result, false>(url, undefined, fetch, false), {
-        ...options,
-        fallbackData: options?.initialData ?? options?.fallbackData,
-    });
+    return useSWRInfinite<Result, Error>(
+        getKey,
+        (key) => {
+            // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
+            const parsedKey = parseQueryKey(key)!;
+            const url = makeUrl(
+                `${endpoint}/${lowerCaseFirst(parsedKey.model)}/${parsedKey.operation}`,
+                parsedKey.args
+            );
+            return fetcher<Result, false>(url, undefined, fetch, false);
+        },
+        {
+            ...options,
+            fallbackData: options?.initialData ?? options?.fallbackData,
+        }
+    );
 }
 
 /**
@@ -155,7 +213,7 @@ export async function post<Result, C extends boolean = boolean>(
         fetch,
         checkReadBack
     );
-    mutate();
+    mutate(getOperationFromUrl(url), data);
     return r;
 }
 
@@ -185,7 +243,7 @@ export async function put<Result, C extends boolean = boolean>(
         fetch,
         checkReadBack
     );
-    mutate();
+    mutate(getOperationFromUrl(url), data);
     return r;
 }
 
@@ -212,29 +270,42 @@ export async function del<Result, C extends boolean = boolean>(
         fetch,
         checkReadBack
     );
-    const path = url.split('/');
-    path.pop();
-    mutate();
+    mutate(getOperationFromUrl(url), args);
     return r;
 }
 
 type Mutator = (
+    operation: string,
     data?: unknown | Promise<unknown> | MutatorCallback,
     opts?: boolean | MutatorOptions
 ) => Promise<unknown[]>;
 
-export function getMutate(prefixes: string[]): Mutator {
+export function useMutate(model: string, modelMeta: ModelMeta, logging?: boolean): Mutator {
     // https://swr.vercel.app/docs/advanced/cache#mutate-multiple-keys-from-regex
     const { cache, mutate } = useSWRConfig();
-    return (data?: unknown | Promise<unknown> | MutatorCallback, opts?: boolean | MutatorOptions) => {
+    return async (operation: string, args: unknown, opts?: boolean | MutatorOptions) => {
         if (!(cache instanceof Map)) {
             throw new Error('mutate requires the cache provider to be a Map instance');
         }
 
-        const keys = Array.from(cache.keys()).filter(
-            (k) => typeof k === 'string' && prefixes.some((prefix) => k.startsWith(prefix))
-        ) as string[];
-        const mutations = keys.map((key) => mutate(key, data, opts));
+        const mutatedModels = await getMutatedModels(model, operation as PrismaWriteActionType, args, modelMeta);
+
+        const keys = Array.from(cache.keys()).filter((key) => {
+            const parsedKey = parseQueryKey(key);
+            if (!parsedKey) {
+                return false;
+            }
+            const modelsRead = getReadModels(parsedKey.model, modelMeta, parsedKey.args);
+            return modelsRead.some((m) => mutatedModels.includes(m));
+        });
+
+        if (logging) {
+            keys.forEach((key) => {
+                console.log(`Invalidating query ${key} due to mutation "${model}.${operation}"`);
+            });
+        }
+
+        const mutations = keys.map((key) => mutate(key, undefined, opts));
         return Promise.all(mutations);
     };
 }
@@ -245,7 +316,7 @@ export async function fetcher<R, C extends boolean>(
     fetch?: FetchFn,
     checkReadBack?: C
 ): Promise<C extends true ? R | undefined : R> {
-    const _fetch = fetch ?? window.fetch;
+    const _fetch = fetch ?? crossFetch.fetch;
     const res = await _fetch(url, options);
     if (!res.ok) {
         const errData = unmarshal(await res.text());
@@ -306,3 +377,13 @@ function makeUrl(url: string, args: unknown) {
     }
     return result;
 }
+
+function getOperationFromUrl(url: string) {
+    const parts = url.split('/');
+    const r = parts.pop();
+    if (!r) {
+        throw new Error(`Invalid URL: ${url}`);
+    } else {
+        return r;
+    }
+}