Implement QuerySnapshotReadWrite transaction mode for query client

Co-authored-by: asmyasnikov <[email protected]>

Author: Copilot

tests/aio/query/test_query_session_pool.py

@@ -69,6 +69,7 @@ async def callee(session: QuerySession):
             (None),
             (ydb.QuerySerializableReadWrite()),
             (ydb.QuerySnapshotReadOnly()),
+            (ydb.QuerySnapshotReadWrite()),
             (ydb.QueryOnlineReadOnly()),
             (ydb.QueryStaleReadOnly()),
         ],

tests/query/test_query_session_pool.py

@@ -68,6 +68,7 @@ def callee(session: QuerySession):
             (None),
             (ydb.QuerySerializableReadWrite()),
             (ydb.QuerySnapshotReadOnly()),
+            (ydb.QuerySnapshotReadWrite()),
             (ydb.QueryOnlineReadOnly()),
             (ydb.QueryStaleReadOnly()),
         ],

ydb/_grpc/grpcwrapper/ydb_query.py

@@ -79,6 +79,8 @@ def to_proto(self) -> ydb_query_pb2.TransactionSettings:
             return ydb_query_pb2.TransactionSettings(online_read_only=self.tx_mode.to_proto())
         if self.tx_mode.name == "stale_read_only":
             return ydb_query_pb2.TransactionSettings(stale_read_only=self.tx_mode.to_proto())
+        if self.tx_mode.name == "snapshot_read_write":
+            return ydb_query_pb2.TransactionSettings(snapshot_read_write=self.tx_mode.to_proto())
 
 
 @dataclass

ydb/_grpc/grpcwrapper/ydb_query_public_types.py

@@ -93,3 +93,21 @@ def name(self):
 
     def to_proto(self) -> ydb_query_pb2.StaleModeSettings:
         return ydb_query_pb2.StaleModeSettings()
+
+
+class QuerySnapshotReadWrite(BaseQueryTxMode):
+    """Snapshot isolation mode with read-write support.
+    All read operations within a transaction access the database snapshot.
+    The snapshot is taken when the transaction begins.
+    Write operations are also allowed in this mode.
+    """
+
+    def __init__(self):
+        self._name = "snapshot_read_write"
+
+    @property
+    def name(self) -> str:
+        return self._name
+
+    def to_proto(self) -> ydb_query_pb2.SnapshotRWModeSettings:
+        return ydb_query_pb2.SnapshotRWModeSettings()

ydb/aio/query/pool.py

@@ -149,7 +149,8 @@ async def retry_tx_async(
           1) QuerySerializableReadWrite() which is default mode;
           2) QueryOnlineReadOnly(allow_inconsistent_reads=False);
           3) QuerySnapshotReadOnly();
-          4) QueryStaleReadOnly().
+          4) QuerySnapshotReadWrite();
+          5) QueryStaleReadOnly().
         :param retry_settings: RetrySettings object.
 
         :return: Result sets or exception in case of execution errors.
@@ -161,7 +162,7 @@ async def retry_tx_async(
         async def wrapped_callee():
             async with self.checkout() as session:
                 async with session.transaction(tx_mode=tx_mode) as tx:
-                    if tx_mode.name in ["serializable_read_write", "snapshot_read_only"]:
+                    if tx_mode.name in ["serializable_read_write", "snapshot_read_only", "snapshot_read_write"]:
                         await tx.begin()
                     result = await callee(tx, *args, **kwargs)
                     await tx.commit()

ydb/aio/query/transaction.py

@@ -33,7 +33,8 @@ def __init__(self, driver, session_state, session, tx_mode):
          1) QuerySerializableReadWrite() which is default mode;
          2) QueryOnlineReadOnly(allow_inconsistent_reads=False);
          3) QuerySnapshotReadOnly();
-         4) QueryStaleReadOnly().
+         4) QuerySnapshotReadWrite();
+         5) QueryStaleReadOnly().
         """
         super().__init__(driver, session_state, session, tx_mode)
         self._init_callback_handler(base.CallbackHandlerMode.ASYNC)

ydb/query/__init__.py

@@ -4,6 +4,7 @@
     "QueryOnlineReadOnly",
     "QuerySerializableReadWrite",
     "QuerySnapshotReadOnly",
+    "QuerySnapshotReadWrite",
     "QueryStaleReadOnly",
     "QuerySessionPool",
     "QueryClientSettings",
@@ -29,6 +30,7 @@
     QueryOnlineReadOnly,
     QuerySerializableReadWrite,
     QuerySnapshotReadOnly,
+    QuerySnapshotReadWrite,
     QueryStaleReadOnly,
 )
 

ydb/query/pool.py

@@ -197,7 +197,8 @@ def retry_tx_sync(
           1) QuerySerializableReadWrite() which is default mode;
           2) QueryOnlineReadOnly(allow_inconsistent_reads=False);
           3) QuerySnapshotReadOnly();
-          4) QueryStaleReadOnly().
+          4) QuerySnapshotReadWrite();
+          5) QueryStaleReadOnly().
         :param retry_settings: RetrySettings object.
 
         :return: Result sets or exception in case of execution errors.
@@ -212,7 +213,7 @@ def retry_tx_sync(
         def wrapped_callee():
             with self.checkout(timeout=retry_settings.max_session_acquire_timeout) as session:
                 with session.transaction(tx_mode=tx_mode) as tx:
-                    if tx_mode.name in ["serializable_read_write", "snapshot_read_only"]:
+                    if tx_mode.name in ["serializable_read_write", "snapshot_read_only", "snapshot_read_write"]:
                         tx.begin()
                     result = callee(tx, *args, **kwargs)
                     tx.commit()

ydb/query/session.py

@@ -310,7 +310,8 @@ def transaction(self, tx_mode: Optional[base.BaseQueryTxMode] = None) -> QueryTx
          1) QuerySerializableReadWrite() which is default mode;
          2) QueryOnlineReadOnly(allow_inconsistent_reads=False);
          3) QuerySnapshotReadOnly();
-         4) QueryStaleReadOnly().
+         4) QuerySnapshotReadWrite();
+         5) QueryStaleReadOnly().
 
         :return transaction context manager.
 

ydb/query/transaction.py

@@ -201,7 +201,8 @@ def __init__(self, driver, session_state, session, tx_mode):
          1) QuerySerializableReadWrite() which is default mode;
          2) QueryOnlineReadOnly(allow_inconsistent_reads=False);
          3) QuerySnapshotReadOnly();
-         4) QueryStaleReadOnly().
+         4) QuerySnapshotReadWrite();
+         5) QueryStaleReadOnly().
         """
 
         self._driver = driver
@@ -351,7 +352,8 @@ def __init__(self, driver, session_state, session, tx_mode):
          1) QuerySerializableReadWrite() which is default mode;
          2) QueryOnlineReadOnly(allow_inconsistent_reads=False);
          3) QuerySnapshotReadOnly();
-         4) QueryStaleReadOnly().
+         4) QuerySnapshotReadWrite();
+         5) QueryStaleReadOnly().
         """
 
         super().__init__(driver, session_state, session, tx_mode)