Skip to content

Commit

Permalink
feat: add setOnLocked for SKIP LOCKED and NO WAIT
Browse files Browse the repository at this point in the history
  • Loading branch information
@taylorhakes
taylorhakes committed Aug 24, 2022
1 parent 15f90e0 commit 54891b0
Show file tree
Hide file tree
Showing 3 changed files with 97 additions and 17 deletions.
6 changes: 6 additions & 0 deletions src/query-builder/QueryExpressionMap.ts
View file
Expand Up @@ -201,6 +201,11 @@ export class QueryExpressionMap {
*/
lockTables?: string[]

/**
* Skip locked rows. "<lockMode> SKIP LOCKED"
*/
onLocked?: "no_wait" | "skip_locked"

/**
* Indicates if soft-deleted rows should be included in entity result.
* By default the soft-deleted rows are not included.
Expand Down Expand Up @@ -492,6 +497,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
Expand Down
89 changes: 72 additions & 17 deletions src/query-builder/SelectQueryBuilder.ts
View file
Expand Up @@ -1515,6 +1515,14 @@ export class SelectQueryBuilder<Entity>
return this
}

/**
* Sets on lock handling NO WAIT or SKIP LOCKED.
*/
setOnLocked(onLocked: "no_wait" | "skip_locked"): this {
this.expressionMap.onLocked = onLocked
return this
}

/**
* Disables the global condition of "non-deleted" for the entity with delete date columns.
*/
Expand Down Expand Up @@ -2505,79 +2513,126 @@ export class SelectQueryBuilder<Entity>
lockTablesClause = " OF " + this.expressionMap.lockTables.join(", ")
}

let lockModeClause
switch (this.expressionMap.lockMode) {
case "pessimistic_read":
if (
DriverUtils.isMySQLFamily(driver) ||
driver.options.type === "aurora-mysql"
) {
return " LOCK IN SHARE MODE"
lockModeClause = " LOCK IN SHARE MODE"
} else if (driver.options.type === "postgres") {
return " FOR SHARE" + lockTablesClause
lockModeClause = " FOR SHARE" + lockTablesClause
} else if (driver.options.type === "oracle") {
return " FOR UPDATE"
lockModeClause = " FOR UPDATE"
} else if (driver.options.type === "mssql") {
return ""
lockModeClause = ""
} else {
throw new LockNotSupportedOnGivenDriverError()
}
break
case "pessimistic_write":
if (
DriverUtils.isMySQLFamily(driver) ||
driver.options.type === "aurora-mysql" ||
driver.options.type === "oracle"
) {
return " FOR UPDATE"
lockModeClause = " FOR UPDATE"
} else if (
driver.options.type === "postgres" ||
driver.options.type === "cockroachdb"
) {
return " FOR UPDATE" + lockTablesClause
lockModeClause = " FOR UPDATE" + lockTablesClause
} else if (driver.options.type === "mssql") {
return ""
lockModeClause = ""
} else {
throw new LockNotSupportedOnGivenDriverError()
}
break
case "pessimistic_partial_write":
if (driver.options.type === "postgres") {
return " FOR UPDATE" + lockTablesClause + " SKIP LOCKED"
lockModeClause =
" FOR UPDATE" + lockTablesClause + " SKIP LOCKED"
} else if (DriverUtils.isMySQLFamily(driver)) {
return " FOR UPDATE SKIP LOCKED"
lockModeClause = " FOR UPDATE SKIP LOCKED"
} else {
throw new LockNotSupportedOnGivenDriverError()
}
break
case "pessimistic_write_or_fail":
if (
driver.options.type === "postgres" ||
driver.options.type === "cockroachdb"
) {
return " FOR UPDATE" + lockTablesClause + " NOWAIT"
lockModeClause =
" FOR UPDATE" + lockTablesClause + " NOWAIT"
} else if (DriverUtils.isMySQLFamily(driver)) {
return " FOR UPDATE NOWAIT"
lockModeClause = " FOR UPDATE NOWAIT"
} else {
throw new LockNotSupportedOnGivenDriverError()
}

break
case "for_no_key_update":
if (
driver.options.type === "postgres" ||
driver.options.type === "cockroachdb"
) {
return " FOR NO KEY UPDATE" + lockTablesClause
lockModeClause = " FOR NO KEY UPDATE" + lockTablesClause
} else {
throw new LockNotSupportedOnGivenDriverError()
}

break
case "for_key_share":
if (driver.options.type === "postgres") {
return " FOR KEY SHARE" + lockTablesClause
lockModeClause = " FOR KEY SHARE" + lockTablesClause
} else {
throw new LockNotSupportedOnGivenDriverError()
}

break
default:
return ""
lockModeClause = ""
}

if (this.expressionMap.onLocked) {
if (this.expressionMap.lockMode === "pessimistic_write_or_fail") {
if (this.expressionMap.onLocked === "skip_locked") {
// User cannot specify both NO WAIT and SKIP Locked
throw new LockNotSupportedOnGivenDriverError()
}

// This lock mode automatically adds NO WAIT. Don't add it twice
return lockModeClause
}

if (this.expressionMap.lockMode === "pessimistic_partial_write") {
if (this.expressionMap.onLocked === "no_wait") {
// User cannot specify both NO WAIT and SKIP Locked
throw new LockNotSupportedOnGivenDriverError()
}

// This lock mode automatically adds SKIP LOCKED. Don't add it twice
return lockModeClause
}

const onLockExpression =
this.expressionMap.onLocked === "no_wait"
? " NO WAIT"
: " SKIP LOCKED"
if (driver.options.type === "postgres") {
return lockModeClause + onLockExpression
} else if (
DriverUtils.isMySQLFamily(driver) &&
["pessimistic_write", "pessimistic_read"].includes(
this.expressionMap.lockMode || "",
)
) {
return lockModeClause + onLockExpression
} else {
throw new LockNotSupportedOnGivenDriverError()
}
}

return lockModeClause
}

/**
Expand Down
19 changes: 19 additions & 0 deletions test/functional/query-builder/locking/query-builder-locking.ts
View file
Expand Up @@ -1084,4 +1084,23 @@ describe("query builder > locking", () => {
})
}),
))

it("should allow locking a relation of a relation", () =>
Promise.all(
connections.map(async (connection) => {
if (
connection.driver.options.type === "postgres" ||
DriverUtils.isMySQLFamily(connection.driver)
) {
const sql = connection
.createQueryBuilder(PostWithVersion, "post")
.setLock("pessimistic_partial_write")
.setOnLocked("skip_locked")
.where("post.id = :id", { id: 1 })
.getSql()

expect(sql.endsWith("FOR UPDATE SKIP LOCKED")).to.be.true
}
}),
))
})

0 comments on commit 54891b0

Please sign in to comment.