WIP: progress

Author: ymc9

packages/schema/src/plugins/enhancer/enhance/checker-type-generator.ts

@@ -0,0 +1,44 @@
+import { getDataModels } from '@zenstackhq/sdk';
+import type { DataModel, DataModelField, Model } from '@zenstackhq/sdk/ast';
+import { lowerCaseFirst } from 'lower-case-first';
+import { P, match } from 'ts-pattern';
+
+/**
+ * Generates a `ModelCheckers` interface that contains a `check` method for each model in the schema.
+ */
+export function generateCheckerType(model: Model) {
+    return `
+type CheckerOperation = 'create' | 'read' | 'update' | 'delete';
+
+export interface ModelCheckers {
+    ${getDataModels(model)
+        .map((dataModel) => `\t${lowerCaseFirst(dataModel.name)}: ${generateDataModelChecker(dataModel)}`)
+        .join(',\n')}
+}
+`;
+}
+
+function generateDataModelChecker(dataModel: DataModel) {
+    return `{
+        check(op: CheckerOperation, args?: ${generateDataModelArgs(dataModel)}): Promise<boolean>
+    }`;
+}
+
+function generateDataModelArgs(dataModel: DataModel) {
+    return `{ ${dataModel.fields
+        .filter((field) => isFieldFilterable(field))
+        .map((field) => `${field.name}?: ${mapFieldType(field)}`)
+        .join('; ')} }`;
+}
+
+function isFieldFilterable(field: DataModelField) {
+    return !!mapFieldType(field);
+}
+
+function mapFieldType(field: DataModelField) {
+    return match(field.type.type)
+        .with('Boolean', () => 'boolean')
+        .with(P.union('BigInt', 'Int', 'Float', 'Decimal'), () => 'number')
+        .with('String', () => 'string')
+        .otherwise(() => undefined);
+}

packages/schema/src/plugins/enhancer/enhance/index.ts

@@ -41,6 +41,7 @@ import { trackPrismaSchemaError } from '../../prisma';
 import { PrismaSchemaGenerator } from '../../prisma/schema-generator';
 import { isDefaultWithAuth } from '../enhancer-utils';
 import { generateAuthType } from './auth-type-generator';
+import { generateCheckerType } from './checker-type-generator';
 
 // information of delegate models and their sub models
 type DelegateInfo = [DataModel, DataModel[]][];
@@ -89,6 +90,8 @@ export class EnhancerGenerator {
         const authTypes = authModel ? generateAuthType(this.model, authModel) : '';
         const authTypeParam = authModel ? `auth.${authModel.name}` : 'AuthUser';
 
+        const checkerTypes = generateCheckerType(this.model);
+
         const enhanceTs = this.project.createSourceFile(
             path.join(this.outDir, 'enhance.ts'),
             `import { type EnhancementContext, type EnhancementOptions, type ZodSchemas, type AuthUser } from '@zenstackhq/runtime';
@@ -105,6 +108,8 @@ ${
 
 ${authTypes}
 
+${checkerTypes}
+
 ${
     logicalPrismaClientDir
         ? this.createLogicalPrismaEnhanceFunction(authTypeParam)
@@ -127,14 +132,14 @@ import type * as _P from '${prismaImport}';
 
     private createSimplePrismaEnhanceFunction(authTypeParam: string) {
         return `
-export function enhance<DbClient extends object>(prisma: DbClient, context?: EnhancementContext<${authTypeParam}>, options?: EnhancementOptions) {
+export function enhance<DbClient extends object>(prisma: DbClient, context?: EnhancementContext<${authTypeParam}>, options?: EnhancementOptions): DbClient & ModelCheckers {
     return createEnhancement(prisma, {
         modelMeta,
         policy,
         zodSchemas: zodSchemas as unknown as (ZodSchemas | undefined),
         prismaModule: Prisma,
         ...options
-    }, context);
+    }, context) as DbClient & ModelCheckers;
 }         
             `;
     }
@@ -157,12 +162,12 @@ import type { Prisma, PrismaClient } from '${logicalPrismaClientDir}/index-fixed
 // overload for plain PrismaClient
 export function enhance<ExtArgs extends Record<string, any> & InternalArgs>(
     prisma: _PrismaClient<any, any, ExtArgs>,
-    context?: EnhancementContext<${authTypeParam}>, options?: EnhancementOptions): PrismaClient;
+    context?: EnhancementContext<${authTypeParam}>, options?: EnhancementOptions): PrismaClient & ModelCheckers;
     
 // overload for extended PrismaClient
 export function enhance<TypeMap extends TypeMapDef, TypeMapCb extends TypeMapCbDef, ExtArgs extends Record<string, any> & InternalArgs>(
     prisma: DynamicClientExtensionThis<TypeMap, TypeMapCb, ExtArgs>,
-    context?: EnhancementContext<${authTypeParam}>, options?: EnhancementOptions): DynamicClientExtensionThis<Prisma.TypeMap, Prisma.TypeMapCb, ExtArgs>;
+    context?: EnhancementContext<${authTypeParam}>, options?: EnhancementOptions): DynamicClientExtensionThis<Prisma.TypeMap, Prisma.TypeMapCb, ExtArgs> & ModelCheckers;
 
 export function enhance(prisma: any, context?: EnhancementContext<${authTypeParam}>, options?: EnhancementOptions): any {
     return createEnhancement(prisma, {

tests/integration/tests/enhancements/with-policy/checker.test.ts

@@ -326,4 +326,33 @@ describe('Permission checker', () => {
         await expect(db.model.check('create')).toResolveTruthy();
         await expect(db.model.check('create', { value: 1 })).toResolveTruthy();
     });
+
+    it('compilation', async () => {
+        await loadSchema(
+            `
+            model Model {
+                id Int @id @default(autoincrement())
+                value Int
+                @@allow('read', value == 1)
+            }
+            `,
+            {
+                compile: true,
+                extraSourceFiles: [
+                    {
+                        name: 'main.ts',
+                        content: `
+                        import { PrismaClient } from '@prisma/client';
+                        import { enhance } from '.zenstack/enhance';
+
+                        const prisma = new PrismaClient();
+                        const db = enhance(prisma);
+                        db.model.check('read');
+                        db.model.check('read', { value: 1 });
+                        `,
+                    },
+                ],
+            }
+        );
+    });
 });