fix: regression with client extension typing (#2182)

Author: ymc9

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

@@ -243,8 +243,14 @@ export function enhance<DbClient extends object>(prisma: DbClient, context?: Enh
 
     private createLogicalPrismaImports(prismaImport: string, prismaClientImport: string, target: string) {
         const prismaTargetImport = target === 'edge' ? `${prismaImport}/edge` : prismaImport;
+        const runtimeLibraryImport = this.isNewPrismaClientGenerator
+            ? // new generator has these typed only in "@prisma/client"
+              '@prisma/client/runtime/library'
+            : // old generator has these types generated with the client
+              `${prismaImport}/runtime/library`;
+
         return `import { Prisma as _Prisma, PrismaClient as _PrismaClient } from '${prismaTargetImport}';
-import type { InternalArgs, DynamicClientExtensionThis } from '@prisma/client/runtime/library';
+import type { InternalArgs, DynamicClientExtensionThis } from '${runtimeLibraryImport}';
 import type * as _P from '${prismaClientImport}';
 import type { Prisma, PrismaClient } from '${prismaClientImport}';
 export type { PrismaClient };

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

@@ -0,0 +1,121 @@
+import { loadSchema } from '@zenstackhq/testtools';
+
+describe('issue 2175', () => {
+    it('regression standard generator', async () => {
+        await loadSchema(
+            `
+        model User {
+            id    Int     @id @default(autoincrement())
+            email String  @unique
+            posts Post[]
+        }   
+
+        model Post {
+            id        Int     @id @default(autoincrement())
+            title     String
+            author    User?   @relation(fields: [authorId], references: [id])
+            authorId  Int?    @default(auth().id)
+        }
+        
+        `,
+            {
+                compile: true,
+                extraSourceFiles: [
+                    {
+                        name: 'main.ts',
+                        content: `
+import { PrismaClient } from "@prisma/client";
+import { enhance } from ".zenstack/enhance";
+
+const prisma = new PrismaClient();
+const prismaExtended = prisma.$extends({
+  model: {
+    user: {
+      async signUp(email: string) {
+        return prisma.user.create({ data: { email } });
+      },
+    },
+  },
+});
+
+const dbExtended = enhance(prismaExtended);
+
+async function main() {
+  const newUser = await dbExtended.user.signUp("[email protected]");
+  console.log(newUser);
+}
+
+main();
+`,
+                    },
+                ],
+            }
+        );
+    });
+
+    it('regression new generator', async () => {
+        await loadSchema(
+            `
+        datasource db {
+            provider = "sqlite"
+            url = "file:./test.db"
+        }
+
+        generator client {
+            provider = "prisma-client"
+            output = "../generated/prisma"
+            moduleFormat = "cjs"
+        }
+
+        model User {
+            id    Int     @id @default(autoincrement())
+            email String  @unique
+            posts Post[]
+        }   
+
+        model Post {
+            id        Int     @id @default(autoincrement())
+            title     String
+            author    User?   @relation(fields: [authorId], references: [id])
+            authorId  Int?    @default(auth().id)
+        }
+        
+        `,
+            {
+                addPrelude: false,
+                compile: true,
+                extraSourceFiles: [
+                    {
+                        name: 'main.ts',
+                        content: `
+import { PrismaClient } from "./generated/prisma/client";
+import { enhance } from "./generated/zenstack/enhance";
+
+const prisma = new PrismaClient();
+const prismaExtended = prisma.$extends({
+  model: {
+    user: {
+      async signUp(email: string) {
+        return prisma.user.create({ data: { email } });
+      },
+    },
+  },
+});
+
+const dbExtended = enhance(prismaExtended);
+
+async function main() {
+  const newUser = await dbExtended.user.signUp("[email protected]");
+  console.log(newUser);
+}
+
+main();
+`,
+                    },
+                ],
+                output: './generated/zenstack',
+                prismaLoadPath: './generated/prisma/client',
+            }
+        );
+    });
+});