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

.github/workflows/build-test.yml

@@ -80,9 +80,9 @@ jobs:
               uses: actions/cache@v4
               with:
                   path: ${{ steps.pnpm-cache.outputs.STORE_PATH }}
-                  key: ${{ runner.os }}-pnpm-store-${{ hashFiles('**/pnpm-lock.yaml') }}
+                  key: ${{ runner.os }}-node-${{ matrix.node-version }}-pnpm-store-${{ hashFiles('**/pnpm-lock.yaml') }}
                   restore-keys: |
-                      ${{ runner.os }}-pnpm-store-
+                      ${{ runner.os }}-node-${{ matrix.node-version }}-pnpm-store-
 
             - name: Install dependencies
               run: pnpm install --frozen-lockfile

packages/cli/src/actions/proxy.ts

@@ -157,7 +157,7 @@ export async function run(options: Options) {
     }
 
     // If a studioAuthKey is provided, create an authDb with the policy plugin
-    const authDb = db.$use(new PolicyPlugin()) as ClientContract<SchemaDef>;
+    const authDb: ClientContract<SchemaDef> = db.$use(new PolicyPlugin());
     if (options.studioAuthKey) {
         console.log(colors.gray('Access policy plugin enabled for authorization.'));
     }
@@ -423,7 +423,7 @@ function resolveClient(
         // SuperUser has full access without policy enforcement, so we return the base client directly.
         return client;
     } else {
-        return authDb.$setAuth(claim.data) as ClientContract<SchemaDef>;
+        return authDb.$setAuth(claim.data);
     }
 }
 

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

@@ -37,6 +37,7 @@ import {
     isComputedField,
     isDataFieldReference,
     isDelegateModel,
+    isEnumFieldReference,
     isRelationshipField,
     mapBuiltinTypeToExpressionType,
     resolved,
@@ -188,6 +189,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) {
@@ -281,6 +284,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/client-impl.ts

@@ -41,6 +41,7 @@ import * as BuiltinFunctions from './functions';
 import { SchemaDbPusher } from './helpers/schema-db-pusher';
 import type { ClientOptions, ProceduresOptions } from './options';
 import type { AnyPlugin } from './plugin';
+import { extractPluginQueryArgs } from './plugin-utils';
 import { createZenStackPromise, type ZenStackPromise } from './promise';
 import { fieldHasDefaultValue, getField, isUnsupportedField, requireModel } from './query-utils';
 import { ResultProcessor } from './result-processor';
@@ -627,7 +628,13 @@ function createModelCrudHandler(
                     ? prepareArgsForExtResult(_args, model, schema, plugins)
                     : _args;
 
-                const _handler = txClient ? handler.withClient(txClient) : handler;
+                const txHandler = txClient ? handler.withClient(txClient) : handler;
+                // make the top-level call context available for embedding into generated queries,
+                // using the final args (after any `onQuery` hook overrides)
+                const _handler = txHandler.withCallContext({
+                    ormOperation: nominalOperation,
+                    pluginArgs: extractPluginQueryArgs(client.$options.plugins, operation, _args),
+                });
                 const r = await _handler.handle(operation, processedArgs);
                 if (!r && throwIfNoResult) {
                     throw createNotFoundError(model);

packages/orm/src/client/constants.ts

@@ -1,7 +1,7 @@
-// /**
-//  * The comment prefix for annotation generated Kysely queries with context information.
-//  */
-// export const CONTEXT_COMMENT_PREFIX = '-- $$context:';
+/**
+ * The comment prefix for annotating generated Kysely queries with context information.
+ */
+export const CONTEXT_COMMENT_PREFIX = '-- $$context:';
 
 /**
  * The types of fields that are numeric.

packages/orm/src/client/crud/dialects/base-dialect.ts

@@ -81,6 +81,20 @@ export abstract class BaseCrudDialect<Schema extends SchemaDef> {
      */
     abstract get insertIgnoreMethod(): 'onConflict' | 'ignore';
 
+    /**
+     * Whether the pre-load SELECT (used to resolve entity IDs before a top-level UPDATE on
+     * non-RETURNING dialects) must bypass the read-policy filter.
+     *
+     * MySQL pre-loads entity IDs before running an UPDATE. If the row is read-denied the
+     * pre-load returns null and the UPDATE never runs, masking update-deny error codes.
+     * Setting this to true tags the pre-load with a `bypassReadPolicy` query context so
+     * the policy plugin skips its read filter for it. Other `onKyselyQuery` interceptors
+     * (e.g. soft-delete) still apply.
+     */
+    get requiresUpdatePreloadBypassReadPolicy(): boolean {
+        return false;
+    }
+
     // #endregion
 
     // #region value transformation

packages/orm/src/client/crud/dialects/mysql.ts

@@ -58,6 +58,10 @@ export class MySqlCrudDialect<Schema extends SchemaDef> extends LateralJoinDiale
         return 'ignore' as const;
     }
 
+    override get requiresUpdatePreloadBypassReadPolicy(): boolean {
+        return true;
+    }
+
     // #endregion
 
     // #region value transformation

packages/orm/src/client/crud/operations/base.ts

@@ -19,7 +19,7 @@ import { ulid } from 'ulid';
 import * as uuid from 'uuid';
 import type { AnyKysely } from '../../../utils/kysely-utils';
 import { extractFields, fieldsToSelectObject, isEmptyObject } from '../../../utils/object-utils';
-import { NUMERIC_FIELD_TYPES } from '../../constants';
+import { CONTEXT_COMMENT_PREFIX, NUMERIC_FIELD_TYPES } from '../../constants';
 import { TransactionIsolationLevel, type ClientContract, type CRUD } from '../../contract';
 import type { FindArgs, SelectIncludeOmit, WhereInput } from '../../crud-types';
 import {
@@ -31,6 +31,7 @@ import {
     ORMError,
     ORMErrorReason,
 } from '../../errors';
+import type { QueryContext } from '../../plugin';
 import type { ToKysely } from '../../query-builder';
 import {
     ensureArray,
@@ -169,6 +170,16 @@ export const AllReadOperations = [...CoreReadOperations, 'findUniqueOrThrow', 'f
  */
 export type AllReadOperations = (typeof AllReadOperations)[number];
 
+/**
+ * List of single-row read operations that throw when no row is found.
+ */
+export const SingleRowOrThrowOperations = ['findUniqueOrThrow', 'findFirstOrThrow'] as const;
+
+/**
+ * List of single-row read operations that throw when no row is found.
+ */
+export type SingleRowOrThrowOperations = (typeof SingleRowOrThrowOperations)[number];
+
 /**
  * List of all write operations - simply an alias of CoreWriteOperations.
  */
@@ -216,6 +227,23 @@ export abstract class BaseOperationHandler<Schema extends SchemaDef> {
 
     abstract handle(operation: CoreCrudOperations, args: any): Promise<unknown>;
 
+    /**
+     * Context of the top-level ORM API call, embedded into generated queries as a
+     * trailing SQL comment (see {@link makeContextComment}).
+     */
+    protected callContext?: { ormOperation: AllCrudOperations; pluginArgs?: Record<string, unknown> };
+
+    /**
+     * Returns a clone carrying the given call context, so concurrent invocations
+     * (e.g. an `onQuery` hook calling `proceed` multiple times in parallel) each
+     * get their own context.
+     */
+    withCallContext(context: { ormOperation: AllCrudOperations; pluginArgs?: Record<string, unknown> }) {
+        const clone = this.withClient(this.client);
+        clone.callContext = context;
+        return clone;
+    }
+
     withClient(client: ClientContract<Schema>) {
         return new (this.constructor as new (...args: any[]) => this)(client, this.model, this.inputValidator);
     }
@@ -313,6 +341,9 @@ export abstract class BaseOperationHandler<Schema extends SchemaDef> {
             const r = await kysely.getExecutor().executeQuery(compiled);
             result = r.rows;
         } catch (err) {
+            // Re-throw ORMErrors (e.g. policy violations with custom error codes) as-is
+            // to avoid wrapping them in a generic DBQueryError and losing their type/code.
+            if (err instanceof ORMError) throw err;
             throw createDBQueryError(`Failed to execute query: ${err}`, err, compiled.sql, compiled.parameters);
         }
 
@@ -1186,11 +1217,22 @@ export abstract class BaseOperationHandler<Schema extends SchemaDef> {
             }
         }
 
+        // For non-RETURNING dialects that require it (e.g. MySQL), the pre-load SELECT must
+        // bypass the read policy so that read-denied rows are still reachable and the UPDATE
+        // can run, allowing its own policy error codes to be surfaced.
+        const bypassReadPolicyForPreload =
+            !this.dialect.supportsReturning && !fromRelation && this.dialect.requiresUpdatePreloadBypassReadPolicy;
+
         // lazily load the entity to be updated
         let thisEntity: any;
         const loadThisEntity = async () => {
             if (thisEntity === undefined) {
-                thisEntity = (await this.getEntityIds(kysely, model, origWhere)) ?? null;
+                thisEntity = bypassReadPolicyForPreload
+                    ? await this.readUniqueBypassingReadPolicy(kysely, model, {
+                          where: origWhere,
+                          select: this.makeIdSelect(model),
+                      } as any)
+                    : ((await this.getEntityIds(kysely, model, origWhere)) ?? null);
                 if (!thisEntity && throwIfNotFound) {
                     throw createNotFoundError(model);
                 }
@@ -1545,9 +1587,19 @@ export abstract class BaseOperationHandler<Schema extends SchemaDef> {
         return NUMERIC_FIELD_TYPES.includes(fieldDef.type) && !fieldDef.array;
     }
 
-    private makeContextComment(_context: { model: string; operation: CRUD }) {
-        return sql``;
-        // return sql.raw(`${CONTEXT_COMMENT_PREFIX}${JSON.stringify(context)}`);
+    private makeContextComment(context: { model: string; operation: CRUD; bypassReadPolicy?: boolean }) {
+        if (!this.options.plugins?.some((plugin) => plugin.onKyselyQuery)) {
+            // the context is only consumed by `onKyselyQuery` hooks, skip the overhead otherwise
+            return sql``;
+        }
+        const payload: QueryContext = {
+            ...context,
+            ...(this.callContext?.ormOperation ? { ormOperation: this.callContext.ormOperation } : {}),
+            ...(this.callContext?.pluginArgs && Object.keys(this.callContext.pluginArgs).length > 0
+                ? { pluginArgs: this.callContext.pluginArgs }
+                : {}),
+        };
+        return sql.raw(`${CONTEXT_COMMENT_PREFIX}${JSON.stringify(payload)}`);
     }
 
     protected async updateMany<
@@ -2528,6 +2580,34 @@ export abstract class BaseOperationHandler<Schema extends SchemaDef> {
         });
     }
 
+    // Like readUnique but tags the query with a `bypassReadPolicy` context so the policy
+    // plugin skips its read filter. Used for the MySQL update pre-load so read-denied rows
+    // are still reachable and the UPDATE's own policy error codes surface. Other query
+    // interceptors (e.g. soft-delete) still apply.
+    private async readUniqueBypassingReadPolicy(
+        kysely: AnyKysely,
+        model: string,
+        args: FindArgs<Schema, GetModels<Schema>, any, true>,
+    ): Promise<any | null> {
+        let query = this.dialect.buildSelectModel(model, model);
+        const argsWithTake = { ...args, take: 1 };
+        query = this.dialect.buildFilterSortTake(model, argsWithTake, query, model);
+        if ('select' in args && args.select) {
+            query = this.buildFieldSelection(model, query, args.select, model);
+        } else {
+            query = this.dialect.buildSelectAllFields(model, query, (args as any)?.omit, model);
+        }
+        query = query.modifyEnd(this.makeContextComment({ model, operation: 'read', bypassReadPolicy: true }));
+        const compiled = kysely.getExecutor().compileQuery(query.toOperationNode(), createQueryId());
+        try {
+            const r = await kysely.getExecutor().executeQuery(compiled);
+            return r.rows[0] ?? null;
+        } catch (err) {
+            if (err instanceof ORMError) throw err;
+            throw createDBQueryError(`Failed to execute query: ${err}`, err, compiled.sql, compiled.parameters);
+        }
+    }
+
     // Given multiple unique filters, load all matching entities and return their id fields in one query
     private getEntitiesIds(kysely: AnyKysely, model: string, uniqueFilters: any[]) {
         return this.read(kysely, model, {

packages/orm/src/client/errors.ts

@@ -92,6 +92,16 @@ 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: surfaced for `create`, `post-update`, `update`, `delete`, and single-row `read`
+     * violations. For `read`, only `findFirst`/`findUnique`-equivalent queries (LIMIT 1)
+     * where a denied row exists will throw; `findMany` uses filter-based enforcement.
+     */
+    public policyCodes?: string[];
+
     /**
      * The SQL query that was executed. Only available when `reason` is `DB_QUERY_ERROR`.
      */

packages/orm/src/client/executor/query-context-utils.ts

@@ -0,0 +1,56 @@
+import { RawNode, SelectModifierNode, type OperationNode, type RootOperationNode } from 'kysely';
+import { CONTEXT_COMMENT_PREFIX } from '../constants';
+import type { QueryContext } from '../plugin';
+
+// `modifyEnd` appends a plain `RawNode` on insert/update/delete nodes, but wraps it in a
+// `SelectModifierNode` on select nodes
+function getRawModifier(modifier: OperationNode): RawNode | undefined {
+    if (RawNode.is(modifier)) {
+        return modifier;
+    }
+    if (SelectModifierNode.is(modifier) && modifier.rawModifier && RawNode.is(modifier.rawModifier)) {
+        return modifier.rawModifier;
+    }
+    return undefined;
+}
+
+/**
+ * Extracts the ORM query context embedded as a trailing `-- $$context:` comment
+ * (a `RawNode` in the root node's `endModifiers`), and returns the node with the
+ * comment stripped so it never reaches plugins or the database driver.
+ */
+export function extractQueryContext(node: RootOperationNode): {
+    queryContext?: QueryContext;
+    strippedNode: RootOperationNode;
+} {
+    const endModifiers = (node as { endModifiers?: readonly OperationNode[] }).endModifiers;
+    if (!endModifiers?.length) {
+        return { strippedNode: node };
+    }
+
+    let queryContext: QueryContext | undefined;
+    const remaining = endModifiers.filter((modifier) => {
+        const rawNode = getRawModifier(modifier);
+        if (!rawNode || !rawNode.sqlFragments[0]?.startsWith(CONTEXT_COMMENT_PREFIX)) {
+            return true;
+        }
+        try {
+            queryContext = JSON.parse(rawNode.sqlFragments[0].slice(CONTEXT_COMMENT_PREFIX.length));
+        } catch {
+            // malformed payload: treat as absent
+        }
+        return false;
+    });
+
+    if (remaining.length === endModifiers.length) {
+        return { strippedNode: node };
+    }
+
+    return {
+        queryContext,
+        strippedNode: {
+            ...node,
+            endModifiers: remaining.length > 0 ? remaining : undefined,
+        } as RootOperationNode,
+    };
+}