query session pool tests

Author: vgvoleg

tests/query/conftest.py

@@ -1,5 +1,6 @@
 import pytest
 from ydb.query.session import QuerySessionSync
+from ydb.query.pool import QuerySessionPool
 
 
 @pytest.fixture
@@ -25,3 +26,9 @@ def tx(session):
         transaction.rollback()
     except BaseException:
         pass
+
+
+@pytest.fixture
+def pool(driver_sync):
+    pool = QuerySessionPool(driver_sync)
+    yield pool

tests/query/test_query_session_pool.py

@@ -0,0 +1,46 @@
+import pytest
+
+from ydb.query.pool import QuerySessionPool
+from ydb.query.session import QuerySessionSync, QuerySessionStateEnum
+
+class TestQuerySessionPool:
+    def test_checkout_provides_created_session(self, pool: QuerySessionPool):
+        with pool.checkout() as session:
+            assert session._state._state == QuerySessionStateEnum.CREATED
+
+        assert session._state._state == QuerySessionStateEnum.CLOSED
+
+    def test_oneshot_query_normal(self, pool: QuerySessionPool):
+        res = pool.execute_with_retries("select 1;")
+        assert len(res) == 1
+
+    def test_oneshot_ddl_query(self, pool: QuerySessionPool):
+        pool.execute_with_retries("create table Queen(key UInt64, PRIMARY KEY (key));")
+        pool.execute_with_retries("drop table Queen;")
+
+    def test_oneshot_query_raises(self, pool: QuerySessionPool):
+        with pytest.raises(Exception):
+            pool.execute_with_retries("Is this the real life? Is this just fantasy?")
+
+    def test_retry_op_uses_created_session(self, pool: QuerySessionPool):
+        def callee(session: QuerySessionSync):
+            assert session._state._state == QuerySessionStateEnum.CREATED
+        pool.retry_operation_sync(callee)
+
+    def test_retry_op_normal(self, pool: QuerySessionPool):
+        def callee(session: QuerySessionSync):
+            with session.transaction() as tx:
+                iterator = tx.execute("select 1;", commit_tx=True)
+                return [result_set for result_set in iterator]
+
+        res = pool.retry_operation_sync(callee)
+        assert len(res) == 1
+
+    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:
+                pass
+
+        with pytest.raises(Exception):
+            pool.retry_operation_sync(callee)

ydb/query/transaction.py

@@ -316,6 +316,7 @@ def _move_to_commited(self):
         if self._tx_state._already_in(QueryTxStateEnum.COMMITTED):
             return
         self._tx_state._change_state(QueryTxStateEnum.COMMITTED)
+        self._tx_state.tx_id = None
 
     def execute(
         self,