Skip to content

Commit

Permalink
fix: allow lock modes for cockroachdb
Browse files Browse the repository at this point in the history
closes #8249
  • Loading branch information
wodka committed Oct 7, 2021
1 parent 8615733 commit 2d79b7a
Show file tree
Hide file tree
Showing 4 changed files with 151 additions and 69 deletions.
1 change: 1 addition & 0 deletions docs/find-options.md
Expand Up @@ -239,6 +239,7 @@ Support of lock modes, and SQL statements they translate to, are listed in the t
| 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 |
```

Expand Down
8 changes: 4 additions & 4 deletions src/query-builder/SelectQueryBuilder.ts
Expand Up @@ -1679,7 +1679,7 @@ export class SelectQueryBuilder<Entity> extends QueryBuilder<Entity> implements
let lockTablesClause = "";

if (this.expressionMap.lockTables) {
if (!(driver instanceof PostgresDriver)) {
if (!(driver instanceof PostgresDriver || driver instanceof CockroachDriver)) {
throw new TypeORMError("Lock tables not supported in selected driver");
}
if (this.expressionMap.lockTables.length < 1) {
Expand Down Expand Up @@ -1710,7 +1710,7 @@ export class SelectQueryBuilder<Entity> extends QueryBuilder<Entity> implements
return " FOR UPDATE";

}
else if (driver instanceof PostgresDriver ) {
else if (driver instanceof PostgresDriver || driver instanceof CockroachDriver) {
return " FOR UPDATE" + lockTablesClause;

} else if (driver instanceof SqlServerDriver) {
Expand All @@ -1730,7 +1730,7 @@ export class SelectQueryBuilder<Entity> extends QueryBuilder<Entity> implements
throw new LockNotSupportedOnGivenDriverError();
}
case "pessimistic_write_or_fail":
if (driver instanceof PostgresDriver) {
if (driver instanceof PostgresDriver || driver instanceof CockroachDriver) {
return " FOR UPDATE" + lockTablesClause + " NOWAIT";

} else if (driver instanceof MysqlDriver) {
Expand All @@ -1741,7 +1741,7 @@ export class SelectQueryBuilder<Entity> extends QueryBuilder<Entity> implements
}

case "for_no_key_update":
if (driver instanceof PostgresDriver) {
if (driver instanceof PostgresDriver || driver instanceof CockroachDriver) {
return " FOR NO KEY UPDATE" + lockTablesClause;
} else {
throw new LockNotSupportedOnGivenDriverError();
Expand Down
125 changes: 81 additions & 44 deletions test/functional/query-builder/locking/query-builder-locking.ts
Expand Up @@ -61,9 +61,20 @@ describe("query builder > locking", () => {
})));

it("should not throw error if pessimistic lock used with transaction", () => Promise.all(connections.map(async connection => {
if (connection.driver instanceof AbstractSqliteDriver || connection.driver instanceof CockroachDriver || connection.driver instanceof SapDriver)
if (connection.driver instanceof AbstractSqliteDriver || connection.driver instanceof SapDriver)
return;

if (connection.driver instanceof CockroachDriver) {
return connection.manager.transaction(entityManager => {
return Promise.all([
entityManager.createQueryBuilder(PostWithVersion, "post")
.setLock("pessimistic_write")
.where("post.id = :id", { id: 1 })
.getOne().should.not.be.rejected
]);
});
}

return connection.manager.transaction(entityManager => {
return Promise.all([
entityManager.createQueryBuilder(PostWithVersion, "post")
Expand All @@ -80,7 +91,7 @@ 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) {
if (connection.driver instanceof PostgresDriver || connection.driver instanceof CockroachDriver) {
return connection.createQueryBuilder(PostWithVersion, "post")
.setLock("for_no_key_update")
.where("post.id = :id", { id: 1 })
Expand All @@ -90,7 +101,7 @@ describe("query builder > locking", () => {
})));

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) {
if (connection.driver instanceof PostgresDriver || connection.driver instanceof CockroachDriver) {
return connection.manager.transaction(entityManager => {
return Promise.all([entityManager.createQueryBuilder(PostWithVersion, "post")
.setLock("for_no_key_update")
Expand Down Expand Up @@ -119,14 +130,14 @@ describe("query builder > locking", () => {
return connection.createQueryBuilder(PostWithVersion, "post")
.setLock("pessimistic_partial_write")
.where("post.id = :id", { id: 1 })
.getOne().should.be.rejectedWith(PessimisticLockTransactionRequiredError);
.getOne().should.be.rejectedWith(PessimisticLockTransactionRequiredError);
}
}
return;
})));
})));

it("should not throw error if pessimistic_partial_write lock used with transaction", () => Promise.all(connections.map(async connection => {
if (connection.driver instanceof PostgresDriver) {
if (connection.driver instanceof PostgresDriver) {
return connection.manager.transaction(entityManager => {
return Promise.all([entityManager.createQueryBuilder(PostWithVersion, "post")
.setLock("pessimistic_partial_write")
Expand All @@ -149,12 +160,12 @@ describe("query builder > locking", () => {
.getOne().should.not.be.rejected]);
});
}
}
}
return;
})));

it("should throw error if pessimistic_write_or_fail lock used without transaction", () => Promise.all(connections.map(async connection => {
if (connection.driver instanceof PostgresDriver) {
if (connection.driver instanceof PostgresDriver || connection.driver instanceof CockroachDriver) {
return connection.createQueryBuilder(PostWithVersion, "post")
.setLock("pessimistic_write_or_fail")
.where("post.id = :id", { id: 1 })
Expand All @@ -173,11 +184,11 @@ describe("query builder > locking", () => {
.getOne().should.be.rejectedWith(PessimisticLockTransactionRequiredError);
}
}
return;
})));
return;
})));

it("should not throw error if pessimistic_write_or_fail lock used with transaction", () => Promise.all(connections.map(async connection => {
if (connection.driver instanceof PostgresDriver) {
if (connection.driver instanceof PostgresDriver || connection.driver instanceof CockroachDriver) {
return connection.manager.transaction(entityManager => {
return Promise.all([entityManager.createQueryBuilder(PostWithVersion, "post")
.setLock("pessimistic_write_or_fail")
Expand All @@ -199,9 +210,9 @@ describe("query builder > locking", () => {
.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)
Expand Down Expand Up @@ -250,15 +261,15 @@ describe("query builder > locking", () => {
})));

it("should attach pessimistic write 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)
if (connection.driver instanceof AbstractSqliteDriver || connection.driver instanceof SapDriver)
return;

const sql = connection.createQueryBuilder(PostWithVersion, "post")
.setLock("pessimistic_write")
.where("post.id = :id", { id: 1 })
.getSql();

if (connection.driver instanceof MysqlDriver || connection.driver instanceof PostgresDriver || connection.driver instanceof OracleDriver) {
if (connection.driver instanceof MysqlDriver || connection.driver instanceof PostgresDriver || connection.driver instanceof CockroachDriver || connection.driver instanceof OracleDriver) {
expect(sql.indexOf("FOR UPDATE") !== -1).to.be.true;

} else if (connection.driver instanceof SqlServerDriver) {
Expand All @@ -268,7 +279,7 @@ 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) {
if (connection.driver instanceof PostgresDriver || connection.driver instanceof CockroachDriver) {
const sql = connection.createQueryBuilder(PostWithVersion, "post")
.where("post.id = :id", { id: 1 })
.getSql();
Expand All @@ -279,7 +290,7 @@ describe("query builder > locking", () => {
})));

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) {
if (connection.driver instanceof PostgresDriver || connection.driver instanceof CockroachDriver) {
const sql = connection.createQueryBuilder(PostWithVersion, "post")
.setLock("for_no_key_update")
.where("post.id = :id", { id: 1 })
Expand All @@ -300,7 +311,7 @@ describe("query builder > locking", () => {
expect(sql.indexOf("FOR UPDATE SKIP LOCKED") === -1).to.be.true;
}
return;
})));
})));

it("should attach pessimistic_partial_write lock statement on query if locking enabled", () => Promise.all(connections.map(async connection => {
if (connection.driver instanceof PostgresDriver || connection.driver instanceof MysqlDriver) {
Expand All @@ -313,21 +324,21 @@ describe("query builder > locking", () => {
}
return;

})));
})));

it("should not attach pessimistic_write_or_fail lock statement on query if locking is not used", () => Promise.all(connections.map(async connection => {
if (connection.driver instanceof PostgresDriver || connection.driver instanceof MysqlDriver) {
if (connection.driver instanceof PostgresDriver || connection.driver instanceof MysqlDriver || connection.driver instanceof CockroachDriver) {
const sql = connection.createQueryBuilder(PostWithVersion, "post")
.where("post.id = :id", { id: 1 })
.getSql();

expect(sql.indexOf("FOR UPDATE NOWAIT") === -1).to.be.true;
}
return;
})));
})));

it("should attach pessimistic_write_or_fail lock statement on query if locking enabled", () => Promise.all(connections.map(async connection => {
if (connection.driver instanceof PostgresDriver || connection.driver instanceof MysqlDriver) {
if (connection.driver instanceof PostgresDriver || connection.driver instanceof MysqlDriver || connection.driver instanceof CockroachDriver) {
const sql = connection.createQueryBuilder(PostWithVersion, "post")
.setLock("pessimistic_write_or_fail")
.where("post.id = :id", { id: 1 })
Expand All @@ -337,7 +348,7 @@ describe("query builder > locking", () => {
}
return;

})));
})));

it("should throw error if optimistic lock used with getMany method", () => Promise.all(connections.map(async connection => {

Expand Down Expand Up @@ -471,7 +482,7 @@ describe("query builder > locking", () => {
})));

it("should throw error if pessimistic locking not supported by given driver", () => Promise.all(connections.map(async connection => {
if (connection.driver instanceof AbstractSqliteDriver || connection.driver instanceof CockroachDriver || connection.driver instanceof SapDriver)
if (connection.driver instanceof AbstractSqliteDriver || connection.driver instanceof SapDriver)
return connection.manager.transaction(entityManager => {
return Promise.all([
entityManager.createQueryBuilder(PostWithVersion, "post")
Expand All @@ -486,11 +497,21 @@ describe("query builder > locking", () => {
]);
});

if (connection.driver instanceof CockroachDriver)
return connection.manager.transaction(entityManager => {
return Promise.all([
entityManager.createQueryBuilder(PostWithVersion, "post")
.setLock("pessimistic_read")
.where("post.id = :id", { id: 1 })
.getOne().should.be.rejectedWith(LockNotSupportedOnGivenDriverError),
]);
});

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)) {
if (!(connection.driver instanceof PostgresDriver || connection.driver instanceof CockroachDriver)) {
return connection.manager.transaction(entityManager => {
return Promise.all([
entityManager.createQueryBuilder(PostWithVersion, "post")
Expand All @@ -505,7 +526,7 @@ describe("query builder > locking", () => {
})));

it("should only specify locked tables in FOR UPDATE OF clause if argument is given", () => Promise.all(connections.map(async connection => {
if (!(connection.driver instanceof PostgresDriver))
if (!(connection.driver instanceof PostgresDriver || connection.driver instanceof CockroachDriver))
return;

const sql = connection.createQueryBuilder(Post, "post")
Expand All @@ -524,7 +545,7 @@ describe("query builder > locking", () => {
})));

it("should not allow empty array for lockTables", () => Promise.all(connections.map(async connection => {
if (!(connection.driver instanceof PostgresDriver))
if (!(connection.driver instanceof PostgresDriver || connection.driver instanceof CockroachDriver))
return;

return connection.manager.transaction(entityManager => {
Expand All @@ -538,7 +559,7 @@ describe("query builder > locking", () => {
})));

it("should throw error when specifying a table that is not part of the query", () => Promise.all(connections.map(async connection => {
if (!(connection.driver instanceof PostgresDriver))
if (!(connection.driver instanceof PostgresDriver || connection.driver instanceof CockroachDriver))
return;

return connection.manager.transaction(entityManager => {
Expand All @@ -552,22 +573,38 @@ describe("query builder > locking", () => {
})));

it("should allow on a left join", () => Promise.all(connections.map(async connection => {
if (!(connection.driver instanceof PostgresDriver))
return;
if (connection.driver instanceof CockroachDriver) {
return connection.manager.transaction(entityManager => {
return Promise.all([
entityManager.createQueryBuilder(Post, "post")
.leftJoin("post.author", "user")
.setLock('pessimistic_write', undefined, ["post"])
.getOne(),
entityManager.createQueryBuilder(Post, "post")
.leftJoin("post.author", "user")
.setLock('pessimistic_write')
.getOne(),
]);
});
}

return connection.manager.transaction(entityManager => {
if (connection.driver instanceof PostgresDriver) {
return connection.manager.transaction(entityManager => {

return Promise.all([
entityManager.createQueryBuilder(Post, "post")
.leftJoin("post.author", "user")
.setLock('pessimistic_write', undefined, ["post"])
.getOne(),
entityManager.createQueryBuilder(Post, "post")
.leftJoin("post.author", "user")
.setLock('pessimistic_write')
.getOne().should.be.rejectedWith('FOR UPDATE cannot be applied to the nullable side of an outer join'),
]);
});
return Promise.all([
entityManager.createQueryBuilder(Post, "post")
.leftJoin("post.author", "user")
.setLock('pessimistic_write', undefined, ["post"])
.getOne(),
entityManager.createQueryBuilder(Post, "post")
.leftJoin("post.author", "user")
.setLock('pessimistic_write')
.getOne().should.be.rejectedWith('FOR UPDATE cannot be applied to the nullable side of an outer join'),
]);
});
}

return;
})));

it("should allow using lockTables on all types of locking", () => Promise.all(connections.map(async connection => {
Expand Down Expand Up @@ -602,7 +639,7 @@ describe("query builder > locking", () => {
})));

it("should allow locking a relation of a relation", () => Promise.all(connections.map(async connection => {
if (!(connection.driver instanceof PostgresDriver))
if (!(connection.driver instanceof PostgresDriver || connection.driver instanceof CockroachDriver))
return;

return connection.manager.transaction(entityManager => {
Expand Down

0 comments on commit 2d79b7a

Please sign in to comment.