fix: issue 961, incorrect policy injection for nested updateMany

ymc9 · ymc9 · commit bf690a072771 · 2024-01-25T21:28:31.000+08:00
diff --git a/packages/runtime/src/cross/model-meta.ts b/packages/runtime/src/cross/model-meta.ts
@@ -66,6 +66,11 @@ export type FieldInfo = {
      * Mapping from foreign key field names to relation field names
      */
     foreignKeyMapping?: Record<string, string>;
+
+    /**
+     * If the field is an auto-increment field
+     */
+    isAutoIncrement?: boolean;
 };
 
 /**
diff --git a/packages/runtime/src/enhancements/policy/handler.ts b/packages/runtime/src/enhancements/policy/handler.ts
@@ -523,29 +523,16 @@ export class PolicyProxyHandler<DbClient extends DbClientContract> implements Pr
         let createResult = await Promise.all(
             enumerate(args.data).map(async (item) => {
                 if (args.skipDuplicates) {
-                    // check unique constraint conflicts
-                    // we can't rely on try/catch/ignore constraint violation error: https://github.com/prisma/prisma/issues/20496
-                    // TODO: for simple cases we should be able to translate it to an `upsert` with empty `update` payload
-
-                    // for each unique constraint, check if the input item has all fields set, and if so, check if
-                    // an entity already exists, and ignore accordingly
-                    const uniqueConstraints = this.utils.getUniqueConstraints(model);
-                    for (const constraint of Object.values(uniqueConstraints)) {
-                        if (constraint.fields.every((f) => item[f] !== undefined)) {
-                            const uniqueFilter = constraint.fields.reduce((acc, f) => ({ ...acc, [f]: item[f] }), {});
-                            const existing = await this.utils.checkExistence(db, model, uniqueFilter);
-                            if (existing) {
-                                if (this.shouldLogQuery) {
-                                    this.logger.info(`[policy] skipping duplicate ${formatObject(item)}`);
-                                }
-                                return undefined;
-                            }
+                    if (await this.hasDuplicatedUniqueConstraint(model, item, db)) {
+                        if (this.shouldLogQuery) {
+                            this.logger.info(`[policy] \`createMany\` skipping duplicate ${formatObject(item)}`);
                         }
+                        return undefined;
                     }
                 }
 
                 if (this.shouldLogQuery) {
-                    this.logger.info(`[policy] \`create\` ${model}: ${formatObject(item)}`);
+                    this.logger.info(`[policy] \`create\` for \`createMany\` ${model}: ${formatObject(item)}`);
                 }
                 return await db[model].create({ select: this.utils.makeIdSelection(model), data: item });
             })
@@ -564,6 +551,26 @@ export class PolicyProxyHandler<DbClient extends DbClientContract> implements Pr
         };
     }
 
+    private async hasDuplicatedUniqueConstraint(model: string, createData: any, db: Record<string, DbOperations>) {
+        // check unique constraint conflicts
+        // we can't rely on try/catch/ignore constraint violation error: https://github.com/prisma/prisma/issues/20496
+        // TODO: for simple cases we should be able to translate it to an `upsert` with empty `update` payload
+
+        // for each unique constraint, check if the input item has all fields set, and if so, check if
+        // an entity already exists, and ignore accordingly
+        const uniqueConstraints = this.utils.getUniqueConstraints(model);
+        for (const constraint of Object.values(uniqueConstraints)) {
+            if (constraint.fields.every((f) => createData[f] !== undefined)) {
+                const uniqueFilter = constraint.fields.reduce((acc, f) => ({ ...acc, [f]: createData[f] }), {});
+                const existing = await this.utils.checkExistence(db, model, uniqueFilter);
+                if (existing) {
+                    return true;
+                }
+            }
+        }
+        return false;
+    }
+
     //#endregion
 
     //#region Update & Upsert
@@ -707,17 +714,22 @@ export class PolicyProxyHandler<DbClient extends DbClientContract> implements Pr
             postWriteChecks.push(...checks);
         };
 
-        const _createMany = async (model: string, args: any, context: NestedWriteVisitorContext) => {
-            if (context.field?.backLink) {
-                // handles the connection to upstream entity
-                const reversedQuery = this.utils.buildReversedQuery(context);
-                for (const item of enumerate(args.data)) {
-                    Object.assign(item, reversedQuery);
+        const _createMany = async (
+            model: string,
+            args: { data: any; skipDuplicates?: boolean },
+            context: NestedWriteVisitorContext
+        ) => {
+            for (const item of enumerate(args.data)) {
+                if (args.skipDuplicates) {
+                    if (await this.hasDuplicatedUniqueConstraint(model, item, db)) {
+                        if (this.shouldLogQuery) {
+                            this.logger.info(`[policy] \`createMany\` skipping duplicate ${formatObject(item)}`);
+                        }
+                        continue;
+                    }
                 }
+                await _create(model, item, context);
             }
-            // proceed with the create and collect post-create checks
-            const { postWriteChecks: checks } = await this.doCreateMany(model, args, db);
-            postWriteChecks.push(...checks);
         };
 
         const _connectDisconnect = async (model: string, args: any, context: NestedWriteVisitorContext) => {
@@ -797,9 +809,6 @@ export class PolicyProxyHandler<DbClient extends DbClientContract> implements Pr
             },
 
             updateMany: async (model, args, context) => {
-                // injects auth guard into where clause
-                this.utils.injectAuthGuard(db, args, model, 'update');
-
                 // prepare for post-update check
                 if (this.utils.hasAuthGuard(model, 'postUpdate') || this.utils.getZodSchema(model)) {
                     let select = this.utils.makeIdSelection(model);
@@ -809,10 +818,12 @@ export class PolicyProxyHandler<DbClient extends DbClientContract> implements Pr
                     }
                     const reversedQuery = this.utils.buildReversedQuery(context);
                     const currentSetQuery = { select, where: reversedQuery };
-                    this.utils.injectAuthGuard(db, currentSetQuery, model, 'read');
+                    this.utils.injectAuthGuardAsWhere(db, currentSetQuery, model, 'read');
 
                     if (this.shouldLogQuery) {
-                        this.logger.info(`[policy] \`findMany\` ${model}:\n${formatObject(currentSetQuery)}`);
+                        this.logger.info(
+                            `[policy] \`findMany\` for post update check ${model}:\n${formatObject(currentSetQuery)}`
+                        );
                     }
                     const currentSet = await db[model].findMany(currentSetQuery);
 
@@ -825,6 +836,27 @@ export class PolicyProxyHandler<DbClient extends DbClientContract> implements Pr
                         }))
                     );
                 }
+
+                const updateGuard = this.utils.getAuthGuard(db, model, 'update');
+                if (this.utils.isTrue(updateGuard) || this.utils.isFalse(updateGuard)) {
+                    // injects simple auth guard into where clause
+                    this.utils.injectAuthGuardAsWhere(db, args, model, 'update');
+                } else {
+                    // we have to process `updateMany` separately because the guard may contain
+                    // filters using relation fields which are not allowed in nested `updateMany`
+                    const reversedQuery = this.utils.buildReversedQuery(context);
+                    const updateWhere = this.utils.and(reversedQuery, updateGuard);
+                    if (this.shouldLogQuery) {
+                        this.logger.info(
+                            `[policy] \`updateMany\` ${model}:\n${formatObject({
+                                where: updateWhere,
+                                data: args.data,
+                            })}`
+                        );
+                    }
+                    await db[model].updateMany({ where: updateWhere, data: args.data });
+                    delete context.parent.updateMany;
+                }
             },
 
             create: async (model, args, context) => {
@@ -931,9 +963,21 @@ export class PolicyProxyHandler<DbClient extends DbClientContract> implements Pr
             },
 
             deleteMany: async (model, args, context) => {
-                // inject delete guard
                 const guard = await this.utils.getAuthGuard(db, model, 'delete');
-                context.parent.deleteMany = this.utils.and(args, guard);
+                if (this.utils.isTrue(guard) || this.utils.isFalse(guard)) {
+                    // inject simple auth guard
+                    context.parent.deleteMany = this.utils.and(args, guard);
+                } else {
+                    // we have to process `deleteMany` separately because the guard may contain
+                    // filters using relation fields which are not allowed in nested `deleteMany`
+                    const reversedQuery = this.utils.buildReversedQuery(context);
+                    const deleteWhere = this.utils.and(reversedQuery, guard);
+                    if (this.shouldLogQuery) {
+                        this.logger.info(`[policy] \`deleteMany\` ${model}:\n${formatObject({ where: deleteWhere })}`);
+                    }
+                    await db[model].deleteMany({ where: deleteWhere });
+                    delete context.parent.deleteMany;
+                }
             },
         });
 
@@ -958,13 +1002,17 @@ export class PolicyProxyHandler<DbClient extends DbClientContract> implements Pr
         }
         for (const k of Object.keys(args)) {
             const field = resolveField(this.modelMeta, model, k);
-            if (field?.isId || field?.isForeignKey) {
+            if (this.isAutoIncrementIdField(field) || field?.isForeignKey) {
                 return true;
             }
         }
         return false;
     }
 
+    private isAutoIncrementIdField(field: FieldInfo) {
+        return field.isId && field.isAutoIncrement;
+    }
+
     async updateMany(args: any) {
         if (!args) {
             throw prismaClientValidationError(this.prisma, 'query argument is required');
@@ -976,7 +1024,7 @@ export class PolicyProxyHandler<DbClient extends DbClientContract> implements Pr
         this.utils.tryReject(this.prisma, this.model, 'update');
 
         args = this.utils.clone(args);
-        this.utils.injectAuthGuard(this.prisma, args, this.model, 'update');
+        this.utils.injectAuthGuardAsWhere(this.prisma, args, this.model, 'update');
 
         if (this.utils.hasAuthGuard(this.model, 'postUpdate') || this.utils.getZodSchema(this.model)) {
             // use a transaction to do post-update checks
@@ -989,7 +1037,7 @@ export class PolicyProxyHandler<DbClient extends DbClientContract> implements Pr
                     select = { ...select, ...preValueSelect };
                 }
                 const currentSetQuery = { select, where: args.where };
-                this.utils.injectAuthGuard(tx, currentSetQuery, this.model, 'read');
+                this.utils.injectAuthGuardAsWhere(tx, currentSetQuery, this.model, 'read');
 
                 if (this.shouldLogQuery) {
                     this.logger.info(`[policy] \`findMany\` ${this.model}: ${formatObject(currentSetQuery)}`);
@@ -1118,7 +1166,7 @@ export class PolicyProxyHandler<DbClient extends DbClientContract> implements Pr
 
         // inject policy conditions
         args = args ?? {};
-        this.utils.injectAuthGuard(this.prisma, args, this.model, 'delete');
+        this.utils.injectAuthGuardAsWhere(this.prisma, args, this.model, 'delete');
 
         // conduct the deletion
         if (this.shouldLogQuery) {
@@ -1139,7 +1187,7 @@ export class PolicyProxyHandler<DbClient extends DbClientContract> implements Pr
         args = this.utils.clone(args);
 
         // inject policy conditions
-        this.utils.injectAuthGuard(this.prisma, args, this.model, 'read');
+        this.utils.injectAuthGuardAsWhere(this.prisma, args, this.model, 'read');
 
         if (this.shouldLogQuery) {
             this.logger.info(`[policy] \`aggregate\` ${this.model}:\n${formatObject(args)}`);
@@ -1155,7 +1203,7 @@ export class PolicyProxyHandler<DbClient extends DbClientContract> implements Pr
         args = this.utils.clone(args);
 
         // inject policy conditions
-        this.utils.injectAuthGuard(this.prisma, args, this.model, 'read');
+        this.utils.injectAuthGuardAsWhere(this.prisma, args, this.model, 'read');
 
         if (this.shouldLogQuery) {
             this.logger.info(`[policy] \`groupBy\` ${this.model}:\n${formatObject(args)}`);
@@ -1166,7 +1214,7 @@ export class PolicyProxyHandler<DbClient extends DbClientContract> implements Pr
     async count(args: any) {
         // inject policy conditions
         args = args ? this.utils.clone(args) : {};
-        this.utils.injectAuthGuard(this.prisma, args, this.model, 'read');
+        this.utils.injectAuthGuardAsWhere(this.prisma, args, this.model, 'read');
 
         if (this.shouldLogQuery) {
             this.logger.info(`[policy] \`count\` ${this.model}:\n${formatObject(args)}`);
diff --git a/packages/runtime/src/enhancements/policy/policy-utils.ts b/packages/runtime/src/enhancements/policy/policy-utils.ts
@@ -355,7 +355,7 @@ export class PolicyUtil {
     /**
      * Injects model auth guard as where clause.
      */
-    injectAuthGuard(db: Record<string, DbOperations>, args: any, model: string, operation: PolicyOperationKind) {
+    injectAuthGuardAsWhere(db: Record<string, DbOperations>, args: any, model: string, operation: PolicyOperationKind) {
         let guard = this.getAuthGuard(db, model, operation);
 
         if (operation === 'update' && args) {
@@ -494,7 +494,7 @@ export class PolicyUtil {
     injectForRead(db: Record<string, DbOperations>, model: string, args: any) {
         // make select and include visible to the injection
         const injected: any = { select: args.select, include: args.include };
-        if (!this.injectAuthGuard(db, injected, model, 'read')) {
+        if (!this.injectAuthGuardAsWhere(db, injected, model, 'read')) {
             return false;
         }
 
@@ -562,7 +562,7 @@ export class PolicyUtil {
     /**
      * Builds a reversed query for the given nested path.
      */
-    buildReversedQuery(context: NestedWriteVisitorContext, mutating = false, unsafeOperation = false) {
+    buildReversedQuery(context: NestedWriteVisitorContext, forMutationPayload = false, unsafeOperation = false) {
         let result, currQuery: any;
         let currField: FieldInfo | undefined;
 
@@ -593,7 +593,7 @@ export class PolicyUtil {
                     throw this.unknownError(`missing backLink field ${currField.backLink} in ${currField.type}`);
                 }
 
-                if (backLinkField.isArray && !mutating) {
+                if (backLinkField.isArray && !forMutationPayload) {
                     // many-side of relationship, wrap with "some" query
                     currQuery[currField.backLink] = { some: { ...visitWhere } };
                     currQuery = currQuery[currField.backLink].some;
@@ -603,7 +603,7 @@ export class PolicyUtil {
                     // calculate if we should preserve the relation condition (e.g., { user: { id: 1 } })
                     const shouldPreserveRelationCondition =
                         // doing a mutation
-                        mutating &&
+                        forMutationPayload &&
                         // and it's a safe mutate
                         !unsafeOperation &&
                         // and the current segment is the direct parent (the last one is the mutate itself),
@@ -666,7 +666,7 @@ export class PolicyUtil {
                     continue;
                 }
                 // inject into the "where" clause inside select
-                this.injectAuthGuard(db, injectTarget._count.select[field], fieldInfo.type, 'read');
+                this.injectAuthGuardAsWhere(db, injectTarget._count.select[field], fieldInfo.type, 'read');
             }
         }
 
@@ -692,7 +692,7 @@ export class PolicyUtil {
                     injectTarget[field] = {};
                 }
                 // inject extra condition for to-many or nullable to-one relation
-                this.injectAuthGuard(db, injectTarget[field], fieldInfo.type, 'read');
+                this.injectAuthGuardAsWhere(db, injectTarget[field], fieldInfo.type, 'read');
 
                 // recurse
                 const subHoisted = this.injectNestedReadConditions(db, fieldInfo.type, injectTarget[field]);
diff --git a/packages/sdk/src/model-meta-generator.ts b/packages/sdk/src/model-meta-generator.ts
@@ -5,6 +5,7 @@ import {
     isArrayExpr,
     isBooleanLiteral,
     isDataModel,
+    isInvocationExpr,
     isNumberLiteral,
     isReferenceExpr,
     isStringLiteral,
@@ -130,6 +131,11 @@ function generateModelMetadata(dataModels: DataModel[], writer: CodeBlockWriter,
                     foreignKeyMapping: ${JSON.stringify(fkMapping)},`);
                         }
 
+                        if (isAutoIncrement(f)) {
+                            writer.write(`
+                    isAutoIncrement: true,`);
+                        }
+
                         writer.write(`
                 },`);
                     }
@@ -338,3 +344,17 @@ function getDeleteCascades(model: DataModel): string[] {
         })
         .map((m) => m.name);
 }
+
+function isAutoIncrement(field: DataModelField) {
+    const defaultAttr = getAttribute(field, '@default');
+    if (!defaultAttr) {
+        return false;
+    }
+
+    const arg = defaultAttr.args[0]?.value;
+    if (!arg) {
+        return false;
+    }
+
+    return isInvocationExpr(arg) && arg.function.$refText === 'autoincrement';
+}
diff --git a/tests/integration/tests/enhancements/with-policy/deep-nested.test.ts b/tests/integration/tests/enhancements/with-policy/deep-nested.test.ts
diff --git a/tests/integration/tests/regression/issue-961.test.ts b/tests/integration/tests/regression/issue-961.test.ts
