feat(policy): surface error codes for single-row read violations

Extend policy error code surfacing to `findFirst`/`findUnique` (and their
`OrThrow` variants): when the query returns 0 rows, a diagnostic query now
checks whether the row exists but was filtered by a deny/allow rule that
carries an `errorCode`. If so, the handler throws `REJECTED_BY_POLICY`
with the relevant codes instead of silently returning `null`.

`findMany` is unaffected — it continues to use filter-based enforcement
and returns an empty array for denied rows. This holds even for
`findMany({ take: 1 })`, which generates LIMIT 1 SQL but remains a
multi-row ORM operation by name.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

Author: Azzerty23

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

@@ -168,6 +168,16 @@ export const AllReadOperations = [...CoreReadOperations, 'findUniqueOrThrow', 'f
  */
 export type AllReadOperations = (typeof AllReadOperations)[number];
 
+/**
+ * List of single-row read operations — `findUnique`/`findFirst` and their 'orThrow' variants.
+ */
+export const SingleRowReadOperations = ['findUnique', 'findFirst', 'findUniqueOrThrow', 'findFirstOrThrow'] as const;
+
+/**
+ * List of single-row read operations.
+ */
+export type SingleRowReadOperations = (typeof SingleRowReadOperations)[number];
+
 /**
  * List of all write operations - simply an alias of CoreWriteOperations.
  */
@@ -312,6 +322,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);
         }
 

packages/orm/src/client/errors.ts

@@ -96,8 +96,9 @@ export class ORMError extends Error {
      * 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`, and `delete` violations.
-     * `read` uses filter-based enforcement and does not throw policy errors.
+     * 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[];
 

packages/orm/src/client/index.ts

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

packages/plugins/policy/src/context.ts

@@ -0,0 +1,14 @@
+import { AsyncLocalStorage } from 'node:async_hooks';
+
+/**
+ * Per-query context shared between the ORM hook (`onQuery`) and the Kysely handler (`onKyselyQuery`).
+ * - `operation`: ORM operation name (e.g. `findUnique`, `create`) — used to distinguish single-row reads
+ *   and to skip the read diagnostic check for nested SELECTs inside mutations
+ * - `fetchPolicyCodes`: per-query override of the plugin-level option
+ */
+export type PolicyContext = {
+    operation?: string;
+    fetchPolicyCodes?: boolean;
+};
+
+export const policyContextStorage = new AsyncLocalStorage<PolicyContext>();

packages/plugins/policy/src/plugin.ts

@@ -1,18 +1,15 @@
-import { AsyncLocalStorage } from 'node:async_hooks';
 import { type OnKyselyQueryArgs, type RuntimePlugin } from '@zenstackhq/orm';
 import type { SchemaDef } from '@zenstackhq/orm/schema';
 import { z } from 'zod';
+import { policyContextStorage } from './context';
 import { check } from './functions';
 import type { PolicyPluginOptions } from './options';
 import { PolicyHandler } from './policy-handler';
 
 export type { PolicyPluginOptions } from './options';
 
-type PolicyQueryContext = Pick<PolicyPluginOptions, 'fetchPolicyCodes'>;
-
-const policyContextStorage = new AsyncLocalStorage<PolicyQueryContext>();
-
 type PolicyExtQueryArgs = {
+    $read: Pick<PolicyPluginOptions, 'fetchPolicyCodes'>;
     $create: Pick<PolicyPluginOptions, 'fetchPolicyCodes'>;
     $update: Pick<PolicyPluginOptions, 'fetchPolicyCodes'>;
     $delete: Pick<PolicyPluginOptions, 'fetchPolicyCodes'>;
@@ -42,23 +39,24 @@ export class PolicyPlugin implements RuntimePlugin<SchemaDef, PolicyExtQueryArgs
     }
 
     readonly queryArgs = {
+        $read: fetchPolicyCodesSchema,
         $create: fetchPolicyCodesSchema,
         $update: fetchPolicyCodesSchema,
         $delete: 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.
+    // AsyncLocalStorage bridges per-query context into the Kysely executor.
     onQuery(ctx: {
+        operation: string;
         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);
+        return policyContextStorage.run(
+            { operation: ctx.operation, fetchPolicyCodes: ctx.args?.['fetchPolicyCodes'] as boolean | undefined },
+            () => ctx.proceed(ctx.args),
+        );
     }
 
     onKyselyQuery({ query, client, proceed }: OnKyselyQueryArgs<SchemaDef>) {

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

@@ -1,6 +1,6 @@
 import { invariant } from '@zenstackhq/common-helpers';
 import type { BaseCrudDialect, ClientContract, CRUD_EXT, ProceedKyselyQueryFunction } from '@zenstackhq/orm';
-import { getCrudDialect, QueryUtils, RejectedByPolicyReason, SchemaUtils } from '@zenstackhq/orm';
+import { CoreWriteOperations, getCrudDialect, QueryUtils, RejectedByPolicyReason, SchemaUtils, SingleRowReadOperations } from '@zenstackhq/orm';
 import {
     ExpressionUtils,
     type BuiltinType,
@@ -42,6 +42,7 @@ import {
 } from 'kysely';
 import { match } from 'ts-pattern';
 import { ColumnCollector } from './column-collector';
+import { policyContextStorage } from './context';
 import { ExpressionTransformer } from './expression-transformer';
 import type { PolicyPluginOptions } from './options';
 import type { Policy, PolicyOperation } from './types';
@@ -59,6 +60,9 @@ import {
     trueNode,
 } from './utils';
 
+const SINGLE_ROW_READ_OPERATIONS = new Set<string>(SingleRowReadOperations);
+const ORM_WRITE_OPERATIONS = new Set<string>(CoreWriteOperations);
+
 export type CrudQueryNode = SelectQueryNode | InsertQueryNode | UpdateQueryNode | DeleteQueryNode;
 
 export type MutationQueryNode = InsertQueryNode | UpdateQueryNode | DeleteQueryNode;
@@ -93,8 +97,13 @@ export class PolicyHandler<Schema extends SchemaDef> extends OperationNodeTransf
         }
 
         if (!this.isMutationQueryNode(node)) {
-            // transform and proceed with read directly
-            return proceed(this.transformNode(node));
+            const selectNode = node as SelectQueryNode;
+            const result = await proceed(this.transformNode(node));
+            // When 0 rows returned on a single-row read, distinguish "not found" from policy denial
+            if (result.rows.length === 0 && SINGLE_ROW_READ_OPERATIONS.has(policyContextStorage.getStore()?.operation ?? '')) {
+                await this.postReadZeroRowsCheck(selectNode, proceed);
+            }
+            return result;
         }
 
         const { mutationModel } = this.getMutationModel(node);
@@ -142,9 +151,9 @@ export class PolicyHandler<Schema extends SchemaDef> extends OperationNodeTransf
         // 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);
+                await this.postZeroRowsCheck(mutationModel, 'delete', node.where?.where, proceed);
             } else if (UpdateQueryNode.is(node)) {
-                await this.postMutationZeroRowsCheck(mutationModel, 'update', node.where?.where, proceed);
+                await this.postZeroRowsCheck(mutationModel, 'update', node.where?.where, proceed);
             }
         }
 
@@ -259,38 +268,47 @@ export class PolicyHandler<Schema extends SchemaDef> extends OperationNodeTransf
         }
     }
 
-    // Checks if any row matching the original WHERE exists without the policy filter.
-    // Called when numAffectedRows == 0 for UPDATE or DELETE.
+    // Called when a single-row read returns 0 rows. Skips internal reads (read-back after mutation).
+    private async postReadZeroRowsCheck(node: SelectQueryNode, proceed: ProceedKyselyQueryFunction): Promise<void> {
+        if (ORM_WRITE_OPERATIONS.has(policyContextStorage.getStore()?.operation ?? '')) return;
+        if (!node.from || node.from.froms.length !== 1) return;
+        const extractedTable = this.extractTableName(node.from.froms[0]!);
+        if (!extractedTable) return;
+        const { model } = extractedTable;
+        if (!QueryUtils.getModel(this.client.$schema, model)) return;
+        return this.postZeroRowsCheck(model, 'read', node.where?.where, proceed);
+    }
+
+    // Checks if any row matching WHERE exists without the policy filter.
     // 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(
+    // If no row matches → returns silently.
+    private async postZeroRowsCheck(
         model: string,
-        operation: 'update' | 'delete',
-        originalWhere: OperationNode | undefined,
+        operation: 'read' | 'update' | 'delete',
+        whereCondition: OperationNode | undefined,
         proceed: ProceedKyselyQueryFunction,
     ) {
-        if (this.isManyToManyJoinTable(model)) return;
         if (this.tryGetConstantPolicy(model, operation) === true) return;
         if (this.options.fetchPolicyCodes === false) return;
         const policiesWithCode = this.getModelPolicies(model, operation).filter((p) => p.code);
-        // Skip if no policies carry an error code — nothing to surface.
         if (policiesWithCode.length === 0) return;
+        if (this.isManyToManyJoinTable(model)) return;
 
-        const whereCondition = originalWhere ?? trueNode(this.dialect);
+        // No WHERE clause means "match all rows" — use a literal TRUE so the existence sub-query is valid SQL.
+        const where = whereCondition ?? trueNode(this.dialect);
 
         const rowExistsInner = this.eb
             .selectFrom(model)
             .select(this.eb.lit(1).as('_'))
-            .where(() => new ExpressionWrapper(whereCondition));
+            .where(() => new ExpressionWrapper(where));
 
         const codeSelections = policiesWithCode.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(model)
                 .select(this.eb.lit(1).as('_'))
-                .where(() => new ExpressionWrapper(conjunction(this.dialect, [whereCondition, violationCondition])));
+                .where(() => new ExpressionWrapper(conjunction(this.dialect, [where, violationCondition])));
             return SelectionNode.create(
                 AliasNode.create(this.eb.exists(inner).toOperationNode(), IdentifierNode.create(`$c${i}`)),
             );
@@ -312,10 +330,7 @@ export class PolicyHandler<Schema extends SchemaDef> extends OperationNodeTransf
         const row = result.rows[0] ?? {};
         if (!row.$exists) return;
 
-        const policyCodes = policiesWithCode
-            ? policiesWithCode.filter((_, i) => row[`$c${i}`]).map((p) => p.code!)
-            : undefined;
-
+        const policyCodes = policiesWithCode.filter((_, i) => row[`$c${i}`]).map((p) => p.code!);
         throw createRejectedByPolicyError(model, RejectedByPolicyReason.NO_ACCESS, undefined, policyCodes);
     }