fix: preserve prisma client extensions's typing

zenstackhq · Mar 17, 2024 · 61438cc · 61438cc
1 parent 01e372e
commit 61438cc
Show file tree

Hide file tree

Showing 2 changed files with 98 additions and 43 deletions.
diff --git a/packages/schema/src/plugins/enhancer/enhance/index.ts b/packages/schema/src/plugins/enhancer/enhance/index.ts
@@ -92,31 +92,39 @@ export class EnhancerGenerator {
             `import { createEnhancement, type EnhancementContext, type EnhancementOptions, type ZodSchemas, type AuthUser } from '@zenstackhq/runtime';
     import modelMeta from './model-meta';
     import policy from './policy';
-    import { Prisma } from '${prismaImport}';
+    import { Prisma as _Prisma, PrismaClient as _PrismaClient } from '${prismaImport}';
+    import type { InternalArgs, TypeMapDef, TypeMapCbDef, DynamicClientExtensionThis } from '${prismaImport}/runtime/library';
     ${
         withLogicalClient
             ? `import type * as _P from '${logicalPrismaClientDir}/index-fixed';
-    import type { PrismaClient } from '${logicalPrismaClientDir}/index-fixed';
+    import type { Prisma, PrismaClient } from '${logicalPrismaClientDir}/index-fixed';
     `
             : `import type * as _P from '${prismaImport}';
-    import type { PrismaClient } from '${prismaImport}';
+    import type { Prisma, PrismaClient } from '${prismaImport}';
     `
     }
     ${this.options.withZodSchemas ? "import * as zodSchemas from './zod';" : 'const zodSchemas = undefined;'}
     
     ${authTypes}
+// overload for plain PrismaClient
+export function enhance<ExtArgs extends Record<string, any> & InternalArgs>(
+    prisma: _PrismaClient<any, any, ExtArgs>,
+    context?: EnhancementContext<${authTypeParam}>, options?: EnhancementOptions): PrismaClient;
     
-    export function enhance<DbClient extends object>(prisma: DbClient, context?: EnhancementContext<${authTypeParam}>, options?: EnhancementOptions)${
-                withLogicalClient ? ': PrismaClient' : ''
-            } {
-        return createEnhancement(prisma, {
-            modelMeta,
-            policy,
-            zodSchemas: zodSchemas as unknown as (ZodSchemas | undefined),
-            prismaModule: Prisma,
-            ...options
-        }, context)${withLogicalClient ? ' as PrismaClient' : ''};
-    }
+// 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>;
+
+export function enhance(prisma: any, context?: EnhancementContext<${authTypeParam}>, options?: EnhancementOptions): any {
+    return createEnhancement(prisma, {
+        modelMeta,
+        policy,
+        zodSchemas: zodSchemas as unknown as (ZodSchemas | undefined),
+        prismaModule: _Prisma,
+        ...options
+    }, context);
+}
     `,
             { overwrite: true }
         );

diff --git a/tests/integration/tests/enhancements/with-delegate/enhanced-client.test.ts b/tests/integration/tests/enhancements/with-delegate/enhanced-client.test.ts
@@ -1,7 +1,5 @@
 import { PrismaErrorCode } from '@zenstackhq/runtime';
-import { loadSchema, run } from '@zenstackhq/testtools';
-import fs from 'fs';
-import path from 'path';
+import { loadSchema } from '@zenstackhq/testtools';
 import { POLYMORPHIC_MANY_TO_MANY_SCHEMA, POLYMORPHIC_SCHEMA } from './utils';
 
 describe('Polymorphism Test', () => {
@@ -1039,16 +1037,14 @@ describe('Polymorphism Test', () => {
         );
     });
 
-    it('typescript compilation', async () => {
-        const { projectDir } = await loadSchema(schema, { enhancements: ['delegate'] });
+    it('typescript compilation plain prisma', async () => {
         const src = `
         import { PrismaClient } from '@prisma/client';
         import { enhance } from '.zenstack/enhance';
         
         const prisma = new PrismaClient();
         
         async function main() {
-            await prisma.user.deleteMany();
             const db = enhance(prisma);
         
             const user1 = await db.user.create({ data: { } });
@@ -1084,31 +1080,82 @@ describe('Polymorphism Test', () => {
             }
         }
         
-        main()
-            .then(async () => {
-                await prisma.$disconnect();
-            })
-            .catch(async (e) => {
-                console.error(e);
-                await prisma.$disconnect();
-                process.exit(1);
-        });        
+        main();     
         `;
-
-        fs.writeFileSync(path.join(projectDir, 'script.ts'), src);
-        fs.writeFileSync(
-            path.join(projectDir, 'tsconfig.json'),
-            JSON.stringify({
-                compilerOptions: {
-                    outDir: 'dist',
-                    strict: true,
-                    lib: ['esnext'],
-                    esModuleInterop: true,
+        await loadSchema(schema, {
+            compile: true,
+            enhancements: ['delegate'],
+            extraSourceFiles: [
+                {
+                    name: 'main.ts',
+                    content: src,
                 },
-            })
-        );
+            ],
+        });
+    });
 
-        run('npm i -D @types/node', undefined, projectDir);
-        run('npx tsc --noEmit --skipLibCheck script.ts', undefined, projectDir);
+    it('typescript compilation extended prisma', async () => {
+        const src = `
+        import { PrismaClient } from '@prisma/client';
+        import { enhance } from '.zenstack/enhance';
+        
+        const prisma = new PrismaClient().$extends({
+            model: {
+                user: {
+                    async signUp() {
+                        return prisma.user.create({ data: {} });
+                    },
+                },
+            },
+        });
+        
+        async function main() {
+            const db = enhance(prisma);
+        
+            const user1 = await db.user.signUp();
+        
+            await db.ratedVideo.create({
+                data: {
+                    owner: { connect: { id: user1.id } },
+                    duration: 100,
+                    url: 'abc',
+                    rating: 10,
+                },
+            });
+        
+            await db.image.create({
+                data: {
+                    owner: { connect: { id: user1.id } },
+                    format: 'webp',
+                },
+            });
+        
+            const video = await db.video.findFirst({ include: { owner: true } });
+            console.log(video?.duration);
+            console.log(video?.viewCount);
+        
+            const asset = await db.asset.findFirstOrThrow();
+            console.log(asset.assetType);
+            console.log(asset.viewCount);
+        
+            if (asset.assetType === 'Video') {
+                console.log('Video: duration', asset.duration);
+            } else {
+                console.log('Image: format', asset.format);
+            }
+        }
+        
+        main();     
+        `;
+        await loadSchema(schema, {
+            compile: true,
+            enhancements: ['delegate'],
+            extraSourceFiles: [
+                {
+                    name: 'main.ts',
+                    content: src,
+                },
+            ],
+        });
     });
 });