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/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/client-impl.ts

@@ -127,10 +127,10 @@ export class ClientImpl {
             });
         }
 
-        if (baseClient?.isTransaction && !executor) {
-            // if we're creating a derived client from a transaction client and not replacing
-            // the executor, reuse the current kysely instance to retain the transaction context
-            this.kysely = baseClient.$qb;
+        if (baseClient?.isTransaction) {
+            // preserve transaction context in derived clients: reuse the kysely instance when
+            // no new executor is provided, or create a Transaction (not plain Kysely) when one is
+            this.kysely = executor ? new Transaction(this.kyselyProps) : baseClient.$qb;
         } else {
             this.kysely = new Kysely(this.kyselyProps);
         }
@@ -611,6 +611,10 @@ function createModelCrudHandler(
         throwIfNoResult = false,
     ) => {
         return createZenStackPromise(async (txClient?: ClientContract<any>) => {
+            // Per-operation context shared between onQuery and onKyselyQuery hooks.
+            // onQuery plugins write here; the context executor passes it to onKyselyQuery.
+            const queryContext = new Map<string, unknown>();
+
             let proceed = async (_args: unknown) => {
                 // prepare args for ext result: strip ext result field names from select/omit,
                 // inject needs fields into select (recursively handles nested relations)
@@ -619,7 +623,32 @@ function createModelCrudHandler(
                     ? prepareArgsForExtResult(_args, model, schema, plugins)
                     : _args;
 
-                const _handler = txClient ? handler.withClient(txClient) : handler;
+                // Bind queryContext to the executor so onKyselyQuery hooks can read it.
+                // Uses txClient's executor (which holds the tx connection) when in a transaction.
+                const baseClient = txClient ?? client;
+                const rawExecutor = (baseClient.$qb as any).getExecutor();
+
+                let contextExecutor: ZenStackQueryExecutor;
+                if (rawExecutor instanceof ZenStackQueryExecutor) {
+                    contextExecutor = rawExecutor.withQueryContext(queryContext);
+                } else {
+                    // Kysely wraps the real executor in NotCommittedOrRolledBackAssertingExecutor
+                    // inside sequential transactions — delegate connection to rawExecutor so
+                    // queries run within the transaction.
+                    const rootZenExecutor = (client as unknown as ClientImpl).kyselyProps
+                        .executor as ZenStackQueryExecutor;
+                    contextExecutor = rootZenExecutor
+                        .withConnectionProvider({
+                            provideConnection: (consumer) => rawExecutor.provideConnection(consumer),
+                        })
+                        .withQueryContext(queryContext);
+                }
+
+                const contextClient = (baseClient as unknown as ClientImpl).withExecutor(
+                    contextExecutor,
+                ) as unknown as ClientContract<any>;
+
+                const _handler = handler.withClient(contextClient);
                 const r = await _handler.handle(operation, processedArgs);
                 if (!r && throwIfNoResult) {
                     throw createNotFoundError(model);
@@ -652,6 +681,7 @@ function createModelCrudHandler(
                             operation: nominalOperation,
                             // reflect the latest override if provided
                             args: _args,
+                            queryContext,
                             // ensure inner overrides are propagated to the previous proceed
                             proceed: (nextArgs: unknown) => _proceed(nextArgs),
                         };

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

@@ -170,9 +170,7 @@ export abstract class BaseCrudDialect<Schema extends SchemaDef> {
                 effectiveOrderBy &&
                 enumerate(effectiveOrderBy).some((ob: any) => typeof ob === 'object' && '_fuzzyRelevance' in ob)
             ) {
-                throw createNotSupportedError(
-                    'cursor pagination cannot be combined with "_fuzzyRelevance" ordering',
-                );
+                throw createNotSupportedError('cursor pagination cannot be combined with "_fuzzyRelevance" ordering');
             }
             result = this.buildCursorFilter(
                 model,
@@ -1670,7 +1668,10 @@ export abstract class BaseCrudDialect<Schema extends SchemaDef> {
             'fuzzy filter must be an object with at least a "search" field',
         );
         const raw = value as Record<string, unknown>;
-        invariant(typeof raw['search'] === 'string' && raw['search'].length > 0, 'fuzzy.search must be a non-empty string');
+        invariant(
+            typeof raw['search'] === 'string' && raw['search'].length > 0,
+            'fuzzy.search must be a non-empty string',
+        );
         const mode = raw['mode'] ?? 'simple';
         invariant(
             mode === 'simple' || mode === 'word' || mode === 'strictWord',

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

@@ -51,6 +51,7 @@ import {
 } from '../../query-utils';
 import { getCrudDialect } from '../dialects';
 import type { BaseCrudDialect } from '../dialects/base-dialect';
+import { ZenStackQueryExecutor } from '../../executor/zenstack-query-executor';
 import { InputValidator } from '../validator';
 
 /**
@@ -168,6 +169,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.
  */
@@ -281,6 +292,7 @@ export abstract class BaseOperationHandler<Schema extends SchemaDef> {
         kysely: AnyKysely,
         model: string,
         args: FindArgs<Schema, GetModels<Schema>, any, true> | undefined,
+        direct = false,
     ): Promise<any[]> {
         // table
         let query = this.dialect.buildSelectModel(model, model);
@@ -306,20 +318,40 @@ export abstract class BaseOperationHandler<Schema extends SchemaDef> {
 
         query = query.modifyEnd(this.makeContextComment({ model, operation: 'read' }));
 
-        let result: any[] = [];
         const compiled = kysely.getExecutor().compileQuery(query.toOperationNode(), createQueryId());
+
+        let result: any[] = [];
         try {
-            const r = await kysely.getExecutor().executeQuery(compiled);
-            result = r.rows;
+            if (direct) {
+                // Bypass onKyselyQuery interceptors (e.g. policy plugin) so read-denied rows
+                // are still reachable. Uses the outer executor for connection acquisition so
+                // the query runs within an active transaction when applicable.
+                const zenExecutor = (this.client as any).kyselyProps.executor as ZenStackQueryExecutor;
+                const r = await kysely
+                    .getExecutor()
+                    .provideConnection((connection) => zenExecutor.executeQueryDirect(compiled, connection));
+                result = r.rows;
+            } else {
+                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);
         }
 
         return result;
     }
 
-    protected async readUnique(kysely: AnyKysely, model: string, args: FindArgs<Schema, GetModels<Schema>, any, true>) {
-        const result = await this.read(kysely, model, { ...args, take: 1 });
+    protected async readUnique(
+        kysely: AnyKysely,
+        model: string,
+        args: FindArgs<Schema, GetModels<Schema>, any, true>,
+        direct = false,
+    ) {
+        const result = await this.read(kysely, model, { ...args, take: 1 }, direct);
         return result[0] ?? null;
     }
 
@@ -1182,11 +1214,16 @@ 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;
+
         // 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 = (await this.getEntityIds(kysely, model, origWhere, bypassReadPolicyForPreload)) ?? null;
                 if (!thisEntity && throwIfNotFound) {
                     throw createNotFoundError(model);
                 }
@@ -2517,11 +2554,16 @@ export abstract class BaseOperationHandler<Schema extends SchemaDef> {
     }
 
     // Given a unique filter of a model, load the entity and return its id fields
-    private getEntityIds(kysely: AnyKysely, model: string, uniqueFilter: any) {
-        return this.readUnique(kysely, model, {
-            where: uniqueFilter,
-            select: this.makeIdSelect(model),
-        });
+    private getEntityIds(kysely: AnyKysely, model: string, uniqueFilter: any, direct = false) {
+        return this.readUnique(
+            kysely,
+            model,
+            {
+                where: uniqueFilter,
+                select: this.makeIdSelect(model),
+            },
+            direct,
+        );
     }
 
     // Given multiple unique filters, load all matching entities and return their id fields in one query

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/zenstack-query-executor.ts

@@ -86,6 +86,7 @@ export class ZenStackQueryExecutor extends DefaultQueryExecutor {
         private readonly connectionProvider: ConnectionProvider,
         plugins: KyselyPlugin[] = [],
         private suppressMutationHooks: boolean = false,
+        private readonly queryContext: Map<string, unknown> = new Map(),
     ) {
         super(compiler, adapter, connectionProvider, plugins);
 
@@ -214,6 +215,7 @@ export class ZenStackQueryExecutor extends DefaultQueryExecutor {
                     schema: this.client.$schema,
                     query,
                     proceed: _p,
+                    queryContext: this.queryContext,
                 });
                 return hookResult;
             };
@@ -673,6 +675,16 @@ In such cases, ZenStack cannot reliably determine the IDs of the mutated entitie
         }
     }
 
+    /**
+     * Execute a compiled query on `connection`, bypassing all `onKyselyQuery` plugin interceptors.
+     */
+    async executeQueryDirect(
+        compiledQuery: CompiledQuery,
+        connection: DatabaseConnection,
+    ): Promise<QueryResult<unknown>> {
+        return this.internalExecuteQuery(compiledQuery.query, connection, compiledQuery.queryId);
+    }
+
     private async internalExecuteQuery(
         query: RootOperationNode,
         connection: DatabaseConnection,
@@ -770,6 +782,7 @@ In such cases, ZenStack cannot reliably determine the IDs of the mutated entitie
             this.connectionProvider,
             [...this.plugins, plugin],
             this.suppressMutationHooks,
+            this.queryContext,
         );
     }
 
@@ -782,6 +795,7 @@ In such cases, ZenStack cannot reliably determine the IDs of the mutated entitie
             this.connectionProvider,
             [...this.plugins, ...plugins],
             this.suppressMutationHooks,
+            this.queryContext,
         );
     }
 
@@ -794,6 +808,7 @@ In such cases, ZenStack cannot reliably determine the IDs of the mutated entitie
             this.connectionProvider,
             [plugin, ...this.plugins],
             this.suppressMutationHooks,
+            this.queryContext,
         );
     }
 
@@ -806,6 +821,7 @@ In such cases, ZenStack cannot reliably determine the IDs of the mutated entitie
             this.connectionProvider,
             [],
             this.suppressMutationHooks,
+            this.queryContext,
         );
     }
 
@@ -818,11 +834,33 @@ In such cases, ZenStack cannot reliably determine the IDs of the mutated entitie
             connectionProvider,
             this.plugins as KyselyPlugin[],
             this.suppressMutationHooks,
+            this.queryContext,
         );
         // replace client with a new one associated with the new executor
         newExecutor.client = this.client.withExecutor(newExecutor);
         return newExecutor;
     }
 
+    /**
+     * Create a new executor carrying the given per-operation query context.
+     * Called once per top-level ORM operation so that onQuery plugins can write
+     * values (e.g. `operation`, `fetchPolicyCodes`) that onKyselyQuery plugins read —
+     * without AsyncLocalStorage.
+     */
+    withQueryContext(queryContext: Map<string, unknown>): ZenStackQueryExecutor {
+        const newExecutor = new ZenStackQueryExecutor(
+            this.client,
+            this.driver,
+            this.compiler,
+            this.adapter,
+            this.connectionProvider,
+            this.plugins as KyselyPlugin[],
+            this.suppressMutationHooks,
+            queryContext,
+        );
+        newExecutor.client = this.client.withExecutor(newExecutor);
+        return newExecutor;
+    }
+
     // #endregion
 }

packages/orm/src/client/index.ts

@@ -13,6 +13,7 @@ export {
     CoreReadOperations,
     CoreUpdateOperations,
     CoreWriteOperations,
+    SingleRowOrThrowOperations,
 } from './crud/operations/base';
 export { InputValidator } from './crud/validator';
 export { ORMError, ORMErrorReason, RejectedByPolicyReason } from './errors';