feat(select/distinctOn): add RawBuilder support (#239)

closes #230

src/operation-node/select-query-node.ts

@@ -15,12 +15,13 @@ import { WithNode } from './with-node.js'
 import { SelectModifierNode } from './select-modifier-node.js'
 import { ExplainNode } from './explain-node.js'
 import { SetOperationNode } from './set-operation-node.js'
+import { SimpleReferenceExpressionNode } from './simple-reference-expression-node.js'
 
 export interface SelectQueryNode extends OperationNode {
   readonly kind: 'SelectQueryNode'
   readonly from: FromNode
   readonly selections?: ReadonlyArray<SelectionNode>
-  readonly distinctOnSelections?: ReadonlyArray<SelectionNode>
+  readonly distinctOnSelections?: ReadonlyArray<SimpleReferenceExpressionNode>
   readonly joins?: ReadonlyArray<JoinNode>
   readonly groupBy?: GroupByNode
   readonly orderBy?: OrderByNode
@@ -68,7 +69,7 @@ export const SelectQueryNode = freeze({
 
   cloneWithDistinctOnSelections(
     select: SelectQueryNode,
-    selections: ReadonlyArray<SelectionNode>
+    selections: ReadonlyArray<SimpleReferenceExpressionNode>
   ): SelectQueryNode {
     return freeze({
       ...select,

src/parser/expression-parser.ts

@@ -6,6 +6,7 @@ import {
 import { AliasNode } from '../operation-node/alias-node.js'
 import { isOperationNodeSource } from '../operation-node/operation-node-source.js'
 import { OperationNode } from '../operation-node/operation-node.js'
+import { SimpleReferenceExpressionNode } from '../operation-node/simple-reference-expression-node.js'
 import { ExpressionBuilder } from '../query-builder/expression-builder.js'
 import { SelectQueryBuilder } from '../query-builder/select-query-builder.js'
 import { isFunction } from '../util/object-utils.js'
@@ -29,11 +30,13 @@ export type AliasedExpressionOrFactory<DB, TB extends keyof DB> =
 
 export function parseExpression(
   exp: ExpressionOrFactory<any, any, any>
-): OperationNode {
+): SimpleReferenceExpressionNode {
   if (isOperationNodeSource(exp)) {
-    return exp.toOperationNode()
+    return exp.toOperationNode() as SimpleReferenceExpressionNode
   } else if (isFunction(exp)) {
-    return exp(createExpressionBuilder()).toOperationNode()
+    return exp(
+      createExpressionBuilder()
+    ).toOperationNode() as SimpleReferenceExpressionNode
   }
 
   throw new Error(`invalid expression: ${JSON.stringify(exp)}`)

src/parser/reference-parser.ts

@@ -88,7 +88,7 @@ export function parseSimpleReferenceExpression(
 
 export function parseReferenceExpressionOrList(
   arg: ReferenceExpressionOrList<any, any>
-): OperationNode[] {
+): SimpleReferenceExpressionNode[] {
   if (isReadonlyArray(arg)) {
     return arg.map((it) => parseReferenceExpression(it))
   } else {
@@ -98,7 +98,7 @@ export function parseReferenceExpressionOrList(
 
 export function parseReferenceExpression(
   exp: ReferenceExpression<any, any>
-): OperationNode {
+): SimpleReferenceExpressionNode {
   if (isExpressionOrFactory(exp)) {
     return parseExpression(exp)
   }

src/query-builder/select-query-builder.ts

@@ -15,7 +15,11 @@ import {
   SelectAllQueryBuilder,
   SelectExpressionOrList,
 } from '../parser/select-parser.js'
-import { ReferenceExpression } from '../parser/reference-parser.js'
+import {
+  parseReferenceExpressionOrList,
+  ReferenceExpression,
+  ReferenceExpressionOrList,
+} from '../parser/reference-parser.js'
 import { SelectQueryNode } from '../operation-node/select-query-node.js'
 import { QueryNode } from '../operation-node/query-node.js'
 import { MergePartial, Nullable, SingleResultType } from '../util/type-utils.js'
@@ -515,20 +519,20 @@ export class SelectQueryBuilder<DB, TB extends keyof DB, O>
    * where "pet"."name" = $1
    * ```
    */
-  distinctOn<SE extends SelectExpression<DB, TB>>(
-    selections: ReadonlyArray<SE>
+  distinctOn<RE extends ReferenceExpression<DB, TB>>(
+    selections: ReadonlyArray<RE>
   ): SelectQueryBuilder<DB, TB, O>
 
-  distinctOn<SE extends SelectExpression<DB, TB>>(
-    selection: SE
+  distinctOn<RE extends ReferenceExpression<DB, TB>>(
+    selection: RE
   ): SelectQueryBuilder<DB, TB, O>
 
-  distinctOn(selection: SelectExpressionOrList<DB, TB>): any {
+  distinctOn(selection: ReferenceExpressionOrList<DB, TB>): any {
     return new SelectQueryBuilder({
       ...this.#props,
       queryNode: SelectQueryNode.cloneWithDistinctOnSelections(
         this.#props.queryNode,
-        parseSelectExpressionOrList(selection)
+        parseReferenceExpressionOrList(selection)
       ),
     })
   }

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

@@ -93,6 +93,7 @@ import { PartitionByItemNode } from '../operation-node/partition-by-item-node.js
 import { SetOperationNode } from '../operation-node/set-operation-node.js'
 import { BinaryOperationNode } from '../operation-node/binary-operation-node.js'
 import { UnaryOperationNode } from '../operation-node/unary-operation-node.js'
+import { SimpleReferenceExpressionNode } from '../operation-node/simple-reference-expression-node.js'
 
 export class DefaultQueryCompiler
   extends OperationNodeVisitor
@@ -225,7 +226,9 @@ export class DefaultQueryCompiler
     this.visitNode(node.column)
   }
 
-  protected compileDistinctOn(selections: ReadonlyArray<SelectionNode>): void {
+  protected compileDistinctOn(
+    selections: ReadonlyArray<SimpleReferenceExpressionNode>
+  ): void {
     this.append('distinct on (')
     this.compileList(selections)
     this.append(')')

test/node/src/select.test.ts

@@ -534,6 +534,30 @@ for (const dialect of BUILT_IN_DIALECTS) {
         ])
       })
 
+      it('should select with distict on that uses a RawBuilder expression', async () => {
+        const query = ctx.db
+          .selectFrom('person')
+          .select(['first_name', 'last_name'])
+          .distinctOn(sql`gender::text`)
+
+        testSql(query, dialect, {
+          postgres: {
+            sql: 'select distinct on (gender::text) "first_name", "last_name" from "person"',
+            parameters: [],
+          },
+          mysql: NOT_SUPPORTED,
+          sqlite: NOT_SUPPORTED,
+        })
+
+        const persons = await query.execute()
+
+        expect(persons).to.have.length(2)
+        expect(persons).to.eql([
+          { first_name: 'Jennifer', last_name: 'Aniston' },
+          { first_name: 'Arnold', last_name: 'Schwarzenegger' },
+        ])
+      })
+
       for (const [methods, sql] of [
         [['forUpdate'], 'for update'],
         [['forShare'], 'for share'],