add modifyEnd to insert, update and delete query builders (#871)

Co-authored-by: Igal Klebanov <igalklebanov@gmail.com>

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

@@ -24,6 +24,7 @@ export interface DeleteQueryNode extends OperationNode {
   readonly orderBy?: OrderByNode
   readonly limit?: LimitNode
   readonly explain?: ExplainNode
+  readonly endModifiers?: ReadonlyArray<OperationNode>
   readonly top?: TopNode
   readonly output?: OutputNode
 }

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

@@ -25,6 +25,7 @@ export interface InsertQueryNode extends OperationNode {
   readonly replace?: boolean
   readonly explain?: ExplainNode
   readonly defaultValues?: boolean
+  readonly endModifiers?: ReadonlyArray<OperationNode>
   readonly top?: TopNode
   readonly output?: OutputNode
 }

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

@@ -16,6 +16,7 @@ export interface MergeQueryNode extends OperationNode {
   readonly with?: WithNode
   readonly top?: TopNode
   readonly output?: OutputNode
+  readonly endModifiers?: ReadonlyArray<OperationNode>
 }
 
 /**

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

@@ -380,6 +380,7 @@ export class OperationNodeTransformer {
       returning: this.transformNode(node.returning),
       onConflict: this.transformNode(node.onConflict),
       onDuplicateKey: this.transformNode(node.onDuplicateKey),
+      endModifiers: this.transformNodeList(node.endModifiers),
       with: this.transformNode(node.with),
       ignore: node.ignore,
       replace: node.replace,
@@ -405,6 +406,7 @@ export class OperationNodeTransformer {
       joins: this.transformNodeList(node.joins),
       where: this.transformNode(node.where),
       returning: this.transformNode(node.returning),
+      endModifiers: this.transformNodeList(node.endModifiers),
       with: this.transformNode(node.with),
       orderBy: this.transformNode(node.orderBy),
       limit: this.transformNode(node.limit),
@@ -514,6 +516,7 @@ export class OperationNodeTransformer {
       where: this.transformNode(node.where),
       updates: this.transformNodeList(node.updates),
       returning: this.transformNode(node.returning),
+      endModifiers: this.transformNodeList(node.endModifiers),
       with: this.transformNode(node.with),
       explain: this.transformNode(node.explain),
       limit: this.transformNode(node.limit),
@@ -1013,6 +1016,7 @@ export class OperationNodeTransformer {
       whens: this.transformNodeList(node.whens),
       with: this.transformNode(node.with),
       top: this.transformNode(node.top),
+      endModifiers: this.transformNodeList(node.endModifiers),
       output: this.transformNode(node.output),
     })
   }

src/operation-node/query-node.ts

@@ -28,6 +28,7 @@ type HasReturning = { returning?: ReturningNode }
 type HasExplain = { explain?: ExplainNode }
 type HasTop = { top?: TopNode }
 type HasOutput = { output?: OutputNode }
+type HasEndModifiers = { endModifiers?: ReadonlyArray<OperationNode> }
 
 /**
  * @internal
@@ -43,6 +44,18 @@ export const QueryNode = freeze({
     )
   },
 
+  cloneWithEndModifier<T extends HasEndModifiers>(
+    node: T,
+    modifier: OperationNode,
+  ): T {
+    return freeze({
+      ...node,
+      endModifiers: node.endModifiers
+        ? freeze([...node.endModifiers, modifier])
+        : freeze([modifier]),
+    })
+  },
+
   cloneWithWhere<T extends HasWhere>(node: T, operation: OperationNode): T {
     return freeze({
       ...node,

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

@@ -101,18 +101,6 @@ export const SelectQueryNode = freeze({
     })
   },
 
-  cloneWithEndModifier(
-    select: SelectQueryNode,
-    modifier: SelectModifierNode,
-  ): SelectQueryNode {
-    return freeze({
-      ...select,
-      endModifiers: select.endModifiers
-        ? freeze([...select.endModifiers, modifier])
-        : freeze([modifier]),
-    })
-  },
-
   cloneWithOrderByItems(
     selectNode: SelectQueryNode,
     items: ReadonlyArray<OrderByItemNode>,

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

@@ -25,6 +25,7 @@ export interface UpdateQueryNode extends OperationNode {
   readonly returning?: ReturningNode
   readonly with?: WithNode
   readonly explain?: ExplainNode
+  readonly endModifiers?: ReadonlyArray<OperationNode>
   readonly limit?: LimitNode
   readonly top?: TopNode
   readonly output?: OutputNode

src/query-builder/delete-query-builder.ts

@@ -823,6 +823,37 @@ export class DeleteQueryBuilder<DB, TB extends keyof DB, O>
     })
   }
 
+
+  /**
+   * This can be used to add any additional SQL to the end of the query.
+   *
+   * ### Examples
+   *
+   * ```ts
+   * await db.deleteFrom('person')
+   * .where('first_name', '=', 'John')
+   * .modifyEnd(sql.raw('-- This is a comment'))
+   * .execute()
+   * ```
+   *
+   * The generated SQL (MySQL):
+   *
+   * ```sql
+   * delete from `person`
+   * where `first_name` = "John" -- This is a comment
+   * ```
+   */
+  modifyEnd(modifier: Expression<any>): DeleteQueryBuilder<DB, TB, O> {
+    return new DeleteQueryBuilder({
+      ...this.#props,
+      queryNode: QueryNode.cloneWithEndModifier(
+        this.#props.queryNode,
+        modifier.toOperationNode(),
+      ),
+    })
+  }
+
+
   /**
    * Simply calls the provided function passing `this` as the only argument. `$call` returns
    * what the provided function returns.

src/query-builder/insert-query-builder.ts

@@ -350,6 +350,36 @@ export class InsertQueryBuilder<DB, TB extends keyof DB, O>
     })
   }
 
+
+  /**
+   * This can be used to add any additional SQL to the end of the query.
+   *
+   * ### Examples
+   *
+   * ```ts
+   * await db.insertInto('person')
+   *   .values(values)
+   *   .modifyEnd(sql.raw('-- This is a comment'))
+   *   .execute()
+   * ```
+   *
+   * The generated SQL (MySQL):
+   *
+   * ```sql
+   * insert into `person` 
+   * values (?, ?, ?) -- This is a comment
+   * ```
+   */
+  modifyEnd(modifier: Expression<any>): InsertQueryBuilder<DB, TB, O> {
+    return new InsertQueryBuilder({
+      ...this.#props,
+      queryNode: QueryNode.cloneWithEndModifier(
+        this.#props.queryNode,
+        modifier.toOperationNode(),
+      ),
+    })
+  }
+
   /**
    * Changes an `insert into` query to an `insert ignore into` query.
    *

src/query-builder/merge-query-builder.ts

@@ -1,4 +1,4 @@
-import { AliasedExpression } from '../expression/expression.js'
+import { AliasedExpression, type Expression } from '../expression/expression.js'
 import { InsertQueryNode } from '../operation-node/insert-query-node.js'
 import { MergeQueryNode } from '../operation-node/merge-query-node.js'
 import { OperationNodeSource } from '../operation-node/operation-node-source.js'
@@ -71,6 +71,37 @@ export class MergeQueryBuilder<DB, TT extends keyof DB, O>
     this.#props = freeze(props)
   }
 
+  /**
+   * This can be used to add any additional SQL to the end of the query.
+   *
+   * ### Examples
+   *
+   * ```ts
+   * const result = await db
+   *   .mergeInto('person')
+   *   .using('pet', 'pet.owner_id', 'person.id')
+   *   .whenMatched()
+   *   .thenDelete()
+   *   .modifyEnd(sql.raw('-- this is a comment'))
+   *   .execute()
+   * ```
+   * 
+   * The generated SQL (PostgreSQL):
+   *
+   * ```sql
+   * merge into "person" using "pet" on "pet"."owner_id" = "person"."id" when matched then delete -- this is a comment
+   * ```
+   */
+  modifyEnd(modifier: Expression<any>): MergeQueryBuilder<DB, TT, O> {
+    return new MergeQueryBuilder({
+      ...this.#props,
+      queryNode: QueryNode.cloneWithEndModifier(
+        this.#props.queryNode,
+        modifier.toOperationNode(),
+      ),
+    })
+  }
+
   /**
    * Changes a `merge into` query to an `merge top into` query.
    *
@@ -256,6 +287,37 @@ export class WheneableMergeQueryBuilder<
     this.#props = freeze(props)
   }
 
+  /**
+   * This can be used to add any additional SQL to the end of the query.
+   *
+   * ### Examples
+   *
+   * ```ts
+   * const result = await db
+   *   .mergeInto('person')
+   *   .using('pet', 'pet.owner_id', 'person.id')
+   *   .whenMatched()
+   *   .thenDelete()
+   *   .modifyEnd(sql.raw('-- this is a comment'))
+   *   .execute()
+   * ```
+   * 
+   * The generated SQL (PostgreSQL):
+   *
+   * ```sql
+   * merge into "person" using "pet" on "pet"."owner_id" = "person"."id" when matched then delete -- this is a comment
+   * ```
+   */
+  modifyEnd(modifier: Expression<any>): WheneableMergeQueryBuilder<DB, TT, ST, O> {
+    return new WheneableMergeQueryBuilder({
+      ...this.#props,
+      queryNode: QueryNode.cloneWithEndModifier(
+        this.#props.queryNode,
+        modifier.toOperationNode(),
+      ),
+    })
+  }
+
   /**
    * See {@link MergeQueryBuilder.top}.
    */

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

@@ -1959,7 +1959,7 @@ class SelectQueryBuilderImpl<DB, TB extends keyof DB, O>
   modifyEnd(modifier: Expression<any>): SelectQueryBuilder<DB, TB, O> {
     return new SelectQueryBuilderImpl({
       ...this.#props,
-      queryNode: SelectQueryNode.cloneWithEndModifier(
+      queryNode: QueryNode.cloneWithEndModifier(
         this.#props.queryNode,
         SelectModifierNode.createWithExpression(modifier.toOperationNode()),
       ),
@@ -1979,7 +1979,7 @@ class SelectQueryBuilderImpl<DB, TB extends keyof DB, O>
   forUpdate(of?: TableOrList<TB>): SelectQueryBuilder<DB, TB, O> {
     return new SelectQueryBuilderImpl({
       ...this.#props,
-      queryNode: SelectQueryNode.cloneWithEndModifier(
+      queryNode: QueryNode.cloneWithEndModifier(
         this.#props.queryNode,
         SelectModifierNode.create(
           'ForUpdate',
@@ -1992,7 +1992,7 @@ class SelectQueryBuilderImpl<DB, TB extends keyof DB, O>
   forShare(of?: TableOrList<TB>): SelectQueryBuilder<DB, TB, O> {
     return new SelectQueryBuilderImpl({
       ...this.#props,
-      queryNode: SelectQueryNode.cloneWithEndModifier(
+      queryNode: QueryNode.cloneWithEndModifier(
         this.#props.queryNode,
         SelectModifierNode.create(
           'ForShare',
@@ -2005,7 +2005,7 @@ class SelectQueryBuilderImpl<DB, TB extends keyof DB, O>
   forKeyShare(of?: TableOrList<TB>): SelectQueryBuilder<DB, TB, O> {
     return new SelectQueryBuilderImpl({
       ...this.#props,
-      queryNode: SelectQueryNode.cloneWithEndModifier(
+      queryNode: QueryNode.cloneWithEndModifier(
         this.#props.queryNode,
         SelectModifierNode.create(
           'ForKeyShare',
@@ -2018,7 +2018,7 @@ class SelectQueryBuilderImpl<DB, TB extends keyof DB, O>
   forNoKeyUpdate(of?: TableOrList<TB>): SelectQueryBuilder<DB, TB, O> {
     return new SelectQueryBuilderImpl({
       ...this.#props,
-      queryNode: SelectQueryNode.cloneWithEndModifier(
+      queryNode: QueryNode.cloneWithEndModifier(
         this.#props.queryNode,
         SelectModifierNode.create(
           'ForNoKeyUpdate',
@@ -2031,7 +2031,7 @@ class SelectQueryBuilderImpl<DB, TB extends keyof DB, O>
   skipLocked(): SelectQueryBuilder<DB, TB, O> {
     return new SelectQueryBuilderImpl({
       ...this.#props,
-      queryNode: SelectQueryNode.cloneWithEndModifier(
+      queryNode: QueryNode.cloneWithEndModifier(
         this.#props.queryNode,
         SelectModifierNode.create('SkipLocked'),
       ),
@@ -2041,7 +2041,7 @@ class SelectQueryBuilderImpl<DB, TB extends keyof DB, O>
   noWait(): SelectQueryBuilder<DB, TB, O> {
     return new SelectQueryBuilderImpl({
       ...this.#props,
-      queryNode: SelectQueryNode.cloneWithEndModifier(
+      queryNode: QueryNode.cloneWithEndModifier(
         this.#props.queryNode,
         SelectModifierNode.create('NoWait'),
       ),

src/query-builder/update-query-builder.ts

@@ -751,6 +751,37 @@ export class UpdateQueryBuilder<DB, UT extends keyof DB, TB extends keyof DB, O>
     })
   }
 
+  /**
+   * This can be used to add any additional SQL to the end of the query.
+   *
+   * ### Examples
+   *
+    * ```ts
+   * await db.updateTable('person')
+   * .set({ age: 39 })
+   * .where('first_name', '=', 'John')
+   * .modifyEnd(sql.raw('-- This is a comment'))
+   * .execute()
+   * ```
+   *
+   * The generated SQL (MySQL):
+   *
+   * ```sql
+   * update `person`
+   * set `age` = 39
+   * where `first_name` = "John" -- This is a comment
+   * ```
+   */
+  modifyEnd(modifier: Expression<any>): UpdateQueryBuilder<DB, UT, TB, O> {
+    return new UpdateQueryBuilder({
+      ...this.#props,
+      queryNode: QueryNode.cloneWithEndModifier(
+        this.#props.queryNode,
+        modifier.toOperationNode(),
+      ),
+    })
+  }
+
   /**
    * Clears all `returning` clauses from the query.
    *