feat: add check hooks generation to tanstack and swr plugins

Author: ymc9

packages/plugins/swr/package.json

@@ -46,6 +46,7 @@
         "lower-case-first": "^2.0.2",
         "semver": "^7.5.2",
         "ts-morph": "^16.0.0",
+        "ts-pattern": "^4.3.0",
         "upper-case-first": "^2.0.2"
     },
     "devDependencies": {

packages/plugins/swr/src/generator.ts

@@ -1,5 +1,6 @@
 import {
     PluginOptions,
+    RUNTIME_PACKAGE,
     createProject,
     ensureEmptyDir,
     generateModelMeta,
@@ -8,11 +9,12 @@ import {
     resolvePath,
     saveProject,
 } from '@zenstackhq/sdk';
-import { DataModel, Model } from '@zenstackhq/sdk/ast';
+import { DataModel, DataModelFieldType, Model, isEnum } from '@zenstackhq/sdk/ast';
 import { getPrismaClientImportSpec, type DMMF } from '@zenstackhq/sdk/prisma';
 import { paramCase } from 'change-case';
 import path from 'path';
 import type { OptionalKind, ParameterDeclarationStructure, Project, SourceFile } from 'ts-morph';
+import { P, match } from 'ts-pattern';
 import { upperCaseFirst } from 'upper-case-first';
 import { name } from '.';
 
@@ -66,6 +68,7 @@ function generateModelHooks(
     });
     sf.addStatements([
         `import { type GetNextArgs, type QueryOptions, type InfiniteQueryOptions, type MutationOptions, type PickEnumerable } from '@zenstackhq/swr/runtime';`,
+        `import type { PolicyCrudKind } from '${RUNTIME_PACKAGE}'`,
         `import metadata from './__model_meta';`,
         `import * as request from '@zenstackhq/swr/runtime';`,
     ]);
@@ -239,6 +242,11 @@ function generateModelHooks(
         const returnType = `T extends { select: any; } ? T['select'] extends true ? number : Prisma.GetScalarType<T['select'], Prisma.${modelNameCap}CountAggregateOutputType> : number`;
         generateQueryHook(sf, model, 'count', argsType, inputType, returnType);
     }
+
+    // extra `check` hook for ZenStack's permission checker API
+    {
+        generateCheckHook(sf, model, prismaImport);
+    }
 }
 
 function makeOptimistic(returnType: string) {
@@ -337,3 +345,47 @@ function generateMutation(
 
     return funcName;
 }
+
+function generateCheckHook(sf: SourceFile, model: DataModel, prismaImport: string) {
+    const mapFilterType = (type: DataModelFieldType) => {
+        return match(type.type)
+            .with(P.union('Int', 'BigInt'), () => 'number')
+            .with('String', () => 'string')
+            .with('Boolean', () => 'boolean')
+            .otherwise(() => undefined);
+    };
+
+    const filterFields: Array<{ name: string; type: string }> = [];
+    const enumsToImport = new Set<string>();
+
+    // collect filterable fields and enums to import
+    model.fields.forEach((f) => {
+        if (isEnum(f.type.reference?.ref)) {
+            enumsToImport.add(f.type.reference.$refText);
+            filterFields.push({ name: f.name, type: f.type.reference.$refText });
+        }
+
+        const mappedType = mapFilterType(f.type);
+        if (mappedType) {
+            filterFields.push({ name: f.name, type: mappedType });
+        }
+    });
+
+    if (enumsToImport.size > 0) {
+        // import enums
+        sf.addStatements(`import type { ${Array.from(enumsToImport).join(', ')} } from '${prismaImport}';`);
+    }
+
+    const whereType = `{ ${filterFields.map(({ name, type }) => `${name}?: ${type}`).join('; ')} }`;
+
+    const func = sf.addFunction({
+        name: `useCheck${model.name}`,
+        isExported: true,
+        parameters: [
+            { name: 'args', type: `{ operation: PolicyCrudKind; where?: ${whereType}; }` },
+            { name: 'options?', type: `QueryOptions<boolean>` },
+        ],
+    });
+
+    func.addStatements(`return request.useModelQuery('${model.name}', 'check', args, options);`);
+}

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

@@ -1,6 +1,7 @@
 import {
     PluginError,
     PluginOptions,
+    RUNTIME_PACKAGE,
     createProject,
     ensureEmptyDir,
     generateModelMeta,
@@ -9,13 +10,13 @@ import {
     resolvePath,
     saveProject,
 } from '@zenstackhq/sdk';
-import { DataModel, Model } from '@zenstackhq/sdk/ast';
+import { DataModel, DataModelFieldType, Model, isEnum } from '@zenstackhq/sdk/ast';
 import { getPrismaClientImportSpec, type DMMF } from '@zenstackhq/sdk/prisma';
 import { paramCase } from 'change-case';
 import { lowerCaseFirst } from 'lower-case-first';
 import path from 'path';
 import { Project, SourceFile, VariableDeclarationKind } from 'ts-morph';
-import { match } from 'ts-pattern';
+import { P, match } from 'ts-pattern';
 import { upperCaseFirst } from 'upper-case-first';
 import { name } from '.';
 
@@ -261,6 +262,62 @@ function generateMutationHook(
     func.addStatements('return mutation;');
 }
 
+function generateCheckHook(
+    target: string,
+    version: TanStackVersion,
+    sf: SourceFile,
+    model: DataModel,
+    prismaImport: string
+) {
+    const mapFilterType = (type: DataModelFieldType) => {
+        return match(type.type)
+            .with(P.union('Int', 'BigInt'), () => 'number')
+            .with('String', () => 'string')
+            .with('Boolean', () => 'boolean')
+            .otherwise(() => undefined);
+    };
+
+    const filterFields: Array<{ name: string; type: string }> = [];
+    const enumsToImport = new Set<string>();
+
+    // collect filterable fields and enums to import
+    model.fields.forEach((f) => {
+        if (isEnum(f.type.reference?.ref)) {
+            enumsToImport.add(f.type.reference.$refText);
+            filterFields.push({ name: f.name, type: f.type.reference.$refText });
+        }
+
+        const mappedType = mapFilterType(f.type);
+        if (mappedType) {
+            filterFields.push({ name: f.name, type: mappedType });
+        }
+    });
+
+    if (enumsToImport.size > 0) {
+        // import enums
+        sf.addStatements(`import type { ${Array.from(enumsToImport).join(', ')} } from '${prismaImport}';`);
+    }
+
+    const whereType = `{ ${filterFields.map(({ name, type }) => `${name}?: ${type}`).join('; ')} }`;
+
+    const func = sf.addFunction({
+        name: `useCheck${model.name}`,
+        isExported: true,
+        typeParameters: ['TError = DefaultError'],
+        parameters: [
+            { name: 'args', type: `{ operation: PolicyCrudKind; where?: ${whereType}; }` },
+            { name: 'options?', type: makeQueryOptions(target, 'boolean', 'boolean', false, false, version) },
+        ],
+    });
+
+    func.addStatements([
+        makeGetContext(target),
+        `return useModelQuery<boolean, boolean, TError>('${model.name}', \`\${endpoint}/${lowerCaseFirst(
+            model.name
+        )}/check\`, args, options, fetch);`,
+    ]);
+}
+
 function generateModelHooks(
     target: TargetFramework,
     version: TanStackVersion,
@@ -494,6 +551,11 @@ function generateModelHooks(
             `TArgs extends { select: any; } ? TArgs['select'] extends true ? number : Prisma.GetScalarType<TArgs['select'], Prisma.${modelNameCap}CountAggregateOutputType> : number`
         );
     }
+
+    {
+        // extra `check` hook for ZenStack's permission checker API
+        generateCheckHook(target, version, sf, model, prismaImport);
+    }
 }
 
 function generateIndex(
@@ -538,6 +600,7 @@ function makeBaseImports(target: TargetFramework, version: TanStackVersion) {
     const shared = [
         `import { useModelQuery, useInfiniteModelQuery, useModelMutation } from '${runtimeImportBase}/${target}';`,
         `import type { PickEnumerable, CheckSelect, QueryError, ExtraQueryOptions, ExtraMutationOptions } from '${runtimeImportBase}';`,
+        `import type { PolicyCrudKind } from '${RUNTIME_PACKAGE}'`,
         `import metadata from './__model_meta';`,
         `type DefaultError = QueryError;`,
     ];