refactor transaction

Author: vgvoleg

docker-compose.yml

@@ -1,7 +1,7 @@
 version: "3.3"
 services:
   ydb:
-    image: ydbplatform/local-ydb:latest
+    image: ydbplatform/local-ydb:trunk
     restart: always
     ports:
       - 2136:2136

tests/query/conftest.py

@@ -14,7 +14,7 @@ def session(driver_sync):
         pass
 
 @pytest.fixture
-def transaction(session):
+def tx(session):
     session.create()
     transaction = session.transaction()
 

tests/query/test_query_transaction.py

@@ -1,17 +1,44 @@
 import pytest
 
-import ydb.query.session
+class TestQueryTransaction:
+    def test_tx_begin(self, tx):
+        assert tx.tx_id == None
 
-class TestQuerySession:
-    def test_transaction_begin(self, driver_sync):
-        session = ydb.query.session.QuerySessionSync(driver_sync)
+        tx.begin()
+        assert tx.tx_id != None
 
-        session.create()
+    def test_tx_allow_double_commit(self, tx):
+        tx.begin()
+        tx.commit()
+        tx.commit()
 
-        tx = session.transaction()
+    def test_tx_allow_double_rollback(self, tx):
+        tx.begin()
+        tx.rollback()
+        tx.rollback()
 
-        assert tx.tx_id == None
+    def test_tx_commit_raises_before_begin(self, tx):
+        with pytest.raises(RuntimeError):
+            tx.commit()
 
+    def test_tx_rollback_raises_before_begin(self, tx):
+        with pytest.raises(RuntimeError):
+            tx.rollback()
+
+    # def test_tx_execute_raises_before_begin(self, tx):
+    #     with pytest.raises(RuntimeError):
+    #         tx.execute("select 1;")
+
+    def text_tx_execute_raises_after_commit(self, tx):
         tx.begin()
+        tx.commit()
+        with pytest.raises(RuntimeError):
+            tx.execute("select 1;")
+
+    def text_tx_execute_raises_after_rollback(self, tx):
+        tx.begin()
+        tx.rollback()
+        with pytest.raises(RuntimeError):
+            tx.execute("select 1;")
+
 
-        assert tx.tx_id != None

ydb/query/base.py

@@ -124,19 +124,32 @@ def session(self) -> IQuerySession:
         pass
 
 
-def create_execute_query_request(query: str, session_id: str, commit_tx: bool):
-    req = ydb_query.ExecuteQueryRequest(
-        session_id=session_id,
-        query_content=ydb_query.QueryContent.from_public(
-            query=query,
-        ),
-        tx_control=ydb_query.TransactionControl(
-            begin_tx=ydb_query.TransactionSettings(
-                tx_mode=QuerySerializableReadWrite(),
+def create_execute_query_request(query: str, session_id: str, tx_id: str = None, commit_tx: bool = False, tx_mode: BaseQueryTxMode = None):
+    if tx_id:
+        req = ydb_query.ExecuteQueryRequest(
+            session_id=session_id,
+            query_content=ydb_query.QueryContent.from_public(
+                query=query,
             ),
-            commit_tx=commit_tx
-        ),
-    )
+            tx_control=ydb_query.TransactionControl(
+                tx_id=tx_id,
+                commit_tx=commit_tx
+            ),
+        )
+    else:
+        tx_mode = tx_mode if tx_mode is not None else QuerySerializableReadWrite()
+        req = ydb_query.ExecuteQueryRequest(
+            session_id=session_id,
+            query_content=ydb_query.QueryContent.from_public(
+                query=query,
+            ),
+            tx_control=ydb_query.TransactionControl(
+                begin_tx=ydb_query.TransactionSettings(
+                    tx_mode=tx_mode,
+                ),
+                commit_tx=commit_tx
+            ),
+        )
 
     return req.to_proto()
 
@@ -148,17 +161,17 @@ def wrap_execute_query_response(rpc_state, response_pb):
 X_YDB_SESSION_CLOSE = "session-close"
 
 
-def _check_session_is_closing(rpc_state, session_state):
-    metadata = rpc_state.trailing_metadata()
-    if X_YDB_SESSION_CLOSE in metadata.get(X_YDB_SERVER_HINTS, []):
-        session_state.set_closing() # TODO: clarify & implement
+# def _check_session_is_closing(rpc_state, session_state):
+#     metadata = rpc_state.trailing_metadata()
+#     if X_YDB_SESSION_CLOSE in metadata.get(X_YDB_SERVER_HINTS, []):
+#         session_state.set_closing() # TODO: clarify & implement
 
 
 def bad_session_handler(func):
     @functools.wraps(func)
     def decorator(rpc_state, response_pb, session_state, *args, **kwargs):
         try:
-            _check_session_is_closing(rpc_state, session_state)
+            # _check_session_is_closing(rpc_state, session_state)
             return func(rpc_state, response_pb, session_state, *args, **kwargs)
         except issues.BadSession:
             session_state.reset()

ydb/query/transaction.py

@@ -15,19 +15,37 @@
 
 logger = logging.getLogger(__name__)
 
-# def patch_table_service_tx_mode_to_query_service(tx_mode: AbstractTransactionModeBuilder):
-#     if tx_mode.name == 'snapshot_read_only':
-#         tx_mode = _ydb_query_public.QuerySnapshotReadOnly()
-#     elif tx_mode.name == 'serializable_read_write':
-#         tx_mode = _ydb_query_public.QuerySerializableReadWrite()
-#     elif tx_mode.name =='online_read_only':
-#         tx_mode = _ydb_query_public.QueryOnlineReadOnly()
-#     elif tx_mode.name == 'stale_read_only':
-#         tx_mode = _ydb_query_public.QueryStaleReadOnly()
-#     else:
-#         raise issues.YDBInvalidArgumentError(f'Unknown transaction mode: {tx_mode.name}')
 
-#     return tx_mode
+class QueryTxStateEnum(enum.Enum):
+    NOT_INITIALIZED = "NOT_INITIALIZED"
+    BEGINED = "BEGINED"
+    COMMITTED = "COMMITTED"
+    ROLLBACKED = "ROLLBACKED"
+    DEAD = "DEAD"
+
+
+class QueryTxStateHelper(abc.ABC):
+    _VALID_TRANSITIONS = {
+        QueryTxStateEnum.NOT_INITIALIZED: [QueryTxStateEnum.BEGINED, QueryTxStateEnum.DEAD],
+        QueryTxStateEnum.BEGINED: [QueryTxStateEnum.COMMITTED, QueryTxStateEnum.ROLLBACKED, QueryTxStateEnum.DEAD],
+        QueryTxStateEnum.COMMITTED: [],
+        QueryTxStateEnum.ROLLBACKED: [],
+        QueryTxStateEnum.DEAD: [],
+    }
+
+    _TERMINAL_STATES = [
+        QueryTxStateEnum.COMMITTED,
+        QueryTxStateEnum.ROLLBACKED,
+        QueryTxStateEnum.DEAD,
+    ]
+
+    @classmethod
+    def valid_transition(cls, before: QueryTxStateEnum, after: QueryTxStateEnum) -> bool:
+        return after in cls._VALID_TRANSITIONS[before]
+
+    @classmethod
+    def terminal(cls, state: QueryTxStateEnum) -> bool:
+        return state in cls._TERMINAL_STATES
 
 
 def reset_tx_id_handler(func):
@@ -36,7 +54,7 @@ def decorator(rpc_state, response_pb, session_state, tx_state, *args, **kwargs):
         try:
             return func(rpc_state, response_pb, session_state, tx_state, *args, **kwargs)
         except issues.Error:
-            tx_state.change_state(base.QueryTxStateEnum.DEAD)
+            tx_state._change_state(QueryTxStateEnum.DEAD)
             tx_state.tx_id = None
             raise
 
@@ -51,16 +69,23 @@ def __init__(self, tx_mode: base.BaseQueryTxMode):
         """
         self.tx_id = None
         self.tx_mode = tx_mode
-        self._state = base.QueryTxStateEnum.NOT_INITIALIZED
+        self._state = QueryTxStateEnum.NOT_INITIALIZED
 
-    def check_invalid_transition(self, target: base.QueryTxStateEnum):
-        if not base.QueryTxStateHelper.is_valid_transition(self._state, target):
+    def _check_invalid_transition(self, target: QueryTxStateEnum):
+        if not QueryTxStateHelper.valid_transition(self._state, target):
             raise RuntimeError(f"Transaction could not be moved from {self._state.value} to {target.value}")
 
-    def change_state(self, target: base.QueryTxStateEnum):
-        self.check_invalid_transition(target)
+    def _change_state(self, target: QueryTxStateEnum):
+        self._check_invalid_transition(target)
         self._state = target
 
+    def _check_tx_not_terminal(self):
+        if QueryTxStateHelper.terminal(self._state):
+            raise RuntimeError(f"Transaction is in terminal state: {self._state.value}")
+
+    def _already_in(self, target: QueryTxStateEnum):
+        return self._state == target
+
 
 def _construct_tx_settings(tx_state):
     tx_settings = _ydb_query.TransactionSettings.from_public(tx_state.tx_mode)
@@ -93,7 +118,7 @@ def _create_rollback_transaction_request(session_state, tx_state):
 def wrap_tx_begin_response(rpc_state, response_pb, session_state, tx_state, tx):
     message = _ydb_query.BeginTransactionResponse.from_proto(response_pb)
     issues._process_response(message.status)
-    tx_state.change_state(base.QueryTxStateEnum.BEGINED)
+    tx_state._change_state(QueryTxStateEnum.BEGINED)
     tx_state.tx_id = message.tx_meta.tx_id
     return tx
 
@@ -104,7 +129,7 @@ def wrap_tx_commit_response(rpc_state, response_pb, session_state, tx_state, tx)
     message = _ydb_query.CommitTransactionResponse(response_pb)
     issues._process_response(message.status)
     tx_state.tx_id = None
-    tx_state.change_state(base.QueryTxStateEnum.COMMITTED)
+    tx_state._change_state(QueryTxStateEnum.COMMITTED)
     return tx
 
 @base.bad_session_handler
@@ -113,7 +138,7 @@ def wrap_tx_rollback_response(rpc_state, response_pb, session_state, tx_state, t
     message = _ydb_query.RollbackTransactionResponse(response_pb)
     issues._process_response(message.status)
     tx_state.tx_id = None
-    tx_state.change_state(base.QueryTxStateEnum.ROLLBACKED)
+    tx_state._change_state(QueryTxStateEnum.ROLLBACKED)
     return tx
 
 
@@ -196,7 +221,7 @@ def begin(self, settings=None):
 
         :return: An open transaction
         """
-        self._tx_state.check_invalid_transition(base.QueryTxStateEnum.BEGINED)
+        self._tx_state._check_invalid_transition(QueryTxStateEnum.BEGINED)
 
         return self._driver(
             _create_begin_transaction_request(self._session_state, self._tx_state),
@@ -216,8 +241,9 @@ def commit(self, settings=None):
 
         :return: A committed transaction or exception if commit is failed
         """
-
-        self._tx_state.check_invalid_transition(base.QueryTxStateEnum.COMMITTED)
+        if self._tx_state._already_in(QueryTxStateEnum.COMMITTED):
+            return
+        self._tx_state._check_invalid_transition(QueryTxStateEnum.COMMITTED)
 
         return self._driver(
             _create_commit_transaction_request(self._session_state, self._tx_state),
@@ -229,7 +255,10 @@ def commit(self, settings=None):
         )
 
     def rollback(self, settings=None):
-        self._tx_state.check_invalid_transition(base.QueryTxStateEnum.ROLLBACKED)
+        if self._tx_state._already_in(QueryTxStateEnum.ROLLBACKED):
+            return
+
+        self._tx_state._check_invalid_transition(QueryTxStateEnum.ROLLBACKED)
 
         return self._driver(
             _create_rollback_transaction_request(self._session_state, self._tx_state),
@@ -240,5 +269,24 @@ def rollback(self, settings=None):
             (self._session_state, self._tx_state, self),
         )
 
+    def _execute_call(self, query: str, commit_tx: bool):
+        request = base.create_execute_query_request(
+            query=query,
+            session_id=self._session_state.session_id,
+            commit_tx=commit_tx
+        )
+        return self._driver(
+            request,
+            _apis.QueryService.Stub,
+            _apis.QueryService.ExecuteQuery,
+        )
+
     def execute(self, query, parameters=None, commit_tx=False, settings=None):
-        pass
+        self._tx_state._check_tx_not_terminal()
+
+        stream_it = self._execute_call(query, commit_tx)
+
+        return _utilities.SyncResponseIterator(
+            stream_it,
+            lambda resp: base.wrap_execute_query_response(rpc_state=None, response_pb=resp),
+        )