Support for update of and friends. Closes #683

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

@@ -811,6 +811,7 @@ export class OperationNodeTransformer {
       kind: 'SelectModifierNode',
       modifier: node.modifier,
       rawModifier: this.transformNode(node.rawModifier),
+      of: this.transformNodeList(node.of),
     })
   }
 

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

@@ -14,6 +14,7 @@ export interface SelectModifierNode extends OperationNode {
   readonly kind: 'SelectModifierNode'
   readonly modifier?: SelectModifier
   readonly rawModifier?: OperationNode
+  readonly of?: ReadonlyArray<OperationNode>
 }
 
 /**
@@ -24,10 +25,14 @@ export const SelectModifierNode = freeze({
     return node.kind === 'SelectModifierNode'
   },
 
-  create(modifier: SelectModifier): SelectModifierNode {
+  create(
+    modifier: SelectModifier,
+    of?: ReadonlyArray<OperationNode>
+  ): SelectModifierNode {
     return freeze({
       kind: 'SelectModifierNode',
       modifier,
+      of,
     })
   },
 

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

@@ -6,7 +6,7 @@ import {
   JoinReferenceExpression,
   parseJoin,
 } from '../parser/join-parser.js'
-import { TableExpression } from '../parser/table-parser.js'
+import { TableExpression, parseTable } from '../parser/table-parser.js'
 import {
   parseSelectArg,
   parseSelectAll,
@@ -46,7 +46,7 @@ import { OffsetNode } from '../operation-node/offset-node.js'
 import { Compilable } from '../util/compilable.js'
 import { QueryExecutor } from '../query-executor/query-executor.js'
 import { QueryId } from '../util/query-id.js'
-import { freeze } from '../util/object-utils.js'
+import { asArray, freeze } from '../util/object-utils.js'
 import { GroupByArg, parseGroupBy } from '../parser/group-by-parser.js'
 import { KyselyPlugin } from '../plugin/kysely-plugin.js'
 import { WhereInterface } from './where-interface.js'
@@ -429,22 +429,22 @@ export interface SelectQueryBuilder<DB, TB extends keyof DB, O>
   /**
    * Adds the `for update` modifier to a select query on supported databases.
    */
-  forUpdate(): SelectQueryBuilder<DB, TB, O>
+  forUpdate(of?: TableOrList<TB>): SelectQueryBuilder<DB, TB, O>
 
   /**
    * Adds the `for share` modifier to a select query on supported databases.
    */
-  forShare(): SelectQueryBuilder<DB, TB, O>
+  forShare(of?: TableOrList<TB>): SelectQueryBuilder<DB, TB, O>
 
   /**
    * Adds the `for key share` modifier to a select query on supported databases.
    */
-  forKeyShare(): SelectQueryBuilder<DB, TB, O>
+  forKeyShare(of?: TableOrList<TB>): SelectQueryBuilder<DB, TB, O>
 
   /**
    * Adds the `for no key update` modifier to a select query on supported databases.
    */
-  forNoKeyUpdate(): SelectQueryBuilder<DB, TB, O>
+  forNoKeyUpdate(of?: TableOrList<TB>): SelectQueryBuilder<DB, TB, O>
 
   /**
    * Adds the `skip locked` modifier to a select query on supported databases.
@@ -1797,42 +1797,54 @@ class SelectQueryBuilderImpl<DB, TB extends keyof DB, O>
     })
   }
 
-  forUpdate(): SelectQueryBuilder<DB, TB, O> {
+  forUpdate(of?: TableOrList<TB>): SelectQueryBuilder<DB, TB, O> {
     return new SelectQueryBuilderImpl({
       ...this.#props,
       queryNode: SelectQueryNode.cloneWithEndModifier(
         this.#props.queryNode,
-        SelectModifierNode.create('ForUpdate')
+        SelectModifierNode.create(
+          'ForUpdate',
+          of ? asArray(of).map(parseTable) : undefined
+        )
       ),
     })
   }
 
-  forShare(): SelectQueryBuilder<DB, TB, O> {
+  forShare(of?: TableOrList<TB>): SelectQueryBuilder<DB, TB, O> {
     return new SelectQueryBuilderImpl({
       ...this.#props,
       queryNode: SelectQueryNode.cloneWithEndModifier(
         this.#props.queryNode,
-        SelectModifierNode.create('ForShare')
+        SelectModifierNode.create(
+          'ForShare',
+          of ? asArray(of).map(parseTable) : undefined
+        )
       ),
     })
   }
 
-  forKeyShare(): SelectQueryBuilder<DB, TB, O> {
+  forKeyShare(of?: TableOrList<TB>): SelectQueryBuilder<DB, TB, O> {
     return new SelectQueryBuilderImpl({
       ...this.#props,
       queryNode: SelectQueryNode.cloneWithEndModifier(
         this.#props.queryNode,
-        SelectModifierNode.create('ForKeyShare')
+        SelectModifierNode.create(
+          'ForKeyShare',
+          of ? asArray(of).map(parseTable) : undefined
+        )
       ),
     })
   }
 
-  forNoKeyUpdate(): SelectQueryBuilder<DB, TB, O> {
+  forNoKeyUpdate(of?: TableOrList<TB>): SelectQueryBuilder<DB, TB, O> {
     return new SelectQueryBuilderImpl({
       ...this.#props,
       queryNode: SelectQueryNode.cloneWithEndModifier(
         this.#props.queryNode,
-        SelectModifierNode.create('ForNoKeyUpdate')
+        SelectModifierNode.create(
+          'ForNoKeyUpdate',
+          of ? asArray(of).map(parseTable) : undefined
+        )
       ),
     })
   }
@@ -2413,3 +2425,7 @@ type OuterJoinedBuilderDB<
     ? DB[C]
     : never
 }>
+
+type TableOrList<TB extends keyof any> =
+  | (TB & string)
+  | ReadonlyArray<TB & string>

src/query-compiler/default-query-compiler.ts

@@ -1219,6 +1219,11 @@ export class DefaultQueryCompiler
     } else {
       this.append(SELECT_MODIFIER_SQL[node.modifier!])
     }
+
+    if (node.of) {
+      this.append(' of ')
+      this.compileList(node.of, ', ')
+    }
   }
 
   protected override visitCreateType(node: CreateTypeNode): void {

src/util/object-utils.ts

@@ -76,8 +76,8 @@ export function freeze<T>(obj: T): Readonly<T> {
   return Object.freeze(obj)
 }
 
-export function asArray<T>(arg: T | T[]): T[] {
-  if (Array.isArray(arg)) {
+export function asArray<T>(arg: T | ReadonlyArray<T>): ReadonlyArray<T> {
+  if (isReadonlyArray(arg)) {
     return arg
   } else {
     return [arg]

test/node/src/select.test.ts

@@ -645,6 +645,32 @@ for (const dialect of DIALECTS) {
         expect(persons).to.eql([{ last_name: 'Aniston' }])
       })
 
+      it('should select a row for update of', async () => {
+        const query = ctx.db
+          .selectFrom('person')
+          .select('last_name')
+          .where('first_name', '=', 'Jennifer')
+          .forUpdate('person')
+
+        testSql(query, dialect, {
+          postgres: {
+            sql: 'select "last_name" from "person" where "first_name" = $1 for update of "person"',
+            parameters: ['Jennifer'],
+          },
+          mysql: {
+            sql: 'select `last_name` from `person` where `first_name` = ? for update of "person"',
+            parameters: ['Jennifer'],
+          },
+          mssql: NOT_SUPPORTED,
+          sqlite: NOT_SUPPORTED,
+        })
+
+        const persons = await query.execute()
+
+        expect(persons).to.have.length(1)
+        expect(persons).to.eql([{ last_name: 'Aniston' }])
+      })
+
       it('should select a row for update with skip locked', async () => {
         const query = ctx.db
           .selectFrom('person')
@@ -672,6 +698,33 @@ for (const dialect of DIALECTS) {
         expect(persons).to.eql([{ last_name: 'Aniston' }])
       })
 
+      it('should select a row for update of with skip locked', async () => {
+        const query = ctx.db
+          .selectFrom('person')
+          .select('last_name')
+          .where('first_name', '=', 'Jennifer')
+          .forUpdate(['person'])
+          .skipLocked()
+
+        testSql(query, dialect, {
+          postgres: {
+            sql: 'select "last_name" from "person" where "first_name" = $1 for update of "person" skip locked',
+            parameters: ['Jennifer'],
+          },
+          mysql: {
+            sql: 'select `last_name` from `person` where `first_name` = ? for update of "person" skip locked',
+            parameters: ['Jennifer'],
+          },
+          mssql: NOT_SUPPORTED,
+          sqlite: NOT_SUPPORTED,
+        })
+
+        const persons = await query.execute()
+
+        expect(persons).to.have.length(1)
+        expect(persons).to.eql([{ last_name: 'Aniston' }])
+      })
+
       it('should select a row for update with skipLocked called before forUpdate', async () => {
         const query = ctx.db
           .selectFrom('person')