fix: make computed fields on base models work (#321)

* fix: make computed fields on base models work

fixes #284

* update

Author: ymc9

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

@@ -19,7 +19,6 @@ import { InternalError, QueryError } from '../../errors';
 import type { ClientOptions } from '../../options';
 import {
     aggregate,
-    buildFieldRef,
     buildJoinPairs,
     ensureArray,
     flattenCompoundUniqueFilters,
@@ -931,15 +930,13 @@ export abstract class BaseCrudDialect<Schema extends SchemaDef> {
         field: string,
     ): SelectQueryBuilder<any, any, any> {
         const fieldDef = requireField(this.schema, model, field);
-        if (fieldDef.computed) {
-            // TODO: computed field from delegate base?
-            return query.select(() => this.fieldRef(model, field, modelAlias).as(field));
-        } else if (!fieldDef.originModel) {
-            // regular field
-            return query.select(this.eb.ref(`${modelAlias}.${field}`).as(field));
-        } else {
-            return this.buildSelectField(query, fieldDef.originModel, fieldDef.originModel, field);
-        }
+
+        // if field is defined on a delegate base, the base model is joined with its
+        // model name from outer query, so we should use it directly as the alias
+        const fieldModel = fieldDef.originModel ?? model;
+        const alias = fieldDef.originModel ?? modelAlias;
+
+        return query.select(() => this.fieldRef(fieldModel, field, alias).as(field));
     }
 
     buildDelegateJoin(
@@ -1071,7 +1068,26 @@ export abstract class BaseCrudDialect<Schema extends SchemaDef> {
     }
 
     fieldRef(model: string, field: string, modelAlias?: string, inlineComputedField = true) {
-        return buildFieldRef(this.schema, model, field, this.options, this.eb, modelAlias, inlineComputedField);
+        const fieldDef = requireField(this.schema, model, field);
+
+        if (!fieldDef.computed) {
+            // regular field
+            return this.eb.ref(modelAlias ? `${modelAlias}.${field}` : field);
+        } else {
+            // computed field
+            if (!inlineComputedField) {
+                return this.eb.ref(modelAlias ? `${modelAlias}.${field}` : field);
+            }
+            let computer: Function | undefined;
+            if ('computedFields' in this.options) {
+                const computedFields = this.options.computedFields as Record<string, any>;
+                computer = computedFields?.[fieldDef.originModel ?? model]?.[field];
+            }
+            if (!computer) {
+                throw new QueryError(`Computed field "${field}" implementation not provided for model "${model}"`);
+            }
+            return computer(this.eb, { modelAlias });
+        }
     }
 
     protected canJoinWithoutNestedSelect(

packages/runtime/src/client/query-utils.ts

@@ -6,7 +6,6 @@ import {
     TableNode,
     type Expression,
     type ExpressionBuilder,
-    type ExpressionWrapper,
     type OperationNode,
 } from 'kysely';
 import { match } from 'ts-pattern';
@@ -15,7 +14,6 @@ import { extractFields } from '../utils/object-utils';
 import type { AGGREGATE_OPERATORS } from './constants';
 import type { OrderBy } from './crud-types';
 import { InternalError, QueryError } from './errors';
-import type { ClientOptions } from './options';
 
 export function hasModel(schema: SchemaDef, model: string) {
     return Object.keys(schema.models)
@@ -180,34 +178,6 @@ export function getIdValues(schema: SchemaDef, model: string, data: any): Record
     return idFields.reduce((acc, field) => ({ ...acc, [field]: data[field] }), {});
 }
 
-export function buildFieldRef<Schema extends SchemaDef>(
-    schema: Schema,
-    model: string,
-    field: string,
-    options: ClientOptions<Schema>,
-    eb: ExpressionBuilder<any, any>,
-    modelAlias?: string,
-    inlineComputedField = true,
-): ExpressionWrapper<any, any, unknown> {
-    const fieldDef = requireField(schema, model, field);
-    if (!fieldDef.computed) {
-        return eb.ref(modelAlias ? `${modelAlias}.${field}` : field);
-    } else {
-        if (!inlineComputedField) {
-            return eb.ref(modelAlias ? `${modelAlias}.${field}` : field);
-        }
-        let computer: Function | undefined;
-        if ('computedFields' in options) {
-            const computedFields = options.computedFields as Record<string, any>;
-            computer = computedFields?.[model]?.[field];
-        }
-        if (!computer) {
-            throw new QueryError(`Computed field "${field}" implementation not provided for model "${model}"`);
-        }
-        return computer(eb, { modelAlias });
-    }
-}
-
 export function fieldHasDefaultValue(fieldDef: FieldDef) {
     return fieldDef.default !== undefined || fieldDef.updatedAt;
 }

samples/blog/zenstack/schema.zmodel

@@ -10,7 +10,8 @@ enum Role {
 }
 
 plugin policy {
-    // due to pnpm layout we can't directly use package name here
+    // due to pnpm layout we can't directly use package name here,
+    // don't do this in your code and use "@zenstackhq/plugin-policy" instead
     provider = '../node_modules/@zenstackhq/plugin-policy/dist/index.js'
 }
 

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

@@ -1,15 +1,9 @@
 import { createTestClient } from '@zenstackhq/testtools';
-import { afterEach, describe, expect, it } from 'vitest';
+import { describe, expect, it } from 'vitest';
 
 describe('Computed fields tests', () => {
-    let db: any;
-
-    afterEach(async () => {
-        await db?.$disconnect();
-    });
-
     it('works with non-optional fields', async () => {
-        db = await createTestClient(
+        const db = await createTestClient(
             `
 model User {
     id Int @id @default(autoincrement())
@@ -97,7 +91,7 @@ model User {
     });
 
     it('is typed correctly for non-optional fields', async () => {
-        db = await createTestClient(
+        await createTestClient(
             `
 model User {
     id Int @id @default(autoincrement())
@@ -137,7 +131,7 @@ main();
     });
 
     it('works with optional fields', async () => {
-        db = await createTestClient(
+        const db = await createTestClient(
             `
 model User {
     id Int @id @default(autoincrement())
@@ -164,7 +158,7 @@ model User {
     });
 
     it('is typed correctly for optional fields', async () => {
-        db = await createTestClient(
+        await createTestClient(
             `
 model User {
     id Int @id @default(autoincrement())
@@ -203,7 +197,7 @@ main();
     });
 
     it('works with read from a relation', async () => {
-        db = await createTestClient(
+        const db = await createTestClient(
             `
 model User {
     id Int @id @default(autoincrement())
@@ -240,4 +234,42 @@ model Post {
             author: expect.objectContaining({ postCount: 1 }),
         });
     });
+
+    it('allows sub models to use computed fields from delegate base', async () => {
+        const db = await createTestClient(
+            `
+model Content {
+    id Int @id @default(autoincrement())
+    title String
+    isNews Boolean @computed
+    contentType String
+    @@delegate(contentType)
+}
+
+model Post extends Content {
+    body String
+}
+`,
+            {
+                computedFields: {
+                    Content: {
+                        isNews: (eb: any) => eb('title', 'like', '%news%'),
+                    },
+                },
+            } as any,
+        );
+
+        const posts = await db.post.createManyAndReturn({
+            data: [
+                { id: 1, title: 'latest news', body: 'some news content' },
+                { id: 2, title: 'random post', body: 'some other content' },
+            ],
+        });
+        expect(posts).toEqual(
+            expect.arrayContaining([
+                expect.objectContaining({ id: 1, isNews: true }),
+                expect.objectContaining({ id: 2, isNews: false }),
+            ]),
+        );
+    });
 });