Lock records on transaction start

Same as c8771e45b861 but when locking only records instead of the table.

Closes #13943

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

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``.

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.

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
 

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)

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):

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
 

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,

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)
 

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):
         """

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