fix(orm): fixed postgres orderBy in nested queries (includes/selects) (#2518)

Co-authored-by: ymc9 <104139426+ymc9@users.noreply.github.com>

Author: eredzik

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

@@ -2,9 +2,10 @@ import { invariant } from '@zenstackhq/common-helpers';
 import { type AliasableExpression, type Expression, type ExpressionBuilder, type SelectQueryBuilder } from 'kysely';
 import type { FieldDef, GetModels, SchemaDef } from '../../../schema';
 import { DELEGATE_JOINED_FIELD_PREFIX } from '../../constants';
-import type { FindArgs } from '../../crud-types';
+import type { FindArgs, NullsOrder, SortOrder } from '../../crud-types';
 import {
     buildJoinPairs,
+    ensureArray,
     getDelegateDescendantModels,
     getManyToManyRelation,
     isRelationField,
@@ -23,7 +24,10 @@ export abstract class LateralJoinDialectBase<Schema extends SchemaDef> extends B
     /**
      * Builds an array aggregation expression.
      */
-    protected abstract buildArrayAgg(arg: Expression<any>): AliasableExpression<any>;
+    protected abstract buildArrayAgg(
+        arg: Expression<any>,
+        orderBy?: { expr: Expression<any>; sort: SortOrder; nulls?: NullsOrder }[],
+    ): AliasableExpression<any>;
 
     override buildRelationSelection(
         query: SelectQueryBuilder<any, any, any>,
@@ -172,7 +176,8 @@ export abstract class LateralJoinDialectBase<Schema extends SchemaDef> extends B
             );
 
             if (relationFieldDef.array) {
-                return this.buildArrayAgg(this.buildJsonObject(objArgs)).as('$data');
+                const orderBy = this.buildRelationOrderByExpressions(relationModel, relationModelAlias, payload);
+                return this.buildArrayAgg(this.buildJsonObject(objArgs), orderBy).as('$data');
             } else {
                 return this.buildJsonObject(objArgs).as('$data');
             }
@@ -181,6 +186,50 @@ export abstract class LateralJoinDialectBase<Schema extends SchemaDef> extends B
         return qb;
     }
 
+    /**
+     * Extracts scalar `orderBy` clauses from the relation payload and maps them to
+     * the array-aggregation ordering format.
+     *
+     * For to-many relations aggregated into a JSON array (via lateral joins), this
+     * lets us preserve a stable ordering by passing `{ expr, sort, nulls? }` into
+     * the dialect's `buildArrayAgg` implementation.
+     */
+    private buildRelationOrderByExpressions(
+        model: string,
+        modelAlias: string,
+        payload: true | FindArgs<Schema, GetModels<Schema>, any, true>,
+    ): { expr: Expression<any>; sort: SortOrder; nulls?: NullsOrder }[] | undefined {
+        if (payload === true || !payload.orderBy) {
+            return undefined;
+        }
+
+        type ScalarSortValue = SortOrder | { sort: SortOrder; nulls?: NullsOrder };
+        const items: { expr: Expression<any>; sort: SortOrder; nulls?: NullsOrder }[] = [];
+
+        for (const orderBy of ensureArray(payload.orderBy)) {
+            for (const [field, value] of Object.entries(orderBy) as [string, ScalarSortValue | undefined][]) {
+                if (!value || requireField(this.schema, model, field).relation) {
+                    continue;
+                }
+
+                const expr = this.fieldRef(model, field, modelAlias);
+                let sort = typeof value === 'string' ? value : value.sort;
+                if (payload.take !== undefined && payload.take < 0) {
+                    // negative `take` requires negated sorting, and the result order
+                    // will be corrected during post-read processing
+                    sort = this.negateSort(sort, true);
+                }
+                if (typeof value === 'string') {
+                    items.push({ expr, sort });
+                } else {
+                    items.push({ expr, sort, nulls: value.nulls });
+                }
+            }
+        }
+
+        return items.length > 0 ? items : undefined;
+    }
+
     private buildRelationObjectArgs(
         relationModel: string,
         relationModelAlias: string,

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

@@ -12,7 +12,7 @@ import {
 } from 'kysely';
 import { AnyNullClass, DbNullClass, JsonNullClass } from '../../../common-types';
 import type { BuiltinType, FieldDef, SchemaDef } from '../../../schema';
-import type { SortOrder } from '../../crud-types';
+import type { NullsOrder, SortOrder } from '../../crud-types';
 import { createInvalidInputError, createNotSupportedError } from '../../errors';
 import type { ClientOptions } from '../../options';
 import { isTypeDef } from '../../query-utils';
@@ -192,7 +192,13 @@ export class MySqlCrudDialect<Schema extends SchemaDef> extends LateralJoinDiale
         return this.eb.exists(this.eb.selectFrom(innerQuery.as('$exists_sub')).select(this.eb.lit(1).as('_')));
     }
 
-    protected buildArrayAgg(arg: Expression<any>): AliasableExpression<any> {
+    protected buildArrayAgg(
+        arg: Expression<any>,
+        _orderBy?: { expr: Expression<any>; sort: SortOrder; nulls?: NullsOrder }[],
+    ): AliasableExpression<any> {
+        // MySQL doesn't support ORDER BY inside JSON_ARRAYAGG.
+        // For relation queries that need deterministic ordering, ordering is applied
+        // by the input subquery before aggregation.
         return this.eb.fn.coalesce(sql`JSON_ARRAYAGG(${arg})`, sql`JSON_ARRAY()`);
     }
 

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

@@ -11,7 +11,7 @@ import {
 import { parse as parsePostgresArray } from 'postgres-array';
 import { AnyNullClass, DbNullClass, JsonNullClass } from '../../../common-types';
 import type { BuiltinType, FieldDef, SchemaDef } from '../../../schema';
-import type { SortOrder } from '../../crud-types';
+import type { NullsOrder, SortOrder } from '../../crud-types';
 import { createInvalidInputError } from '../../errors';
 import type { ClientOptions } from '../../options';
 import { isEnum, isTypeDef } from '../../query-utils';
@@ -272,8 +272,24 @@ export class PostgresCrudDialect<Schema extends SchemaDef> extends LateralJoinDi
 
     // #region other overrides
 
-    protected buildArrayAgg(arg: Expression<any>) {
-        return this.eb.fn.coalesce(sql`jsonb_agg(${arg})`, sql`'[]'::jsonb`);
+    protected buildArrayAgg(
+        arg: Expression<any>,
+        orderBy?: { expr: Expression<any>; sort: SortOrder; nulls?: NullsOrder }[],
+    ) {
+        if (!orderBy || orderBy.length === 0) {
+            return this.eb.fn.coalesce(sql`jsonb_agg(${arg})`, sql`'[]'::jsonb`);
+        }
+
+        const orderBySql = sql.join(
+            orderBy.map(({ expr, sort, nulls }) => {
+                const dir = sql.raw(sort.toUpperCase());
+                const nullsSql = nulls ? sql` NULLS ${sql.raw(nulls.toUpperCase())}` : sql``;
+                return sql`${expr} ${dir}${nullsSql}`;
+            }),
+            sql.raw(', '),
+        );
+
+        return this.eb.fn.coalesce(sql`jsonb_agg(${arg} ORDER BY ${orderBySql})`, sql`'[]'::jsonb`);
     }
 
     override buildSkipTake(

tests/e2e/orm/client-api/relation/order-by-nested-includes.test.ts

@@ -0,0 +1,145 @@
+import { createTestClient } from '@zenstackhq/testtools';
+import { afterEach, describe, expect, it } from 'vitest';
+
+const schema = `
+model User {
+    id String @id
+    email String @unique
+    posts Post[]
+    comments Comment[]
+}
+
+model Post {
+    id String @id
+    sequence Int
+    title String
+    author User @relation(fields: [authorId], references: [id])
+    authorId String
+    comments Comment[]
+}
+
+model Comment {
+    id String @id
+    content String
+    post Post @relation(fields: [postId], references: [id])
+    postId String
+    author User? @relation(fields: [authorId], references: [id])
+    authorId String?
+}
+`;
+
+function makePostsData(count: number) {
+    return Array.from({ length: count }, (_, i) => {
+        const sequence = count - i; // insert descending
+        return {
+            id: `p${sequence}`,
+            sequence,
+            title: `P${sequence}`,
+            // Keep outer relation (User -> posts) required.
+            authorId: 'u1',
+        };
+    });
+}
+
+function makeCommentsData(count: number) {
+    return Array.from({ length: count }, (_, i) => {
+        const sequence = count - i;
+        return {
+            id: `c${sequence}`,
+            postId: `p${sequence}`,
+            content: `C${sequence}`,
+            // Make nested to-one include nullable to vary lateral join execution.
+            authorId: sequence % 11 === 0 ? null : 'u1',
+        };
+    });
+}
+
+describe('Relation orderBy with nested includes', () => {
+    let db: any;
+
+    afterEach(async () => {
+        await db?.$disconnect();
+    });
+
+    it('keeps stable order for to-many include with nested includes', async () => {
+        const count = 2000;
+
+        db = await createTestClient(schema);
+
+        await db.user.create({ data: { id: 'u1', email: 'u1@example.com' } });
+        await db.post.createMany({ data: makePostsData(count) });
+        await db.comment.createMany({ data: makeCommentsData(count) });
+
+        const user = await db.user.findFirst({
+            where: { id: 'u1' },
+            include: {
+                posts: {
+                    orderBy: { sequence: 'asc' },
+                    include: { author: true, comments: { include: { author: true } } },
+                },
+            },
+        });
+
+        const ascSequences = user.posts.map((p: any) => p.sequence);
+        expect(ascSequences).toEqual(Array.from({ length: count }, (_, i) => i + 1));
+
+        const userDesc = await db.user.findFirst({
+            where: { id: 'u1' },
+            include: {
+                posts: {
+                    orderBy: { sequence: 'desc' },
+                    include: { author: true, comments: { include: { author: true } } },
+                },
+            },
+        });
+
+        const descSequences = userDesc.posts.map((p: any) => p.sequence);
+        expect(descSequences).toEqual(Array.from({ length: count }, (_, i) => count - i));
+    });
+
+    it('keeps stable order for to-many select with nested selects', async () => {
+        const count = 2000;
+
+        db = await createTestClient(schema);
+
+        await db.user.create({ data: { id: 'u1', email: 'u1@example.com' } });
+        await db.post.createMany({ data: makePostsData(count) });
+        await db.comment.createMany({ data: makeCommentsData(count) });
+
+        const user = await db.user.findFirst({
+            where: { id: 'u1' },
+            select: {
+                id: true,
+                posts: {
+                    orderBy: { sequence: 'asc' },
+                    select: {
+                        sequence: true,
+                        author: { select: { id: true } },
+                        comments: { select: { author: { select: { id: true } } } },
+                    },
+                },
+            },
+        });
+
+        const ascSequences = user.posts.map((p: any) => p.sequence);
+        expect(ascSequences).toEqual(Array.from({ length: count }, (_, i) => i + 1));
+
+        const userDesc = await db.user.findFirst({
+            where: { id: 'u1' },
+            select: {
+                id: true,
+                posts: {
+                    orderBy: { sequence: 'desc' },
+                    select: {
+                        sequence: true,
+                        author: { select: { id: true } },
+                        comments: { select: { author: { select: { id: true } } } },
+                    },
+                },
+            },
+        });
+
+        const descSequences = userDesc.posts.map((p: any) => p.sequence);
+        expect(descSequences).toEqual(Array.from({ length: count }, (_, i) => count - i));
+    });
+});