feat: add FOR NO KEY UPDATE lock mode for postgresql (typeorm#5971)

* [Add] FOR NO KEY UPDATE lock mode for postgresql * [Add] for no key update lock test * [Fix] lint * [Fix] test Co-authored-by: JeyongOh <[email protected]>
zaro · Jan 12, 2021 · e6ae38f · e6ae38f
1 parent b701a23
commit e6ae38f
Show file tree

Hide file tree

Showing 3 changed files with 71 additions and 4 deletions.
diff --git a/src/query-builder/QueryExpressionMap.ts b/src/query-builder/QueryExpressionMap.ts
@@ -150,7 +150,7 @@ export class QueryExpressionMap {
     /**
      * Locking mode.
      */
-    lockMode?: "optimistic"|"pessimistic_read"|"pessimistic_write"|"dirty_read";
+    lockMode?: "optimistic"|"pessimistic_read"|"pessimistic_write"|"dirty_read"|"for_no_key_update";
 
     /**
      * Current version of the entity, used for locking.

diff --git a/src/query-builder/SelectQueryBuilder.ts b/src/query-builder/SelectQueryBuilder.ts
@@ -967,12 +967,12 @@ export class SelectQueryBuilder<Entity> extends QueryBuilder<Entity> implements
     /**
      * Sets locking mode.
      */
-    setLock(lockMode: "pessimistic_read"|"pessimistic_write"|"dirty_read"): this;
+    setLock(lockMode: "pessimistic_read"|"pessimistic_write"|"dirty_read"|"for_no_key_update"): this;
 
     /**
      * Sets locking mode.
      */
-    setLock(lockMode: "optimistic"|"pessimistic_read"|"pessimistic_write"|"dirty_read", lockVersion?: number|Date): this {
+    setLock(lockMode: "optimistic"|"pessimistic_read"|"pessimistic_write"|"dirty_read"|"for_no_key_update", lockVersion?: number|Date): this {
         this.expressionMap.lockMode = lockMode;
         this.expressionMap.lockVersion = lockVersion;
         return this;
@@ -1672,6 +1672,12 @@ export class SelectQueryBuilder<Entity> extends QueryBuilder<Entity> implements
                 } else {
                     throw new LockNotSupportedOnGivenDriverError();
                 }
+            case "for_no_key_update":
+                if (driver instanceof PostgresDriver) {
+                    return " FOR NO KEY UPDATE";
+                } else {
+                    throw new LockNotSupportedOnGivenDriverError();
+                }
             default:
                 return "";
         }
@@ -1806,7 +1812,7 @@ export class SelectQueryBuilder<Entity> extends QueryBuilder<Entity> implements
         if (!this.expressionMap.mainAlias)
             throw new Error(`Alias is not set. Use "from" method to set an alias.`);
 
-        if ((this.expressionMap.lockMode === "pessimistic_read" || this.expressionMap.lockMode === "pessimistic_write") && !queryRunner.isTransactionActive)
+        if ((this.expressionMap.lockMode === "pessimistic_read" || this.expressionMap.lockMode === "pessimistic_write" || this.expressionMap.lockMode === "for_no_key_update") && !queryRunner.isTransactionActive)
             throw new PessimisticLockTransactionRequiredError();
 
         if (this.expressionMap.lockMode === "optimistic") {

diff --git a/test/functional/query-builder/locking/query-builder-locking.ts b/test/functional/query-builder/locking/query-builder-locking.ts
@@ -77,6 +77,28 @@ describe("query builder > locking", () => {
         });
     })));
 
+    it("should throw error if for no key update lock used without transaction", () => Promise.all(connections.map(async connection => {
+        if (connection.driver instanceof PostgresDriver) {
+            return connection.createQueryBuilder(PostWithVersion, "post")
+                .setLock("for_no_key_update")
+                .where("post.id = :id", { id: 1 })
+                .getOne().should.be.rejectedWith(PessimisticLockTransactionRequiredError);
+        }
+        return;
+    })));
+
+    it("should not throw error if for no key update lock used with transaction", () => Promise.all(connections.map(async connection => {
+        if (connection.driver instanceof PostgresDriver) {
+            return connection.manager.transaction(entityManager => {
+                return Promise.all([entityManager.createQueryBuilder(PostWithVersion, "post")
+                    .setLock("for_no_key_update")
+                    .where("post.id = :id", { id: 1})
+                    .getOne().should.not.be.rejected]);
+            });
+        }
+        return;
+    })));
+
     it("should attach pessimistic read lock statement on query if locking enabled", () => Promise.all(connections.map(async connection => {
         if (connection.driver instanceof AbstractSqliteDriver || connection.driver instanceof CockroachDriver || connection.driver instanceof SapDriver)
             return;
@@ -141,6 +163,30 @@ describe("query builder > locking", () => {
 
     })));
 
+    it("should not attach for no key update lock statement on query if locking is not used", () => Promise.all(connections.map(async connection => {
+        if (connection.driver instanceof PostgresDriver) {
+            const sql = connection.createQueryBuilder(PostWithVersion, "post")
+                .where("post.id = :id", { id: 1 })
+                .getSql();
+
+                expect(sql.indexOf("FOR NO KEY UPDATE") === -1).to.be.true;
+            }
+        return;
+    })));
+
+    it("should attach for no key update lock statement on query if locking enabled", () => Promise.all(connections.map(async connection => {
+        if (connection.driver instanceof PostgresDriver) {
+            const sql = connection.createQueryBuilder(PostWithVersion, "post")
+                .setLock("for_no_key_update")
+                .where("post.id = :id", { id: 1 })
+                .getSql();
+
+            expect(sql.indexOf("FOR NO KEY UPDATE") !== -1).to.be.true;
+        }
+        return;
+
+    })));
+
     it("should throw error if optimistic lock used with getMany method", () => Promise.all(connections.map(async connection => {
 
        return connection.createQueryBuilder(PostWithVersion, "post")
@@ -291,4 +337,19 @@ describe("query builder > locking", () => {
         return;
     })));
 
+    it("should throw error if for no key update locking not supported by given driver", () => Promise.all(connections.map(async connection => {
+        if (!(connection.driver instanceof PostgresDriver)) {
+            return connection.manager.transaction(entityManager => {
+                return Promise.all([
+                    entityManager.createQueryBuilder(PostWithVersion, "post")
+                        .setLock("for_no_key_update")
+                        .where("post.id = :id", { id: 1 })
+                        .getOne().should.be.rejectedWith(LockNotSupportedOnGivenDriverError),
+                ]);
+            });
+        }
+
+        return;
+    })));
+
 });