fix: compatibility with Prisma's "omit" feature

ymc9 · ymc9 · commit b2e001ccadc8 · 2024-05-11T15:35:27.000+08:00
diff --git a/packages/plugins/openapi/src/rpc-generator.ts b/packages/plugins/openapi/src/rpc-generator.ts
@@ -1,22 +1,22 @@
 // Inspired by: https://github.com/omar-dulaimi/prisma-trpc-generator
 
-import { analyzePolicies, PluginError, requireOption, resolvePath } from '@zenstackhq/sdk';
+import { PluginError, analyzePolicies, requireOption, resolvePath, supportCreateMany } from '@zenstackhq/sdk';
 import { DataModel, isDataModel } from '@zenstackhq/sdk/ast';
 import {
+    AggregateOperationSupport,
     addMissingInputObjectTypesForAggregate,
     addMissingInputObjectTypesForInclude,
     addMissingInputObjectTypesForModelArgs,
     addMissingInputObjectTypesForSelect,
-    AggregateOperationSupport,
     resolveAggregateOperationSupport,
 } from '@zenstackhq/sdk/dmmf-helpers';
-import type { DMMF } from '@zenstackhq/sdk/prisma';
+import { type DMMF } from '@zenstackhq/sdk/prisma';
 import * as fs from 'fs';
 import { lowerCaseFirst } from 'lower-case-first';
 import type { OpenAPIV3_1 as OAPI } from 'openapi-types';
 import * as path from 'path';
 import invariant from 'tiny-invariant';
-import { match, P } from 'ts-pattern';
+import { P, match } from 'ts-pattern';
 import { upperCaseFirst } from 'upper-case-first';
 import YAML from 'yaml';
 import { name } from '.';
@@ -166,7 +166,7 @@ export class RPCOpenAPIGenerator extends OpenAPIGeneratorBase {
             });
         }
 
-        if (ops['createMany']) {
+        if (ops['createMany'] && supportCreateMany(zmodel.$container)) {
             definitions.push({
                 method: 'post',
                 operation: 'createMany',
@@ -704,7 +704,7 @@ export class RPCOpenAPIGenerator extends OpenAPIGeneratorBase {
     private generateEnumComponent(_enum: DMMF.SchemaEnum): OAPI.SchemaObject {
         const schema: OAPI.SchemaObject = {
             type: 'string',
-            enum: _enum.values,
+            enum: _enum.values as string[],
         };
         return schema;
     }
@@ -793,17 +793,14 @@ export class RPCOpenAPIGenerator extends OpenAPIGeneratorBase {
         return result;
     }
 
-    private setInputRequired(fields: { name: string; isRequired: boolean }[], result: OAPI.NonArraySchemaObject) {
+    private setInputRequired(fields: readonly DMMF.SchemaArg[], result: OAPI.NonArraySchemaObject) {
         const required = fields.filter((f) => f.isRequired).map((f) => f.name);
         if (required.length > 0) {
             result.required = required;
         }
     }
 
-    private setOutputRequired(
-        fields: { name: string; isNullable?: boolean; outputType: DMMF.OutputTypeRef }[],
-        result: OAPI.NonArraySchemaObject
-    ) {
+    private setOutputRequired(fields: readonly DMMF.SchemaField[], result: OAPI.NonArraySchemaObject) {
         const required = fields.filter((f) => f.isNullable !== true).map((f) => f.name);
         if (required.length > 0) {
             result.required = required;
diff --git a/packages/plugins/swr/src/generator.ts b/packages/plugins/swr/src/generator.ts
@@ -8,6 +8,7 @@ import {
     requireOption,
     resolvePath,
     saveProject,
+    supportCreateMany,
 } from '@zenstackhq/sdk';
 import { DataModel, DataModelFieldType, Model, isEnum } from '@zenstackhq/sdk/ast';
 import { getPrismaClientImportSpec, type DMMF } from '@zenstackhq/sdk/prisma';
@@ -85,7 +86,7 @@ function generateModelHooks(
     }
 
     // createMany
-    if (mapping.createMany) {
+    if (mapping.createMany && supportCreateMany(model.$container)) {
         const argsType = `Prisma.${model.name}CreateManyArgs`;
         mutationFuncs.push(generateMutation(sf, model, 'POST', 'createMany', argsType, true));
     }
diff --git a/packages/plugins/tanstack-query/src/generator.ts b/packages/plugins/tanstack-query/src/generator.ts
@@ -9,6 +9,7 @@ import {
     requireOption,
     resolvePath,
     saveProject,
+    supportCreateMany,
 } from '@zenstackhq/sdk';
 import { DataModel, DataModelFieldType, Model, isEnum } from '@zenstackhq/sdk/ast';
 import { getPrismaClientImportSpec, type DMMF } from '@zenstackhq/sdk/prisma';
@@ -348,7 +349,7 @@ function generateModelHooks(
     }
 
     // createMany
-    if (mapping.createMany) {
+    if (mapping.createMany && supportCreateMany(model.$container)) {
         generateMutationHook(target, sf, model.name, 'createMany', 'post', false, 'Prisma.BatchPayload');
     }
 
diff --git a/packages/plugins/trpc/src/generator.ts b/packages/plugins/trpc/src/generator.ts
@@ -7,6 +7,7 @@ import {
     requireOption,
     resolvePath,
     saveProject,
+    supportCreateMany,
     type PluginOptions,
 } from '@zenstackhq/sdk';
 import { Model } from '@zenstackhq/sdk/ast';
@@ -79,11 +80,11 @@ export async function generate(model: Model, options: PluginOptions, dmmf: DMMF.
 
 function createAppRouter(
     outDir: string,
-    modelOperations: DMMF.ModelMapping[],
+    modelOperations: readonly DMMF.ModelMapping[],
     hiddenModels: string[],
     generateModelActions: string[] | undefined,
     generateClientHelpers: string[] | undefined,
-    _zmodel: Model,
+    zmodel: Model,
     zodSchemasImport: string,
     options: PluginOptions
 ) {
@@ -171,7 +172,8 @@ function createAppRouter(
                         generateModelActions,
                         generateClientHelpers,
                         zodSchemasImport,
-                        options
+                        options,
+                        zmodel
                     );
 
                     appRouter.addImportDeclaration({
@@ -241,7 +243,8 @@ function generateModelCreateRouter(
     generateModelActions: string[] | undefined,
     generateClientHelpers: string[] | undefined,
     zodSchemasImport: string,
-    options: PluginOptions
+    options: PluginOptions,
+    zmodel: Model
 ) {
     const modelRouter = project.createSourceFile(path.resolve(outputDir, 'routers', `${model}.router.ts`), undefined, {
         overwrite: true,
@@ -298,6 +301,10 @@ function generateModelCreateRouter(
                     inputType &&
                     (!generateModelActions || generateModelActions.includes(generateOpName))
                 ) {
+                    if (generateOpName === 'createMany' && !supportCreateMany(zmodel)) {
+                        continue;
+                    }
+
                     generateProcedure(funcWriter, generateOpName, upperCaseFirst(inputType), model, baseOpType);
 
                     if (routerTypingStructure) {
diff --git a/packages/plugins/trpc/src/helpers.ts b/packages/plugins/trpc/src/helpers.ts
@@ -327,7 +327,7 @@ export const getProcedureTypeByOpName = (opName: string) => {
     return procType;
 };
 
-export function resolveModelsComments(models: DMMF.Model[], hiddenModels: string[]) {
+export function resolveModelsComments(models: readonly DMMF.Model[], hiddenModels: string[]) {
     const modelAttributeRegex = /(@@Gen\.)+([A-z])+(\()+(.+)+(\))+/;
     const attributeNameRegex = /(?:\.)+([A-Za-z])+(?:\()+/;
     const attributeArgsRegex = /(?:\()+([A-Za-z])+:+(.+)+(?:\))+/;
diff --git a/packages/runtime/src/enhancements/policy/policy-utils.ts b/packages/runtime/src/enhancements/policy/policy-utils.ts
@@ -982,6 +982,11 @@ export class PolicyUtil extends QueryUtils {
     }
 
     private doInjectReadCheckSelect(model: string, args: any, input: any) {
+        // omit should be ignored to avoid interfering with field selection
+        if (args.omit) {
+            delete args.omit;
+        }
+
         if (!input?.select) {
             return;
         }
@@ -1178,6 +1183,12 @@ export class PolicyUtil extends QueryUtils {
                     continue;
                 }
 
+                if (queryArgs?.omit?.[field] === true) {
+                    // respect `{ omit: { [field]: true } }`
+                    delete entityData[field];
+                    continue;
+                }
+
                 if (hasFieldLevelPolicy) {
                     // 1. remove fields selected for checking field-level policies but not selected by the original query args
                     // 2. evaluate field-level policies and remove fields that are not readable
diff --git a/packages/schema/src/plugins/enhancer/enhance/index.ts b/packages/schema/src/plugins/enhancer/enhance/index.ts
@@ -56,7 +56,7 @@ export class EnhancerGenerator {
         private readonly outDir: string
     ) {}
 
-    async generate() {
+    async generate(): Promise<{ dmmf: DMMF.Document | undefined }> {
         let logicalPrismaClientDir: string | undefined;
         let dmmf: DMMF.Document | undefined;
 
diff --git a/packages/schema/src/plugins/zod/generator.ts b/packages/schema/src/plugins/zod/generator.ts
@@ -107,7 +107,7 @@ export class ZodSchemaGenerator {
                 project: this.project,
                 inputObjectTypes,
             });
-            await transformer.generateInputSchemas(this.options);
+            await transformer.generateInputSchemas(this.options, this.model);
             this.sourceFiles.push(...transformer.sourceFiles);
         }
 
@@ -189,7 +189,10 @@ export class ZodSchemaGenerator {
         );
     }
 
-    private async generateEnumSchemas(prismaSchemaEnum: DMMF.SchemaEnum[], modelSchemaEnum: DMMF.SchemaEnum[]) {
+    private async generateEnumSchemas(
+        prismaSchemaEnum: readonly DMMF.SchemaEnum[],
+        modelSchemaEnum: readonly DMMF.SchemaEnum[]
+    ) {
         const enumTypes = [...prismaSchemaEnum, ...modelSchemaEnum];
         const enumNames = enumTypes.map((enumItem) => upperCaseFirst(enumItem.name));
         Transformer.enumNames = enumNames ?? [];
diff --git a/packages/schema/src/plugins/zod/transformer.ts b/packages/schema/src/plugins/zod/transformer.ts
@@ -1,5 +1,6 @@
 /* eslint-disable @typescript-eslint/ban-ts-comment */
-import { indentString, type PluginOptions } from '@zenstackhq/sdk';
+import { indentString, supportCreateMany, type PluginOptions } from '@zenstackhq/sdk';
+import type { Model } from '@zenstackhq/sdk/ast';
 import { checkModelHasModelRelation, findModelByName, isAggregateInputType } from '@zenstackhq/sdk/dmmf-helpers';
 import { type DMMF as PrismaDMMF } from '@zenstackhq/sdk/prisma';
 import path from 'path';
@@ -11,12 +12,12 @@ import { AggregateOperationSupport, TransformerParams } from './types';
 export default class Transformer {
     name: string;
     originalName: string;
-    fields: PrismaDMMF.SchemaArg[];
+    fields: readonly PrismaDMMF.SchemaArg[];
     schemaImports = new Set<string>();
-    models: PrismaDMMF.Model[];
+    models: readonly PrismaDMMF.Model[];
     modelOperations: PrismaDMMF.ModelMapping[];
     aggregateOperationSupport: AggregateOperationSupport;
-    enumTypes: PrismaDMMF.SchemaEnum[];
+    enumTypes: readonly PrismaDMMF.SchemaEnum[];
 
     static enumNames: string[] = [];
     static rawOpsMap: { [name: string]: string } = {};
@@ -389,7 +390,7 @@ export const ${this.name}ObjectSchema: SchemaType = ${schema} as SchemaType;`;
         return wrapped;
     }
 
-    async generateInputSchemas(options: PluginOptions) {
+    async generateInputSchemas(options: PluginOptions, zmodel: Model) {
         const globalExports: string[] = [];
 
         // whether Prisma's Unchecked* series of input types should be generated
@@ -489,7 +490,7 @@ export const ${this.name}ObjectSchema: SchemaType = ${schema} as SchemaType;`;
                 operations.push(['create', origModelName]);
             }
 
-            if (createMany) {
+            if (createMany && supportCreateMany(zmodel)) {
                 imports.push(
                     `import { ${modelName}CreateManyInputObjectSchema } from '../objects/${modelName}CreateManyInput.schema'`
                 );
diff --git a/packages/schema/src/plugins/zod/types.ts b/packages/schema/src/plugins/zod/types.ts
@@ -2,10 +2,10 @@ import type { DMMF as PrismaDMMF } from '@zenstackhq/sdk/prisma';
 import { Project } from 'ts-morph';
 
 export type TransformerParams = {
-    enumTypes?: PrismaDMMF.SchemaEnum[];
-    fields?: PrismaDMMF.SchemaArg[];
+    enumTypes?: readonly PrismaDMMF.SchemaEnum[];
+    fields?: readonly PrismaDMMF.SchemaArg[];
     name?: string;
-    models?: PrismaDMMF.Model[];
+    models?: readonly PrismaDMMF.Model[];
     modelOperations?: PrismaDMMF.ModelMapping[];
     aggregateOperationSupport?: AggregateOperationSupport;
     isDefaultPrismaClientOutput?: boolean;
diff --git a/packages/sdk/package.json b/packages/sdk/package.json
@@ -18,8 +18,8 @@
     "author": "",
     "license": "MIT",
     "dependencies": {
-        "@prisma/generator-helper": "5.7.0",
-        "@prisma/internals": "5.7.0",
+        "@prisma/generator-helper": "^5.13.0",
+        "@prisma/internals": "^5.13.0",
         "@zenstackhq/language": "workspace:*",
         "@zenstackhq/runtime": "workspace:*",
         "langium": "1.3.1",
diff --git a/packages/sdk/src/dmmf-helpers/include-helpers.ts b/packages/sdk/src/dmmf-helpers/include-helpers.ts
@@ -1,15 +1,18 @@
 import type { DMMF } from '../prisma';
 import { checkIsModelRelationField, checkModelHasManyModelRelation, checkModelHasModelRelation } from './model-helpers';
 
-export function addMissingInputObjectTypesForInclude(inputObjectTypes: DMMF.InputType[], models: DMMF.Model[]) {
+export function addMissingInputObjectTypesForInclude(
+    inputObjectTypes: DMMF.InputType[],
+    models: readonly DMMF.Model[]
+) {
     // generate input object types necessary to support ModelInclude with relation support
     const generatedIncludeInputObjectTypes = generateModelIncludeInputObjectTypes(models);
 
     for (const includeInputObjectType of generatedIncludeInputObjectTypes) {
         inputObjectTypes.push(includeInputObjectType);
     }
 }
-function generateModelIncludeInputObjectTypes(models: DMMF.Model[]) {
+function generateModelIncludeInputObjectTypes(models: readonly DMMF.Model[]) {
     const modelIncludeInputObjectTypes: DMMF.InputType[] = [];
     for (const model of models) {
         const { name: modelName, fields: modelFields } = model;
diff --git a/packages/sdk/src/dmmf-helpers/model-helpers.ts b/packages/sdk/src/dmmf-helpers/model-helpers.ts
@@ -31,6 +31,6 @@ export function checkIsManyModelRelationField(modelField: DMMF.Field) {
     return checkIsModelRelationField(modelField) && modelField.isList;
 }
 
-export function findModelByName(models: DMMF.Model[], modelName: string) {
+export function findModelByName(models: readonly DMMF.Model[], modelName: string) {
     return models.find(({ name }) => name === modelName);
 }
diff --git a/packages/sdk/src/dmmf-helpers/modelArgs-helpers.ts b/packages/sdk/src/dmmf-helpers/modelArgs-helpers.ts
@@ -1,14 +1,17 @@
 import type { DMMF } from '../prisma';
 import { checkModelHasModelRelation } from './model-helpers';
 
-export function addMissingInputObjectTypesForModelArgs(inputObjectTypes: DMMF.InputType[], models: DMMF.Model[]) {
+export function addMissingInputObjectTypesForModelArgs(
+    inputObjectTypes: DMMF.InputType[],
+    models: readonly DMMF.Model[]
+) {
     const modelArgsInputObjectTypes = generateModelArgsInputObjectTypes(models);
 
     for (const modelArgsInputObjectType of modelArgsInputObjectTypes) {
         inputObjectTypes.push(modelArgsInputObjectType);
     }
 }
-function generateModelArgsInputObjectTypes(models: DMMF.Model[]) {
+function generateModelArgsInputObjectTypes(models: readonly DMMF.Model[]) {
     const modelArgsInputObjectTypes: DMMF.InputType[] = [];
     for (const model of models) {
         const { name: modelName } = model;
diff --git a/packages/sdk/src/dmmf-helpers/select-helpers.ts b/packages/sdk/src/dmmf-helpers/select-helpers.ts
@@ -4,7 +4,7 @@ import { checkIsModelRelationField, checkModelHasManyModelRelation } from './mod
 export function addMissingInputObjectTypesForSelect(
     inputObjectTypes: DMMF.InputType[],
     outputObjectTypes: DMMF.OutputType[],
-    models: DMMF.Model[]
+    models: readonly DMMF.Model[]
 ) {
     // generate input object types necessary to support ModelSelect._count
     const modelCountOutputTypes = getModelCountOutputTypes(outputObjectTypes);
@@ -89,34 +89,33 @@ function generateModelCountOutputTypeArgsInputObjectTypes(modelCountOutputTypes:
     return modelCountOutputTypeArgsInputObjectTypes;
 }
 
-function generateModelSelectInputObjectTypes(models: DMMF.Model[]) {
+function generateModelSelectInputObjectTypes(models: readonly DMMF.Model[]) {
     const modelSelectInputObjectTypes: DMMF.InputType[] = [];
     for (const model of models) {
         const { name: modelName, fields: modelFields } = model;
         const fields: DMMF.SchemaArg[] = [];
 
         for (const modelField of modelFields) {
             const { name: modelFieldName, isList, type } = modelField;
+            const inputTypes = [{ isList: false, type: 'Boolean', location: 'scalar' }];
 
             const isRelationField = checkIsModelRelationField(modelField);
-
-            const field: DMMF.SchemaArg = {
-                name: modelFieldName,
-                isRequired: false,
-                isNullable: false,
-                inputTypes: [{ isList: false, type: 'Boolean', location: 'scalar' }],
-            };
-
             if (isRelationField) {
                 const schemaArgInputType: DMMF.InputTypeRef = {
                     isList: false,
                     type: isList ? `${type}FindManyArgs` : `${type}Args`,
                     location: 'inputObjectTypes',
                     namespace: 'prisma',
                 };
-                field.inputTypes.push(schemaArgInputType);
+                inputTypes.push(schemaArgInputType);
             }
 
+            const field: DMMF.SchemaArg = {
+                name: modelFieldName,
+                isRequired: false,
+                isNullable: false,
+                inputTypes: inputTypes as DMMF.InputTypeRef[],
+            };
             fields.push(field);
         }
 
diff --git a/packages/sdk/src/utils.ts b/packages/sdk/src/utils.ts
diff --git a/packages/testtools/src/schema.ts b/packages/testtools/src/schema.ts
diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml
diff --git a/tests/integration/tests/enhancements/with-policy/prisma-omit.test.ts b/tests/integration/tests/enhancements/with-policy/prisma-omit.test.ts

Original file line number	Diff line number	Diff line change
`@@ -9,6 +9,7 @@ import {`
`9`	`9`	`requireOption,`
`10`	`10`	`resolvePath,`
`11`	`11`	`saveProject,`
	`12`	`+ supportCreateMany,`
`12`	`13`	`} from '@zenstackhq/sdk';`
`13`	`14`	`import { DataModel, DataModelFieldType, Model, isEnum } from '@zenstackhq/sdk/ast';`
`14`	`15`	`import { getPrismaClientImportSpec, type DMMF } from '@zenstackhq/sdk/prisma';`
`@@ -348,7 +349,7 @@ function generateModelHooks(`
`348`	`349`	`}`
`349`	`350`
`350`	`351`	`// createMany`
`351`		`- if (mapping.createMany) {`
	`352`	`+ if (mapping.createMany && supportCreateMany(model.$container)) {`
`352`	`353`	`generateMutationHook(target, sf, model.name, 'createMany', 'post', false, 'Prisma.BatchPayload');`
`353`	`354`	`}`
`354`	`355`