From 7686058beccf88ad3e3e1e055b54696ce105ebb2 Mon Sep 17 00:00:00 2001
From: igalklebanov <igalklebanov@gmail.com>
Date: Mon, 7 Oct 2024 03:04:13 +0300
Subject: [PATCH] add returning/returningAll @ MergeQueryBuilder.

---
 src/query-builder/merge-query-builder.ts | 118 ++++++++++++++++++++++-
 1 file changed, 114 insertions(+), 4 deletions(-)

diff --git a/src/query-builder/merge-query-builder.ts b/src/query-builder/merge-query-builder.ts
index 8bab1b6c6..dfdcddb42 100644
--- a/src/query-builder/merge-query-builder.ts
+++ b/src/query-builder/merge-query-builder.ts
@@ -22,8 +22,17 @@ import {
 } from '../parser/join-parser.js'
 import { parseMergeThen, parseMergeWhen } from '../parser/merge-parser.js'
 import { ReferenceExpression } from '../parser/reference-parser.js'
-import { ReturningAllRow, ReturningRow } from '../parser/returning-parser.js'
-import { parseSelectAll, parseSelectArg } from '../parser/select-parser.js'
+import {
+  ReturningAllRow,
+  ReturningCallbackRow,
+  ReturningRow,
+} from '../parser/returning-parser.js'
+import {
+  parseSelectAll,
+  parseSelectArg,
+  SelectCallback,
+  SelectExpression,
+} from '../parser/select-parser.js'
 import { TableExpression } from '../parser/table-parser.js'
 import { parseTop } from '../parser/top-parser.js'
 import {
@@ -59,10 +68,13 @@ import {
   SelectExpressionFromOutputCallback,
   SelectExpressionFromOutputExpression,
 } from './output-interface.js'
+import { MultiTableReturningInterface } from './returning-interface.js'
 import { UpdateQueryBuilder } from './update-query-builder.js'
 
 export class MergeQueryBuilder<DB, TT extends keyof DB, O>
-  implements OutputInterface<DB, TT, O>
+  implements
+    MultiTableReturningInterface<DB, TT, O>,
+    OutputInterface<DB, TT, O>
 {
   readonly #props: MergeQueryBuilderProps
 
@@ -214,6 +226,48 @@ export class MergeQueryBuilder<DB, TT extends keyof DB, O>
     })
   }
 
+  returning<SE extends SelectExpression<DB, TT>>(
+    selections: ReadonlyArray<SE>,
+  ): MergeQueryBuilder<DB, TT, ReturningRow<DB, TT, O, SE>>
+
+  returning<CB extends SelectCallback<DB, TT>>(
+    callback: CB,
+  ): MergeQueryBuilder<DB, TT, ReturningCallbackRow<DB, TT, O, CB>>
+
+  returning<SE extends SelectExpression<DB, TT>>(
+    selection: SE,
+  ): MergeQueryBuilder<DB, TT, ReturningRow<DB, TT, O, SE>>
+
+  returning(args: any): any {
+    return new MergeQueryBuilder({
+      ...this.#props,
+      queryNode: QueryNode.cloneWithReturning(
+        this.#props.queryNode,
+        parseSelectArg(args),
+      ),
+    })
+  }
+
+  returningAll<T extends TT>(
+    tables: ReadonlyArray<T>,
+  ): MergeQueryBuilder<DB, TT, ReturningAllRow<DB, T, O>>
+
+  returningAll<T extends TT>(
+    table: T,
+  ): MergeQueryBuilder<DB, TT, ReturningAllRow<DB, T, O>>
+
+  returningAll(): MergeQueryBuilder<DB, TT, ReturningAllRow<DB, TT, O>>
+
+  returningAll(table?: any): any {
+    return new MergeQueryBuilder({
+      ...this.#props,
+      queryNode: QueryNode.cloneWithReturning(
+        this.#props.queryNode,
+        parseSelectAll(table),
+      ),
+    })
+  }
+
   output<OE extends OutputExpression<DB, TT>>(
     selections: readonly OE[],
   ): MergeQueryBuilder<
@@ -273,7 +327,11 @@ export class WheneableMergeQueryBuilder<
     ST extends keyof DB,
     O,
   >
-  implements Compilable<O>, OutputInterface<DB, TT, O>, OperationNodeSource
+  implements
+    Compilable<O>,
+    MultiTableReturningInterface<DB, TT | ST, O>,
+    OutputInterface<DB, TT, O>,
+    OperationNodeSource
 {
   readonly #props: MergeQueryBuilderProps
 
@@ -605,6 +663,58 @@ export class WheneableMergeQueryBuilder<
     return this.#whenNotMatched([lhs, op, rhs], true, true)
   }
 
+  returning<SE extends SelectExpression<DB, TT | ST>>(
+    selections: ReadonlyArray<SE>,
+  ): WheneableMergeQueryBuilder<DB, TT, ST, ReturningRow<DB, TT | ST, O, SE>>
+
+  returning<CB extends SelectCallback<DB, TT | ST>>(
+    callback: CB,
+  ): WheneableMergeQueryBuilder<
+    DB,
+    TT,
+    ST,
+    ReturningCallbackRow<DB, TT | ST, O, CB>
+  >
+
+  returning<SE extends SelectExpression<DB, TT | ST>>(
+    selection: SE,
+  ): WheneableMergeQueryBuilder<DB, TT, ST, ReturningRow<DB, TT | ST, O, SE>>
+
+  returning(args: any): any {
+    return new WheneableMergeQueryBuilder({
+      ...this.#props,
+      queryNode: QueryNode.cloneWithReturning(
+        this.#props.queryNode,
+        parseSelectArg(args),
+      ),
+    })
+  }
+
+  returningAll<T extends TT | ST>(
+    tables: ReadonlyArray<T>,
+  ): WheneableMergeQueryBuilder<DB, TT, ST, ReturningAllRow<DB, T, O>>
+
+  returningAll<T extends TT>(
+    table: T,
+  ): WheneableMergeQueryBuilder<DB, TT, ST, ReturningAllRow<DB, T, O>>
+
+  returningAll(): WheneableMergeQueryBuilder<
+    DB,
+    TT,
+    ST,
+    ReturningAllRow<DB, TT | ST, O>
+  >
+
+  returningAll(table?: any): any {
+    return new WheneableMergeQueryBuilder({
+      ...this.#props,
+      queryNode: QueryNode.cloneWithReturning(
+        this.#props.queryNode,
+        parseSelectAll(table),
+      ),
+    })
+  }
+
   output<OE extends OutputExpression<DB, TT>>(
     selections: readonly OE[],
   ): WheneableMergeQueryBuilder<