From 17fb3a2eda04af54d1240520643de5e122253e7d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sami=20Koskim=C3=A4ki?= <koskomi@gmail.com>
Date: Sat, 16 Apr 2022 10:57:10 +0300
Subject: [PATCH] add update from queries

---
 package.json                                  |   2 +-
 .../operation-node-transformer.ts             |   1 +
 src/operation-node/update-query-node.ts       |  18 ++-
 src/parser/update-set-parser.ts               |   6 +-
 src/parser/with-parser.ts                     |   4 +-
 src/query-builder/insert-query-builder.ts     |   2 +-
 src/query-builder/on-conflict-builder.ts      |   6 +-
 src/query-builder/update-query-builder.ts     | 112 ++++++++++++++----
 src/query-compiler/default-query-compiler.ts  |   5 +
 src/query-creator.ts                          |   1 +
 test/node/src/update.test.ts                  |  24 ++++
 11 files changed, 151 insertions(+), 30 deletions(-)

diff --git a/package.json b/package.json
index d8e1b8304..3698e4d2e 100644
--- a/package.json
+++ b/package.json
@@ -1,6 +1,6 @@
 {
   "name": "kysely",
-  "version": "0.17.3",
+  "version": "0.18.0",
   "description": "Type safe SQL query builder",
   "repository": {
     "type": "git",
diff --git a/src/operation-node/operation-node-transformer.ts b/src/operation-node/operation-node-transformer.ts
index 6a7ad54d1..14edf28ce 100644
--- a/src/operation-node/operation-node-transformer.ts
+++ b/src/operation-node/operation-node-transformer.ts
@@ -439,6 +439,7 @@ export class OperationNodeTransformer {
     return {
       kind: 'UpdateQueryNode',
       table: this.transformNode(node.table),
+      from: this.transformNode(node.from),
       joins: this.transformNodeList(node.joins),
       where: this.transformNode(node.where),
       updates: this.transformNodeList(node.updates),
diff --git a/src/operation-node/update-query-node.ts b/src/operation-node/update-query-node.ts
index 8d5597a5a..3fe57b496 100644
--- a/src/operation-node/update-query-node.ts
+++ b/src/operation-node/update-query-node.ts
@@ -8,12 +8,14 @@ import { ReturningNode } from './returning-node.js'
 import { ValueListNode } from './value-list-node.js'
 import { WhereNode } from './where-node.js'
 import { WithNode } from './with-node.js'
+import { FromNode } from './from-node.js'
 
 export type UpdateValuesNode = ValueListNode | PrimitiveValueListNode
 
 export interface UpdateQueryNode extends OperationNode {
   readonly kind: 'UpdateQueryNode'
   readonly table: TableExpressionNode
+  readonly from?: FromNode
   readonly joins?: ReadonlyArray<JoinNode>
   readonly where?: WhereNode
   readonly updates?: ReadonlyArray<ColumnUpdateNode>
@@ -30,11 +32,23 @@ export const UpdateQueryNode = freeze({
   },
 
   create(table: TableExpressionNode, withNode?: WithNode): UpdateQueryNode {
-    return {
+    return freeze({
       kind: 'UpdateQueryNode',
       table,
       ...(withNode && { with: withNode }),
-    }
+    })
+  },
+
+  cloneWithFromItems(
+    updateQuery: UpdateQueryNode,
+    fromItems: ReadonlyArray<TableExpressionNode>
+  ): UpdateQueryNode {
+    return freeze({
+      ...updateQuery,
+      from: updateQuery.from
+        ? FromNode.cloneWithFroms(updateQuery.from, fromItems)
+        : FromNode.create(fromItems),
+    })
   },
 
   cloneWithUpdates(
diff --git a/src/parser/update-set-parser.ts b/src/parser/update-set-parser.ts
index f45309d11..cd00624ca 100644
--- a/src/parser/update-set-parser.ts
+++ b/src/parser/update-set-parser.ts
@@ -3,12 +3,12 @@ import { ColumnUpdateNode } from '../operation-node/column-update-node.js'
 import { UpdateKeys, UpdateType } from '../util/column-type.js'
 import { parseValueExpression, ValueExpression } from './value-parser.js'
 
-export type MutationObject<DB, TB extends keyof DB> = {
-  [C in UpdateKeys<DB[TB]>]?: ValueExpression<DB, TB, UpdateType<DB[TB][C]>>
+export type MutationObject<DB, TB extends keyof DB, TM extends keyof DB> = {
+  [C in UpdateKeys<DB[TM]>]?: ValueExpression<DB, TB, UpdateType<DB[TM][C]>>
 }
 
 export function parseUpdateObject(
-  row: MutationObject<any, any>
+  row: MutationObject<any, any, any>
 ): ReadonlyArray<ColumnUpdateNode> {
   return Object.entries(row)
     .filter(([_, value]) => value !== undefined)
diff --git a/src/parser/with-parser.ts b/src/parser/with-parser.ts
index 9a483a7b3..e838456db 100644
--- a/src/parser/with-parser.ts
+++ b/src/parser/with-parser.ts
@@ -38,7 +38,7 @@ export type QueryCreatorWithCommonTableExpression<
 type CommonTableExpressionOutput<DB, CN extends string> =
   | SelectQueryBuilder<DB, any, ExtractRowFromCommonTableExpressionName<CN>>
   | InsertQueryBuilder<DB, any, ExtractRowFromCommonTableExpressionName<CN>>
-  | UpdateQueryBuilder<DB, any, ExtractRowFromCommonTableExpressionName<CN>>
+  | UpdateQueryBuilder<DB, any, any, ExtractRowFromCommonTableExpressionName<CN>>
   | DeleteQueryBuilder<DB, any, ExtractRowFromCommonTableExpressionName<CN>>
   | RawBuilder<ExtractRowFromCommonTableExpressionName<CN>>
 
@@ -55,7 +55,7 @@ type ExtractRowFromCommonTableExpression<CTE> = CTE extends (
     ? QO
     : Q extends InsertQueryBuilder<any, any, infer QO>
     ? QO
-    : Q extends UpdateQueryBuilder<any, any, infer QO>
+    : Q extends UpdateQueryBuilder<any, any, any, infer QO>
     ? QO
     : Q extends DeleteQueryBuilder<any, any, infer QO>
     ? QO
diff --git a/src/query-builder/insert-query-builder.ts b/src/query-builder/insert-query-builder.ts
index ac098a88f..c39a0878a 100644
--- a/src/query-builder/insert-query-builder.ts
+++ b/src/query-builder/insert-query-builder.ts
@@ -475,7 +475,7 @@ export class InsertQueryBuilder<DB, TB extends keyof DB, O>
    * ```
    */
   onDuplicateKeyUpdate(
-    updates: MutationObject<DB, TB>
+    updates: MutationObject<DB, TB, TB>
   ): InsertQueryBuilder<DB, TB, O> {
     return new InsertQueryBuilder({
       ...this.#props,
diff --git a/src/query-builder/on-conflict-builder.ts b/src/query-builder/on-conflict-builder.ts
index a8bf6cef9..831570f75 100644
--- a/src/query-builder/on-conflict-builder.ts
+++ b/src/query-builder/on-conflict-builder.ts
@@ -305,7 +305,11 @@ export class OnConflictBuilder<DB, TB extends keyof DB>
    * ```
    */
   doUpdateSet(
-    updates: MutationObject<OnConflictDatabase<DB, TB>, OnConflictTables<TB>>
+    updates: MutationObject<
+      OnConflictDatabase<DB, TB>,
+      OnConflictTables<TB>,
+      OnConflictTables<TB>
+    >
   ): OnConflictUpdateBuilder<OnConflictDatabase<DB, TB>, OnConflictTables<TB>> {
     return new OnConflictUpdateBuilder({
       ...this.#props,
diff --git a/src/query-builder/update-query-builder.ts b/src/query-builder/update-query-builder.ts
index b0e476bce..bb555a03e 100644
--- a/src/query-builder/update-query-builder.ts
+++ b/src/query-builder/update-query-builder.ts
@@ -12,6 +12,8 @@ import {
   LeftJoinTableExpressionDatabase,
   RightJoinTableExpressionDatabase,
   FullJoinTableExpressionDatabase,
+  parseTableExpressionOrList,
+  TableExpressionOrList,
 } from '../parser/table-parser.js'
 import {
   parseSelectExpressionOrList,
@@ -55,7 +57,7 @@ import { ReturningInterface } from './returning-interface.js'
 import { NoResultError, NoResultErrorConstructor } from './no-result-error.js'
 import { Selectable } from '../util/column-type.js'
 
-export class UpdateQueryBuilder<DB, TB extends keyof DB, O>
+export class UpdateQueryBuilder<DB, UT extends keyof DB, TB extends keyof DB, O>
   implements
     WhereInterface<DB, TB>,
     JoinInterface<DB, TB>,
@@ -73,10 +75,10 @@ export class UpdateQueryBuilder<DB, TB extends keyof DB, O>
     lhs: RE,
     op: FilterOperator,
     rhs: FilterValueExpressionOrList<DB, TB, RE>
-  ): UpdateQueryBuilder<DB, TB, O>
+  ): UpdateQueryBuilder<DB, UT, TB, O>
 
-  where(grouper: WhereGrouper<DB, TB>): UpdateQueryBuilder<DB, TB, O>
-  where(raw: AnyRawBuilder): UpdateQueryBuilder<DB, TB, O>
+  where(grouper: WhereGrouper<DB, TB>): UpdateQueryBuilder<DB, UT, TB, O>
+  where(raw: AnyRawBuilder): UpdateQueryBuilder<DB, UT, TB, O>
 
   where(...args: any[]): any {
     return new UpdateQueryBuilder({
@@ -92,7 +94,7 @@ export class UpdateQueryBuilder<DB, TB extends keyof DB, O>
     lhs: ReferenceExpression<DB, TB>,
     op: FilterOperator,
     rhs: ReferenceExpression<DB, TB>
-  ): UpdateQueryBuilder<DB, TB, O> {
+  ): UpdateQueryBuilder<DB, UT, TB, O> {
     return new UpdateQueryBuilder({
       ...this.#props,
       queryNode: QueryNode.cloneWithWhere(
@@ -106,10 +108,10 @@ export class UpdateQueryBuilder<DB, TB extends keyof DB, O>
     lhs: RE,
     op: FilterOperator,
     rhs: FilterValueExpressionOrList<DB, TB, RE>
-  ): UpdateQueryBuilder<DB, TB, O>
+  ): UpdateQueryBuilder<DB, UT, TB, O>
 
-  orWhere(grouper: WhereGrouper<DB, TB>): UpdateQueryBuilder<DB, TB, O>
-  orWhere(raw: AnyRawBuilder): UpdateQueryBuilder<DB, TB, O>
+  orWhere(grouper: WhereGrouper<DB, TB>): UpdateQueryBuilder<DB, UT, TB, O>
+  orWhere(raw: AnyRawBuilder): UpdateQueryBuilder<DB, UT, TB, O>
 
   orWhere(...args: any[]): any {
     return new UpdateQueryBuilder({
@@ -125,7 +127,7 @@ export class UpdateQueryBuilder<DB, TB extends keyof DB, O>
     lhs: ReferenceExpression<DB, TB>,
     op: FilterOperator,
     rhs: ReferenceExpression<DB, TB>
-  ): UpdateQueryBuilder<DB, TB, O> {
+  ): UpdateQueryBuilder<DB, UT, TB, O> {
     return new UpdateQueryBuilder({
       ...this.#props,
       queryNode: QueryNode.cloneWithOrWhere(
@@ -135,7 +137,9 @@ export class UpdateQueryBuilder<DB, TB extends keyof DB, O>
     })
   }
 
-  whereExists(arg: ExistsExpression<DB, TB>): UpdateQueryBuilder<DB, TB, O> {
+  whereExists(
+    arg: ExistsExpression<DB, TB>
+  ): UpdateQueryBuilder<DB, UT, TB, O> {
     return new UpdateQueryBuilder({
       ...this.#props,
       queryNode: QueryNode.cloneWithWhere(
@@ -145,7 +149,9 @@ export class UpdateQueryBuilder<DB, TB extends keyof DB, O>
     })
   }
 
-  whereNotExists(arg: ExistsExpression<DB, TB>): UpdateQueryBuilder<DB, TB, O> {
+  whereNotExists(
+    arg: ExistsExpression<DB, TB>
+  ): UpdateQueryBuilder<DB, UT, TB, O> {
     return new UpdateQueryBuilder({
       ...this.#props,
       queryNode: QueryNode.cloneWithWhere(
@@ -155,7 +161,9 @@ export class UpdateQueryBuilder<DB, TB extends keyof DB, O>
     })
   }
 
-  orWhereExists(arg: ExistsExpression<DB, TB>): UpdateQueryBuilder<DB, TB, O> {
+  orWhereExists(
+    arg: ExistsExpression<DB, TB>
+  ): UpdateQueryBuilder<DB, UT, TB, O> {
     return new UpdateQueryBuilder({
       ...this.#props,
       queryNode: QueryNode.cloneWithOrWhere(
@@ -167,7 +175,7 @@ export class UpdateQueryBuilder<DB, TB extends keyof DB, O>
 
   orWhereNotExists(
     arg: ExistsExpression<DB, TB>
-  ): UpdateQueryBuilder<DB, TB, O> {
+  ): UpdateQueryBuilder<DB, UT, TB, O> {
     return new UpdateQueryBuilder({
       ...this.#props,
       queryNode: QueryNode.cloneWithOrWhere(
@@ -177,6 +185,61 @@ export class UpdateQueryBuilder<DB, TB extends keyof DB, O>
     })
   }
 
+  /**
+   * Adds a from clause to the update query.
+   *
+   * This is supported only on some databases like PostgreSQL.
+   *
+   * The API is the same as {@link QueryCreator.selectFrom}.
+   *
+   * ### Examples
+   *
+   * ```ts
+   * db.updateTable('person')
+   *   .from('pet')
+   *   .set({
+   *     first_name: (eb) => eb.ref('pet.name')
+   *   })
+   *   .whereRef('pet.owner_id', '=', 'person.id')
+   * ```
+   *
+   * The generated SQL (PostgreSQL):
+   *
+   * ```sql
+   * update "person"
+   * set "first_name" = "pet"."name"
+   * from "pet"
+   * where "pet"."owner_id" = "person"."id"
+   * ```
+   */
+  from<TE extends TableExpression<DB, TB>>(
+    table: TE
+  ): UpdateQueryBuilder<
+    TableExpressionDatabase<DB, TE>,
+    UT,
+    TableExpressionTables<DB, TB, TE>,
+    O
+  >
+
+  from<TE extends TableExpression<DB, TB>>(
+    table: TE[]
+  ): UpdateQueryBuilder<
+    TableExpressionDatabase<DB, TE>,
+    UT,
+    TableExpressionTables<DB, TB, TE>,
+    O
+  >
+
+  from(from: TableExpressionOrList<any, any>): any {
+    return new UpdateQueryBuilder({
+      ...this.#props,
+      queryNode: UpdateQueryNode.cloneWithFromItems(
+        this.#props.queryNode,
+        parseTableExpressionOrList(from)
+      ),
+    })
+  }
+
   innerJoin<
     TE extends TableExpression<DB, TB>,
     K1 extends JoinReferenceExpression<DB, TB, TE>,
@@ -187,6 +250,7 @@ export class UpdateQueryBuilder<DB, TB extends keyof DB, O>
     k2: K2
   ): UpdateQueryBuilder<
     TableExpressionDatabase<DB, TE>,
+    UT,
     TableExpressionTables<DB, TB, TE>,
     O
   >
@@ -199,6 +263,7 @@ export class UpdateQueryBuilder<DB, TB extends keyof DB, O>
     callback: FN
   ): UpdateQueryBuilder<
     TableExpressionDatabase<DB, TE>,
+    UT,
     TableExpressionTables<DB, TB, TE>,
     O
   >
@@ -223,6 +288,7 @@ export class UpdateQueryBuilder<DB, TB extends keyof DB, O>
     k2: K2
   ): UpdateQueryBuilder<
     LeftJoinTableExpressionDatabase<DB, TE>,
+    UT,
     TableExpressionTables<DB, TB, TE>,
     O
   >
@@ -235,6 +301,7 @@ export class UpdateQueryBuilder<DB, TB extends keyof DB, O>
     callback: FN
   ): UpdateQueryBuilder<
     LeftJoinTableExpressionDatabase<DB, TE>,
+    UT,
     TableExpressionTables<DB, TB, TE>,
     O
   >
@@ -259,6 +326,7 @@ export class UpdateQueryBuilder<DB, TB extends keyof DB, O>
     k2: K2
   ): UpdateQueryBuilder<
     RightJoinTableExpressionDatabase<DB, TB, TE>,
+    UT,
     TableExpressionTables<DB, TB, TE>,
     O
   >
@@ -271,6 +339,7 @@ export class UpdateQueryBuilder<DB, TB extends keyof DB, O>
     callback: FN
   ): UpdateQueryBuilder<
     RightJoinTableExpressionDatabase<DB, TB, TE>,
+    UT,
     TableExpressionTables<DB, TB, TE>,
     O
   >
@@ -295,6 +364,7 @@ export class UpdateQueryBuilder<DB, TB extends keyof DB, O>
     k2: K2
   ): UpdateQueryBuilder<
     FullJoinTableExpressionDatabase<DB, TB, TE>,
+    UT,
     TableExpressionTables<DB, TB, TE>,
     O
   >
@@ -307,6 +377,7 @@ export class UpdateQueryBuilder<DB, TB extends keyof DB, O>
     callback: FN
   ): UpdateQueryBuilder<
     FullJoinTableExpressionDatabase<DB, TB, TE>,
+    UT,
     TableExpressionTables<DB, TB, TE>,
     O
   >
@@ -408,7 +479,7 @@ export class UpdateQueryBuilder<DB, TB extends keyof DB, O>
    * where "id" = $4
    * ```
    */
-  set(row: MutationObject<DB, TB>): UpdateQueryBuilder<DB, TB, O> {
+  set(row: MutationObject<DB, TB, UT>): UpdateQueryBuilder<DB, UT, TB, O> {
     return new UpdateQueryBuilder({
       ...this.#props,
       queryNode: UpdateQueryNode.cloneWithUpdates(
@@ -420,11 +491,11 @@ export class UpdateQueryBuilder<DB, TB extends keyof DB, O>
 
   returning<SE extends SelectExpression<DB, TB>>(
     selections: ReadonlyArray<SE>
-  ): UpdateQueryBuilder<DB, TB, ReturningRow<DB, TB, O, SE>>
+  ): UpdateQueryBuilder<DB, UT, TB, ReturningRow<DB, TB, O, SE>>
 
   returning<SE extends SelectExpression<DB, TB>>(
     selection: SE
-  ): UpdateQueryBuilder<DB, TB, ReturningRow<DB, TB, O, SE>>
+  ): UpdateQueryBuilder<DB, UT, TB, ReturningRow<DB, TB, O, SE>>
 
   returning(selection: SelectExpressionOrList<DB, TB>): any {
     return new UpdateQueryBuilder({
@@ -436,7 +507,7 @@ export class UpdateQueryBuilder<DB, TB extends keyof DB, O>
     })
   }
 
-  returningAll(): UpdateQueryBuilder<DB, TB, Selectable<DB[TB]>> {
+  returningAll(): UpdateQueryBuilder<DB, UT, TB, Selectable<DB[TB]>> {
     return new UpdateQueryBuilder({
       ...this.#props,
       queryNode: QueryNode.cloneWithReturning(
@@ -510,9 +581,10 @@ export class UpdateQueryBuilder<DB, TB extends keyof DB, O>
    */
   if<O2>(
     condition: boolean,
-    func: (qb: this) => UpdateQueryBuilder<DB, TB, O2>
+    func: (qb: this) => UpdateQueryBuilder<DB, UT, TB, O2>
   ): UpdateQueryBuilder<
     DB,
+    UT,
     TB,
     O2 extends UpdateResult
       ? UpdateResult
@@ -535,14 +607,14 @@ export class UpdateQueryBuilder<DB, TB extends keyof DB, O>
    * You should only use this method as the last resort if the types
    * don't support your use case.
    */
-  castTo<T>(): UpdateQueryBuilder<DB, TB, T> {
+  castTo<T>(): UpdateQueryBuilder<DB, UT, TB, T> {
     return new UpdateQueryBuilder(this.#props)
   }
 
   /**
    * Returns a copy of this UpdateQueryBuilder instance with the given plugin installed.
    */
-  withPlugin(plugin: KyselyPlugin): UpdateQueryBuilder<DB, TB, O> {
+  withPlugin(plugin: KyselyPlugin): UpdateQueryBuilder<DB, UT, TB, O> {
     return new UpdateQueryBuilder({
       ...this.#props,
       executor: this.#props.executor.withPlugin(plugin),
diff --git a/src/query-compiler/default-query-compiler.ts b/src/query-compiler/default-query-compiler.ts
index c31fae0e4..4b382a29c 100644
--- a/src/query-compiler/default-query-compiler.ts
+++ b/src/query-compiler/default-query-compiler.ts
@@ -627,6 +627,11 @@ export class DefaultQueryCompiler
       this.compileList(node.updates)
     }
 
+    if (node.from) {
+      this.append(' ')
+      this.visitNode(node.from)
+    }
+
     if (node.joins) {
       this.append(' ')
       this.compileList(node.joins, ' ')
diff --git a/src/query-creator.ts b/src/query-creator.ts
index caaa91e11..9f8ac676c 100644
--- a/src/query-creator.ts
+++ b/src/query-creator.ts
@@ -287,6 +287,7 @@ export class QueryCreator<DB> {
   ): UpdateQueryBuilder<
     TableExpressionDatabase<DB, TR>,
     TableExpressionTables<DB, never, TR>,
+    TableExpressionTables<DB, never, TR>,
     UpdateResult
   > {
     return new UpdateQueryBuilder({
diff --git a/test/node/src/update.test.ts b/test/node/src/update.test.ts
index 48b9943e6..5626c40e7 100644
--- a/test/node/src/update.test.ts
+++ b/test/node/src/update.test.ts
@@ -207,6 +207,30 @@ for (const dialect of BUILT_IN_DIALECTS) {
         const result = await query.executeTakeFirstOrThrow()
         expect(result.last_name).to.equal('Barson')
       })
+
+      it('should join a table when `from` is called', async () => {
+        const query = ctx.db
+          .updateTable('person')
+          .from('pet')
+          .set({
+            first_name: (eb) => eb.ref('pet.name'),
+          })
+          .whereRef('pet.owner_id', '=', 'person.id')
+          .where('person.first_name', '=', 'Arnold')
+          .returning('first_name')
+
+        testSql(query, dialect, {
+          postgres: {
+            sql: 'update "person" set "first_name" = "pet"."name" from "pet" where "pet"."owner_id" = "person"."id" and "person"."first_name" = $1 returning "first_name"',
+            parameters: ['Arnold'],
+          },
+          mysql: NOT_SUPPORTED,
+          sqlite: NOT_SUPPORTED,
+        })
+
+        const result = await query.execute()
+        expect(result[0].first_name).to.equal('Doggo')
+      })
     }
   })
 }