diff --git a/trytond/CHANGELOG b/trytond/CHANGELOG index 501f45d438..8e737a1781 100644 --- a/trytond/CHANGELOG +++ b/trytond/CHANGELOG @@ -1,3 +1,4 @@ +* Lock records on transaction start * Add DecimalNull as tools * Add chat backend * Transaction.check_warnings is now a directory mapping the key to the warnings diff --git a/trytond/doc/ref/backend.rst b/trytond/doc/ref/backend.rst index 4a541cda3f..41ac6c0db1 100644 --- a/trytond/doc/ref/backend.rst +++ b/trytond/doc/ref/backend.rst @@ -89,6 +89,10 @@ Database Lock the ``table`` using the ``connection``. +.. classmethod:: Database.lock_records(connection, table, ids) + + Lock the record ``ids`` of the ``table`` using the ``connection``. + .. method:: Database.lock_id(id[, timeout]) Return the SQL expression to lock the ``id``. diff --git a/trytond/doc/ref/transaction.rst b/trytond/doc/ref/transaction.rst index c50346f9cf..37a7961942 100644 --- a/trytond/doc/ref/transaction.rst +++ b/trytond/doc/ref/transaction.rst @@ -112,6 +112,11 @@ Transaction Raise a :exc:`TransactionError` to retry the transaction if the table has not been locked at the start. +.. method:: Transaction.lock_records(table, ids) + + Raise a :exc:`TransactionError` to retry the transaction if the ``ids`` of + the ``table`` has not been locked at the start. + .. method:: Transaction.set_current_transaction(transaction) Add a specific ``transaction`` on the top of the transaction stack. diff --git a/trytond/trytond/backend/database.py b/trytond/trytond/backend/database.py index fd2b88016a..976bc765c7 100644 --- a/trytond/trytond/backend/database.py +++ b/trytond/trytond/backend/database.py @@ -73,6 +73,10 @@ class DatabaseInterface(object): def lock(cls, connection, table): raise NotImplementedError + @classmethod + def lock_records(cls, connection, table, ids): + raise NotImplementedError + def lock_id(self, id, timeout=None): raise NotImplementedError diff --git a/trytond/trytond/backend/postgresql/database.py b/trytond/trytond/backend/postgresql/database.py index 9af0e2b0df..4e6cfd2e0c 100644 --- a/trytond/trytond/backend/postgresql/database.py +++ b/trytond/trytond/backend/postgresql/database.py @@ -40,6 +40,7 @@ from trytond import __series__ from trytond.backend.database import DatabaseInterface, SQLType from trytond.config import config, parse_uri from trytond.sql.operators import RangeOperator +from trytond.tools import grouped_slice, reduce_ids from trytond.tools.gevent import is_gevent_monkey_patched from .table import index_method @@ -474,11 +475,22 @@ class Database(DatabaseInterface): cursor.execute(*from_item.select(Count(Literal('*')))) return cursor.fetchone()[0] - def lock(self, connection, table): + @classmethod + def lock(cls, connection, table): cursor = connection.cursor() cursor.execute(SQL('LOCK {} IN EXCLUSIVE MODE NOWAIT').format( Identifier(table))) + @classmethod + def lock_records(cls, connection, table, ids): + table = Table(table) + cursor = connection.cursor() + for sub_ids in grouped_slice(ids): + where = reduce_ids(table.id, sub_ids) + query = table.select( + Literal(1), where=where, for_=For('UPDATE', nowait=True)) + cursor.execute(*query) + def lock_id(self, id, timeout=None): if not timeout: return TryAdvisoryLock(id) diff --git a/trytond/trytond/backend/sqlite/database.py b/trytond/trytond/backend/sqlite/database.py index ee2b64613e..6a0928eb3a 100644 --- a/trytond/trytond/backend/sqlite/database.py +++ b/trytond/trytond/backend/sqlite/database.py @@ -728,7 +728,12 @@ class Database(DatabaseInterface): cursor.execute(*table.select(Count(Literal('*')))) return cursor.fetchone()[0] - def lock(self, connection, table): + @classmethod + def lock(cls, connection, table): + pass + + @classmethod + def lock_records(cls, connection, table, ids): pass def lock_id(self, id, timeout=None): diff --git a/trytond/trytond/commandline.py b/trytond/trytond/commandline.py index 944101e712..054ba6a087 100644 --- a/trytond/trytond/commandline.py +++ b/trytond/trytond/commandline.py @@ -160,6 +160,12 @@ def get_parser_admin(): return parser +def _convert_records(s): + table, ids = s.split('=', 1) + ids = list(map(int, ids.split(','))) + return table, ids + + def get_parser_console(): parser = get_base_parser() parser.add_argument( @@ -173,6 +179,10 @@ def get_parser_console(): parser.add_argument( "--lock-table", dest="lock_tables", nargs='+', default=[], metavar='TABLE', help="lock tables") + parser.register('type', 'records', _convert_records) + parser.add_argument( + "--lock-records", dest="lock_records", type='records', nargs='+', + default=[], metavar='TABLE=id,id', help="lock record ids from TABLE") parser.epilog = "To store changes, `transaction.commit()` must be called." return parser diff --git a/trytond/trytond/console.py b/trytond/trytond/console.py index 8243f40644..1966d82359 100644 --- a/trytond/trytond/console.py +++ b/trytond/trytond/console.py @@ -47,7 +47,8 @@ def run(options): with Transaction().start( db_name, 0, readonly=options.readonly, - _lock_tables=options.lock_tables) as transaction: + _lock_tables=options.lock_tables, + _lock_records=dict(options.lock_records)) as transaction: local = { 'pool': pool, 'transaction': transaction, diff --git a/trytond/trytond/model/modelsql.py b/trytond/trytond/model/modelsql.py index c04d5a339e..47116b3df5 100644 --- a/trytond/trytond/model/modelsql.py +++ b/trytond/trytond/model/modelsql.py @@ -6,8 +6,8 @@ from functools import wraps from itertools import chain, groupby, islice, product, repeat from sql import ( - Asc, Column, Desc, Expression, For, Literal, Null, NullsFirst, NullsLast, - Table, Union, Window, With) + Asc, Column, Desc, Expression, Literal, Null, NullsFirst, NullsLast, Table, + Union, Window, With) from sql.aggregate import Count, Max from sql.conditionals import Coalesce from sql.functions import CurrentTimestamp, Extract, RowNumber, Substring @@ -2248,17 +2248,10 @@ class ModelSQL(ModelStorage): @dualmethod def lock(cls, records=None): transaction = Transaction() - database = transaction.database - connection = transaction.connection - table = cls.__table__() - - if records is not None and database.has_select_for(): - for sub_records in grouped_slice(records): - where = reduce_ids(table.id, sub_records) - query = table.select( - Literal(1), where=where, for_=For('UPDATE', nowait=True)) - with connection.cursor() as cursor: - cursor.execute(*query) + if records is not None: + new_ids = transaction.create_records[cls.__name__] + ids = [id_ for id_ in map(int, records) if id_ not in new_ids] + transaction.lock_records(cls._table, ids) else: transaction.lock_table(cls._table) diff --git a/trytond/trytond/tests/test_modelsql.py b/trytond/trytond/tests/test_modelsql.py index a1d5860a01..b8d12a153e 100644 --- a/trytond/trytond/tests/test_modelsql.py +++ b/trytond/trytond/tests/test_modelsql.py @@ -923,13 +923,28 @@ class ModelSQLTestCase(TestCase): record_id = Model.create([{}])[0].id transaction.commit() - with transaction.new_transaction(): - record = Model(record_id) - record.lock() - with transaction.new_transaction(): + args1 = {} + while True: + with transaction.new_transaction(**args1): record = Model(record_id) - with self.assertRaises(backend.DatabaseOperationalError): + try: record.lock() + except TransactionError as e: + e.fix(args1) + continue + with self.assertRaises( + backend.DatabaseOperationalError): + args2 = {} + while True: + with transaction.new_transaction(**args2): + record = Model(record_id) + try: + record.lock() + except TransactionError as e: + e.fix(args2) + continue + break + break @unittest.skipIf(backend.name == 'sqlite', 'SQLite does not have lock at table level but on file') @@ -961,6 +976,19 @@ class ModelSQLTestCase(TestCase): break break + @with_transaction() + def test_record_lock_with_table_lock(self): + "Test not locking record when table is locked" + pool = Pool() + Model = pool.get('test.modelsql.lock') + transaction = Transaction() + record_id = Model.create([{}])[0].id + transaction.commit() + + with transaction.new_transaction(_lock_tables=[Model._table]): + record = Model(record_id) + record.lock() + @with_transaction() def test_search_or_to_union(self): """ diff --git a/trytond/trytond/transaction.py b/trytond/trytond/transaction.py index 03cdbf5bc2..52f63d747b 100644 --- a/trytond/trytond/transaction.py +++ b/trytond/trytond/transaction.py @@ -39,6 +39,18 @@ class _TransactionLockError(TransactionError): extras.setdefault('_lock_tables', []).append(self._table) +class _TransactionLockRecordsError(TransactionError): + def __init__(self, table, ids): + super().__init__() + self._table = table + self._ids = ids + + def fix(self, extras): + super().fix(extras) + extras.setdefault('_lock_records', {}).setdefault( + self._table, []).extend(self._ids) + + def record_cache_size(transaction): return transaction.context.get('_record_cache_size', _cache_record) @@ -204,6 +216,11 @@ class Transaction(object): for table in lock_tables: self.database.lock(self.connection, table) self._locked_tables = set(lock_tables) + locked_records = defaultdict(set) + for table, ids in extras.get('_lock_records', {}).items(): + self.database.lock_records(self.connection, table, ids) + locked_records[table].update(ids) + self._locked_records = locked_records except backend.DatabaseOperationalError: if count < _retry: self.connection.rollback() @@ -298,6 +315,11 @@ class Transaction(object): if table not in self._locked_tables: raise _TransactionLockError(table) + def lock_records(self, table, ids): + if table not in self._locked_tables: + if missing := (set(ids) - self._locked_records[table]): + raise _TransactionLockRecordsError(table, missing) + def set_current_transaction(self, transaction): self._local.transactions.append(transaction) return transaction