fix review comments

Author: vgvoleg

tests/query/test_query_session_pool.py

@@ -1,5 +1,5 @@
 import pytest
-
+import ydb
 from ydb.query.pool import QuerySessionPool
 from ydb.query.session import QuerySessionSync, QuerySessionStateEnum
 
@@ -20,7 +20,7 @@ def test_oneshot_ddl_query(self, pool: QuerySessionPool):
         pool.execute_with_retries("drop table Queen;")
 
     def test_oneshot_query_raises(self, pool: QuerySessionPool):
-        with pytest.raises(Exception):
+        with pytest.raises(ydb.GenericError):
             pool.execute_with_retries("Is this the real life? Is this just fantasy?")
 
     def test_retry_op_uses_created_session(self, pool: QuerySessionPool):
@@ -40,9 +40,8 @@ def callee(session: QuerySessionSync):
 
     def test_retry_op_raises(self, pool: QuerySessionPool):
         def callee(session: QuerySessionSync):
-            res = session.execute("Caught in a landslide, no escape from reality")
-            for _ in res:
+            with session.execute("Caught in a landslide, no escape from reality"):
                 pass
 
-        with pytest.raises(Exception):
+        with pytest.raises(ydb.GenericError):
             pool.retry_operation_sync(callee)

tests/query/test_query_transaction.py

@@ -21,10 +21,12 @@ def test_tx_allow_double_rollback(self, tx: BaseTxContext):
         tx.rollback()
         tx.rollback()
 
+    @pytest.mark.skip(reason="Not sure should we keep this behavior or not")
     def test_tx_commit_raises_before_begin(self, tx: BaseTxContext):
         with pytest.raises(RuntimeError):
             tx.commit()
 
+    @pytest.mark.skip(reason="Not sure should we keep this behavior or not")
     def test_tx_rollback_raises_before_begin(self, tx: BaseTxContext):
         with pytest.raises(RuntimeError):
             tx.rollback()

ydb/query/base.py

@@ -328,5 +328,5 @@ def __enter__(self):
         return self
 
     def __exit__(self, exc_type, exc_val, exc_tb):
-        for _ in self.it:
+        for _ in self:
             pass

ydb/query/transaction.py

@@ -25,7 +25,12 @@ class QueryTxStateEnum(enum.Enum):
 
 class QueryTxStateHelper(abc.ABC):
     _VALID_TRANSITIONS = {
-        QueryTxStateEnum.NOT_INITIALIZED: [QueryTxStateEnum.BEGINED, QueryTxStateEnum.DEAD],
+        QueryTxStateEnum.NOT_INITIALIZED: [
+            QueryTxStateEnum.BEGINED,
+            QueryTxStateEnum.DEAD,
+            QueryTxStateEnum.COMMITTED,
+            QueryTxStateEnum.ROLLBACKED,
+        ],
         QueryTxStateEnum.BEGINED: [QueryTxStateEnum.COMMITTED, QueryTxStateEnum.ROLLBACKED, QueryTxStateEnum.DEAD],
         QueryTxStateEnum.COMMITTED: [],
         QueryTxStateEnum.ROLLBACKED: [],
@@ -246,6 +251,13 @@ def commit(self, settings=None):
         if self._tx_state._already_in(QueryTxStateEnum.COMMITTED):
             return
         self._ensure_prev_stream_finished()
+
+        if self._tx_state._state == QueryTxStateEnum.NOT_INITIALIZED:
+            # TODO(vgvoleg): Discuss should we raise before begin or not
+            self._tx_state._change_state(QueryTxStateEnum.COMMITTED)
+            self._tx_state.tx_id = None
+            return
+
         self._tx_state._check_invalid_transition(QueryTxStateEnum.COMMITTED)
 
         return self._driver(
@@ -262,6 +274,13 @@ def rollback(self, settings=None):
             return
 
         self._ensure_prev_stream_finished()
+
+        if self._tx_state._state == QueryTxStateEnum.NOT_INITIALIZED:
+            # TODO(vgvoleg): Discuss should we raise before begin or not
+            self._tx_state._change_state(QueryTxStateEnum.ROLLBACKED)
+            self._tx_state.tx_id = None
+            return
+
         self._tx_state._check_invalid_transition(QueryTxStateEnum.ROLLBACKED)
 
         return self._driver(