Query Service AIO

ydb/aio/query/base.py

@@ -6,5 +6,6 @@ async def __aenter__(self) -> "AsyncResponseContextIterator":
         return self
 
     async def __aexit__(self, exc_type, exc_val, exc_tb):
+        #  To close stream on YDB it is necessary to scroll through it to the end
         async for _ in self:
             pass

ydb/aio/query/session.py

@@ -58,7 +58,9 @@ async def _check_session_status_loop(self) -> None:
                     self._state.reset()
                     self._state._change_state(QuerySessionStateEnum.CLOSED)
         except Exception:
-            pass
+            if not self._state._already_in(QuerySessionStateEnum.CLOSED):
+                self._state.reset()
+                self._state._change_state(QuerySessionStateEnum.CLOSED)
 
     async def delete(self) -> None:
         """WARNING: This API is experimental and could be changed.

ydb/aio/query/transaction.py

@@ -28,7 +28,7 @@ async def __aexit__(self, *args, **kwargs):
         Closes a transaction context manager and rollbacks transaction if
         it is not finished explicitly
         """
-        self._ensure_prev_stream_finished()
+        await self._ensure_prev_stream_finished()
         if self._tx_state._state == QueryTxStateEnum.BEGINED:
             # It's strictly recommended to close transactions directly
             # by using commit_tx=True flag while executing statement or by
@@ -42,7 +42,7 @@ async def __aexit__(self, *args, **kwargs):
 
     async def _ensure_prev_stream_finished(self) -> None:
         if self._prev_stream is not None:
-            async for _ in self._prev_stream:
+            async with self._prev_stream:
                 pass
             self._prev_stream = None
 

ydb/query/base.py

@@ -48,6 +48,7 @@ def __enter__(self) -> "SyncResponseContextIterator":
         return self
 
     def __exit__(self, exc_type, exc_val, exc_tb):
+        #  To close stream on YDB it is necessary to scroll through it to the end
         for _ in self:
             pass
 

ydb/query/session.py

@@ -224,7 +224,9 @@ def _check_session_status_loop(self, status_stream: _utilities.SyncResponseItera
                     self._state.reset()
                     self._state._change_state(QuerySessionStateEnum.CLOSED)
         except Exception:
-            pass
+            if not self._state._already_in(QuerySessionStateEnum.CLOSED):
+                self._state.reset()
+                self._state._change_state(QuerySessionStateEnum.CLOSED)
 
     def delete(self) -> None:
         """WARNING: This API is experimental and could be changed.

ydb/query/transaction.py

@@ -196,31 +196,6 @@ def __init__(self, driver, session_state, session, tx_mode):
         self.session = session
         self._prev_stream = None
 
-    def __enter__(self) -> "BaseQueryTxContext":
-        """
-        Enters a context manager and returns a transaction
-
-        :return: A transaction instance
-        """
-        return self
-
-    def __exit__(self, *args, **kwargs):
-        """
-        Closes a transaction context manager and rollbacks transaction if
-        it is not finished explicitly
-        """
-        self._ensure_prev_stream_finished()
-        if self._tx_state._state == QueryTxStateEnum.BEGINED:
-            # It's strictly recommended to close transactions directly
-            # by using commit_tx=True flag while executing statement or by
-            # .commit() or .rollback() methods, but here we trying to do best
-            # effort to avoid useless open transactions
-            logger.warning("Potentially leaked tx: %s", self._tx_state.tx_id)
-            try:
-                self.rollback()
-            except issues.Error:
-                logger.warning("Failed to rollback leaked tx: %s", self._tx_state.tx_id)
-
     @property
     def session_id(self) -> str:
         """
@@ -304,12 +279,6 @@ def _execute_call(
             _apis.QueryService.ExecuteQuery,
         )
 
-    def _ensure_prev_stream_finished(self) -> None:
-        if self._prev_stream is not None:
-            for _ in self._prev_stream:
-                pass
-            self._prev_stream = None
-
     def _move_to_beginned(self, tx_id: str) -> None:
         if self._tx_state._already_in(QueryTxStateEnum.BEGINED):
             return
@@ -323,6 +292,37 @@ def _move_to_commited(self) -> None:
 
 
 class QueryTxContextSync(BaseQueryTxContext):
+    def __enter__(self) -> "BaseQueryTxContext":
+        """
+        Enters a context manager and returns a transaction
+
+        :return: A transaction instance
+        """
+        return self
+
+    def __exit__(self, *args, **kwargs):
+        """
+        Closes a transaction context manager and rollbacks transaction if
+        it is not finished explicitly
+        """
+        self._ensure_prev_stream_finished()
+        if self._tx_state._state == QueryTxStateEnum.BEGINED:
+            # It's strictly recommended to close transactions directly
+            # by using commit_tx=True flag while executing statement or by
+            # .commit() or .rollback() methods, but here we trying to do best
+            # effort to avoid useless open transactions
+            logger.warning("Potentially leaked tx: %s", self._tx_state.tx_id)
+            try:
+                self.rollback()
+            except issues.Error:
+                logger.warning("Failed to rollback leaked tx: %s", self._tx_state.tx_id)
+
+    def _ensure_prev_stream_finished(self) -> None:
+        if self._prev_stream is not None:
+            with self._prev_stream:
+                pass
+            self._prev_stream = None
+
     def begin(self, settings: Optional[base.QueryClientSettings] = None) -> "QueryTxContextSync":
         """WARNING: This API is experimental and could be changed.