fix: type with @@auth cannot be properly compiled with new "prisma-client" generator

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

@@ -1,10 +1,11 @@
-import { getIdFields, isAuthInvocation, isDataModelFieldReference } from '@zenstackhq/sdk';
+import { getIdFields, getPrismaClientGenerator, isAuthInvocation, isDataModelFieldReference } from '@zenstackhq/sdk';
 import {
     DataModel,
     DataModelField,
     Expression,
     isDataModel,
     isMemberAccessExpr,
+    isTypeDef,
     TypeDef,
     type Model,
 } from '@zenstackhq/sdk/ast';
@@ -19,27 +20,39 @@ export function generateAuthType(model: Model, authDecl: DataModel | TypeDef) {
     const types = new Map<
         string,
         {
+            isTypeDef: boolean;
             // relation fields to require
             requiredRelations: { name: string; type: string }[];
         }
     >();
 
-    types.set(authDecl.name, { requiredRelations: [] });
+    types.set(authDecl.name, { isTypeDef: isTypeDef(authDecl), requiredRelations: [] });
 
-    const ensureType = (model: string) => {
-        if (!types.has(model)) {
-            types.set(model, { requiredRelations: [] });
+    const findType = (name: string) =>
+        model.declarations.find((d) => (isDataModel(d) || isTypeDef(d)) && d.name === name);
+
+    const ensureType = (name: string) => {
+        if (!types.has(name)) {
+            const decl = findType(name);
+            if (!decl) {
+                return;
+            }
+            types.set(name, { isTypeDef: isTypeDef(decl), requiredRelations: [] });
         }
     };
 
-    const addAddField = (model: string, name: string, type: string, array: boolean) => {
-        let fields = types.get(model);
-        if (!fields) {
-            fields = { requiredRelations: [] };
-            types.set(model, fields);
+    const addTypeField = (typeName: string, fieldName: string, fieldType: string, array: boolean) => {
+        let typeInfo = types.get(typeName);
+        if (!typeInfo) {
+            const decl = findType(typeName);
+            if (!decl) {
+                return;
+            }
+            typeInfo = { isTypeDef: isTypeDef(decl), requiredRelations: [] };
+            types.set(typeName, typeInfo);
         }
-        if (!fields.requiredRelations.find((f) => f.name === name)) {
-            fields.requiredRelations.push({ name, type: array ? `${type}[]` : type });
+        if (!typeInfo.requiredRelations.find((f) => f.name === fieldName)) {
+            typeInfo.requiredRelations.push({ name: fieldName, type: array ? `${fieldType}[]` : fieldType });
         }
     };
 
@@ -57,7 +70,7 @@ export function generateAuthType(model: Model, authDecl: DataModel | TypeDef) {
                         // member is a relation
                         const fieldType = memberDecl.type.reference.ref.name;
                         ensureType(fieldType);
-                        addAddField(exprType.name, memberDecl.name, fieldType, memberDecl.type.array);
+                        addTypeField(exprType.name, memberDecl.name, fieldType, memberDecl.type.array);
                     }
                 }
             }
@@ -69,12 +82,15 @@ export function generateAuthType(model: Model, authDecl: DataModel | TypeDef) {
                 if (isDataModel(fieldType)) {
                     // field is a relation
                     ensureType(fieldType.name);
-                    addAddField(fieldDecl.$container.name, node.target.$refText, fieldType.name, fieldDecl.type.array);
+                    addTypeField(fieldDecl.$container.name, node.target.$refText, fieldType.name, fieldDecl.type.array);
                 }
             }
         });
     });
 
+    const prismaGenerator = getPrismaClientGenerator(model);
+    const isNewGenerator = !!prismaGenerator?.isNewGenerator;
+
     // generate:
     // `
     // namespace auth {
@@ -86,25 +102,27 @@ export function generateAuthType(model: Model, authDecl: DataModel | TypeDef) {
     return `export namespace auth {
     type WithRequired<T, K extends keyof T> = T & { [P in K]-?: T[P] };
 ${Array.from(types.entries())
-    .map(([model, fields]) => {
-        let result = `Partial<_P.${model}>`;
+    .map(([type, typeInfo]) => {
+        // TypeDef types are generated in "json-types.ts" for the new "prisma-client" generator
+        const typeRef = isNewGenerator && typeInfo.isTypeDef ? `$TypeDefs.${type}` : `_P.${type}`;
+        let result = `Partial<${typeRef}>`;
 
-        if (model === authDecl.name) {
+        if (type === authDecl.name) {
             // auth model's id fields are always required
             const idFields = getIdFields(authDecl).map((f) => f.name);
             if (idFields.length > 0) {
                 result = `WithRequired<${result}, ${idFields.map((f) => `'${f}'`).join('|')}>`;
             }
         }
 
-        if (fields.requiredRelations.length > 0) {
+        if (typeInfo.requiredRelations.length > 0) {
             // merge required relation fields
-            result = `${result} & { ${fields.requiredRelations.map((f) => `${f.name}: ${f.type}`).join('; ')} }`;
+            result = `${result} & { ${typeInfo.requiredRelations.map((f) => `${f.name}: ${f.type}`).join('; ')} }`;
         }
 
         result = `${result} & Record<string, unknown>`;
 
-        return `    export type ${model} = ${result};`;
+        return `    export type ${type} = ${result};`;
     })
     .join('\n')}
 }`;

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

@@ -322,9 +322,13 @@ export function enhance<DbClient extends object>(prisma: DbClient, context?: Enh
             : // old generator has these types generated with the client
               `${prismaImport}/runtime/library`;
 
+        const hasTypeDef = this.model.declarations.some(isTypeDef);
+
         return `import { Prisma as _Prisma, PrismaClient as _PrismaClient } from '${prismaTargetImport}';
 import type { InternalArgs, DynamicClientExtensionThis } from '${runtimeLibraryImport}';
-import type * as _P from '${prismaClientImport}';
+import type * as _P from '${prismaClientImport}';${
+            hasTypeDef && this.isNewPrismaClientGenerator ? `\nimport type * as $TypeDefs from './json-types';` : ''
+        }
 import type { Prisma, PrismaClient } from '${prismaClientImport}';
 export type { PrismaClient };
 `;

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

@@ -865,7 +865,7 @@ describe('auth() compile-time test', () => {
         );
     });
 
-    it('"User" type as auth', async () => {
+    it('"User" type as auth legacy generator', async () => {
         const { enhance } = await loadSchema(
             `
         type Profile {
@@ -920,4 +920,75 @@ describe('auth() compile-time test', () => {
             })
         ).toResolveTruthy();
     });
+
+    it('"User" type as auth new generator', async () => {
+        const { enhance } = await loadSchema(
+            `
+        datasource db {
+            provider = "sqlite"
+            url = "file:./dev.db"
+        }
+
+        generator js {
+            provider = "prisma-client"
+            output = "./generated/client"
+            moduleFormat = "cjs"
+        }
+
+        type Profile {
+            age Int
+        }
+
+        type Role {
+            name String
+            permissions String[]
+        }
+
+        type User {
+            myId Int @id
+            banned Boolean
+            profile Profile
+            roles Role[]
+        }
+
+        model Foo {
+            id Int @id @default(autoincrement())
+            @@allow('read', true)
+            @@allow('create', auth().myId == 1 && !auth().banned)
+            @@allow('delete', auth().roles?['DELETE' in permissions])
+            @@deny('all', auth().profile.age < 18)
+        }
+        `,
+            {
+                addPrelude: false,
+                output: './zenstack',
+                preserveTsFiles: true,
+                prismaLoadPath: './prisma/generated/client/client',
+                compile: true,
+                extraSourceFiles: [
+                    {
+                        name: 'main.ts',
+                        content: `
+                import { enhance } from "./zenstack/enhance";
+                import { PrismaClient } from '@prisma/client';
+                enhance(new PrismaClient(), { user: { myId: 1, profile: { age: 20 } } });
+                `,
+                    },
+                ],
+            }
+        );
+
+        await expect(enhance().foo.create({ data: {} })).toBeRejectedByPolicy();
+        await expect(enhance({ myId: 1, banned: true }).foo.create({ data: {} })).toBeRejectedByPolicy();
+        await expect(enhance({ myId: 1, profile: { age: 16 } }).foo.create({ data: {} })).toBeRejectedByPolicy();
+        const r = await enhance({ myId: 1, profile: { age: 20 } }).foo.create({ data: {} });
+        await expect(
+            enhance({ myId: 1, profile: { age: 20 } }).foo.delete({ where: { id: r.id } })
+        ).toBeRejectedByPolicy();
+        await expect(
+            enhance({ myId: 1, profile: { age: 20 }, roles: [{ name: 'ADMIN', permissions: ['DELETE'] }] }).foo.delete({
+                where: { id: r.id },
+            })
+        ).toResolveTruthy();
+    });
 });

tests/regression/tests/issue-2294.test.ts

@@ -0,0 +1,48 @@
+import { loadSchema } from '@zenstackhq/testtools';
+
+describe('Issue 2294', () => {
+    it('should work', async () => {
+        await loadSchema(
+            `
+datasource db {
+    provider = "sqlite"
+    url = "file:./dev.db"
+}
+
+generator js {
+    provider = "prisma-client"
+    output = "./generated/client"
+    moduleFormat = "cjs"
+}
+
+type AuthUser {
+    id     Int @id
+    name   String?
+
+    @@auth
+}
+
+model Foo {
+    id Int @id
+    @@allow('all', auth().name == 'admin')
+}
+`,
+            {
+                addPrelude: false,
+                output: './zenstack',
+                prismaLoadPath: './prisma/generated/client/client',
+                compile: true,
+                extraSourceFiles: [
+                    {
+                        name: 'main.ts',
+                        content: `
+import { enhance } from "./zenstack/enhance";
+import { PrismaClient } from './prisma/generated/client/client';
+enhance(new PrismaClient(), { user: { id: 1, name: 'admin' } });
+                `,
+                    },
+                ],
+            }
+        );
+    });
+});