interactive tx support

Author: vgvoleg

tests/query/test_query_transaction.py

@@ -26,9 +26,14 @@ 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 test_tx_first_execute_begins_tx(self, tx):
+        tx.execute("select 1;")
+        tx.commit()
+
+    def test_interactive_tx_commit(self, tx):
+        tx.execute("select 1;", commit_tx=True)
+        with pytest.raises(RuntimeError):
+            tx.execute("select 1;")
 
     def text_tx_execute_raises_after_commit(self, tx):
         tx.begin()

ydb/query/base.py

@@ -192,26 +192,23 @@ def create_execute_query_request(
     return req.to_proto()
 
 
-def wrap_execute_query_response(rpc_state, response_pb):
+def wrap_execute_query_response(rpc_state, response_pb, tx, commit_tx=False):
     issues._process_response(response_pb)
+    if response_pb.tx_meta and not tx.tx_id:
+        tx._handle_tx_meta(response_pb.tx_meta)
+    if commit_tx:
+        tx._move_to_commited()
     return convert.ResultSet.from_message(response_pb.result_set)
 
 
 X_YDB_SERVER_HINTS = "x-ydb-server-hints"
 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 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)
             return func(rpc_state, response_pb, session_state, *args, **kwargs)
         except issues.BadSession:
             session_state.reset()

ydb/query/transaction.py

@@ -245,9 +245,8 @@ def commit(self, settings=None):
         """
         if self._tx_state._already_in(QueryTxStateEnum.COMMITTED):
             return
-        self._tx_state._check_invalid_transition(QueryTxStateEnum.COMMITTED)
-
         self._ensure_prev_stream_finished()
+        self._tx_state._check_invalid_transition(QueryTxStateEnum.COMMITTED)
 
         return self._driver(
             _create_commit_transaction_request(self._session_state, self._tx_state),
@@ -262,9 +261,8 @@ def rollback(self, settings=None):
         if self._tx_state._already_in(QueryTxStateEnum.ROLLBACKED):
             return
 
-        self._tx_state._check_invalid_transition(QueryTxStateEnum.ROLLBACKED)
-
         self._ensure_prev_stream_finished()
+        self._tx_state._check_invalid_transition(QueryTxStateEnum.ROLLBACKED)
 
         return self._driver(
             _create_rollback_transaction_request(self._session_state, self._tx_state),
@@ -309,6 +307,16 @@ def _ensure_prev_stream_finished(self):
                 pass
             self._prev_stream = None
 
+    def _handle_tx_meta(self, tx_meta=None):
+        if not self.tx_id:
+            self._tx_state._change_state(QueryTxStateEnum.BEGINED)
+            self._tx_state.tx_id = tx_meta.id
+
+    def _move_to_commited(self):
+        if self._tx_state._already_in(QueryTxStateEnum.COMMITTED):
+            return
+        self._tx_state._change_state(QueryTxStateEnum.COMMITTED)
+
     def execute(
         self,
         query: str,
@@ -319,8 +327,8 @@ def execute(
         parameters: dict = None,
         concurrent_result_sets: bool = False,
     ):
-        self._tx_state._check_tx_not_terminal()
         self._ensure_prev_stream_finished()
+        self._tx_state._check_tx_not_terminal()
 
         stream_it = self._execute_call(
             query=query,
@@ -333,6 +341,11 @@ def execute(
         )
         self._prev_stream = _utilities.SyncResponseIterator(
             stream_it,
-            lambda resp: base.wrap_execute_query_response(rpc_state=None, response_pb=resp),
+            lambda resp: base.wrap_execute_query_response(
+                rpc_state=None,
+                response_pb=resp,
+                tx=self,
+                commit_tx=commit_tx,
+            ),
         )
         return self._prev_stream