Refactor session state reset logic

Author: vgvoleg

tests/query/test_query_transaction.py

@@ -26,6 +26,12 @@ def test_tx_commit_before_begin(self, tx: QueryTxContext):
         tx.commit()
         assert tx._tx_state._state == QueryTxStateEnum.COMMITTED
 
+    def test_tx_rollback_after_session_detached(self, tx: QueryTxContext):
+        with tx:
+            tx.begin()
+            tx.session.delete()
+
+
     def test_tx_rollback_before_begin(self, tx: QueryTxContext):
         tx.rollback()
         assert tx._tx_state._state == QueryTxStateEnum.ROLLBACKED

ydb/aio/query/session.py

@@ -57,7 +57,7 @@ async def _attach(self) -> None:
             if first_response.status != issues.StatusCode.SUCCESS:
                 raise RuntimeError("Failed to attach session")
         except Exception as e:
-            self._state.reset()
+            self._state.set_attached(False)
             self._status_stream.cancel()
             raise e
 
@@ -70,11 +70,11 @@ async def _check_session_status_loop(self) -> None:
         try:
             async for status in self._status_stream:
                 if status.status != issues.StatusCode.SUCCESS:
-                    self._state.reset()
+                    self._state.set_attached(False)
                     self._state._change_state(QuerySessionStateEnum.CLOSED)
         except Exception:
             if not self._state._already_in(QuerySessionStateEnum.CLOSED):
-                self._state.reset()
+                self._state.set_attached(False)
                 self._state._change_state(QuerySessionStateEnum.CLOSED)
 
     async def delete(self, settings: Optional[BaseRequestSettings] = None) -> None:

ydb/query/base.py

@@ -111,10 +111,6 @@ class IQuerySessionState(abc.ABC):
     def __init__(self, settings: Optional[QueryClientSettings] = None):
         pass
 
-    @abc.abstractmethod
-    def reset(self) -> None:
-        pass
-
     @property
     @abc.abstractmethod
     def session_id(self) -> Optional[str]:
@@ -210,7 +206,7 @@ def decorator(rpc_state, response_pb, session_state: IQuerySessionState, *args,
         try:
             return func(rpc_state, response_pb, session_state, *args, **kwargs)
         except issues.BadSession:
-            session_state.reset()
+            session_state.set_attached(False)
             raise
 
     return decorator

ydb/query/session.py

@@ -65,11 +65,6 @@ class QuerySessionState(base.IQuerySessionState):
     def __init__(self, settings: base.QueryClientSettings = None):
         self._settings = settings
 
-    def reset(self) -> None:
-        self._session_id = None
-        self._node_id = None
-        self._attached = False
-
     @property
     def session_id(self) -> Optional[str]:
         return self._session_id
@@ -129,7 +124,7 @@ def wrapper_delete_session(
 ) -> "BaseQuerySession":
     message = _ydb_query.DeleteSessionResponse.from_proto(response_pb)
     issues._process_response(message.status)
-    session_state.reset()
+    session_state.set_attached(False)
     session_state._change_state(QuerySessionStateEnum.CLOSED)
     return session
 
@@ -257,7 +252,7 @@ def _attach(self, first_resp_timeout: int = DEFAULT_INITIAL_RESPONSE_TIMEOUT) ->
             if first_response.status != issues.StatusCode.SUCCESS:
                 raise RuntimeError("Failed to attach session")
         except Exception as e:
-            self._state.reset()
+            self._state.set_attached(False)
             status_stream.cancel()
             raise e
 
@@ -275,11 +270,11 @@ def _check_session_status_loop(self, status_stream: _utilities.SyncResponseItera
         try:
             for status in status_stream:
                 if status.status != issues.StatusCode.SUCCESS:
-                    self._state.reset()
+                    self._state.set_attached(False)
                     self._state._change_state(QuerySessionStateEnum.CLOSED)
         except Exception:
             if not self._state._already_in(QuerySessionStateEnum.CLOSED):
-                self._state.reset()
+                self._state.set_attached(False)
                 self._state._change_state(QuerySessionStateEnum.CLOSED)
 
     def delete(self, settings: Optional[BaseRequestSettings] = None) -> None:

ydb/query/transaction.py

@@ -377,6 +377,10 @@ def __exit__(self, *args, **kwargs):
         it is not finished explicitly
         """
         self._ensure_prev_stream_finished()
+
+        if not self._session_state.attached:
+            return
+
         if self._tx_state._state == QueryTxStateEnum.BEGINED and self._external_error is None:
             # It's strictly recommended to close transactions directly
             # by using commit_tx=True flag while executing statement or by