feat: export Enhanced type to infer typeof enhanced PrismaClient (#2010)

Author: ymc9

packages/runtime/res/enhance.d.ts

@@ -1 +1 @@
-export { auth, enhance, type PrismaClient } from '.zenstack/enhance';
+export { auth, enhance, type PrismaClient, type Enhanced } from '.zenstack/enhance';

packages/runtime/src/enhance.d.ts

@@ -1,2 +1,2 @@
 // @ts-expect-error stub for re-exporting generated code
-export { auth, enhance } from '.zenstack/enhance';
+export { auth, enhance, type PrismaClient, type Enhanced } from '.zenstack/enhance';

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

@@ -266,6 +266,19 @@ export function enhance(prisma: any, context?: EnhancementContext<${authTypePara
         ...options
     }, context);
 }
+
+/**
+ * Infers the type of PrismaClient with ZenStack's enhancements.
+ * @example
+ * type EnhancedPrismaClient = Enhanced<typeof prisma>;
+ */
+export type Enhanced<Client> = 
+    Client extends _PrismaClient<infer _ClientOptions, infer _U, infer ExtArgs> ? PrismaClient :
+    Client extends DynamicClientExtensionThis<infer _TypeMap, infer _TypeMapCb, infer ExtArgs${
+        hasClientOptions ? ', infer ClientOptions' : ''
+    }> ? DynamicClientExtensionThis<Prisma.TypeMap<ExtArgs>, Prisma.TypeMapCb, ExtArgs${
+            hasClientOptions ? ', ClientOptions' : ''
+        }> : Client;
 `;
     }
 

tests/integration/tests/enhancements/typing/enhancement-typing.test.ts

@@ -0,0 +1,60 @@
+import { loadSchema } from '@zenstackhq/testtools';
+
+describe('Enhancement typing tests', () => {
+    it('infers correct typing', async () => {
+        await loadSchema(
+            `
+            model User {
+                id Int @id @default(autoincrement())
+                posts Post[]
+            }
+            
+            model Post {
+                id Int @id @default(autoincrement())
+                title String
+                author User @relation(fields: [authorId], references: [id])
+                authorId Int @default(auth().id)
+            }
+            `,
+            {
+                pushDb: false,
+                compile: true,
+                extraSourceFiles: [
+                    {
+                        name: 'main.ts',
+                        content: `
+import { PrismaClient } from '@prisma/client';
+import type { Enhanced } from '.zenstack/enhance';
+
+async function withoutClientExtension() {
+    const prisma = new PrismaClient();
+    const db = {} as any as Enhanced<typeof prisma>;
+    // note that "author" becomes optional
+    const r = await db.post.create({ data: { title: 'Post1' }});
+    console.log(r);
+}
+
+async function withClientExtension() {
+    const prisma = (new PrismaClient())
+        .$extends({
+            client: {
+                $log: (message: string) => {
+                    console.log(message);
+                },
+            },
+        });
+    const db = {} as any as Enhanced<typeof prisma>;
+    // note that "author" becomes optional
+    const r = await db.post.create({ data: { title: 'Post1' }});
+    console.log(r);
+
+    // note that "$log" is preserved
+    db.$log('hello');
+}
+                        `,
+                    },
+                ],
+            }
+        );
+    });
+});