feat: add setOnLocked for SKIP LOCKED and NOWAIT (typeorm#9317)

nordinh · Aug 29, 2022 · b499d85 · b499d85
1 parent 2780140
commit b499d85
Show file tree

Hide file tree

Showing 11 changed files with 402 additions and 33 deletions.
diff --git a/docs/find-options.md b/docs/find-options.md
@@ -219,10 +219,19 @@ or
     mode: "pessimistic_read" |
         "pessimistic_write" |
         "dirty_read" |
+        /*
+            "pessimistic_partial_write" and "pessimistic_write_or_fail" are deprecated and
+            will be removed in a future version.
+
+            Use onLocked instead.
+         */
         "pessimistic_partial_write" |
         "pessimistic_write_or_fail" |
         "for_no_key_update" |
-        "for_key_share"
+        "for_key_share",
+
+    tables: string[],
+    onLocked: "nowait" | "skip_locked"
 }
 ```
 
@@ -237,19 +246,7 @@ userRepository.findOne({
 })
 ```
 
-Support of lock modes, and SQL statements they translate to, are listed in the table below (blank cell denotes unsupported). When specified lock mode is not supported, a `LockNotSupportedOnGivenDriverError` error will be thrown.
-
-```text
-|                 | pessimistic_read         | pessimistic_write       | dirty_read    | pessimistic_partial_write   | pessimistic_write_or_fail   | for_no_key_update   | for_key_share |
-| --------------- | --------------------     | ----------------------- | ------------- | --------------------------- | --------------------------- | ------------------- | ------------- |
-| MySQL           | LOCK IN SHARE MODE       | FOR UPDATE              | (nothing)     | FOR UPDATE SKIP LOCKED      | FOR UPDATE NOWAIT           |                     |               |
-| Postgres        | FOR SHARE                | FOR UPDATE              | (nothing)     | FOR UPDATE SKIP LOCKED      | FOR UPDATE NOWAIT           | FOR NO KEY UPDATE   | FOR KEY SHARE |
-| Oracle          | FOR UPDATE               | FOR UPDATE              | (nothing)     |                             |                             |                     |               |
-| SQL Server      | WITH (HOLDLOCK, ROWLOCK) | WITH (UPDLOCK, ROWLOCK) | WITH (NOLOCK) |                             |                             |                     |               |
-| AuroraDataApi   | LOCK IN SHARE MODE       | FOR UPDATE              | (nothing)     |                             |                             |                     |               |
-| CockroachDB     |                          | FOR UPDATE              | (nothing)     |                             | FOR UPDATE NOWAIT           | FOR NO KEY UPDATE   |               |
-
-```
+See [lock modes](./select-query-builder.md#lock-modes) for more information
 
 Complete example of find options:
 

diff --git a/docs/select-query-builder.md b/docs/select-query-builder.md
@@ -901,6 +901,22 @@ Using `take` and `skip` will prevent those issues.
 ## Set locking
 
 QueryBuilder supports both optimistic and pessimistic locking.
+
+#### Lock modes
+Support of lock modes, and SQL statements they translate to, are listed in the table below (blank cell denotes unsupported). When specified lock mode is not supported, a `LockNotSupportedOnGivenDriverError` error will be thrown.
+
+```text
+|                 | pessimistic_read                  | pessimistic_write       | dirty_read    | pessimistic_partial_write (Deprecated, use onLocked instead)   | pessimistic_write_or_fail (Deprecated, use onLocked instead)   | for_no_key_update   | for_key_share |
+| --------------- | --------------------------------- | ----------------------- | ------------- | -------------------------------------------------------------- | -------------------------------------------------------------- | ------------------- | ------------- |
+| MySQL           | FOR SHARE (8+)/LOCK IN SHARE MODE | FOR UPDATE              | (nothing)     | FOR UPDATE SKIP LOCKED                                         | FOR UPDATE NOWAIT                                              |                     |               |
+| Postgres        | FOR SHARE                         | FOR UPDATE              | (nothing)     | FOR UPDATE SKIP LOCKED                                         | FOR UPDATE NOWAIT                                              | FOR NO KEY UPDATE   | FOR KEY SHARE |
+| Oracle          | FOR UPDATE                        | FOR UPDATE              | (nothing)     |                                                                |                                                                |                     |               |
+| SQL Server      | WITH (HOLDLOCK, ROWLOCK)          | WITH (UPDLOCK, ROWLOCK) | WITH (NOLOCK) |                                                                |                                                                |                     |               |
+| AuroraDataApi   | LOCK IN SHARE MODE                | FOR UPDATE              | (nothing)     |                                                                |                                                                |                     |               |
+| CockroachDB     |                                   | FOR UPDATE              | (nothing)     |                                                                | FOR UPDATE NOWAIT                                              | FOR NO KEY UPDATE   |               |
+
+```
+
 To use pessimistic read locking use the following method:
 
 ```typescript
@@ -943,6 +959,39 @@ const users = await dataSource
 
 Optimistic locking works in conjunction with both `@Version` and `@UpdatedDate` decorators.
 
+### setOnLock
+Allows you to control what happens when a row is locked. By default, the database will wait for the lock.
+You can control that behavior by using `setOnLocked`
+
+
+To not wait:
+
+```typescript
+const users = await dataSource
+    .getRepository(User)
+    .createQueryBuilder("user")
+    .setLock("pessimistic_write")
+    .setOnLocked("nowait")
+    .getMany()
+```
+
+To skip the row:
+
+```typescript
+const users = await dataSource
+    .getRepository(User)
+    .createQueryBuilder("user")
+    .setLock("pessimistic_write")
+    .setOnLocked("skip_locked")
+    .getMany()
+```
+
+Database support for `setOnLocked` based on [lock mode](#lock-modes):
+- Postgres: `pessimistic_read`, `pessimistic_write`, `for_no_key_update`, `for_key_share`
+- MySQL 8+: `pessimistic_read`, `pessimistic_write`
+- MySQL < 8, Maria DB: `pessimistic_write`
+- Cockroach: `pessimistic_write` (`nowait` only)
+
 ## Use custom index
 
 You can provide a certain index for database server to use in some cases. This feature is only supported in MySQL.

diff --git a/src/driver/Driver.ts b/src/driver/Driver.ts
@@ -26,6 +26,11 @@ export interface Driver {
      */
     options: BaseDataSourceOptions
 
+    /**
+     * Database version/release. Often requires a SQL query to the DB, so it is not always set
+     */
+    version?: string
+
     /**
      * Database name used to perform all write queries.
      *

diff --git a/src/driver/DriverUtils.ts b/src/driver/DriverUtils.ts
@@ -1,5 +1,6 @@
 import { Driver } from "./Driver"
 import { hash, shorten } from "../util/StringUtils"
+import { VersionUtils } from "../util/VersionUtils"
 
 /**
  * Common driver utility functions.
@@ -32,6 +33,13 @@ export class DriverUtils {
         return ["mysql", "mariadb"].includes(driver.options.type)
     }
 
+    static isReleaseVersionOrGreater(driver: Driver, version: string): boolean {
+        return (
+            driver.version != null &&
+            VersionUtils.isGreaterOrEqual(driver.version, version)
+        )
+    }
+
     static isPostgresFamily(driver: Driver): boolean {
         return ["postgres", "aurora-postgres"].includes(driver.options.type)
     }

diff --git a/src/driver/mysql/MysqlDriver.ts b/src/driver/mysql/MysqlDriver.ts
@@ -65,6 +65,11 @@ export class MysqlDriver implements Driver {
      */
     options: MysqlConnectionOptions
 
+    /**
+     * Version of MySQL. Requires a SQL query to the DB, so it is not always set
+     */
+    version?: string
+
     /**
      * Master database used to perform all write queries.
      */
@@ -402,6 +407,7 @@ export class MysqlDriver implements Driver {
             version: string
         }[] = await queryRunner.query(`SELECT VERSION() AS \`version\``)
         const dbVersion = result[0].version
+        this.version = dbVersion
         await queryRunner.release()
 
         if (this.options.type === "mariadb") {

diff --git a/src/driver/postgres/PostgresDriver.ts b/src/driver/postgres/PostgresDriver.ts
@@ -71,6 +71,11 @@ export class PostgresDriver implements Driver {
      */
     options: PostgresConnectionOptions
 
+    /**
+     * Version of Postgres. Requires a SQL query to the DB, so it is not always set
+     */
+    version?: string
+
     /**
      * Database name used to perform all write queries.
      */
@@ -385,6 +390,7 @@ export class PostgresDriver implements Driver {
             /^PostgreSQL ([\d\.]+) .*$/,
             "$1",
         )
+        this.version = versionString
         this.isGeneratedColumnsSupported = VersionUtils.isGreaterOrEqual(
             versionString,
             "12.0",

diff --git a/src/find-options/FindOneOptions.ts b/src/find-options/FindOneOptions.ts
@@ -74,11 +74,18 @@ export interface FindOneOptions<Entity = any> {
                   | "pessimistic_read"
                   | "pessimistic_write"
                   | "dirty_read"
+                  /*
+                    "pessimistic_partial_write" and "pessimistic_write_or_fail" are deprecated and
+                    will be removed in a future version.
+
+                    Use onLocked instead.
+                 */
                   | "pessimistic_partial_write"
                   | "pessimistic_write_or_fail"
                   | "for_no_key_update"
                   | "for_key_share"
               tables?: string[]
+              onLocked?: "nowait" | "skip_locked"
           }
 
     /**

diff --git a/src/query-builder/QueryExpressionMap.ts b/src/query-builder/QueryExpressionMap.ts
@@ -186,6 +186,12 @@ export class QueryExpressionMap {
         | "pessimistic_read"
         | "pessimistic_write"
         | "dirty_read"
+        /*
+            "pessimistic_partial_write" and "pessimistic_write_or_fail" are deprecated and
+            will be removed in a future version.
+
+            Use onLocked instead.
+         */
         | "pessimistic_partial_write"
         | "pessimistic_write_or_fail"
         | "for_no_key_update"
@@ -201,6 +207,11 @@ export class QueryExpressionMap {
      */
     lockTables?: string[]
 
+    /**
+     * Modify behavior when encountering locked rows. NOWAIT or SKIP LOCKED
+     */
+    onLocked?: "nowait" | "skip_locked"
+
     /**
      * Indicates if soft-deleted rows should be included in entity result.
      * By default the soft-deleted rows are not included.
@@ -492,6 +503,7 @@ export class QueryExpressionMap {
         map.skip = this.skip
         map.take = this.take
         map.lockMode = this.lockMode
+        map.onLocked = this.onLocked
         map.lockVersion = this.lockVersion
         map.lockTables = this.lockTables
         map.withDeleted = this.withDeleted