Lock records on transaction start

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

Closes #13943
(grafted from 00790a58e6a364c46dd7178de02bd6842e6dc0c6)

--HG--
branch : 7.0

trytond/CHANGELOG

@@ -1,3 +1,4 @@
+* Lock records on transaction start
 
 Version 7.0.37 - 2025-10-02
 ---------------------------

trytond/doc/ref/backend.rst

@@ -83,6 +83,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

@@ -111,6 +111,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

@@ -68,6 +68,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

@@ -27,13 +27,14 @@ from psycopg2 import IntegrityError as DatabaseIntegrityError
 from psycopg2 import OperationalError as DatabaseOperationalError
 from psycopg2 import ProgrammingError
 from psycopg2.extras import register_default_json, register_default_jsonb
-from sql import Cast, Flavor, For, Table
+from sql import Cast, Flavor, For, Literal, Table
 from sql.conditionals import Coalesce
 from sql.functions import Function
 from sql.operators import BinaryOperator, Concat
 
 from trytond.backend.database import DatabaseInterface, SQLType
 from trytond.config import config, parse_uri
+from trytond.tools import grouped_slice, reduce_ids
 from trytond.tools.gevent import is_gevent_monkey_patched
 
 __all__ = [
@@ -440,11 +441,22 @@ class Database(DatabaseInterface):
         cursor.execute(f"SELECT last_value FROM {sequence_name}")
         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

@@ -588,7 +588,12 @@ class Database(DatabaseInterface):
         # This call is not thread safe
         return cursor.lastrowid
 
-    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

@@ -144,6 +144,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(
@@ -157,6 +163,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, With)
+    Asc, Column, Desc, Expression, Literal, Null, NullsFirst, NullsLast, Table,
+    Union, With)
 from sql.aggregate import Count, Max
 from sql.conditionals import Coalesce
 from sql.functions import CurrentTimestamp, Extract, Substring
@@ -2180,17 +2180,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

@@ -856,13 +856,28 @@ class ModelSQLTestCase(unittest.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')
@@ -894,6 +909,19 @@ class ModelSQLTestCase(unittest.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

@@ -38,6 +38,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)
 
@@ -203,6 +215,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()
@@ -296,6 +313,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