feat: provide the client in the computed field context

packages/orm/src/client/client-impl.ts

@@ -253,7 +253,7 @@ export class ClientImpl {
     ): Promise<any> {
         if (this.kysely.isTransaction) {
             // proceed directly if already in a transaction
-            return callback(this as unknown as ClientContract<SchemaDef>);
+            return callback(this.$contract);
         } else {
             // otherwise, create a new transaction, clone the client, and execute the callback
             let txBuilder = this.kysely.transaction();
@@ -263,7 +263,7 @@ export class ClientImpl {
             return txBuilder.execute((tx) => {
                 const txClient = new ClientImpl(this.schema, this.$options, this);
                 txClient.kysely = tx;
-                return callback(txClient as unknown as ClientContract<SchemaDef>);
+                return callback(txClient.$contract);
             });
         }
     }
@@ -285,7 +285,7 @@ export class ClientImpl {
             const result: any[] = [];
             for (const promise of arg) {
                 const cb = this.getPromiseCallback(promise);
-                result.push(await cb(txClient as unknown as ClientContract<SchemaDef>));
+                result.push(await cb(txClient.$contract));
             }
             return result;
         };
@@ -446,6 +446,17 @@ export class ClientImpl {
         return this.auth;
     }
 
+    /**
+     * This client viewed through its public typed contract. `ClientImpl` is intentionally
+     * untyped internally — the model accessors are added by the runtime proxy — so this
+     * getter is the single sanctioned bridge to `ClientContract`. The proxy invokes it
+     * with the proxy as `this` (`Reflect.get` with receiver), so the returned reference
+     * keeps the model accessors.
+     */
+    get $contract(): ClientContract<SchemaDef> {
+        return this as unknown as ClientContract<SchemaDef>;
+    }
+
     $setOptions<Options extends ClientOptions<SchemaDef>>(options: Options): ClientContract<SchemaDef, Options> {
         const newClient = new ClientImpl(this.schema, options as ClientOptions<SchemaDef>, this);
         // create a new validator to have a fresh schema cache, because options may change validation settings

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

@@ -15,6 +15,7 @@ import type {
     SortOrder,
     StringFilter,
 } from '../../crud-types';
+import type { ClientContract } from '../../contract';
 import { createConfigError, createInvalidInputError, createNotSupportedError } from '../../errors';
 import type { ClientOptions } from '../../options';
 import {
@@ -42,6 +43,9 @@ export abstract class BaseCrudDialect<Schema extends SchemaDef> {
     constructor(
         protected readonly schema: Schema,
         protected readonly options: ClientOptions<Schema>,
+        // the client executing the query; optional so the dialect can still be
+        // constructed standalone (e.g. for output transformation only)
+        protected readonly client?: ClientContract<Schema>,
     ) {}
 
     // #region capability flags1
@@ -1662,7 +1666,7 @@ export abstract class BaseCrudDialect<Schema extends SchemaDef> {
             }
             // `computedArgs` is the query-time args object for a parameterized computed
             // field (undefined otherwise); forwarded as the implementation's 3rd argument.
-            return computer(this.eb, { modelAlias }, computedArgs);
+            return computer(this.eb, { modelAlias, client: this.client }, computedArgs);
         }
     }
 

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

@@ -1,5 +1,6 @@
 import type { SchemaDef } from '@zenstackhq/schema';
 import { match } from 'ts-pattern';
+import type { ClientContract } from '../../contract';
 import type { ClientOptions } from '../../options';
 import type { BaseCrudDialect } from './base-dialect';
 import { MySqlCrudDialect } from './mysql';
@@ -9,10 +10,11 @@ import { SqliteCrudDialect } from './sqlite';
 export function getCrudDialect<Schema extends SchemaDef>(
     schema: Schema,
     options: ClientOptions<Schema>,
+    client?: ClientContract<Schema>,
 ): BaseCrudDialect<Schema> {
     return match(schema.provider.type)
-        .with('sqlite', () => new SqliteCrudDialect(schema, options))
-        .with('postgresql', () => new PostgresCrudDialect(schema, options))
-        .with('mysql', () => new MySqlCrudDialect(schema, options))
+        .with('sqlite', () => new SqliteCrudDialect(schema, options, client))
+        .with('postgresql', () => new PostgresCrudDialect(schema, options, client))
+        .with('mysql', () => new MySqlCrudDialect(schema, options, client))
         .exhaustive();
 }

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

@@ -12,6 +12,7 @@ import {
     type SqlBool,
 } from 'kysely';
 import { AnyNullClass, DbNullClass, JsonNullClass } from '../../../common-types';
+import type { ClientContract } from '../../contract';
 import type { NullsOrder, SortOrder } from '../../crud-types';
 import { createInvalidInputError, createNotSupportedError } from '../../errors';
 import type { ClientOptions } from '../../options';
@@ -20,8 +21,8 @@ import type { FuzzyFilterOptions } from './base-dialect';
 import { LateralJoinDialectBase } from './lateral-join-dialect-base';
 
 export class MySqlCrudDialect<Schema extends SchemaDef> extends LateralJoinDialectBase<Schema> {
-    constructor(schema: Schema, options: ClientOptions<Schema>) {
-        super(schema, options);
+    constructor(schema: Schema, options: ClientOptions<Schema>, client?: ClientContract<Schema>) {
+        super(schema, options, client);
     }
 
     override get provider() {

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

@@ -11,6 +11,7 @@ import {
 } from 'kysely';
 import { parse as parsePostgresArray } from 'postgres-array';
 import { AnyNullClass, DbNullClass, JsonNullClass } from '../../../common-types';
+import type { ClientContract } from '../../contract';
 import type { NullsOrder, SortOrder } from '../../crud-types';
 import { createInvalidInputError } from '../../errors';
 import type { ClientOptions } from '../../options';
@@ -73,8 +74,8 @@ export class PostgresCrudDialect<Schema extends SchemaDef> extends LateralJoinDi
         '@db.Boolean': 'boolean',
     };
 
-    constructor(schema: Schema, options: ClientOptions<Schema>) {
-        super(schema, options);
+    constructor(schema: Schema, options: ClientOptions<Schema>, client?: ClientContract<Schema>) {
+        super(schema, options, client);
         this.overrideTypeParsers();
     }
 

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

@@ -199,7 +199,7 @@ export abstract class BaseOperationHandler<Schema extends SchemaDef> {
         protected readonly model: GetModels<Schema>,
         protected readonly inputValidator: InputValidator<Schema>,
     ) {
-        this.dialect = getCrudDialect(this.schema, this.client.$options);
+        this.dialect = getCrudDialect(this.schema, this.client.$options, this.client);
     }
 
     protected get schema() {

packages/orm/src/client/executor/name-mapper.ts

@@ -62,7 +62,7 @@ export class QueryNameMapper extends OperationNodeTransformer {
 
     constructor(private readonly client: ClientContract<SchemaDef>) {
         super();
-        this.dialect = getCrudDialect(client.$schema, client.$options);
+        this.dialect = getCrudDialect(client.$schema, client.$options, client);
         for (const [modelName, modelDef] of Object.entries(client.$schema.models)) {
             const mappedName = this.getMappedName(modelDef);
             if (mappedName) {

packages/orm/src/client/executor/zenstack-query-executor.ts

@@ -93,10 +93,10 @@ export class ZenStackQueryExecutor extends DefaultQueryExecutor {
             client.$schema.provider.type === 'postgresql' || // postgres queries need to be schema-qualified
             this.schemaHasMappedNames(client.$schema)
         ) {
-            this.nameMapper = new QueryNameMapper(client as unknown as ClientContract<SchemaDef>);
+            this.nameMapper = new QueryNameMapper(client.$contract);
         }
 
-        this.dialect = getCrudDialect(client.$schema, client.$options);
+        this.dialect = getCrudDialect(client.$schema, client.$options, client.$contract);
     }
 
     private schemaHasMappedNames(schema: SchemaDef) {
@@ -210,7 +210,7 @@ export class ZenStackQueryExecutor extends DefaultQueryExecutor {
             proceed = async (query: RootOperationNode) => {
                 const _p = (q: RootOperationNode) => _proceed(q);
                 const hookResult = await hook!({
-                    client: this.client as unknown as ClientContract<SchemaDef>,
+                    client: this.client.$contract,
                     schema: this.client.$schema,
                     query,
                     proceed: _p,
@@ -660,7 +660,7 @@ In such cases, ZenStack cannot reliably determine the IDs of the mutated entitie
         if (inTx) {
             innerClient.forceTransaction();
         }
-        return innerClient as unknown as ClientContract<SchemaDef>;
+        return innerClient.$contract;
     }
 
     private andNodes(condition1: WhereNode | undefined, condition2: WhereNode | undefined) {

packages/orm/src/client/options.ts

@@ -283,16 +283,38 @@ export type OmitConfig<Schema extends SchemaDef> = {
     };
 };
 
+/**
+ * Context object passed to computed field implementations.
+ */
+export type ComputedFieldContext<Schema extends SchemaDef> = {
+    /**
+     * The alias name that can be used to refer to the containing model
+     */
+    modelAlias: string;
+
+    /**
+     * The ZenStack client executing the query. Useful for reading per-client state,
+     * e.g. the auth context set via `$setAuth`. Undefined only when the CRUD dialect
+     * is constructed standalone rather than by the ORM runtime — queries issued
+     * through the client API always have it.
+     */
+    client?: ClientContract<Schema>;
+};
+
 export type ComputedFieldsOptions<Schema extends SchemaDef> = {
     [Model in GetModels<Schema> as 'computedFields' extends keyof GetModel<Schema, Model>
         ? Uncapitalize<Model>
         : never]: {
         [Field in keyof Schema['models'][Model]['computedFields']]: Schema['models'][Model]['computedFields'][Field] extends infer Func
-            ? Func extends (...args: any[]) => infer R
+            ? Func extends (...args: infer Params) => infer R
                 ? (
                       // inject a first parameter for expression builder
                       p: ExpressionBuilder<ToKyselySchema<Schema>, Model>,
-                      ...args: Parameters<Func>
+                      // runtime-provided context (the generated stub only declares
+                      // `modelAlias`; the runtime passes the full context)
+                      context: ComputedFieldContext<Schema>,
+                      // query-time args of a parameterized field, from the stub
+                      ...args: Params extends [any, ...infer Rest] ? Rest : []
                   ) => OperandExpression<R> // wrap the return type with Kysely `OperandExpression`
                 : never
             : never;

tests/e2e/orm/client-api/computed-fields.test.ts

@@ -989,4 +989,41 @@ model User {
             }),
         ).toBeRejectedByValidation(['upperName']);
     });
+
+    it('provides the client in the computed field context', async () => {
+        const db = await createTestClient(
+            `
+model Post {
+    id Int @id @default(autoincrement())
+    authorId Int
+    isMine Boolean @computed
+}
+`,
+            {
+                computedFields: {
+                    Post: {
+                        // parenthesized: the expression gets embedded into larger ones
+                        // (e.g. `where: { isMine: true }` wraps it with `= true`), and an
+                        // unparenthesized chained comparison is a syntax error on postgres
+                        isMine: (eb: any, { client }: any) => eb.parens(eb('authorId', '=', client?.$auth?.id ?? -1)),
+                    },
+                },
+            } as any,
+        );
+
+        await db.post.create({ data: { id: 1, authorId: 1 } });
+        await db.post.create({ data: { id: 2, authorId: 2 } });
+
+        // no auth set: nothing is mine
+        await expect(db.post.findUnique({ where: { id: 1 } })).resolves.toMatchObject({ isMine: false });
+
+        // the client derived with $setAuth carries its auth into the computed field
+        const authedDb = db.$setAuth({ id: 1 });
+        await expect(authedDb.post.findUnique({ where: { id: 1 } })).resolves.toMatchObject({ isMine: true });
+        await expect(authedDb.post.findUnique({ where: { id: 2 } })).resolves.toMatchObject({ isMine: false });
+        await expect(authedDb.post.findMany({ where: { isMine: true } })).resolves.toHaveLength(1);
+
+        // the original client is unaffected
+        await expect(db.post.findUnique({ where: { id: 1 } })).resolves.toMatchObject({ isMine: false });
+    });
 });