feat(query-builder): improve typing of qb.execute()

Also forbids calling getResult() etc on non-select QBs. Closes #2396
mikro-orm · Nov 14, 2021 · c4cfedb · c4cfedb
1 parent 89d63b3
commit c4cfedb
Show file tree

Hide file tree

Showing 2 changed files with 58 additions and 7 deletions.
diff --git a/packages/knex/src/query/QueryBuilder.ts b/packages/knex/src/query/QueryBuilder.ts
@@ -11,6 +11,7 @@ import type {
   PopulateOptions,
   QBFilterQuery,
   QueryOrderMap,
+  QueryResult,
 } from '@mikro-orm/core';
 import {
   LoadStrategy,
@@ -103,7 +104,7 @@ export class QueryBuilder<T extends AnyEntity<T> = AnyEntity> {
     this.helper = new QueryBuilderHelper(this.entityName, this.alias, this._aliasMap, this.subQueries, this.metadata, this.knex, this.platform);
   }
 
-  select(fields: Field<T> | Field<T>[], distinct = false): this {
+  select(fields: Field<T> | Field<T>[], distinct = false): SelectQueryBuilder<T> {
     this._fields = Utils.asArray(fields);
 
     if (distinct) {
@@ -113,31 +114,31 @@ export class QueryBuilder<T extends AnyEntity<T> = AnyEntity> {
     return this.init(QueryType.SELECT);
   }
 
-  addSelect(fields: Field<T> | Field<T>[]): this {
+  addSelect(fields: Field<T> | Field<T>[]): SelectQueryBuilder<T> {
     if (this.type && this.type !== QueryType.SELECT) {
       return this;
     }
 
     return this.select([...Utils.asArray(this._fields), ...Utils.asArray(fields)]);
   }
 
-  insert(data: EntityData<T> | EntityData<T>[]): this {
+  insert(data: EntityData<T> | EntityData<T>[]): InsertQueryBuilder<T> {
     return this.init(QueryType.INSERT, data);
   }
 
-  update(data: EntityData<T>): this {
+  update(data: EntityData<T>): UpdateQueryBuilder<T> {
     return this.init(QueryType.UPDATE, data);
   }
 
-  delete(cond?: QBFilterQuery): this {
+  delete(cond?: QBFilterQuery): DeleteQueryBuilder<T> {
     return this.init(QueryType.DELETE, undefined, cond);
   }
 
-  truncate(): this {
+  truncate(): TruncateQueryBuilder<T> {
     return this.init(QueryType.TRUNCATE);
   }
 
-  count(field?: string | string[], distinct = false): this {
+  count(field?: string | string[], distinct = false): CountQueryBuilder<T> {
     this._fields = [...(field ? Utils.asArray(field) : this.metadata.find(this.entityName)!.primaryKeys)];
 
     if (distinct) {
@@ -838,3 +839,24 @@ export class QueryBuilder<T extends AnyEntity<T> = AnyEntity> {
   }
 
 }
+
+export interface RunQueryBuilder<T> extends Omit<QueryBuilder<T>, 'getResult' | 'getSingleResult' | 'getResultList'> {
+  execute<U = QueryResult<T>>(method?: 'all' | 'get' | 'run', mapResults?: boolean): Promise<U>;
+}
+
+export interface SelectQueryBuilder<T> extends QueryBuilder<T> {
+  execute<U = T[]>(method?: 'all' | 'get' | 'run', mapResults?: boolean): Promise<U>;
+  execute<U = T[]>(method: 'all', mapResults?: boolean): Promise<U>;
+  execute<U = T>(method: 'get', mapResults?: boolean): Promise<U>;
+  execute<U = QueryResult<T>>(method: 'run', mapResults?: boolean): Promise<U>;
+}
+
+export interface CountQueryBuilder<T> extends SelectQueryBuilder<T> {}
+
+export interface InsertQueryBuilder<T> extends RunQueryBuilder<T> {}
+
+export interface UpdateQueryBuilder<T> extends RunQueryBuilder<T> {}
+
+export interface DeleteQueryBuilder<T> extends RunQueryBuilder<T> {}
+
+export interface TruncateQueryBuilder<T> extends RunQueryBuilder<T> {}
diff --git a/tests/QueryBuilder.test.ts b/tests/QueryBuilder.test.ts
@@ -8,6 +8,7 @@ import { initORMMySql } from './bootstrap';
 import { BaseEntity2 } from './entities-sql/BaseEntity2';
 import { performance } from 'perf_hooks';
 import { BaseEntity22 } from './entities-sql/BaseEntity22';
+import { QueryBuilder } from '@mikro-orm/postgresql';
 
 describe('QueryBuilder', () => {
 
@@ -2143,6 +2144,34 @@ describe('QueryBuilder', () => {
     expect(sql2).toBe(expected);
   });
 
+  test('execute return type works based on qb.select/insert/update/delete() being used', async () => {
+    const spy = jest.spyOn(QueryBuilder.prototype, 'execute');
+
+    spy.mockResolvedValueOnce([]);
+    const res1 = await orm.em.createQueryBuilder(Book2).select('*').execute();
+    expect(res1).toEqual([]);
+
+    spy.mockResolvedValue({ insertId: 123 });
+    const res2 = await orm.em.createQueryBuilder(Book2).insert({}).execute();
+    expect(res2.insertId).toBe(123);
+    const res3 = await orm.em.createQueryBuilder(Book2).update({}).execute();
+    expect(res3.insertId).toBe(123);
+    const res4 = await orm.em.createQueryBuilder(Book2).delete().execute();
+    expect(res4.insertId).toBe(123);
+
+    spy.mockResolvedValue([]);
+    // @ts-expect-error
+    await orm.em.createQueryBuilder(Book2).insert({}).getResultList();
+    // @ts-expect-error
+    await orm.em.createQueryBuilder(Book2).update({}).getResultList();
+    // @ts-expect-error
+    await orm.em.createQueryBuilder(Book2).delete().getResultList();
+    // @ts-expect-error
+    await orm.em.createQueryBuilder(Book2).truncate().getResultList();
+
+    spy.mockRestore();
+  });
+
   afterAll(async () => orm.close(true));
 
 });