refactor(policy): replace postModelLevelCheck with postMutationZeroRowsCheck

Azzerty23 · claude · Azzerty23 · commit 0ffadd402530 · 2026-05-02T12:41:08.000+02:00
Rename and tighten the zero-rows diagnostic helper for update/delete:
- Skip the diagnostic query entirely when no policies carry an error code,
  removing an unnecessary round-trip for the common case.
- Fix BigInt comparison: use `&gt; 0` negation instead of `=== 0` so
  numAffectedRows is handled correctly across drivers.
- Reorder delete/update branches so the more specific guard (delete has
  no `using` restriction) runs first.
- Inline the fetchPolicyCodes guard into a single early-exit path.

Tests: add "does not throw for nonexistent row" cases for update and delete,
and a focused test asserting the opt-in behaviour — rules without an errorCode
keep the existing NotFound error, rules with an errorCode surface
RejectedByPolicy instead.

Co-Authored-By: Claude Sonnet 4.6 &lt;noreply@anthropic.com&gt;
diff --git a/packages/plugins/policy/src/policy-handler.ts b/packages/plugins/policy/src/policy-handler.ts
@@ -138,12 +138,13 @@ export class PolicyHandler<Schema extends SchemaDef> extends OperationNodeTransf
 
         // #region Post mutation work
 
-        // post-check: detect policy denial when 0 rows affected (replaces pre-check for update/delete)
-        if (!(result.numAffectedRows ?? 0n)) {
-            if (UpdateQueryNode.is(node)) {
-                await this.postModelLevelCheck(mutationModel, 'update', node.where?.where ?? trueNode(this.dialect), proceed);
-            } else if (DeleteQueryNode.is(node) && !node.using) {
-                await this.postModelLevelCheck(mutationModel, 'delete', node.where?.where ?? trueNode(this.dialect), proceed);
+        // When 0 rows affected, distinguish "row not found" from "row denied by policy"
+        // Use > 0 negation (not === 0) because numAffectedRows is BigInt in some drivers
+        if (!((result.numAffectedRows ?? 0) > 0)) {
+            if (DeleteQueryNode.is(node)) {
+                await this.postMutationZeroRowsCheck(mutationModel, 'delete', node.where?.where, proceed);
+            } else if (UpdateQueryNode.is(node)) {
+                await this.postMutationZeroRowsCheck(mutationModel, 'update', node.where?.where, proceed);
             }
         }
 
@@ -258,62 +259,63 @@ export class PolicyHandler<Schema extends SchemaDef> extends OperationNodeTransf
         }
     }
 
-    // Post-check for model-level update/delete policy enforcement.
-    // Runs only when 0 rows were affected: distinguishes "row not found" from "row denied by policy".
-    // Combines existence check and diagnostic into a single query.
-    private async postModelLevelCheck(
-        mutationModel: string,
+    // Checks if any row matching the original WHERE exists without the policy filter.
+    // Called when numAffectedRows == 0 for UPDATE or DELETE.
+    // If a row exists but was filtered by policy → throws REJECTED_BY_POLICY with codes.
+    // If no row matches → returns silently (ORM layer handles "not found").
+    // Combines existence check and code diagnostics into a single query.
+    private async postMutationZeroRowsCheck(
+        model: string,
         operation: 'update' | 'delete',
-        userWhere: OperationNode,
+        originalWhere: OperationNode | undefined,
         proceed: ProceedKyselyQueryFunction,
     ) {
-        const modelLevelFilter = this.buildPolicyFilter(mutationModel, undefined, operation);
-        if (isTrueNode(modelLevelFilter)) {
-            return;
-        }
+        if (this.isManyToManyJoinTable(model)) return;
+        if (this.tryGetConstantPolicy(model, operation) === true) return;
+        const codedPolicies = this.getModelPolicies(model, operation).filter((p) => p.code);
+        // Skip if no policies carry an error code — nothing to surface.
+        if (codedPolicies.length === 0) return;
 
-        const codedPolicies =
-            this.options.fetchPolicyCodes !== false
-                ? this.getModelPolicies(mutationModel, operation).filter((p) => p.code)
-                : [];
+        const whereCondition = originalWhere ?? trueNode(this.dialect);
 
-        // $exists: does the row exist at all (without policy filter)?
-        const existsInner = this.eb
-            .selectFrom(mutationModel)
+        const rowExistsInner = this.eb
+            .selectFrom(model)
             .select(this.eb.lit(1).as('_'))
-            .where(() => new ExpressionWrapper(userWhere));
+            .where(() => new ExpressionWrapper(whereCondition));
 
-        const selections: SelectionNode[] = [
-            SelectionNode.create(
-                AliasNode.create(this.eb.exists(existsInner).toOperationNode(), IdentifierNode.create('$exists')),
-            ),
-        ];
+        const fetchCodes = this.options.fetchPolicyCodes !== false;
+        const selectedPolicies = fetchCodes ? codedPolicies : [];
 
-        // one EXISTS column per coded policy, folded into the same query
-        for (const [i, policy] of codedPolicies.entries()) {
-            const condition = this.compilePolicyCondition(mutationModel, undefined, operation, policy);
-            const existsCondition = policy.kind === 'allow' ? logicalNot(this.dialect, condition) : condition;
+        const codeSelections = selectedPolicies.map((policy, i) => {
+            const condition = this.compilePolicyCondition(model, undefined, operation, policy);
+            const violationCondition = policy.kind === 'allow' ? logicalNot(this.dialect, condition) : condition;
             const inner = this.eb
-                .selectFrom(mutationModel)
+                .selectFrom(model)
                 .select(this.eb.lit(1).as('_'))
-                .where(() => new ExpressionWrapper(conjunction(this.dialect, [userWhere, existsCondition])));
-            selections.push(
+                .where(() => new ExpressionWrapper(conjunction(this.dialect, [whereCondition, violationCondition])));
+            return SelectionNode.create(
+                AliasNode.create(this.eb.exists(inner).toOperationNode(), IdentifierNode.create(`$c${i}`)),
+            );
+        });
+
+        const result = await proceed({
+            kind: 'SelectQueryNode',
+            selections: [
                 SelectionNode.create(
-                    AliasNode.create(this.eb.exists(inner).toOperationNode(), IdentifierNode.create(`$c${i}`)),
+                    AliasNode.create(this.eb.exists(rowExistsInner).toOperationNode(), IdentifierNode.create('$exists')),
                 ),
-            );
-        }
+                ...codeSelections,
+            ],
+        } satisfies SelectQueryNode);
 
-        const checkResult = await proceed({ kind: 'SelectQueryNode', selections } satisfies SelectQueryNode);
-        const row = checkResult.rows[0] ?? {};
+        const row = result.rows[0] ?? {};
+        if (!row.$exists) return;
 
-        if (row.$exists) {
-            const policyCodes =
-                this.options.fetchPolicyCodes !== false
-                    ? codedPolicies.filter((_, i) => row[`$c${i}`]).map((p) => p.code!)
-                    : undefined;
-            throw createRejectedByPolicyError(mutationModel, RejectedByPolicyReason.NO_ACCESS, undefined, policyCodes);
-        }
+        const policyCodes = fetchCodes
+            ? selectedPolicies.filter((_, i) => row[`$c${i}`]).map((p) => p.code!)
+            : undefined;
+
+        throw createRejectedByPolicyError(model, RejectedByPolicyReason.NO_ACCESS, undefined, policyCodes);
     }
 
     private async postUpdateCheck(
diff --git a/tests/e2e/orm/policy/crud/delete.test.ts b/tests/e2e/orm/policy/crud/delete.test.ts
@@ -48,4 +48,24 @@ model Foo {
         await expect(db.$qb.deleteFrom('Foo').executeTakeFirst()).resolves.toMatchObject({ numDeletedRows: 1n });
         await expect(db.foo.count()).resolves.toBe(1);
     });
+
+    it('does not throw for nonexistent row', async () => {
+        const db = await createPolicyTestClient(
+            `
+model Foo {
+    id Int @id
+    x  Int
+    @@allow('create,read', true)
+    @@allow('delete', x > 0)
+}
+`,
+        );
+        await db.foo.create({ data: { id: 1, x: 1 } });
+
+        // nonexistent row — row does not exist at all, so postModelLevelCheck must NOT throw
+        await expect(db.$qb.deleteFrom('Foo').where('id', '=', 999).executeTakeFirst()).resolves.toMatchObject({
+            numDeletedRows: 0n,
+        });
+        await expect(db.foo.count()).resolves.toBe(1);
+    });
 });
diff --git a/tests/e2e/orm/policy/crud/error-codes.test.ts b/tests/e2e/orm/policy/crud/error-codes.test.ts
@@ -10,6 +10,7 @@ describe('Policy error code tests', () => {
     // │ allow rule fails (string code)          │   ✓    │      ✓      │   ✓    │   ✓    │
     // │ constant deny (true condition)          │   ✓    │             │        │        │
     // │ no errorCode on rule                    │   ✓    │             │        │        │
+    // │ code opt-in: NotFound vs RejectedByPol. │        │             │   ✓    │   ✓    │
     // │ multiple deny rules fire                │   ✓    │      ✓      │        │        │
     // │ deny + allow conflict                   │   ✓    │      ✓      │        │        │
     // │ multiple allow rules all fail           │   ✓    │      ✓      │        │        │
@@ -72,6 +73,33 @@ describe('Policy error code tests', () => {
         await expect(db.foo.create({ data: { x: 0 } })).toBeRejectedByPolicy(undefined, []);
     });
 
+    // ── opt-in: adding a code changes error type from NotFound to RejectedByPolicy ──
+
+    it('blocked update/delete yields NotFound without code, RejectedByPolicy with code', async () => {
+        const schema = (withCode: boolean) => `
+model Foo {
+    id Int @id
+    x  Int
+    @@allow('create,read', true)
+    @@allow('update', x > 0${withCode ? ", 'NEED_POSITIVE_X'" : ''})
+    @@allow('delete', x > 0${withCode ? ", 'NEED_POSITIVE_X'" : ''})
+}
+`;
+        // Without error code: the ORM sees 0 affected rows and raises NotFound
+        const dbNoCode = await createPolicyTestClient(schema(false));
+        await dbNoCode.foo.create({ data: { id: 1, x: 0 } });
+        await expect(dbNoCode.foo.update({ where: { id: 1 }, data: { x: -1 } })).toBeRejectedNotFound();
+        await expect(dbNoCode.foo.delete({ where: { id: 1 } })).toBeRejectedNotFound();
+
+        // With error code: the plugin detects the policy block and raises RejectedByPolicy
+        const dbWithCode = await createPolicyTestClient(schema(true));
+        await dbWithCode.foo.create({ data: { id: 1, x: 0 } });
+        await expect(dbWithCode.foo.update({ where: { id: 1 }, data: { x: -1 } })).toBeRejectedByPolicy(undefined, [
+            'NEED_POSITIVE_X',
+        ]);
+        await expect(dbWithCode.foo.delete({ where: { id: 1 } })).toBeRejectedByPolicy(undefined, ['NEED_POSITIVE_X']);
+    });
+
     // ── post-update: single rule, single code ─────────────────────────────────
 
     it('surfaces code from deny/allow rule on post-update violation', async () => {
diff --git a/tests/e2e/orm/policy/crud/update.test.ts b/tests/e2e/orm/policy/crud/update.test.ts
@@ -1208,6 +1208,27 @@ model Foo {
             await expect(db.foo.findUnique({ where: { id: 1 } })).resolves.toMatchObject({ x: 1 });
         });
 
+        it('does not throw for nonexistent row', async () => {
+            const db = await createPolicyTestClient(
+                `
+model Foo {
+    id Int @id
+    x  Int
+    @@allow('create', true)
+    @@allow('update', x > 1)
+    @@allow('read', true)
+}
+`,
+            );
+
+            await db.foo.createMany({ data: [{ id: 1, x: 2 }] });
+
+            // nonexistent row — row does not exist at all, so postModelLevelCheck must NOT throw
+            await expect(
+                db.$qb.updateTable('Foo').set({ x: 5 }).where('id', '=', 999).executeTakeFirst(),
+            ).resolves.toMatchObject({ numUpdatedRows: 0n });
+        });
+
         it('works with insert on conflict do update', async () => {
             const db = await createPolicyTestClient(
                 `
