From 61438ccdf9ba500cab1bb6eaeba26b5486121b18 Mon Sep 17 00:00:00 2001 From: ymc9 <104139426+ymc9@users.noreply.github.com> Date: Sat, 16 Mar 2024 22:26:49 -0700 Subject: [PATCH 1/3] fix: preserve prisma client extensions's typing --- .../src/plugins/enhancer/enhance/index.ts | 36 +++--- .../with-delegate/enhanced-client.test.ts | 105 +++++++++++++----- 2 files changed, 98 insertions(+), 43 deletions(-) diff --git a/packages/schema/src/plugins/enhancer/enhance/index.ts b/packages/schema/src/plugins/enhancer/enhance/index.ts index 35bc246b4..fbb8a442a 100644 --- 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 & InternalArgs>( + prisma: _PrismaClient, + context?: EnhancementContext<${authTypeParam}>, options?: EnhancementOptions): PrismaClient; - export function enhance(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 & InternalArgs>( + prisma: DynamicClientExtensionThis, + context?: EnhancementContext<${authTypeParam}>, options?: EnhancementOptions): DynamicClientExtensionThis; + +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 index 6ec03486f..076202553 100644 --- 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,8 +1037,7 @@ 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'; @@ -1048,7 +1045,6 @@ describe('Polymorphism Test', () => { 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, + }, + ], + }); }); }); From 6671675db692a2f67b41d68e94cddb95ab038627 Mon Sep 17 00:00:00 2001 From: ymc9 <104139426+ymc9@users.noreply.github.com> Date: Sun, 17 Mar 2024 09:51:16 -0700 Subject: [PATCH 2/3] fix tests --- tests/integration/tests/cli/plugins.test.ts | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/integration/tests/cli/plugins.test.ts b/tests/integration/tests/cli/plugins.test.ts index ee32ce9e9..b5dbf9fb8 100644 --- a/tests/integration/tests/cli/plugins.test.ts +++ b/tests/integration/tests/cli/plugins.test.ts @@ -73,7 +73,7 @@ describe('CLI Plugins Tests', () => { 'swr', '@tanstack/react-query@^5.0.0', '@trpc/server', - '@prisma/client@^4.0.0', + '@prisma/client@^5.0.0', `${path.join(__dirname, '../../../../.build/zenstackhq-language-' + ver + '.tgz')}`, `${path.join(__dirname, '../../../../.build/zenstackhq-sdk-' + ver + '.tgz')}`, `${path.join(__dirname, '../../../../.build/zenstackhq-runtime-' + ver + '.tgz')}`, @@ -83,7 +83,7 @@ describe('CLI Plugins Tests', () => { const devDepPkgs = [ 'typescript', '@types/react', - 'prisma@^4.0.0', + 'prisma@^5.0.0', `${path.join(__dirname, '../../../../.build/zenstack-' + ver + '.tgz')}`, `${path.join(__dirname, '../../../../.build/zenstackhq-tanstack-query-' + ver + '.tgz')}`, `${path.join(__dirname, '../../../../.build/zenstackhq-swr-' + ver + '.tgz')}`, From a9603dedd09228f86f7c1b8de620a0a5aa98a62e Mon Sep 17 00:00:00 2001 From: ymc9 <104139426+ymc9@users.noreply.github.com> Date: Sun, 17 Mar 2024 11:15:52 -0700 Subject: [PATCH 3/3] chore: improve code generation --- .../src/plugins/enhancer/enhance/index.ts | 86 +++++++++++++------ 1 file changed, 62 insertions(+), 24 deletions(-) diff --git a/packages/schema/src/plugins/enhancer/enhance/index.ts b/packages/schema/src/plugins/enhancer/enhance/index.ts index fbb8a442a..7b042f05a 100644 --- a/packages/schema/src/plugins/enhancer/enhance/index.ts +++ b/packages/schema/src/plugins/enhancer/enhance/index.ts @@ -56,10 +56,9 @@ export class EnhancerGenerator { let logicalPrismaClientDir: string | undefined; let dmmf: DMMF.Document | undefined; - const withLogicalClient = this.needsLogicalClient(); const prismaImport = getPrismaClientImportSpec(this.outDir, this.options); - if (withLogicalClient) { + if (this.needsLogicalClient()) { // schema contains delegate models, need to generate a logical prisma schema const result = await this.generateLogicalPrisma(); @@ -90,22 +89,67 @@ export class EnhancerGenerator { const enhanceTs = this.project.createSourceFile( path.join(this.outDir, 'enhance.ts'), `import { createEnhancement, type EnhancementContext, type EnhancementOptions, type ZodSchemas, type AuthUser } from '@zenstackhq/runtime'; - import modelMeta from './model-meta'; - import policy from './policy'; - 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 { Prisma, PrismaClient } from '${logicalPrismaClientDir}/index-fixed'; - ` - : `import type * as _P from '${prismaImport}'; - import type { Prisma, PrismaClient } from '${prismaImport}'; - ` +import modelMeta from './model-meta'; +import policy from './policy'; +${this.options.withZodSchemas ? "import * as zodSchemas from './zod';" : 'const zodSchemas = undefined;'} + +${ + logicalPrismaClientDir + ? this.createLogicalPrismaImports(prismaImport, logicalPrismaClientDir) + : this.createSimplePrismaImports(prismaImport) +} + +${authTypes} + +${ + logicalPrismaClientDir + ? this.createLogicalPrismaEnhanceFunction(authTypeParam) + : this.createSimplePrismaEnhanceFunction(authTypeParam) +} + `, + { overwrite: true } + ); + + await this.saveSourceFile(enhanceTs); + + return { dmmf }; } - ${this.options.withZodSchemas ? "import * as zodSchemas from './zod';" : 'const zodSchemas = undefined;'} - - ${authTypes} + + private createSimplePrismaImports(prismaImport: string) { + return `import { Prisma } from '${prismaImport}'; +import type * as _P from '${prismaImport}'; + `; + } + + private createSimplePrismaEnhanceFunction(authTypeParam: string) { + return ` +export function enhance(prisma: DbClient, context?: EnhancementContext<${authTypeParam}>, options?: EnhancementOptions) { + return createEnhancement(prisma, { + modelMeta, + policy, + zodSchemas: zodSchemas as unknown as (ZodSchemas | undefined), + prismaModule: Prisma, + ...options + }, context); +} + `; + } + + private createLogicalPrismaImports(prismaImport: string, logicalPrismaClientDir: string) { + return `import { Prisma as _Prisma, PrismaClient as _PrismaClient } from '${prismaImport}'; +import type { + InternalArgs, + TypeMapDef, + TypeMapCbDef, + DynamicClientExtensionThis, +} from '${prismaImport}/runtime/library'; +import type * as _P from '${logicalPrismaClientDir}/index-fixed'; +import type { Prisma, PrismaClient } from '${logicalPrismaClientDir}/index-fixed'; +`; + } + + private createLogicalPrismaEnhanceFunction(authTypeParam: string) { + return ` // overload for plain PrismaClient export function enhance & InternalArgs>( prisma: _PrismaClient, @@ -125,13 +169,7 @@ export function enhance(prisma: any, context?: EnhancementContext<${authTypePara ...options }, context); } - `, - { overwrite: true } - ); - - await this.saveSourceFile(enhanceTs); - - return { dmmf }; +`; } private needsLogicalClient() {