From 074da5f86d0e828e0a76f4c123527bb28a61863d Mon Sep 17 00:00:00 2001
From: Gary Mathews <contact@garymathews.com>
Date: Tue, 9 Jan 2024 22:53:03 -0800
Subject: [PATCH] implement type predicates for query builders

---
 src/index.ts                              |   1 +
 src/query-builder/delete-query-builder.ts |   4 +
 src/query-builder/insert-query-builder.ts |   4 +
 src/query-builder/merge-query-builder.ts  |  16 +++
 src/query-builder/update-query-builder.ts |   4 +
 src/util/query-utils.ts                   | 129 ++++++++++++++++++++
 test/node/src/query-utils.test.ts         | 136 ++++++++++++++++++++++
 test/typings/test-d/query-utils.test-d.ts |  98 ++++++++++++++++
 8 files changed, 392 insertions(+)
 create mode 100644 src/util/query-utils.ts
 create mode 100644 test/node/src/query-utils.test.ts
 create mode 100644 test/typings/test-d/query-utils.test-d.ts

diff --git a/src/index.ts b/src/index.ts
index 7213c01d2..196c695b4 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -207,6 +207,7 @@ export * from './util/compilable.js'
 export * from './util/explainable.js'
 export * from './util/streamable.js'
 export * from './util/log.js'
+export * from './util/query-utils.js'
 export {
   AnyAliasedColumn,
   AnyAliasedColumnWithTable,
diff --git a/src/query-builder/delete-query-builder.ts b/src/query-builder/delete-query-builder.ts
index c0daf3b5c..2ef727988 100644
--- a/src/query-builder/delete-query-builder.ts
+++ b/src/query-builder/delete-query-builder.ts
@@ -87,6 +87,10 @@ export class DeleteQueryBuilder<DB, TB extends keyof DB, O>
     this.#props = freeze(props)
   }
 
+  get isDeleteQueryBuilder(): true {
+    return true
+  }
+
   where<
     RE extends ReferenceExpression<DB, TB>,
     VE extends OperandValueExpressionOrList<DB, TB, RE>,
diff --git a/src/query-builder/insert-query-builder.ts b/src/query-builder/insert-query-builder.ts
index 05168695a..3be7ad97b 100644
--- a/src/query-builder/insert-query-builder.ts
+++ b/src/query-builder/insert-query-builder.ts
@@ -73,6 +73,10 @@ export class InsertQueryBuilder<DB, TB extends keyof DB, O>
     this.#props = freeze(props)
   }
 
+  get isInsertQueryBuilder(): true {
+    return true
+  }
+
   /**
    * Sets the values to insert for an {@link Kysely.insertInto | insert} query.
    *
diff --git a/src/query-builder/merge-query-builder.ts b/src/query-builder/merge-query-builder.ts
index c00951997..8b8b9a286 100644
--- a/src/query-builder/merge-query-builder.ts
+++ b/src/query-builder/merge-query-builder.ts
@@ -58,6 +58,10 @@ export class MergeQueryBuilder<DB, TT extends keyof DB, O> {
     this.#props = freeze(props)
   }
 
+  get isMergeQueryBuilder(): true {
+    return true
+  }
+
   /**
    * Adds the `using` clause to the query.
    *
@@ -137,6 +141,10 @@ export class WheneableMergeQueryBuilder<
     this.#props = freeze(props)
   }
 
+  get isMergeQueryBuilder(): true {
+    return true
+  }
+
   /**
    * Adds a simple `when matched` clause to the query.
    *
@@ -596,6 +604,10 @@ export class MatchedThenableMergeQueryBuilder<
     this.#props = freeze(props)
   }
 
+  get isMergeQueryBuilder(): true {
+    return true
+  }
+
   /**
    * Performs the `delete` action.
    *
@@ -796,6 +808,10 @@ export class NotMatchedThenableMergeQueryBuilder<
     this.#props = freeze(props)
   }
 
+  get isMergeQueryBuilder(): true {
+    return true
+  }
+
   /**
    * Performs the `do nothing` action.
    *
diff --git a/src/query-builder/update-query-builder.ts b/src/query-builder/update-query-builder.ts
index 1651fa18a..ba45eb31c 100644
--- a/src/query-builder/update-query-builder.ts
+++ b/src/query-builder/update-query-builder.ts
@@ -83,6 +83,10 @@ export class UpdateQueryBuilder<DB, UT extends keyof DB, TB extends keyof DB, O>
     this.#props = freeze(props)
   }
 
+  get isUpdateQueryBuilder(): true {
+    return true
+  }
+
   where<
     RE extends ReferenceExpression<DB, TB>,
     VE extends OperandValueExpressionOrList<DB, TB, RE>,
diff --git a/src/util/query-utils.ts b/src/util/query-utils.ts
new file mode 100644
index 000000000..f6a92bb6b
--- /dev/null
+++ b/src/util/query-utils.ts
@@ -0,0 +1,129 @@
+import { DeleteQueryBuilder } from '../query-builder/delete-query-builder.js'
+import { InsertQueryBuilder } from '../query-builder/insert-query-builder.js'
+import {
+  MatchedThenableMergeQueryBuilder,
+  MergeQueryBuilder,
+  NotMatchedThenableMergeQueryBuilder,
+  WheneableMergeQueryBuilder,
+} from '../query-builder/merge-query-builder.js'
+import { SelectQueryBuilder } from '../query-builder/select-query-builder.js'
+import { UpdateQueryBuilder } from '../query-builder/update-query-builder.js'
+
+type AnyMergeQueryBuilder<
+  DB,
+  TT extends keyof DB,
+  ST extends keyof DB,
+  UT extends TT | ST,
+  O
+> =
+  | MergeQueryBuilder<DB, TT, O>
+  | WheneableMergeQueryBuilder<DB, TT, ST, O>
+  | MatchedThenableMergeQueryBuilder<DB, TT, ST, UT, O>
+  | NotMatchedThenableMergeQueryBuilder<DB, TT, ST, O>
+
+/**
+ * A helper type that can accept any query builder object.
+ *
+ * You can use helper methods to determine the type of query builder object.
+ *
+ * See {@link isSelectQueryBuilder}, {@link isInsertQueryBuilder}, {@link isUpdateQueryBuilder},
+ * {@link isDeleteQueryBuilder}, {@link isMergeQueryBuilder}.
+ *
+ * ### Example
+ * ```ts
+ * import { AnyQueryBuilder, isSelectQueryBuilder } from 'kysely'
+ *
+ * function alwaysSelectAll<DB, TB extends keyof DB, O>(qb: AnyQueryBuilder<DB, TB, O>) {
+ *   // Here `qb` could be any query builder object.
+ *   // We can use the helper method `isSelectQueryBuilder` to determine the type.
+ *   if (isSelectQueryBuilder(qb)) {
+ *     // We can now use `SelectQueryBuilder` methods and properties.
+ *     return qb.clearSelect().selectAll()
+ *   }
+ *   return qb
+ * }
+ * ```
+ */
+export type AnyQueryBuilder<
+  DB,
+  TB extends keyof DB,
+  O,
+  UT extends keyof DB = any,
+  ST extends keyof DB = any
+> =
+  | SelectQueryBuilder<DB, TB, O>
+  | InsertQueryBuilder<DB, TB, O>
+  | UpdateQueryBuilder<DB, UT, TB, O>
+  | DeleteQueryBuilder<DB, TB, O>
+  | AnyMergeQueryBuilder<DB, TB, ST, TB | ST, O>
+
+/**
+ * A helper method to determine if a query builder object is of type {@link SelectQueryBuilder}.
+ *
+ * Useful when using the {@link AnyQueryBuilder} type.
+ */
+export function isSelectQueryBuilder<DB, TB extends keyof DB, O>(
+  qb: AnyQueryBuilder<DB, TB, O>
+): qb is SelectQueryBuilder<DB, TB, O> {
+  return !!(qb as SelectQueryBuilder<DB, TB, O>).isSelectQueryBuilder
+}
+
+/**
+ * A helper method to determine if a query builder object is of type {@link InsertQueryBuilder}.
+ *
+ * Useful when using the {@link AnyQueryBuilder} type.
+ */
+export function isInsertQueryBuilder<DB, TB extends keyof DB, O>(
+  qb: AnyQueryBuilder<DB, TB, O>
+): qb is InsertQueryBuilder<DB, TB, O> {
+  return !!(qb as InsertQueryBuilder<DB, TB, O>).isInsertQueryBuilder
+}
+
+/**
+ * A helper method to determine if a query builder object is of type {@link UpdateQueryBuilder}.
+ *
+ * Useful when using the {@link AnyQueryBuilder} type.
+ */
+export function isUpdateQueryBuilder<
+  DB,
+  UT extends keyof DB,
+  TB extends keyof DB,
+  O
+>(qb: AnyQueryBuilder<DB, TB, O, UT>): qb is UpdateQueryBuilder<DB, UT, TB, O> {
+  return !!(qb as UpdateQueryBuilder<DB, UT, TB, O>).isUpdateQueryBuilder
+}
+
+/**
+ * A helper method to determine if a query builder object is of type {@link DeleteQueryBuilder}.
+ *
+ * Useful when using the {@link AnyQueryBuilder} type.
+ */
+export function isDeleteQueryBuilder<DB, TB extends keyof DB, O>(
+  qb: AnyQueryBuilder<DB, TB, O>
+): qb is DeleteQueryBuilder<DB, TB, O> {
+  return !!(qb as DeleteQueryBuilder<DB, TB, O>).isDeleteQueryBuilder
+}
+
+/**
+ * A helper method to determine if a query builder object is of type {@link MergeQueryBuilder}.
+ *
+ * Useful when using the {@link AnyQueryBuilder} type.
+ */
+export function isMergeQueryBuilder<
+  DB,
+  TB extends keyof DB,
+  O,
+  ST extends keyof DB = any
+>(
+  qb: AnyQueryBuilder<DB, TB, O, ST>
+): qb is typeof qb extends MergeQueryBuilder<DB, TB, O>
+  ? MergeQueryBuilder<DB, TB, O>
+  : typeof qb extends WheneableMergeQueryBuilder<DB, TB, ST, O>
+  ? WheneableMergeQueryBuilder<DB, TB, ST, O>
+  : typeof qb extends MatchedThenableMergeQueryBuilder<DB, TB, ST, TB | ST, O>
+  ? MatchedThenableMergeQueryBuilder<DB, TB, ST, TB | ST, O>
+  : typeof qb extends NotMatchedThenableMergeQueryBuilder<DB, TB, ST, O>
+  ? NotMatchedThenableMergeQueryBuilder<DB, TB, ST, O>
+  : never {
+  return !!(qb as any).isMergeQueryBuilder
+}
diff --git a/test/node/src/query-utils.test.ts b/test/node/src/query-utils.test.ts
new file mode 100644
index 000000000..17cb5581d
--- /dev/null
+++ b/test/node/src/query-utils.test.ts
@@ -0,0 +1,136 @@
+import {
+  AnyQueryBuilder,
+  MergeResult,
+  isDeleteQueryBuilder,
+  isInsertQueryBuilder,
+  isMergeQueryBuilder,
+  isSelectQueryBuilder,
+  isUpdateQueryBuilder,
+} from '../../../'
+
+import {
+  destroyTest,
+  initTest,
+  TestContext,
+  expect,
+  DIALECTS,
+  Database,
+} from './test-setup.js'
+
+for (const dialect of DIALECTS) {
+  describe(`${dialect}: query-utils`, () => {
+    let ctx: TestContext
+
+    before(async function () {
+      ctx = await initTest(this, dialect)
+    })
+
+    after(async () => {
+      await destroyTest(ctx)
+    })
+
+    it('should isSelectQueryBuilder be true', async () => {
+      const query = ctx.db
+        .selectFrom('person')
+        .where('first_name', '=', 'Jennifer')
+
+      expect(query.isSelectQueryBuilder).to.be.true
+      expect(isSelectQueryBuilder(query)).to.be.true
+      expect(isInsertQueryBuilder(query)).to.be.false
+      expect(isUpdateQueryBuilder(query)).to.be.false
+      expect(isDeleteQueryBuilder(query)).to.be.false
+      expect(isMergeQueryBuilder(query)).to.be.false
+    })
+
+    it('should isInsertQueryBuilder be true', async () => {
+      const query = ctx.db.insertInto('person').values({
+        first_name: 'David',
+        last_name: 'Bowie',
+        gender: 'male',
+      })
+
+      expect(query.isInsertQueryBuilder).to.be.true
+      expect(isSelectQueryBuilder(query)).to.be.false
+      expect(isInsertQueryBuilder(query)).to.be.true
+      expect(isUpdateQueryBuilder(query)).to.be.false
+      expect(isDeleteQueryBuilder(query)).to.be.false
+      expect(isMergeQueryBuilder(query)).to.be.false
+    })
+
+    it('should isUpdateQueryBuilder be true', async () => {
+      const query = ctx.db
+        .updateTable('person')
+        .where('first_name', '=', 'John')
+        .set({ last_name: 'Wick' })
+
+      expect(query.isUpdateQueryBuilder).to.be.true
+      expect(isSelectQueryBuilder(query)).to.be.false
+      expect(isInsertQueryBuilder(query)).to.be.false
+      expect(isUpdateQueryBuilder(query)).to.be.true
+      expect(isDeleteQueryBuilder(query)).to.be.false
+      expect(isMergeQueryBuilder(query)).to.be.false
+    })
+
+    it('should isDeleteQueryBuilder be true', async () => {
+      const query = ctx.db.deleteFrom('person').where('first_name', '=', 'John')
+
+      expect(query.isDeleteQueryBuilder).to.be.true
+      expect(isSelectQueryBuilder(query)).to.be.false
+      expect(isInsertQueryBuilder(query)).to.be.false
+      expect(isUpdateQueryBuilder(query)).to.be.false
+      expect(isDeleteQueryBuilder(query)).to.be.true
+      expect(isMergeQueryBuilder(query)).to.be.false
+    })
+
+    it('should isMergeQueryBuilder be true', async () => {
+      // MergeQueryBuilder
+      let query: AnyQueryBuilder<Database, 'person', MergeResult> =
+        ctx.db.mergeInto('person')
+
+      expect(query.isMergeQueryBuilder).to.be.true
+      expect(isSelectQueryBuilder(query)).to.be.false
+      expect(isInsertQueryBuilder(query)).to.be.false
+      expect(isUpdateQueryBuilder(query)).to.be.false
+      expect(isDeleteQueryBuilder(query)).to.be.false
+      expect(isMergeQueryBuilder(query)).to.be.true
+
+      // WheneableMergeQueryBuilder
+      query = ctx.db
+        .mergeInto('person')
+        .using('pet', 'person.id', 'pet.owner_id')
+
+      expect(query.isMergeQueryBuilder).to.be.true
+      expect(isSelectQueryBuilder(query)).to.be.false
+      expect(isInsertQueryBuilder(query)).to.be.false
+      expect(isUpdateQueryBuilder(query)).to.be.false
+      expect(isDeleteQueryBuilder(query)).to.be.false
+      expect(isMergeQueryBuilder(query)).to.be.true
+
+      // MatchedThenableMergeQueryBuilder
+      query = ctx.db
+        .mergeInto('person')
+        .using('pet', 'person.id', 'pet.owner_id')
+        .whenMatched()
+
+      expect(query.isMergeQueryBuilder).to.be.true
+      expect(isSelectQueryBuilder(query)).to.be.false
+      expect(isInsertQueryBuilder(query)).to.be.false
+      expect(isUpdateQueryBuilder(query)).to.be.false
+      expect(isDeleteQueryBuilder(query)).to.be.false
+      expect(isMergeQueryBuilder(query)).to.be.true
+
+      // NotMatchedThenableMergeQueryBuilder
+      query = ctx.db
+        .mergeInto('person')
+        .using('pet', 'person.id', 'pet.owner_id')
+        .whenNotMatched()
+
+      expect(query.isMergeQueryBuilder).to.be.true
+      expect(isSelectQueryBuilder(query)).to.be.false
+      expect(isInsertQueryBuilder(query)).to.be.false
+      expect(isUpdateQueryBuilder(query)).to.be.false
+      expect(isDeleteQueryBuilder(query)).to.be.false
+      expect(isMergeQueryBuilder(query)).to.be.true
+    })
+  })
+}
diff --git a/test/typings/test-d/query-utils.test-d.ts b/test/typings/test-d/query-utils.test-d.ts
new file mode 100644
index 000000000..8fb5eface
--- /dev/null
+++ b/test/typings/test-d/query-utils.test-d.ts
@@ -0,0 +1,98 @@
+import {
+  AnyQueryBuilder,
+  DeleteResult,
+  InsertResult,
+  Kysely,
+  MatchedThenableMergeQueryBuilder,
+  MergeQueryBuilder,
+  MergeResult,
+  NotMatchedThenableMergeQueryBuilder,
+  UpdateResult,
+  WheneableMergeQueryBuilder,
+  isDeleteQueryBuilder,
+  isInsertQueryBuilder,
+  isMergeQueryBuilder,
+  isSelectQueryBuilder,
+  isUpdateQueryBuilder,
+} from '..'
+import { Database } from '../shared'
+import { expectType } from 'tsd'
+
+function testIsSelectQueryBuilder(db: Kysely<Database>) {
+  const query: AnyQueryBuilder<Database, 'person', { age: number }> = db
+    .selectFrom('person')
+    .select('age')
+  expectType<true>(query.isSelectQueryBuilder)
+  expectType<boolean>(isSelectQueryBuilder(query))
+}
+
+function testIsInsertQueryBuilder(db: Kysely<Database>) {
+  const query: AnyQueryBuilder<Database, 'person', InsertResult> = db
+    .insertInto('person')
+    .values({
+      first_name: 'David',
+      last_name: 'Bowie',
+      gender: 'male',
+      age: 69,
+    })
+  expectType<true>(query.isInsertQueryBuilder)
+  expectType<boolean>(isInsertQueryBuilder(query))
+}
+
+function testIsUpdateQueryBuilder(db: Kysely<Database>) {
+  const query: AnyQueryBuilder<Database, 'person', UpdateResult> = db
+    .updateTable('pet')
+    .where('name', '=', 'Max')
+    .set({ name: 'Dana' })
+  expectType<true>(query.isUpdateQueryBuilder)
+  expectType<boolean>(isUpdateQueryBuilder(query))
+}
+
+function testIsDeleteQueryBuilder(db: Kysely<Database>) {
+  const query: AnyQueryBuilder<Database, 'person', DeleteResult> = db
+    .deleteFrom('person')
+    .where('gender', '=', 'male')
+  expectType<true>(query.isDeleteQueryBuilder)
+  expectType<boolean>(isDeleteQueryBuilder(query))
+}
+
+function testIsMergeQueryBuilder(db: Kysely<Database>) {
+  let query: AnyQueryBuilder<Database, 'person', MergeResult> =
+    db.mergeInto('person')
+  expectType<MergeQueryBuilder<Database, 'person', MergeResult>>(query)
+  expectType<true>(query.isMergeQueryBuilder)
+  expectType<boolean>(isMergeQueryBuilder(query))
+
+  query = db.mergeInto('person').using('pet', 'pet.owner_id', 'person.id')
+  expectType<
+    WheneableMergeQueryBuilder<Database, 'person', 'pet', MergeResult>
+  >(query)
+  expectType<true>(query.isMergeQueryBuilder)
+  expectType<boolean>(isMergeQueryBuilder(query))
+
+  query = db
+    .mergeInto('person')
+    .using('pet', 'pet.owner_id', 'person.id')
+    .whenMatched()
+  expectType<
+    MatchedThenableMergeQueryBuilder<
+      Database,
+      'person',
+      'pet',
+      'person' | 'pet',
+      MergeResult
+    >
+  >(query)
+  expectType<true>(query.isMergeQueryBuilder)
+  expectType<boolean>(isMergeQueryBuilder(query))
+
+  query = db
+    .mergeInto('person')
+    .using('pet', 'pet.owner_id', 'person.id')
+    .whenNotMatched()
+  expectType<
+    NotMatchedThenableMergeQueryBuilder<Database, 'person', 'pet', MergeResult>
+  >(query)
+  expectType<true>(query.isMergeQueryBuilder)
+  expectType<boolean>(isMergeQueryBuilder(query))
+}