From 02de413ca6850fa65086a99c22a3d5306d9bc5ae Mon Sep 17 00:00:00 2001 From: Dominik Krejcik Date: Sat, 23 Oct 2021 14:01:36 +0100 Subject: [PATCH] Support SELECT .. FOR NO KEY UPDATE / KEY SHARE row level locking clauses in Postgres (#4755) --- .../postgres/query/pg-querycompiler.js | 22 +++--- .../redshift/query/redshift-querycompiler.js | 12 ++++ .../sqlite3/query/sqlite-querycompiler.js | 2 + lib/query/constants.js | 2 + lib/query/querybuilder.js | 21 +++++- .../integration2/query/select/selects.spec.js | 72 +++++++++++++++++++ test/unit/query/builder.js | 21 ++++++ test/util/tableCreatorHelper.js | 13 ++++ types/index.d.ts | 6 ++ 9 files changed, 162 insertions(+), 9 deletions(-) diff --git a/lib/dialects/postgres/query/pg-querycompiler.js b/lib/dialects/postgres/query/pg-querycompiler.js index 6b249d645e..872a538ddb 100644 --- a/lib/dialects/postgres/query/pg-querycompiler.js +++ b/lib/dialects/postgres/query/pg-querycompiler.js @@ -147,20 +147,26 @@ class QueryCompiler_PG extends QueryCompiler { return sql.join(', '); } - forUpdate() { + _lockingClause(lockMode) { const tables = this.single.lockTables || []; - return ( - 'for update' + (tables.length ? ' of ' + this._tableNames(tables) : '') - ); + return lockMode + (tables.length ? ' of ' + this._tableNames(tables) : ''); + } + + forUpdate() { + return this._lockingClause('for update'); } forShare() { - const tables = this.single.lockTables || []; + return this._lockingClause('for share'); + } + + forNoKeyUpdate() { + return this._lockingClause('for no key update'); + } - return ( - 'for share' + (tables.length ? ' of ' + this._tableNames(tables) : '') - ); + forKeyShare() { + return this._lockingClause('for key share'); } skipLocked() { diff --git a/lib/dialects/redshift/query/redshift-querycompiler.js b/lib/dialects/redshift/query/redshift-querycompiler.js index b345fb8e8b..b1044bf49c 100644 --- a/lib/dialects/redshift/query/redshift-querycompiler.js +++ b/lib/dialects/redshift/query/redshift-querycompiler.js @@ -61,6 +61,18 @@ class QueryCompiler_Redshift extends QueryCompiler_PG { return ''; } + forNoKeyUpdate() { + this.client.logger.warn('table lock is not supported by redshift dialect'); + return ''; + } + + forKeyShare() { + this.client.logger.warn( + 'lock for share is not supported by redshift dialect' + ); + return ''; + } + // Compiles a columnInfo query columnInfo() { const column = this.single.columnInfo; diff --git a/lib/dialects/sqlite3/query/sqlite-querycompiler.js b/lib/dialects/sqlite3/query/sqlite-querycompiler.js index 0f8210d0ca..39c0f934e3 100644 --- a/lib/dialects/sqlite3/query/sqlite-querycompiler.js +++ b/lib/dialects/sqlite3/query/sqlite-querycompiler.js @@ -27,7 +27,9 @@ class QueryCompiler_SQLite3 extends QueryCompiler { // The locks are not applicable in SQLite3 this.forShare = emptyStr; + this.forKeyShare = emptyStr; this.forUpdate = emptyStr; + this.forNoKeyUpdate = emptyStr; } // SQLite requires us to build the multi-row insert as a listing of select with diff --git a/lib/query/constants.js b/lib/query/constants.js index 0b080715dd..125001d9c4 100644 --- a/lib/query/constants.js +++ b/lib/query/constants.js @@ -5,6 +5,8 @@ module.exports = { lockMode: { forShare: 'forShare', forUpdate: 'forUpdate', + forNoKeyUpdate: 'forNoKeyUpdate', + forKeyShare: 'forKeyShare', }, waitMode: { skipLocked: 'skipLocked', diff --git a/lib/query/querybuilder.js b/lib/query/querybuilder.js index 19c9d0807e..e997ab9bc9 100644 --- a/lib/query/querybuilder.js +++ b/lib/query/querybuilder.js @@ -46,7 +46,12 @@ const CLEARABLE_STATEMENTS = new Set([ 'counter', 'counters', ]); -const LOCK_MODES = new Set([lockMode.forShare, lockMode.forUpdate]); +const LOCK_MODES = new Set([ + lockMode.forShare, + lockMode.forUpdate, + lockMode.forNoKeyUpdate, + lockMode.forKeyShare, +]); // Typically called from `knex.builder`, // start a new query building chain. @@ -1198,6 +1203,20 @@ class Builder extends EventEmitter { return this; } + // Set a lock for no key update constraint. + forNoKeyUpdate(...tables) { + this._single.lock = lockMode.forNoKeyUpdate; + this._single.lockTables = tables; + return this; + } + + // Set a lock for key share constraint. + forKeyShare(...tables) { + this._single.lock = lockMode.forKeyShare; + this._single.lockTables = tables; + return this; + } + // Skips locked rows when using a lock constraint. skipLocked() { if (!this._isSelectQuery()) { diff --git a/test/integration2/query/select/selects.spec.js b/test/integration2/query/select/selects.spec.js index cd4ccdd1fc..553952889d 100644 --- a/test/integration2/query/select/selects.spec.js +++ b/test/integration2/query/select/selects.spec.js @@ -23,6 +23,7 @@ const { createTestTableTwo, dropTables, createDefaultTable, + createParentAndChildTables, } = require('../../../util/tableCreatorHelper'); const { insertAccounts } = require('../../../util/dataInsertHelper'); const { assertNumberArrayStrict } = require('../../../util/assertHelper'); @@ -1773,6 +1774,77 @@ describe('Selects', function () { }); }); + it('select for no key update doesnt stop other transactions from inserting into tables that have a foreign key relationship', async function () { + if (!isPostgreSQL(knex)) { + return this.skip(); + } + + await createParentAndChildTables(knex); + + return knex('parent') + .insert({ + id: 1, + }) + .then(() => { + return knex('child') + .insert({ + id: 1, + parent_id: 1, + }) + .then(() => { + return knex.transaction((trx) => { + // select all from the parent table in the for no key update mode + return trx('parent') + .forNoKeyUpdate() + .then((res) => { + // Insert should into the child table not hang + return knex('child') + .insert({ + id: 2, + parent_id: 1, + }) + .timeout(150); + }); + }); + }); + }); + }); + + it('select for key share blocks select for update but not select for no key update', async function () { + if (!isPostgreSQL(knex)) { + return this.skip(); + } + + return knex('test_default_table') + .insert({ string: 'making sure there is a row to lock' }) + .then(() => { + return knex + .transaction((trx) => { + // select all from test table and lock + return trx('test_default_table') + .forKeyShare() + .then((res) => { + // trying to select stuff from table in other connection should succeed with for no key update + return knex('test_default_table') + .forNoKeyUpdate() + .timeout(200); + }) + .then((res) => { + // trying to select stuff from table in other connection should hang with for update + return knex('test_default_table').forUpdate().timeout(100); + }); + }) + .then((res) => { + expect('Second query should have timed out').to.be.false; + }) + .catch((err) => { + expect(err.message).to.be.contain( + 'Defined query timeout of 100ms exceeded when running query' + ); + }); + }); + }); + it('select for share prevents updating in other transaction', function () { // Query cancellation is not yet implemented for CockroachDB if (isSQLite(knex) || isOracle(knex) || isCockroachDB(knex)) { diff --git a/test/unit/query/builder.js b/test/unit/query/builder.js index 223708fad6..99655e9500 100644 --- a/test/unit/query/builder.js +++ b/test/unit/query/builder.js @@ -6507,6 +6507,27 @@ describe('QueryBuilder', () => { }); }); + it('lock for no key update', () => { + testsql( + qb().select('*').from('foo').where('bar', '=', 'baz').forNoKeyUpdate(), + { + pg: { + sql: 'select * from "foo" where "bar" = ? for no key update', + bindings: ['baz'], + }, + } + ); + }); + + it('lock for key share', () => { + testsql(qb().select('*').from('foo').where('bar', '=', 'baz').forShare(), { + pg: { + sql: 'select * from "foo" where "bar" = ? for share', + bindings: ['baz'], + }, + }); + }); + it('should allow lock (such as forUpdate) outside of a transaction', () => { testsql(qb().select('*').from('foo').where('bar', '=', 'baz').forUpdate(), { mysql: { diff --git a/test/util/tableCreatorHelper.js b/test/util/tableCreatorHelper.js index bd1f4062e3..64359881e8 100644 --- a/test/util/tableCreatorHelper.js +++ b/test/util/tableCreatorHelper.js @@ -85,6 +85,16 @@ async function createCompositeKeyTable(knex) { }); } +async function createParentAndChildTables(knex) { + await knex.schema.createTable('parent', (table) => { + table.integer('id').primary(); + }); + await knex.schema.createTable('child', (table) => { + table.integer('id').primary(); + table.integer('parent_id').references('parent.id'); + }); +} + async function dropTables(knex) { await knex.schema.dropTableIfExists('accounts'); await knex.schema.dropTableIfExists('users'); @@ -95,6 +105,8 @@ async function dropTables(knex) { await knex.schema.dropTableIfExists('datatype_test'); await knex.schema.dropTableIfExists('test_default_table'); await knex.schema.dropTableIfExists('test_default_table2'); + await knex.schema.dropTableIfExists('child'); + await knex.schema.dropTableIfExists('parent'); } module.exports = { @@ -104,5 +116,6 @@ module.exports = { createDefaultTable, createUsers, createTestTableTwo, + createParentAndChildTables, dropTables, }; diff --git a/types/index.d.ts b/types/index.d.ts index 334594931e..1515a1463b 100644 --- a/types/index.d.ts +++ b/types/index.d.ts @@ -1740,6 +1740,12 @@ export declare namespace Knex { forShare(...tableNames: string[]): QueryBuilder; forShare(tableNames: readonly string[]): QueryBuilder; + forNoKeyUpdate(...tableNames: string[]): QueryBuilder; + forNoKeyUpdate(tableNames: readonly string[]): QueryBuilder; + + forKeyShare(...tableNames: string[]): QueryBuilder; + forKeyShare(tableNames: readonly string[]): QueryBuilder; + skipLocked(): QueryBuilder; noWait(): QueryBuilder;