Add orderBy clause support to aggregate function builder (#896)

Co-authored-by: Igal Klebanov <igalklebanov@gmail.com>
Co-authored-by: Dev K0te <dev.k0te42@gmail.com>

src/operation-node/aggregate-function-node.ts

@@ -2,12 +2,15 @@ import { freeze } from '../util/object-utils.js'
 import { OperationNode } from './operation-node.js'
 import { OverNode } from './over-node.js'
 import { WhereNode } from './where-node.js'
+import { OrderByNode } from './order-by-node.js'
+import { OrderByItemNode } from './order-by-item-node.js'
 
 export interface AggregateFunctionNode extends OperationNode {
   readonly kind: 'AggregateFunctionNode'
   readonly func: string
   readonly aggregated: readonly OperationNode[]
   readonly distinct?: boolean
+  readonly orderBy?: OrderByNode
   readonly filter?: WhereNode
   readonly over?: OverNode
 }
@@ -40,6 +43,18 @@ export const AggregateFunctionNode = freeze({
     })
   },
 
+  cloneWithOrderBy(
+    aggregateFunctionNode: AggregateFunctionNode,
+    orderItems: ReadonlyArray<OrderByItemNode>,
+  ): AggregateFunctionNode {
+    return freeze({
+      ...aggregateFunctionNode,
+      orderBy: aggregateFunctionNode.orderBy
+        ? OrderByNode.cloneWithItems(aggregateFunctionNode.orderBy, orderItems)
+        : OrderByNode.create(orderItems),
+    })
+  },
+
   cloneWithFilter(
     aggregateFunctionNode: AggregateFunctionNode,
     filter: OperationNode,

src/operation-node/operation-node-transformer.ts

@@ -883,6 +883,7 @@ export class OperationNodeTransformer {
       kind: 'AggregateFunctionNode',
       aggregated: this.transformNodeList(node.aggregated),
       distinct: node.distinct,
+      orderBy: this.transformNode(node.orderBy),
       filter: this.transformNode(node.filter),
       func: node.func,
       over: this.transformNode(node.over),

src/query-builder/aggregate-function-builder.ts

@@ -10,7 +10,10 @@ import {
   AliasedExpression,
   Expression,
 } from '../expression/expression.js'
-import { ReferenceExpression } from '../parser/reference-parser.js'
+import {
+  ReferenceExpression,
+  StringReference,
+} from '../parser/reference-parser.js'
 import {
   ComparisonOperatorExpression,
   OperandValueExpressionOrList,
@@ -19,6 +22,11 @@ import {
 } from '../parser/binary-operation-parser.js'
 import { SqlBool } from '../util/type-utils.js'
 import { ExpressionOrFactory } from '../parser/expression-parser.js'
+import { DynamicReferenceBuilder } from '../dynamic/dynamic-reference-builder.js'
+import {
+  OrderByDirectionExpression,
+  parseOrderBy,
+} from '../parser/order-by-parser.js'
 
 export class AggregateFunctionBuilder<DB, TB extends keyof DB, O = unknown>
   implements AliasableExpression<O>
@@ -95,6 +103,42 @@ export class AggregateFunctionBuilder<DB, TB extends keyof DB, O = unknown>
     })
   }
 
+  /**
+   * Adds an `order by` clause inside the aggregate function.
+   *
+   * ### Examples
+   *
+   * ```ts
+   * const result = await db
+   *   .selectFrom('person')
+   *   .innerJoin('pet', 'pet.owner_id', 'person.id')
+   *   .select((eb) =>
+   *     eb.fn.jsonAgg('pet.name').orderBy('pet.name').as('person_pets')
+   *   )
+   *   .executeTakeFirstOrThrow()
+   * ```
+   *
+   * The generated SQL (PostgreSQL):
+   *
+   * ```sql
+   * select json_agg("pet"."name" order by "pet"."name") as "person_pets"
+   * from "person"
+   * inner join "pet" ON "pet"."owner_id" = "person"."id"
+   * ```
+   */
+  orderBy<OE extends StringReference<DB, TB> | DynamicReferenceBuilder<any>>(
+    orderBy: OE,
+    direction?: OrderByDirectionExpression,
+  ): AggregateFunctionBuilder<DB, TB, O> {
+    return new AggregateFunctionBuilder({
+      ...this.#props,
+      aggregateFunctionNode: AggregateFunctionNode.cloneWithOrderBy(
+        this.#props.aggregateFunctionNode,
+        parseOrderBy([orderBy, direction]),
+      ),
+    })
+  }
+
   /**
    * Adds a `filter` clause with a nested `where` clause after the function.
    *

src/query-compiler/default-query-compiler.ts

@@ -1374,6 +1374,12 @@ export class DefaultQueryCompiler
     }
 
     this.compileList(node.aggregated)
+
+    if (node.orderBy) {
+      this.append(' ')
+      this.visitNode(node.orderBy)
+    }
+
     this.append(')')
 
     if (node.filter) {

test/node/src/aggregate-function.test.ts

@@ -2,6 +2,7 @@ import {
   AggregateFunctionBuilder,
   ExpressionBuilder,
   SimpleReferenceExpression,
+  ReferenceExpression,
   sql,
 } from '../../../'
 import {
@@ -1106,6 +1107,57 @@ for (const dialect of DIALECTS) {
 
       await query.execute()
     })
+
+    describe(`should execute order-sensitive aggregate functions`, () => {
+      if (dialect === 'postgres' || dialect === 'mysql' || dialect === 'sqlite') {
+        const isMySql = dialect === 'mysql'
+        const funcName = isMySql ? 'group_concat' : 'string_agg'
+        const funcArgs: Array<ReferenceExpression<Database, 'person'>> = [
+          'first_name',
+        ]
+        if (!isMySql) {
+          funcArgs.push(sql.lit(','))
+        }
+
+        it(`should execute a query with ${funcName}(column order by column) in select clause`, async () => {
+          const query = ctx.db
+            .selectFrom('person')
+            .select((eb) =>
+              eb.fn
+                .agg(funcName, funcArgs)
+                .orderBy('first_name', 'desc')
+                .as('first_names'),
+            )
+
+          testSql(query, dialect, {
+            postgres: {
+              sql: [
+                `select ${funcName}("first_name", ',' order by "first_name" desc) as "first_names"`,
+                `from "person"`,
+              ],
+              parameters: [],
+            },
+            mysql: {
+              sql: [
+                `select ${funcName}(\`first_name\` order by \`first_name\` desc) as \`first_names\``,
+                `from \`person\``,
+              ],
+              parameters: [],
+            },
+            mssql: NOT_SUPPORTED,
+            sqlite: {
+              sql: [
+                `select ${funcName}("first_name", ',' order by "first_name" desc) as "first_names"`,
+                `from "person"`,
+              ],
+              parameters: [],
+            },
+          })
+
+          await query.execute()
+        })
+      }
+    })
   })
 }
 

test/typings/test-d/aggregate-function.test-d.ts

@@ -4,8 +4,8 @@ import {
   expectNotAssignable,
   expectType,
 } from 'tsd'
-import { Generated, Kysely, sql } from '..'
-import { Database } from '../shared'
+import { Generated, Kysely, Selectable, sql } from '..'
+import { Database, Pet } from '../shared'
 
 async function testSelectWithoutAs(db: Kysely<Database>) {
   const { avg, count, countAll, max, min, sum } = db.fn