feat(policy): add optional error code to @@allow/@@deny attributes

packages/language/src/validators/attribute-application-validator.ts

@@ -36,6 +36,7 @@ import {
     isComputedField,
     isDataFieldReference,
     isDelegateModel,
+    isEnumFieldReference,
     isRelationshipField,
     mapBuiltinTypeToExpressionType,
     resolved,
@@ -187,6 +188,8 @@ export default class AttributeApplicationValidator implements AstValidator<Attri
                 accept('error', `"before()" is only allowed in "post-update" policy rules`, { node: beforeCall });
             }
         }
+
+        this.validateCustomErrorCode(attr.args[2], accept);
     }
 
     private rejectNonOwnedRelationInExpression(expr: Expression, accept: ValidationAcceptor) {
@@ -280,6 +283,25 @@ export default class AttributeApplicationValidator implements AstValidator<Attri
         }
     }
 
+    private validateCustomErrorCode(codeArg: AttributeArg | undefined, accept: ValidationAcceptor) {
+        if (codeArg === undefined) return;
+
+        if (isEnumFieldReference(codeArg.value)) {
+            // enum field references are always valid as error codes
+            return;
+        }
+
+        const codeValue = getStringLiteral(codeArg.value);
+        if (codeValue === undefined) {
+            accept('error', 'Custom error code must be a string literal or an enum value', { node: codeArg });
+            return;
+        }
+        if (codeValue.trim().length === 0) {
+            accept('error', 'Custom error code cannot be empty', { node: codeArg });
+            return;
+        }
+    }
+
     @check('@@validate')
     private _checkValidate(attr: AttributeApplication, accept: ValidationAcceptor) {
         const condition = attr.args[0]?.value;

packages/orm/src/client/errors.ts

@@ -92,6 +92,15 @@ export class ORMError extends Error {
      */
     public rejectedByPolicyReason?: RejectedByPolicyReason;
 
+    /**
+     * Custom error codes from every policy rule that contributed to this rejection.
+     * Set via the optional third argument of `@@allow` / `@@deny`. Only available when
+     * `reason` is `REJECTED_BY_POLICY` and at least one matching rule carries a code.
+     * Note: only surfaced for `create` and `post-update` violations; `update`, `delete`,
+     * and `read` use filter-based enforcement and do not throw policy errors.
+     */
+    public policyCodes?: string[];
+
     /**
      * The SQL query that was executed. Only available when `reason` is `DB_QUERY_ERROR`.
      */

packages/plugins/policy/package.json

@@ -48,13 +48,15 @@
     "dependencies": {
         "@zenstackhq/common-helpers": "workspace:*",
         "@zenstackhq/orm": "workspace:*",
-        "ts-pattern": "catalog:"
+        "ts-pattern": "catalog:",
+        "zod": "catalog:"
     },
     "peerDependencies": {
         "kysely": "catalog:"
     },
     "devDependencies": {
         "@types/better-sqlite3": "catalog:",
+        "@types/node": "catalog:",
         "@types/pg": "^8.0.0",
         "@zenstackhq/eslint-config": "workspace:*",
         "@zenstackhq/tsdown-config": "workspace:*",

packages/plugins/policy/plugin.zmodel

@@ -3,8 +3,11 @@
  *
  * @param operation: comma-separated list of "create", "read", "update", "post-update", "delete". Use "all" to denote all operations.
  * @param condition: a boolean expression that controls if the operation should be allowed.
+ * @param errorCode: an optional code attached to the error thrown when this policy rejects an operation.
+ *   Accepts a string literal or an enum value. Only surfaced for "create" and "post-update" violations
+ *   (other operations use filter-based enforcement).
  */
-attribute @@allow(_ operation: String @@@completionHint(["'create'", "'read'", "'update'", "'post-update'","'delete'", "'all'"]), _ condition: Boolean)
+attribute @@allow(_ operation: String @@@completionHint(["'create'", "'read'", "'update'", "'post-update'","'delete'", "'all'"]), _ condition: Boolean, _ errorCode: Any?)
 
 /**
  * Defines an access policy that allows the annotated field to be read or updated.
@@ -20,8 +23,11 @@ attribute @allow(_ operation: String @@@completionHint(["'read'", "'update'", "'
  *
  * @param operation: comma-separated list of "create", "read", "update", "post-update", "delete". Use "all" to denote all operations.
  * @param condition: a boolean expression that controls if the operation should be denied.
+ * @param errorCode: an optional code attached to the error thrown when this policy rejects an operation.
+ *   Accepts a string literal or an enum value. Only surfaced for "create" and "post-update" violations
+ *   (other operations use filter-based enforcement).
  */
-attribute @@deny(_ operation: String @@@completionHint(["'create'", "'read'", "'update'", "'post-update'","'delete'", "'all'"]), _ condition: Boolean)
+attribute @@deny(_ operation: String @@@completionHint(["'create'", "'read'", "'update'", "'post-update'","'delete'", "'all'"]), _ condition: Boolean, _ errorCode: Any?)
 
 /**
  * Defines an access policy that denies the annotated field to be read or updated.

packages/plugins/policy/src/options.ts

@@ -5,4 +5,11 @@ export type PolicyPluginOptions = {
      * not inspect or reject them.
      */
     dangerouslyAllowRawSql?: boolean;
+
+    /**
+     * Whether to run the diagnostic query to determine which policy rule was violated when
+     * a write is rejected. Defaults to `true`. Set to `false` to skip it globally for
+     * performance. Can be overridden per-query with the `fetchPolicyCodes` option.
+     */
+    fetchPolicyCodes?: boolean;
 };

packages/plugins/policy/src/plugin.ts

@@ -1,12 +1,25 @@
+import { AsyncLocalStorage } from 'node:async_hooks';
 import { type OnKyselyQueryArgs, type RuntimePlugin } from '@zenstackhq/orm';
 import type { SchemaDef } from '@zenstackhq/orm/schema';
-import type { PolicyPluginOptions } from './options';
+import { z } from 'zod';
 import { check } from './functions';
+import type { PolicyPluginOptions } from './options';
 import { PolicyHandler } from './policy-handler';
 
 export type { PolicyPluginOptions } from './options';
 
-export class PolicyPlugin implements RuntimePlugin<SchemaDef, {}, {}, {}> {
+type PolicyQueryContext = Pick<PolicyPluginOptions, 'fetchPolicyCodes'>;
+
+const policyContextStorage = new AsyncLocalStorage<PolicyQueryContext>();
+
+type PolicyExtQueryArgs = {
+    $create: Pick<PolicyPluginOptions, 'fetchPolicyCodes'>;
+    $update: Pick<PolicyPluginOptions, 'fetchPolicyCodes'>;
+};
+
+const fetchPolicyCodesSchema = z.object({ fetchPolicyCodes: z.boolean().optional() });
+
+export class PolicyPlugin implements RuntimePlugin<SchemaDef, PolicyExtQueryArgs, {}, {}> {
     constructor(private readonly options: PolicyPluginOptions = {}) {}
 
     get id() {
@@ -27,8 +40,32 @@ export class PolicyPlugin implements RuntimePlugin<SchemaDef, {}, {}, {}> {
         };
     }
 
+    readonly queryArgs = {
+        $create: fetchPolicyCodesSchema,
+        $update: fetchPolicyCodesSchema,
+    };
+
+    // onQuery and onKyselyQuery are decoupled hook call sites with no shared argument path;
+    // AsyncLocalStorage bridges the per-query fetchPolicyCodes arg into the Kysely executor.
+    onQuery(ctx: {
+        args: Record<string, unknown> | undefined;
+        proceed: (args: Record<string, unknown> | undefined) => Promise<unknown>;
+        [key: string]: unknown;
+    }) {
+        const fetchPolicyCodes = ctx.args?.['fetchPolicyCodes'] as boolean | undefined;
+        if (fetchPolicyCodes !== undefined) {
+            return policyContextStorage.run({ fetchPolicyCodes }, () => ctx.proceed(ctx.args));
+        }
+        return ctx.proceed(ctx.args);
+    }
+
     onKyselyQuery({ query, client, proceed }: OnKyselyQueryArgs<SchemaDef>) {
-        const handler = new PolicyHandler<SchemaDef>(client, this.options);
+        const ctx = policyContextStorage.getStore();
+        const effectiveOptions: PolicyPluginOptions =
+            ctx?.fetchPolicyCodes !== undefined
+                ? { ...this.options, fetchPolicyCodes: ctx.fetchPolicyCodes }
+                : this.options;
+        const handler = new PolicyHandler<SchemaDef>(client, effectiveOptions);
         return handler.handle(query, proceed);
     }
 }

packages/plugins/policy/src/policy-handler.ts

@@ -175,7 +175,16 @@ export class PolicyHandler<Schema extends SchemaDef> extends OperationNodeTransf
             if (constCondition === true) {
                 needCheckPreCreate = false;
             } else if (constCondition === false) {
-                throw createRejectedByPolicyError(mutationModel, RejectedByPolicyReason.NO_ACCESS);
+                const policies = this.getModelPolicies(mutationModel, 'create');
+                const constantDenyCodes = policies
+                    .filter((p) => p.kind === 'deny' && this.isTrueExpr(p.condition) && p.code)
+                    .map((p) => p.code!);
+                throw createRejectedByPolicyError(
+                    mutationModel,
+                    RejectedByPolicyReason.NO_ACCESS,
+                    undefined,
+                    constantDenyCodes,
+                );
             }
         }
 
@@ -331,10 +340,15 @@ export class PolicyHandler<Schema extends SchemaDef> extends OperationNodeTransf
 
         const postUpdateResult = await proceed(postUpdateQuery.toOperationNode());
         if (!postUpdateResult.rows[0]?.$condition) {
+            const policyCodes =
+                this.options.fetchPolicyCodes !== false
+                    ? await this.findViolatingPostUpdatePolicyCodes(model, idConditions, beforeUpdateInfo, proceed)
+                    : undefined;
             throw createRejectedByPolicyError(
                 model,
                 RejectedByPolicyReason.NO_ACCESS,
                 'some or all updated rows failed to pass post-update policy check',
+                policyCodes,
             );
         }
     }
@@ -950,7 +964,11 @@ export class PolicyHandler<Schema extends SchemaDef> extends OperationNodeTransf
             } satisfies SelectQueryNode,
         );
         if (!result.rows[0]?.$condition) {
-            throw createRejectedByPolicyError(model, RejectedByPolicyReason.NO_ACCESS);
+            const policyCodes =
+                this.options.fetchPolicyCodes !== false
+                    ? await this.findViolatingCreatePolicyCodes(model, valuesTable, proceed)
+                    : undefined;
+            throw createRejectedByPolicyError(model, RejectedByPolicyReason.NO_ACCESS, undefined, policyCodes);
         }
     }
 
@@ -1047,6 +1065,98 @@ export class PolicyHandler<Schema extends SchemaDef> extends OperationNodeTransf
         return ExpressionUtils.isLiteral(expr) && expr.value === true;
     }
 
+    private async findViolatingCreatePolicyCodes(
+        model: string,
+        valuesTable: ReturnType<BaseCrudDialect<Schema>['buildValuesTableSelect']>,
+        proceed: ProceedKyselyQueryFunction,
+    ): Promise<string[]> {
+        const codedPolicies = this.getModelPolicies(model, 'create').filter((p) => p.code);
+        if (codedPolicies.length === 0) {
+            return [];
+        }
+
+        const selections = codedPolicies.map((policy, i) => {
+            const condition = this.compilePolicyCondition(model, undefined, 'create', policy);
+            // For allow rules, negate: EXISTS(NOT condition) = true when any proposed row violates allow.
+            // For deny rules, keep as-is: EXISTS(condition) = true when deny fires.
+            const existsCondition = policy.kind === 'allow' ? logicalNot(this.dialect, condition) : condition;
+            const inner = this.eb
+                .selectFrom(valuesTable.as(model))
+                .select(this.eb.lit(1).as('_'))
+                .where(() => new ExpressionWrapper(existsCondition));
+            return SelectionNode.create(
+                AliasNode.create(this.eb.exists(inner).toOperationNode(), IdentifierNode.create(`$c${i}`)),
+            );
+        });
+
+        return this.evaluatePolicyDiagnostics(codedPolicies, selections, proceed);
+    }
+
+    private async findViolatingPostUpdatePolicyCodes(
+        model: string,
+        idConditions: OperationNode,
+        beforeUpdateInfo: Awaited<ReturnType<typeof this.loadBeforeUpdateEntities>>,
+        proceed: ProceedKyselyQueryFunction,
+    ): Promise<string[]> {
+        const codedPolicies = this.getModelPolicies(model, 'post-update').filter((p) => p.code);
+        if (codedPolicies.length === 0) {
+            return [];
+        }
+
+        const needsBeforeUpdateJoin = !!beforeUpdateInfo?.fields;
+        let beforeUpdateTable: SelectQueryNode | undefined;
+        if (needsBeforeUpdateJoin) {
+            const fieldDefs = beforeUpdateInfo!.fields!.map((name) =>
+                QueryUtils.requireField(this.client.$schema, model, name),
+            );
+            const rows = beforeUpdateInfo!.rows.map((r) => beforeUpdateInfo!.fields!.map((f) => r[f]));
+            beforeUpdateTable = this.dialect.buildValuesTableSelect(fieldDefs, rows).toOperationNode();
+        }
+
+        const eb = expressionBuilder<any, any>();
+
+        const buildInnerExists = (condition: OperationNode) => {
+            const inner = eb
+                .selectFrom(model)
+                .select(eb.lit(1).as('_'))
+                .where(() => new ExpressionWrapper(conjunction(this.dialect, [idConditions, condition])))
+                .$if(needsBeforeUpdateJoin, (qb) =>
+                    qb.leftJoin(
+                        () => new ExpressionWrapper(beforeUpdateTable!).as('$before'),
+                        (join) => {
+                            const idFields = QueryUtils.requireIdFields(this.client.$schema, model);
+                            return idFields.reduce((acc, f) => acc.onRef(`${model}.${f}`, '=', `$before.${f}`), join);
+                        },
+                    ),
+                );
+            return eb.exists(inner).toOperationNode();
+        };
+
+        const selections = codedPolicies.map((policy, i) => {
+            const condition = this.compilePolicyCondition(model, undefined, 'post-update', policy);
+            // For allow rules, negate: EXISTS(NOT condition) = true when any updated row violates allow.
+            // For deny rules, keep as-is: EXISTS(condition) = true when deny fires.
+            const existsCondition = policy.kind === 'allow' ? logicalNot(this.dialect, condition) : condition;
+            return SelectionNode.create(
+                AliasNode.create(buildInnerExists(existsCondition), IdentifierNode.create(`$c${i}`)),
+            );
+        });
+
+        return this.evaluatePolicyDiagnostics(codedPolicies, selections, proceed);
+    }
+
+    // Single diagnostic query: one EXISTS column per coded policy.
+    // EXISTS=true means a violation: deny condition fired, or allow condition wasn't met (negated in caller).
+    private async evaluatePolicyDiagnostics(
+        codedPolicies: Policy[],
+        selections: SelectionNode[],
+        proceed: ProceedKyselyQueryFunction,
+    ): Promise<string[]> {
+        const result = await proceed({ kind: 'SelectQueryNode', selections } satisfies SelectQueryNode);
+        const row = result.rows[0] ?? {};
+        return codedPolicies.filter((_, i) => row[`$c${i}`]).map((p) => p.code!);
+    }
+
     private async processReadBack(node: CrudQueryNode, result: QueryResult<any>, proceed: ProceedKyselyQueryFunction) {
         if (result.rows.length === 0) {
             return result;
@@ -1240,14 +1350,18 @@ export class PolicyHandler<Schema extends SchemaDef> extends OperationNodeTransf
             result.push(
                 ...modelDef.attributes
                     .filter((attr) => attr.name === '@@allow' || attr.name === '@@deny')
-                    .map(
-                        (attr) =>
-                            ({
-                                kind: attr.name === '@@allow' ? 'allow' : 'deny',
-                                operations: extractOperations(attr.args![0]!.value),
-                                condition: attr.args![1]!.value,
-                            }) as const,
-                    )
+                    .map((attr) => {
+                        const codeExpr = attr.args?.[2]?.value;
+                        return {
+                            kind: attr.name === '@@allow' ? 'allow' : 'deny',
+                            operations: extractOperations(attr.args![0]!.value),
+                            condition: attr.args![1]!.value,
+                            code:
+                                ExpressionUtils.isLiteral(codeExpr) && typeof codeExpr.value === 'string'
+                                    ? codeExpr.value
+                                    : undefined,
+                        } as const;
+                    })
                     .filter(
                         (policy) =>
                             (operation !== 'post-update' && policy.operations.includes('all')) ||
@@ -1275,14 +1389,18 @@ export class PolicyHandler<Schema extends SchemaDef> extends OperationNodeTransf
             result.push(
                 ...fieldDef.attributes
                     .filter((attr) => attr.name === '@allow' || attr.name === '@deny')
-                    .map(
-                        (attr) =>
-                            ({
-                                kind: attr.name === '@allow' ? 'allow' : 'deny',
-                                operations: extractOperations(attr.args![0]!.value),
-                                condition: attr.args![1]!.value,
-                            }) as const,
-                    )
+                    .map((attr) => {
+                        const codeExpr = attr.args?.[2]?.value;
+                        return {
+                            kind: attr.name === '@allow' ? 'allow' : 'deny',
+                            operations: extractOperations(attr.args![0]!.value),
+                            condition: attr.args![1]!.value,
+                            code:
+                                ExpressionUtils.isLiteral(codeExpr) && typeof codeExpr.value === 'string'
+                                    ? codeExpr.value
+                                    : undefined,
+                        } as const;
+                    })
                     .filter((policy) => policy.operations.includes('all') || policy.operations.includes(operation)),
             );
         }

packages/plugins/policy/src/types.ts

@@ -18,6 +18,7 @@ export type Policy = {
     kind: PolicyKind;
     operations: readonly PolicyOperation[];
     condition: Expression;
+    code?: string;
 };
 
 /**

packages/plugins/policy/src/utils.ts

@@ -181,10 +181,12 @@ export function createRejectedByPolicyError(
     model: string | undefined,
     reason: RejectedByPolicyReason,
     message?: string,
+    policyCodes?: string[],
 ) {
     const err = new ORMError(ORMErrorReason.REJECTED_BY_POLICY, message ?? 'operation is rejected by access policies');
     err.rejectedByPolicyReason = reason;
     err.model = model;
+    err.policyCodes = policyCodes?.length ? policyCodes : undefined;
     return err;
 }