fix async transactions

Author: vgvoleg

tests/aio/query/test_query_transaction.py

@@ -89,6 +89,6 @@ 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]
+            res = [result_set async for result_set in results]
 
         assert len(res) == 1

ydb/aio/query/session.py

@@ -26,9 +26,14 @@ class QuerySessionAsync(BaseQuerySession):
     _loop: asyncio.AbstractEventLoop
     _status_stream: _utilities.AsyncResponseIterator = None
 
-    def __init__(self, driver: base.SupportedDriverType, settings: Optional[base.QueryClientSettings] = None):
+    def __init__(
+        self,
+        driver: base.SupportedDriverType,
+        settings: Optional[base.QueryClientSettings] = None,
+        loop: asyncio.AbstractEventLoop = None
+    ):
         super(QuerySessionAsync, self).__init__(driver, settings)
-        self._loop = asyncio.get_running_loop()
+        self._loop = loop if loop is not None else asyncio.get_running_loop()
 
     async def _attach(self) -> None:
         self._stream = await self._attach_call()
@@ -94,6 +99,7 @@ def transaction(self, tx_mode=None) -> base.IQueryTxContext:
             self._state,
             self,
             tx_mode,
+            self._loop,
         )
 
     async def execute(

ydb/aio/query/transaction.py

@@ -1,3 +1,4 @@
+import asyncio
 import logging
 from typing import (
     Optional,
@@ -15,6 +16,36 @@
 
 
 class QueryTxContextAsync(BaseQueryTxContext):
+    _loop: asyncio.AbstractEventLoop
+
+    def __init__(self, driver, session_state, session, tx_mode, loop):
+        """
+        An object that provides a simple transaction context manager that allows statements execution
+        in a transaction. You don't have to open transaction explicitly, because context manager encapsulates
+        transaction control logic, and opens new transaction if:
+
+        1) By explicit .begin() method;
+        2) On execution of a first statement, which is strictly recommended method, because that avoids useless round trip
+
+        This context manager is not thread-safe, so you should not manipulate on it concurrently.
+
+        :param driver: A driver instance
+        :param session_state: A state of session
+        :param tx_mode: Transaction mode, which is a one from the following choises:
+         1) QuerySerializableReadWrite() which is default mode;
+         2) QueryOnlineReadOnly(allow_inconsistent_reads=False);
+         3) QuerySnapshotReadOnly();
+         4) QueryStaleReadOnly().
+        """
+
+        super(QueryTxContextAsync, self).__init__(
+            driver,
+            session_state,
+            session,
+            tx_mode,
+        )
+        self._loop = loop
+
     async def __aenter__(self) -> "QueryTxContextAsync":
         """
         Enters a context manager and returns a transaction
@@ -40,6 +71,12 @@ async def __aexit__(self, *args, **kwargs):
             except issues.Error:
                 logger.warning("Failed to rollback leaked tx: %s", self._tx_state.tx_id)
 
+    async def _ensure_prev_stream_finished(self) -> None:
+        if self._prev_stream is not None:
+            async for _ in self._prev_stream:
+                pass
+            self._prev_stream = None
+
     async def begin(self, settings: Optional[base.QueryClientSettings] = None) -> None:
         """WARNING: This API is experimental and could be changed.
 
@@ -61,9 +98,27 @@ async def commit(self, settings: Optional[base.QueryClientSettings] = None) -> N
 
         :return: A committed transaction or exception if commit is failed
         """
+        if self._tx_state._already_in(QueryTxStateEnum.COMMITTED):
+            return
+
+        if self._tx_state._state == QueryTxStateEnum.NOT_INITIALIZED:
+            self._tx_state._change_state(QueryTxStateEnum.COMMITTED)
+            return
+
+        await self._ensure_prev_stream_finished()
+
         await self._commit_call(settings)
 
     async def rollback(self, settings: Optional[base.QueryClientSettings] = None) -> None:
+        if self._tx_state._already_in(QueryTxStateEnum.ROLLBACKED):
+            return
+
+        if self._tx_state._state == QueryTxStateEnum.NOT_INITIALIZED:
+            self._tx_state._change_state(QueryTxStateEnum.ROLLBACKED)
+            return
+
+        await self._ensure_prev_stream_finished()
+
         await self._rollback_call(settings)
 
     async def execute(
@@ -93,6 +148,8 @@ async def execute(
 
         :return: Iterator with result sets
         """
+        await self._ensure_prev_stream_finished()
+
         stream_it = await self._execute_call(
             query=query,
             commit_tx=commit_tx,

ydb/query/transaction.py

@@ -252,14 +252,6 @@ def _begin_call(self, settings: Optional[base.QueryClientSettings]) -> "BaseQuer
         )
 
     def _commit_call(self, settings: Optional[base.QueryClientSettings]) -> "BaseQueryTxContext":
-        if self._tx_state._already_in(QueryTxStateEnum.COMMITTED):
-            return
-        self._ensure_prev_stream_finished()
-
-        if self._tx_state._state == QueryTxStateEnum.NOT_INITIALIZED:
-            self._tx_state._change_state(QueryTxStateEnum.COMMITTED)
-            return
-
         self._tx_state._check_invalid_transition(QueryTxStateEnum.COMMITTED)
 
         return self._driver(
@@ -272,15 +264,6 @@ def _commit_call(self, settings: Optional[base.QueryClientSettings]) -> "BaseQue
         )
 
     def _rollback_call(self, settings: Optional[base.QueryClientSettings]) -> "BaseQueryTxContext":
-        if self._tx_state._already_in(QueryTxStateEnum.ROLLBACKED):
-            return
-
-        self._ensure_prev_stream_finished()
-
-        if self._tx_state._state == QueryTxStateEnum.NOT_INITIALIZED:
-            self._tx_state._change_state(QueryTxStateEnum.ROLLBACKED)
-            return
-
         self._tx_state._check_invalid_transition(QueryTxStateEnum.ROLLBACKED)
 
         return self._driver(
@@ -301,7 +284,6 @@ def _execute_call(
         parameters: dict = None,
         concurrent_result_sets: bool = False,
     ) -> Iterable[_apis.ydb_query.ExecuteQueryResponsePart]:
-        self._ensure_prev_stream_finished()
         self._tx_state._check_tx_ready_to_use()
 
         request = base.create_execute_query_request(
@@ -362,9 +344,27 @@ def commit(self, settings: Optional[base.QueryClientSettings] = None) -> None:
 
         :return: A committed transaction or exception if commit is failed
         """
+        if self._tx_state._already_in(QueryTxStateEnum.COMMITTED):
+            return
+
+        if self._tx_state._state == QueryTxStateEnum.NOT_INITIALIZED:
+            self._tx_state._change_state(QueryTxStateEnum.COMMITTED)
+            return
+
+        self._ensure_prev_stream_finished()
+
         self._commit_call(settings)
 
     def rollback(self, settings: Optional[base.QueryClientSettings] = None) -> None:
+        if self._tx_state._already_in(QueryTxStateEnum.ROLLBACKED):
+            return
+
+        if self._tx_state._state == QueryTxStateEnum.NOT_INITIALIZED:
+            self._tx_state._change_state(QueryTxStateEnum.ROLLBACKED)
+            return
+
+        self._ensure_prev_stream_finished()
+
         self._rollback_call(settings)
 
     def execute(
@@ -396,6 +396,8 @@ def execute(
 
         :return: Iterator with result sets
         """
+        self._ensure_prev_stream_finished()
+
         stream_it = self._execute_call(
             query=query,
             commit_tx=commit_tx,