From eb6fbca539cbf4ef4d81d367dd2b9e035cac98af Mon Sep 17 00:00:00 2001
From: ymc9 <104139426+ymc9@users.noreply.github.com>
Date: Mon, 24 Feb 2025 10:38:05 -0800
Subject: [PATCH] feat: export `Enhanced` type to infer typeof enhanced
 PrismaClient

---
 packages/runtime/res/enhance.d.ts             |  2 +-
 packages/runtime/src/enhance.d.ts             |  2 +-
 .../src/plugins/enhancer/enhance/index.ts     | 13 ++++
 .../typing/enhancement-typing.test.ts         | 60 +++++++++++++++++++
 4 files changed, 75 insertions(+), 2 deletions(-)
 create mode 100644 tests/integration/tests/enhancements/typing/enhancement-typing.test.ts

diff --git a/packages/runtime/res/enhance.d.ts b/packages/runtime/res/enhance.d.ts
index e6ca800d2..3058f736e 100644
--- a/packages/runtime/res/enhance.d.ts
+++ b/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';
diff --git a/packages/runtime/src/enhance.d.ts b/packages/runtime/src/enhance.d.ts
index 38a519830..9a4fe97cd 100644
--- a/packages/runtime/src/enhance.d.ts
+++ b/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';
diff --git a/packages/schema/src/plugins/enhancer/enhance/index.ts b/packages/schema/src/plugins/enhancer/enhance/index.ts
index 00addd409..1cd49c0a4 100644
--- a/packages/schema/src/plugins/enhancer/enhance/index.ts
+++ b/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;
 `;
     }
 
diff --git a/tests/integration/tests/enhancements/typing/enhancement-typing.test.ts b/tests/integration/tests/enhancements/typing/enhancement-typing.test.ts
new file mode 100644
index 000000000..d29897a71
--- /dev/null
+++ b/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');
+}
+                        `,
+                    },
+                ],
+            }
+        );
+    });
+});