fix(delegate): relation selection

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

@@ -28,11 +28,13 @@ import {
     makeDefaultOrderBy,
     requireField,
 } from '../../query-utils';
+import type { BaseOperationHandler } from '../operations/base';
 
 export abstract class BaseCrudDialect<Schema extends SchemaDef> {
     constructor(
         protected readonly schema: Schema,
         protected readonly options: ClientOptions<Schema>,
+        protected readonly handler: BaseOperationHandler<Schema>,
     ) {}
 
     abstract get provider(): DataSourceProviderType;

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

@@ -1,16 +1,18 @@
 import { match } from 'ts-pattern';
 import type { SchemaDef } from '../../../schema';
 import type { ClientOptions } from '../../options';
+import type { BaseOperationHandler } from '../operations/base';
 import type { BaseCrudDialect } from './base';
 import { PostgresCrudDialect } from './postgresql';
 import { SqliteCrudDialect } from './sqlite';
 
 export function getCrudDialect<Schema extends SchemaDef>(
     schema: Schema,
     options: ClientOptions<Schema>,
+    handler: BaseOperationHandler<Schema>,
 ): BaseCrudDialect<Schema> {
     return match(schema.provider.type)
-        .with('sqlite', () => new SqliteCrudDialect(schema, options))
-        .with('postgresql', () => new PostgresCrudDialect(schema, options))
+        .with('sqlite', () => new SqliteCrudDialect(schema, options, handler))
+        .with('postgresql', () => new PostgresCrudDialect(schema, options, handler))
         .exhaustive();
 }

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

@@ -9,10 +9,12 @@ import {
 } from 'kysely';
 import { match } from 'ts-pattern';
 import type { BuiltinType, FieldDef, GetModels, SchemaDef } from '../../../schema';
+import { DELEGATE_JOINED_FIELD_PREFIX } from '../../constants';
 import type { FindArgs } from '../../crud-types';
 import {
     buildFieldRef,
     buildJoinPairs,
+    getDelegateDescendantModels,
     getIdFields,
     getManyToManyRelation,
     isRelationField,
@@ -79,10 +81,18 @@ export class PostgresCrudDialect<Schema extends SchemaDef> extends BaseCrudDiale
                 // simple select by default
                 let result = eb.selectFrom(`${relationModel} as ${joinTableName}`);
 
+                const joinBases: string[] = [];
+
                 // however if there're filter/orderBy/take/skip,
                 // we need to build a subquery to handle them before aggregation
                 result = eb.selectFrom(() => {
-                    let subQuery = eb.selectFrom(`${relationModel}`).selectAll();
+                    let subQuery = eb.selectFrom(relationModel);
+                    subQuery = this.handler.buildSelectAllFields(
+                        relationModel,
+                        subQuery,
+                        typeof payload === 'object' ? payload?.omit : undefined,
+                        joinBases,
+                    );
 
                     if (payload && typeof payload === 'object') {
                         if (payload.where) {
@@ -200,6 +210,20 @@ export class PostgresCrudDialect<Schema extends SchemaDef> extends BaseCrudDiale
             string | ExpressionWrapper<any, any, any> | SelectQueryBuilder<any, any, any> | RawBuilder<any>
         > = [];
 
+        // TODO: descendant JSON shouldn't be joined and selected if none of its fields are selected
+        const descendantModels = getDelegateDescendantModels(this.schema, relationModel);
+        if (descendantModels.length > 0) {
+            // select all JSONs built from delegate descendants
+            objArgs.push(
+                ...descendantModels
+                    .map((subModel) => [
+                        sql.lit(`${DELEGATE_JOINED_FIELD_PREFIX}${subModel.name}`),
+                        eb.ref(`${DELEGATE_JOINED_FIELD_PREFIX}${subModel.name}`),
+                    ])
+                    .flatMap((v) => v),
+            );
+        }
+
         if (payload === true || !payload.select) {
             // select all scalar fields
             objArgs.push(

packages/runtime/src/client/crud/dialects/sqlite.ts

@@ -10,9 +10,11 @@ import {
 } from 'kysely';
 import { match } from 'ts-pattern';
 import type { BuiltinType, GetModels, SchemaDef } from '../../../schema';
+import { DELEGATE_JOINED_FIELD_PREFIX } from '../../constants';
 import type { FindArgs } from '../../crud-types';
 import {
     buildFieldRef,
+    getDelegateDescendantModels,
     getIdFields,
     getManyToManyRelation,
     getRelationForeignKeyFieldPairs,
@@ -75,7 +77,15 @@ export class SqliteCrudDialect<Schema extends SchemaDef> extends BaseCrudDialect
         const subQueryName = `${parentName}$${relationField}`;
 
         let tbl = eb.selectFrom(() => {
-            let subQuery = eb.selectFrom(relationModel).selectAll();
+            let subQuery = eb.selectFrom(relationModel);
+
+            const joinBases: string[] = [];
+            subQuery = this.handler.buildSelectAllFields(
+                relationModel,
+                subQuery,
+                typeof payload === 'object' ? payload?.omit : undefined,
+                joinBases,
+            );
 
             if (payload && typeof payload === 'object') {
                 if (payload.where) {
@@ -143,6 +153,20 @@ export class SqliteCrudDialect<Schema extends SchemaDef> extends BaseCrudDialect
             type ArgsType = Expression<any> | RawBuilder<any> | SelectQueryBuilder<any, any, any>;
             const objArgs: ArgsType[] = [];
 
+            // TODO: descendant JSON shouldn't be joined and selected if none of its fields are selected
+            const descendantModels = getDelegateDescendantModels(this.schema, relationModel);
+            if (descendantModels.length > 0) {
+                // select all JSONs built from delegate descendants
+                objArgs.push(
+                    ...descendantModels
+                        .map((subModel) => [
+                            sql.lit(`${DELEGATE_JOINED_FIELD_PREFIX}${subModel.name}`),
+                            eb.ref(`${DELEGATE_JOINED_FIELD_PREFIX}${subModel.name}`),
+                        ])
+                        .flatMap((v) => v),
+                );
+            }
+
             if (payload === true || !payload.select) {
                 // select all scalar fields
                 objArgs.push(

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

@@ -33,6 +33,7 @@ import {
     ensureArray,
     extractIdFields,
     flattenCompoundUniqueFilters,
+    getDelegateDescendantModels,
     getDiscriminatorField,
     getField,
     getIdFields,
@@ -84,7 +85,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);
     }
 
     protected get schema() {
@@ -183,19 +184,22 @@ export abstract class BaseOperationHandler<Schema extends SchemaDef> {
             }
         }
 
+        // for deduplicating base joins
+        const joinedBases: string[] = [];
+
         // select
         if (args && 'select' in args && args.select) {
             // select is mutually exclusive with omit
-            query = this.buildFieldSelection(model, query, args.select, model);
+            query = this.buildFieldSelection(model, query, args.select, model, joinedBases);
         } else {
             // include all scalar fields except those in omit
-            query = this.buildSelectAllScalarFields(model, query, (args as any)?.omit);
+            query = this.buildSelectAllFields(model, query, (args as any)?.omit, joinedBases);
         }
 
         // include
         if (args && 'include' in args && args.include) {
             // note that 'omit' is handled above already
-            query = this.buildFieldSelection(model, query, args.include, model);
+            query = this.buildFieldSelection(model, query, args.include, model, joinedBases);
         }
 
         if (args?.cursor) {
@@ -246,9 +250,9 @@ export abstract class BaseOperationHandler<Schema extends SchemaDef> {
         query: SelectQueryBuilder<any, any, any>,
         selectOrInclude: Record<string, any>,
         parentAlias: string,
+        joinedBases: string[],
     ) {
         let result = query;
-        const joinedBases: string[] = [];
 
         for (const [field, payload] of Object.entries(selectOrInclude)) {
             if (!payload) {
@@ -262,12 +266,29 @@ export abstract class BaseOperationHandler<Schema extends SchemaDef> {
 
             const fieldDef = this.requireField(model, field);
             if (!fieldDef.relation) {
+                // scalar field
                 result = this.selectField(result, model, parentAlias, field, joinedBases);
             } else {
                 if (!fieldDef.array && !fieldDef.optional && payload.where) {
                     throw new QueryError(`Field "${field}" doesn't support filtering`);
                 }
-                result = this.dialect.buildRelationSelection(result, model, field, parentAlias, payload);
+                if (fieldDef.originModel) {
+                    // relation is inherited from a delegate base model, need to build a join
+                    if (!joinedBases.includes(fieldDef.originModel)) {
+                        joinedBases.push(fieldDef.originModel);
+                        result = this.buildDelegateJoin(parentAlias, fieldDef.originModel, result);
+                    }
+                    result = this.dialect.buildRelationSelection(
+                        result,
+                        fieldDef.originModel,
+                        field,
+                        fieldDef.originModel,
+                        payload,
+                    );
+                } else {
+                    //  regular relation
+                    result = this.dialect.buildRelationSelection(result, model, field, parentAlias, payload);
+                }
             }
         }
 
@@ -332,14 +353,14 @@ export abstract class BaseOperationHandler<Schema extends SchemaDef> {
         return query;
     }
 
-    private buildSelectAllScalarFields(
+    buildSelectAllFields(
         model: string,
         query: SelectQueryBuilder<any, any, any>,
         omit?: Record<string, boolean | undefined>,
+        joinedBases: string[] = [],
     ) {
         const modelDef = this.requireModel(model);
         let result = query;
-        const joinedBases: string[] = [];
 
         for (const field of Object.keys(modelDef.fields)) {
             if (isRelationField(this.schema, model, field)) {
@@ -352,7 +373,7 @@ export abstract class BaseOperationHandler<Schema extends SchemaDef> {
         }
 
         // select all fields from delegate descendants and pack into a JSON field `$delegate$Model`
-        const descendants = this.getDelegateDescendantModels(model);
+        const descendants = getDelegateDescendantModels(this.schema, model);
         for (const subModel of descendants) {
             if (!joinedBases.includes(subModel.name)) {
                 joinedBases.push(subModel.name);
@@ -378,17 +399,6 @@ export abstract class BaseOperationHandler<Schema extends SchemaDef> {
         return result;
     }
 
-    private getDelegateDescendantModels(model: string, collected: Set<ModelDef> = new Set<ModelDef>()): ModelDef[] {
-        const subModels = Object.values(this.schema.models).filter((m) => m.baseModel === model);
-        subModels.forEach((def) => {
-            if (!collected.has(def)) {
-                collected.add(def);
-                this.getDelegateDescendantModels(def.name, collected);
-            }
-        });
-        return [...collected];
-    }
-
     private selectField(
         query: SelectQueryBuilder<any, any, any>,
         model: string,

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

@@ -1,5 +1,5 @@
 import type { ExpressionBuilder, ExpressionWrapper } from 'kysely';
-import { ExpressionUtils, type FieldDef, type GetModels, type SchemaDef } from '../schema';
+import { ExpressionUtils, type FieldDef, type GetModels, type ModelDef, type SchemaDef } from '../schema';
 import type { OrderBy } from './crud-types';
 import { InternalError, QueryError } from './errors';
 import type { ClientOptions } from './options';
@@ -313,3 +313,18 @@ export function getDiscriminatorField(schema: SchemaDef, model: string) {
     }
     return discriminator.value.field;
 }
+
+export function getDelegateDescendantModels(
+    schema: SchemaDef,
+    model: string,
+    collected: Set<ModelDef> = new Set<ModelDef>(),
+): ModelDef[] {
+    const subModels = Object.values(schema.models).filter((m) => m.baseModel === model);
+    subModels.forEach((def) => {
+        if (!collected.has(def)) {
+            collected.add(def);
+            getDelegateDescendantModels(schema, def.name, collected);
+        }
+    });
+    return [...collected];
+}