async transactions

Author: vgvoleg

tests/aio/query/conftest.py

@@ -10,4 +10,17 @@ async def session(driver):
     try:
         await session.delete()
     except BaseException:
-        pass
+        pass
+
+
+@pytest.fixture
+async def tx(session):
+    await session.create()
+    transaction = session.transaction()
+
+    yield transaction
+
+    try:
+        await transaction.rollback()
+    except BaseException:
+        pass

tests/aio/query/test_query_session.py

@@ -58,21 +58,23 @@ async def test_create_after_delete_not_possible(self, session: QuerySessionAsync
         with pytest.raises(RuntimeError):
             await session.create()
 
-    # def test_transaction_before_create_raises(self, session: QuerySessionAsync):
-    #     with pytest.raises(RuntimeError):
-    #         session.transaction()
+    def test_transaction_before_create_raises(self, session: QuerySessionAsync):
+        with pytest.raises(RuntimeError):
+            session.transaction()
 
-    # def test_transaction_after_delete_raises(self, session: QuerySessionAsync):
-    #     session.create()
+    @pytest.mark.asyncio
+    async def test_transaction_after_delete_raises(self, session: QuerySessionAsync):
+        await session.create()
 
-    #     session.delete()
+        await session.delete()
 
-    #     with pytest.raises(RuntimeError):
-    #         session.transaction()
+        with pytest.raises(RuntimeError):
+            session.transaction()
 
-    # def test_transaction_after_create_not_raises(self, session: QuerySessionAsync):
-    #     session.create()
-    #     session.transaction()
+    @pytest.mark.asyncio
+    async def test_transaction_after_create_not_raises(self, session: QuerySessionAsync):
+        await session.create()
+        session.transaction()
 
     @pytest.mark.asyncio
     async def test_execute_before_create_raises(self, session: QuerySessionAsync):

tests/aio/query/test_query_transaction.py

@@ -0,0 +1,94 @@
+import pytest
+
+from ydb.aio.query.transaction import QueryTxContextAsync
+from ydb.query.transaction import QueryTxStateEnum
+
+
+class TestAsyncQueryTransaction:
+    @pytest.mark.asyncio
+    async def test_tx_begin(self, tx: QueryTxContextAsync):
+        assert tx.tx_id is None
+
+        await tx.begin()
+        assert tx.tx_id is not None
+
+    @pytest.mark.asyncio
+    async def test_tx_allow_double_commit(self, tx: QueryTxContextAsync):
+        await tx.begin()
+        await tx.commit()
+        await tx.commit()
+
+    @pytest.mark.asyncio
+    async def test_tx_allow_double_rollback(self, tx: QueryTxContextAsync):
+        await tx.begin()
+        await tx.rollback()
+        await tx.rollback()
+
+    @pytest.mark.asyncio
+    async def test_tx_commit_before_begin(self, tx: QueryTxContextAsync):
+        await tx.commit()
+        assert tx._tx_state._state == QueryTxStateEnum.COMMITTED
+
+    @pytest.mark.asyncio
+    async def test_tx_rollback_before_begin(self, tx: QueryTxContextAsync):
+        await tx.rollback()
+        assert tx._tx_state._state == QueryTxStateEnum.ROLLBACKED
+
+    @pytest.mark.asyncio
+    async def test_tx_first_execute_begins_tx(self, tx: QueryTxContextAsync):
+        await tx.execute("select 1;")
+        await tx.commit()
+
+    @pytest.mark.asyncio
+    async def test_interactive_tx_commit(self, tx: QueryTxContextAsync):
+        await tx.execute("select 1;", commit_tx=True)
+        with pytest.raises(RuntimeError):
+            await tx.execute("select 1;")
+
+    @pytest.mark.asyncio
+    async def test_tx_execute_raises_after_commit(self, tx: QueryTxContextAsync):
+        await tx.begin()
+        await tx.commit()
+        with pytest.raises(RuntimeError):
+            await tx.execute("select 1;")
+
+    @pytest.mark.asyncio
+    async def test_tx_execute_raises_after_rollback(self, tx: QueryTxContextAsync):
+        await tx.begin()
+        await tx.rollback()
+        with pytest.raises(RuntimeError):
+            await tx.execute("select 1;")
+
+    @pytest.mark.asyncio
+    async def test_context_manager_rollbacks_tx(self, tx: QueryTxContextAsync):
+        async with tx:
+            await tx.begin()
+
+        assert tx._tx_state._state == QueryTxStateEnum.ROLLBACKED
+
+    @pytest.mark.asyncio
+    async def test_context_manager_normal_flow(self, tx: QueryTxContextAsync):
+        async with tx:
+            await tx.begin()
+            await tx.execute("select 1;")
+            await tx.commit()
+
+        assert tx._tx_state._state == QueryTxStateEnum.COMMITTED
+
+    @pytest.mark.asyncio
+    async def test_context_manager_does_not_hide_exceptions(self, tx: QueryTxContextAsync):
+        class CustomException(Exception):
+            pass
+
+        with pytest.raises(CustomException):
+            async with tx:
+                raise CustomException()
+
+    @pytest.mark.asyncio
+    async def test_execute_as_context_manager(self, tx: QueryTxContextAsync):
+        await tx.begin()
+
+        async with await tx.execute("select 1;") as results:
+            res = [result_set for result_set in results]
+
+        assert len(res) == 1

tests/query/test_query_transaction.py

@@ -1,78 +1,78 @@
 import pytest
 
-from ydb.query.transaction import BaseQueryTxContext
+from ydb.query.transaction import QueryTxContextSync
 from ydb.query.transaction import QueryTxStateEnum
 
 
 class TestQueryTransaction:
-    def test_tx_begin(self, tx: BaseQueryTxContext):
+    def test_tx_begin(self, tx: QueryTxContextSync):
         assert tx.tx_id is None
 
         tx.begin()
         assert tx.tx_id is not None
 
-    def test_tx_allow_double_commit(self, tx: BaseQueryTxContext):
+    def test_tx_allow_double_commit(self, tx: QueryTxContextSync):
         tx.begin()
         tx.commit()
         tx.commit()
 
-    def test_tx_allow_double_rollback(self, tx: BaseQueryTxContext):
+    def test_tx_allow_double_rollback(self, tx: QueryTxContextSync):
         tx.begin()
         tx.rollback()
         tx.rollback()
 
-    def test_tx_commit_before_begin(self, tx: BaseQueryTxContext):
+    def test_tx_commit_before_begin(self, tx: QueryTxContextSync):
         tx.commit()
         assert tx._tx_state._state == QueryTxStateEnum.COMMITTED
 
-    def test_tx_rollback_before_begin(self, tx: BaseQueryTxContext):
+    def test_tx_rollback_before_begin(self, tx: QueryTxContextSync):
         tx.rollback()
         assert tx._tx_state._state == QueryTxStateEnum.ROLLBACKED
 
-    def test_tx_first_execute_begins_tx(self, tx: BaseQueryTxContext):
+    def test_tx_first_execute_begins_tx(self, tx: QueryTxContextSync):
         tx.execute("select 1;")
         tx.commit()
 
-    def test_interactive_tx_commit(self, tx: BaseQueryTxContext):
+    def test_interactive_tx_commit(self, tx: QueryTxContextSync):
         tx.execute("select 1;", commit_tx=True)
         with pytest.raises(RuntimeError):
             tx.execute("select 1;")
 
-    def test_tx_execute_raises_after_commit(self, tx: BaseQueryTxContext):
+    def test_tx_execute_raises_after_commit(self, tx: QueryTxContextSync):
         tx.begin()
         tx.commit()
         with pytest.raises(RuntimeError):
             tx.execute("select 1;")
 
-    def test_tx_execute_raises_after_rollback(self, tx: BaseQueryTxContext):
+    def test_tx_execute_raises_after_rollback(self, tx: QueryTxContextSync):
         tx.begin()
         tx.rollback()
         with pytest.raises(RuntimeError):
             tx.execute("select 1;")
 
-    def test_context_manager_rollbacks_tx(self, tx: BaseQueryTxContext):
+    def test_context_manager_rollbacks_tx(self, tx: QueryTxContextSync):
         with tx:
             tx.begin()
 
         assert tx._tx_state._state == QueryTxStateEnum.ROLLBACKED
 
-    def test_context_manager_normal_flow(self, tx: BaseQueryTxContext):
+    def test_context_manager_normal_flow(self, tx: QueryTxContextSync):
         with tx:
             tx.begin()
             tx.execute("select 1;")
             tx.commit()
 
         assert tx._tx_state._state == QueryTxStateEnum.COMMITTED
 
-    def test_context_manager_does_not_hide_exceptions(self, tx: BaseQueryTxContext):
+    def test_context_manager_does_not_hide_exceptions(self, tx: QueryTxContextSync):
         class CustomException(Exception):
             pass
 
         with pytest.raises(CustomException):
             with tx:
                 raise CustomException()
 
-    def test_execute_as_context_manager(self, tx: BaseQueryTxContext):
+    def test_execute_as_context_manager(self, tx: QueryTxContextSync):
         tx.begin()
 
         with tx.execute("select 1;") as results:

ydb/aio/_utilities.py

@@ -7,6 +7,9 @@ def cancel(self):
         self.it.cancel()
         return self
 
+    def __iter__(self):
+        return self
+
     def __aiter__(self):
         return self
 

ydb/aio/query/base.py

@@ -0,0 +1,9 @@
+from .. import _utilities
+
+class AsyncResponseContextIterator(_utilities.AsyncResponseIterator):
+    async def __aenter__(self) -> "AsyncResponseContextIterator":
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        async for _ in self:
+            pass

ydb/aio/query/session.py

@@ -4,25 +4,19 @@
     Optional,
 )
 
+from .base import AsyncResponseContextIterator
+from .transaction import QueryTxContextAsync
 from .. import _utilities
 from ... import issues
 from ..._grpc.grpcwrapper import common_utils
+from ..._grpc.grpcwrapper import ydb_query_public_types as _ydb_query_public
+
 from ...query import base
 from ...query.session import (
     BaseQuerySession,
     QuerySessionStateEnum,
 )
 
-
-class AsyncResponseContextIterator(_utilities.AsyncResponseIterator):
-    async def __aenter__(self) -> "AsyncResponseContextIterator":
-        return self
-
-    async def __aexit__(self, exc_type, exc_val, exc_tb):
-        async for _ in self:
-            pass
-
-
 class QuerySessionAsync(BaseQuerySession):
     """Session object for Query Service. It is not recommended to control
     session's lifecycle manually - use a QuerySessionPool is always a better choise.
@@ -91,14 +85,21 @@ async def create(self) -> "QuerySessionAsync":
 
         return self
 
-    async def transaction(self, tx_mode) -> base.IQueryTxContext:
-        return super().transaction(tx_mode)
+    def transaction(self, tx_mode = None) -> base.IQueryTxContext:
+        self._state._check_session_ready_to_use()
+        tx_mode = tx_mode if tx_mode else _ydb_query_public.QuerySerializableReadWrite()
+
+        return QueryTxContextAsync(
+            self._driver,
+            self._state,
+            self,
+            tx_mode,
+        )
 
     async def execute(
         self,
         query: str,
         parameters: dict = None,
-        commit_tx: bool = False,
         syntax: base.QuerySyntax = None,
         exec_mode: base.QueryExecMode = None,
         concurrent_result_sets: bool = False,