WL10974: Add Row locking methods to find and select operations

This patch add two methods `lock_shared` and `lock_exclusive` to the
Colleciton.find() and Table.select() command chains.

Multiple calls to these methods are allowed, but only the last call is
considered as final locking mode.

Tests have been added for regression.

Author: amitab

lib/mysqlx/__init__.py

@@ -42,7 +42,8 @@
                         DeleteStatement, UpdateStatement,
                         CreateCollectionIndexStatement, CreateTableStatement,
                         CreateViewStatement, AlterViewStatement, ColumnDef,
-                        GeneratedColumnDef, ForeignKeyDef, Expr)
+                        GeneratedColumnDef, ForeignKeyDef, Expr,
+                        ReadStatement)
 
 _SPLIT = re.compile(r',(?![^\(\)]*\))')
 _PRIORITY = re.compile(r'^\(address=(.+),priority=(\d+)\)$', re.VERBOSE)

lib/mysqlx/protocol.py

@@ -167,6 +167,14 @@ def send_find(self, stmt):
         if stmt._has_projection:
             msg["projection"] = stmt._projection_expr
         self._apply_filter(msg, stmt)
+
+        if stmt._lock_exclusive:
+            msg["locking"] = \
+                mysqlxpb_enum("Mysqlx.Crud.Find.RowLock.EXCLUSIVE_LOCK")
+        elif stmt._lock_shared:
+            msg["locking"] = \
+                mysqlxpb_enum("Mysqlx.Crud.Find.RowLock.SHARED_LOCK")
+
         self._writer.write_message(
             mysqlxpb_enum("Mysqlx.ClientMessages.Type.CRUD_FIND"), msg)
 

lib/mysqlx/statement.py

@@ -477,30 +477,37 @@ def execute(self):
         return self._connection.update(self)
 
 
-class FindStatement(FilterableStatement):
-    """A statement document selection on a Collection.
+class ReadStatement(FilterableStatement):
+    """Provide base functionality for Read operations
 
     Args:
-        collection (mysqlx.Collection): The Collection object.
-        condition (Optional[str]): An optional expression to identify the
-                                   documents to be retrieved. If not specified
-                                   all the documents will be included on the
-                                   result unless a limit is set.
+        target (object): The target database object, it can be
+                         :class:`mysqlx.Collection` or :class:`mysqlx.Table`.
+        doc_based (Optional[bool]): `True` if it is document based
+                                    (default: `True`).
+        condition (Optional[str]): Sets the search condition to filter
+                                   documents or records.
     """
-    def __init__(self, collection, condition=None):
-        super(FindStatement, self).__init__(collection, True, condition)
-
-    def fields(self, *fields):
-        """Sets a document field filter.
+    def __init__(self, target, doc_based=True, condition=None):
+        super(ReadStatement, self).__init__(target, doc_based, condition)
+        self._lock_exclusive = False
+        self._lock_shared = False
 
-        Args:
-            *fields: The string expressions identifying the fields to be
-                     extracted.
+    def lock_shared(self):
+        """Execute a read operation with SHARED LOCK. Only one lock can be
+           active at a time.
+        """
+        self._lock_exclusive = False
+        self._lock_shared = True
+        return self
 
-        Returns:
-            mysqlx.FindStatement: FindStatement object.
+    def lock_exclusive(self):
+        """Execute a read operation with EXCLUSIVE LOCK. Only one lock can be
+           active at a time.
         """
-        return self._projection(*fields)
+        self._lock_exclusive = True
+        self._lock_shared = False
+        return self
 
     def group_by(self, *fields):
         """Sets a grouping criteria for the resultset.
@@ -509,7 +516,7 @@ def group_by(self, *fields):
             *fields: The string expressions identifying the grouping criteria.
 
         Returns:
-            mysqlx.FindStatement: FindStatement object.
+            mysqlx.ReadStatement: ReadStatement object.
         """
         self._group_by(*fields)
         return self
@@ -523,7 +530,7 @@ def having(self, condition):
                                 the grouping criteria.
 
         Returns:
-            mysqlx.FindStatement: FindStatement object.
+            mysqlx.ReadStatement: ReadStatement object.
         """
         self._having(condition)
         return self
@@ -532,12 +539,38 @@ def execute(self):
         """Execute the statement.
 
         Returns:
-            mysqlx.DocResult: DocResult object.
+            mysqlx.Result: Result object.
         """
         return self._connection.find(self)
 
 
-class SelectStatement(FilterableStatement):
+class FindStatement(ReadStatement):
+    """A statement document selection on a Collection.
+
+    Args:
+        collection (mysqlx.Collection): The Collection object.
+        condition (Optional[str]): An optional expression to identify the
+                                   documents to be retrieved. If not specified
+                                   all the documents will be included on the
+                                   result unless a limit is set.
+    """
+    def __init__(self, collection, condition=None):
+        super(FindStatement, self).__init__(collection, True, condition)
+
+    def fields(self, *fields):
+        """Sets a document field filter.
+
+        Args:
+            *fields: The string expressions identifying the fields to be
+                     extracted.
+
+        Returns:
+            mysqlx.FindStatement: FindStatement object.
+        """
+        return self._projection(*fields)
+
+
+class SelectStatement(ReadStatement):
     """A statement for record retrieval operations on a Table.
 
     Args:
@@ -560,40 +593,6 @@ def order_by(self, *clauses):
         self.sort(*clauses)
         return self
 
-    def group_by(self, *fields):
-        """Sets a grouping criteria for the resultset.
-
-        Args:
-            *fields: The fields identifying the grouping criteria.
-
-        Returns:
-            mysqlx.SelectStatement: SelectStatement object.
-        """
-        self._group_by(*fields)
-        return self
-
-    def having(self, condition):
-        """Sets a condition for records to be considered in agregate function
-        operations.
-
-        Args:
-            condition (str): A condition on the agregate functions used on the
-                             grouping criteria.
-
-        Returns:
-            mysqlx.SelectStatement: SelectStatement object.
-        """
-        self._having(condition)
-        return self
-
-    def execute(self):
-        """Execute the statement.
-
-        Returns:
-            mysqlx.RowResult: RowResult object.
-        """
-        return self._connection.find(self)
-
     def get_sql(self):
         where = " WHERE {0}".format(self._where) if self._has_where else ""
         group_by = " GROUP BY {0}".format(self._grouping_str) if \
@@ -611,6 +610,7 @@ def get_sql(self):
 
         return stmt
 
+
 class InsertStatement(Statement):
     """A statement for insert operations on Table.
 

src/mysqlxpb/mysqlx/protocol/mysqlx_crud.proto

@@ -122,8 +122,14 @@ message UpdateOperation {
 // :param order: sort-order in which the rows/document shall be returned in
 // :param grouping: column expression list for aggregation (GROUP BY)
 // :param grouping_criteria: filter criteria for aggregated groups
+// :param locking: perform row locking on matches
 // :Returns: :protobuf:msg:`Mysqlx.Resultset::`
 message Find {
+  enum RowLock {
+    SHARED_LOCK = 1;   // Lock matching rows against updates
+    EXCLUSIVE_LOCK = 2;  // Lock matching rows so no other transaction can read or write to it
+  };
+
   required Collection collection = 2;
 
   optional DataModel data_model = 3;
@@ -134,6 +140,7 @@ message Find {
   repeated Order order = 7;
   repeated Mysqlx.Expr.Expr grouping = 8;
   optional Mysqlx.Expr.Expr grouping_criteria = 9;
+  optional RowLock locking = 12;
 };
 
 // Insert documents/rows into a collection/table

tests/test_mysqlx_crud.py

@@ -27,6 +27,8 @@
 
 import logging
 import unittest
+import threading
+import time
 
 import tests
 import mysqlx
@@ -443,6 +445,171 @@ def test_exists_in_database(self):
         self.assertTrue(collection.exists_in_database())
         self.schema.drop_collection(collection_name)
 
+    @unittest.skipIf(tests.MYSQL_VERSION < (8, 0, 3), "Row locks unavailable.")
+    def test_lock_shared(self):
+        collection_name = "collection_test"
+        collection = self.schema.create_collection(collection_name)
+        result = collection.add({"name": "Fred", "age": 21}).execute()
+
+        pause = threading.Event()
+        locking = threading.Event()
+        waiting = threading.Event()
+
+        errors = []
+
+        def client_a(pause, locking, waiting):
+            sess1 = mysqlx.get_session(self.connect_kwargs)
+            schema = sess1.get_schema(self.schema_name)
+            collection = schema.get_collection(collection_name)
+
+            sess1.start_transaction()
+            result = collection.find("name = 'Fred'").lock_shared().execute()
+            locking.set()
+            time.sleep(2)
+            locking.clear()
+            if waiting.is_set():
+                errors.append("S-S lock test failure.")
+                sess1.commit()
+                return
+            sess1.commit()
+
+            pause.set()
+
+            sess1.start_transaction()
+            result = collection.find("name = 'Fred'").lock_shared().execute()
+            locking.set()
+            time.sleep(2)
+            locking.clear()
+            if not waiting.is_set():
+                errors.append("S-X lock test failure.")
+                sess1.commit()
+                return
+            sess1.commit()
+
+        def client_b(pause, locking, waiting):
+            sess1 = mysqlx.get_session(self.connect_kwargs)
+            schema = sess1.get_schema(self.schema_name)
+            collection = schema.get_collection(collection_name)
+
+            if not locking.wait(2):
+                return
+            sess1.start_transaction()
+
+            waiting.set()
+            result = collection.find("name = 'Fred'").lock_shared().execute()
+            waiting.clear()
+
+            sess1.commit()
+
+            if not pause.wait(2):
+                return
+
+            if not locking.wait(2):
+                return
+            sess1.start_transaction()
+            waiting.set()
+            result = collection.find("name = 'Fred'").lock_exclusive().execute()
+            waiting.clear()
+            sess1.commit()
+
+        client1 = threading.Thread(target=client_a,
+                                   args=(pause, locking, waiting,))
+        client2 = threading.Thread(target=client_b,
+                                   args=(pause, locking, waiting,))
+
+        client1.start()
+        client2.start()
+
+        client1.join()
+        client2.join()
+
+        self.schema.drop_collection(collection_name)
+        if errors:
+            self.fail(errors[0])
+
+    @unittest.skipIf(tests.MYSQL_VERSION < (8, 0, 3), "Row locks unavailable.")
+    def test_lock_exclusive(self):
+        collection_name = "collection_test"
+        collection = self.schema.create_collection(collection_name)
+        result = collection.add({"name": "Fred", "age": 21}).execute()
+        event = threading.Event()
+
+        pause = threading.Event()
+        locking = threading.Event()
+        waiting = threading.Event()
+
+        errors = []
+
+        def client_a(pause, locking, waiting):
+            sess1 = mysqlx.get_session(self.connect_kwargs)
+            schema = sess1.get_schema(self.schema_name)
+            collection = schema.get_collection(collection_name)
+
+            sess1.start_transaction()
+            result = collection.find("name = 'Fred'").lock_exclusive().execute()
+            locking.set()
+            time.sleep(2)
+            locking.clear()
+            if not waiting.is_set():
+                sess1.commit()
+                errors.append("X-X lock test failure.")
+                return
+            sess1.commit()
+
+            pause.set()
+
+            sess1.start_transaction()
+            result = collection.find("name = 'Fred'").lock_exclusive().execute()
+            locking.set()
+            time.sleep(2)
+            locking.clear()
+            if not waiting.is_set():
+                errors.append("X-S lock test failure.")
+                sess1.commit()
+                return
+            sess1.commit()
+
+        def client_b(pause, locking, waiting):
+            sess1 = mysqlx.get_session(self.connect_kwargs)
+            schema = sess1.get_schema(self.schema_name)
+            collection = schema.get_collection(collection_name)
+
+            if not locking.wait(2):
+                return
+            sess1.start_transaction()
+
+            waiting.set()
+            result = collection.find("name = 'Fred'").lock_exclusive().execute()
+            waiting.clear()
+
+            sess1.commit()
+
+            if not pause.wait(2):
+                return
+
+            if not locking.wait(2):
+                return
+            sess1.start_transaction()
+            waiting.set()
+            result = collection.find("name = 'Fred'").lock_shared().execute()
+            waiting.clear()
+            sess1.commit()
+
+        client1 = threading.Thread(target=client_a,
+                                   args=(pause, locking, waiting,))
+        client2 = threading.Thread(target=client_b,
+                                   args=(pause, locking, waiting,))
+
+        client1.start()
+        client2.start()
+
+        client1.join()
+        client2.join()
+
+        self.schema.drop_collection(collection_name)
+        if errors:
+            self.fail(errors[0])
+
     def test_add(self):
         collection_name = "collection_test"
         collection = self.schema.create_collection(collection_name)