feat(postgres): add qb.distinctOn() support

mikro-orm · Oct 2, 2022 · 307d3a1 · 307d3a1
1 parent 659c2de
commit 307d3a1
Show file tree

Hide file tree

Showing 2 changed files with 52 additions and 9 deletions.
diff --git a/packages/knex/src/query/QueryBuilder.ts b/packages/knex/src/query/QueryBuilder.ts
@@ -5,22 +5,32 @@ import type {
   Dictionary,
   EntityData,
   EntityMetadata,
+  EntityName,
   EntityProperty,
+  FilterQuery,
   FlatQueryOrderMap,
-  RequiredEntityData,
-  ObjectQuery,
+  FlushMode,
   GroupOperator,
   MetadataStorage,
+  ObjectQuery,
   PopulateOptions,
   QBFilterQuery,
+  QBQueryOrderMap,
   QueryOrderMap,
   QueryResult,
-  FlushMode,
-  FilterQuery,
-  QBQueryOrderMap,
-  EntityName,
+  RequiredEntityData,
+} from '@mikro-orm/core';
+import {
+  helper,
+  LoadStrategy,
+  LockMode,
+  PopulateHint,
+  QueryFlag,
+  QueryHelper,
+  ReferenceType,
+  Utils,
+  ValidationError,
 } from '@mikro-orm/core';
-import { LoadStrategy, LockMode, PopulateHint, QueryFlag, QueryHelper, ReferenceType, Utils, ValidationError, helper } from '@mikro-orm/core';
 import { QueryType } from './enums';
 import type { AbstractSqlDriver } from '../AbstractSqlDriver';
 import { QueryBuilderHelper } from './QueryBuilderHelper';
@@ -89,6 +99,7 @@ export class QueryBuilder<T extends object = AnyEntity> {
   private _onConflict?: { fields: string[]; ignore?: boolean; merge?: EntityData<T> | Field<T>[]; where?: QBFilterQuery<T> }[];
   private _limit?: number;
   private _offset?: number;
+  private _distinctOn?: string[];
   private _joinedProps = new Map<string, PopulateOptions<any>>();
   private _cache?: boolean | number | [string, number];
   private _indexHint?: string;
@@ -140,6 +151,16 @@ export class QueryBuilder<T extends object = AnyEntity> {
     return this.select([...Utils.asArray(this._fields), ...Utils.asArray(fields)]);
   }
 
+  distinct(): SelectQueryBuilder<T> {
+    return this.setFlag(QueryFlag.DISTINCT) as SelectQueryBuilder<T>;
+  }
+
+  /** postgres only */
+  distinctOn(fields: string | string[]): SelectQueryBuilder<T> {
+    this._distinctOn = Utils.asArray(fields);
+    return this as SelectQueryBuilder<T>;
+  }
+
   insert(data: RequiredEntityData<T> | RequiredEntityData<T>[]): InsertQueryBuilder<T> {
     return this.init(QueryType.INSERT, data) as InsertQueryBuilder<T>;
   }
@@ -866,7 +887,9 @@ export class QueryBuilder<T extends object = AnyEntity> {
       case QueryType.SELECT:
         qb.select(this.prepareFields(this._fields!));
 
-        if (this.flags.has(QueryFlag.DISTINCT)) {
+        if (this._distinctOn) {
+          qb.distinctOn(this._distinctOn as string[]);
+        } else if (this.flags.has(QueryFlag.DISTINCT)) {
           qb.distinct();
         }
 

diff --git a/tests/QueryBuilder.test.ts b/tests/QueryBuilder.test.ts
@@ -2366,7 +2366,27 @@ describe('QueryBuilder', () => {
       dbName: `mikro_orm_test`,
       type: 'postgresql',
     });
-    await pg.getSchemaGenerator().ensureDatabase();
+    await pg.schema.ensureDatabase();
+
+    {
+      const qb = pg.em.createQueryBuilder(FooBar2, 'fb1');
+      qb.select('*')
+        .distinctOn('fb1.id')
+        .joinAndSelect('fb1.baz', 'fz')
+        .leftJoinAndSelect('fz.bar', 'fb2')
+        .where({ 'fz.name': 'baz' })
+        .limit(1);
+      const sql = 'select distinct on ("fb1"."id") "fb1".*, ' +
+        '"fz"."id" as "fz__id", "fz"."name" as "fz__name", "fz"."version" as "fz__version", ' +
+        '"fb2"."id" as "fb2__id", "fb2"."name" as "fb2__name", "fb2"."name with space" as "fb2__name with space", "fb2"."baz_id" as "fb2__baz_id", "fb2"."foo_bar_id" as "fb2__foo_bar_id", "fb2"."version" as "fb2__version", "fb2"."blob" as "fb2__blob", "fb2"."array" as "fb2__array", "fb2"."object_property" as "fb2__object_property", (select 123) as "fb2__random", ' +
+        '(select 123) as "random" from "foo_bar2" as "fb1" ' +
+        'inner join "foo_baz2" as "fz" on "fb1"."baz_id" = "fz"."id" ' +
+        'left join "foo_bar2" as "fb2" on "fz"."id" = "fb2"."baz_id" ' +
+        'where "fz"."name" = $1 ' +
+        'limit $2';
+      expect(qb.getQuery()).toEqual(sql);
+      expect(qb.getParams()).toEqual(['baz', 1]);
+    }
 
     const qb01 = pg.em.createQueryBuilder(FooBar2);
     qb01.insert({ array: [] } as any);