session attach tests

Author: vgvoleg

tests/query/test_query_session.py

@@ -1,4 +1,9 @@
 import pytest
+import threading
+import time
+from concurrent.futures import _base as b
+from unittest import mock
+
 
 from ydb.query.session import QuerySession
 
@@ -100,3 +105,38 @@ def test_two_results(self, session: QuerySession):
                     res.append(list(result_set.rows[0].values()))
 
         assert res == [[1], [2]]
+
+    def test_thread_leaks(self, session: QuerySession):
+        session.create()
+        thread_names = [t.name for t in threading.enumerate()]
+        assert "first response attach stream thread" not in thread_names
+        assert "attach stream thread" in thread_names
+
+    def test_first_resp_timeout(self, session: QuerySession):
+        class FakeStream:
+            def __iter__(self):
+                return self
+
+            def __next__(self):
+                time.sleep(10)
+                return 1
+
+            def cancel(self):
+                pass
+
+        fake_stream = mock.Mock(spec=FakeStream)
+
+        session._attach_call = mock.MagicMock(return_value=fake_stream)
+        assert session._attach_call() == fake_stream
+
+        session._create_call()
+        with pytest.raises(b.TimeoutError):
+            session._attach(0.1)
+
+        fake_stream.cancel.assert_called()
+
+        thread_names = [t.name for t in threading.enumerate()]
+        assert "first response attach stream thread" not in thread_names
+        assert "attach stream thread" not in thread_names
+
+        _check_session_state_empty(session)

ydb/aio/query/session.py

@@ -47,12 +47,12 @@ async def _attach(self) -> None:
         async def get_first_response():
             first_response = await self._status_stream.next()
             if first_response.status != issues.StatusCode.SUCCESS:
-                self._state.reset()
                 raise RuntimeError("Failed to attach session")
 
         try:
             await asyncio.wait_for(get_first_response(), DEFAULT_ATTACH_FIRST_RESP_TIMEOUT)
         except Exception as e:
+            self._state.reset()
             self._status_stream.cancel()
             raise e
 

ydb/query/session.py

@@ -23,6 +23,7 @@
 
 
 DEFAULT_ATTACH_FIRST_RESP_TIMEOUT = 600
+DEFAULT_ATTACH_LONG_TIMEOUT = 31536000  # year
 
 
 class QuerySessionStateEnum(enum.Enum):
@@ -141,9 +142,9 @@ def __init__(self, driver: common_utils.SupportedDriverType, settings: Optional[
         self._state = QuerySessionState(settings)
         self._attach_settings: BaseRequestSettings = (
             BaseRequestSettings()
-            .with_operation_timeout(31536000)  # year
-            .with_cancel_after(31536000)  # year
-            .with_timeout(31536000)  # year
+            .with_operation_timeout(DEFAULT_ATTACH_LONG_TIMEOUT)
+            .with_cancel_after(DEFAULT_ATTACH_LONG_TIMEOUT)
+            .with_timeout(DEFAULT_ATTACH_LONG_TIMEOUT)
         )
 
     def _get_client_settings(
@@ -222,7 +223,7 @@ class QuerySession(BaseQuerySession):
 
     _stream = None
 
-    def _attach(self) -> None:
+    def _attach(self, first_resp_timeout: int = DEFAULT_ATTACH_FIRST_RESP_TIMEOUT) -> None:
         self._stream = self._attach_call()
         status_stream = _utilities.SyncResponseIterator(
             self._stream,
@@ -234,7 +235,6 @@ def _attach(self) -> None:
         def get_first_response(waiter):
             first_response = next(status_stream)
             if first_response.status != issues.StatusCode.SUCCESS:
-                self._state.reset()
                 raise RuntimeError("Failed to attach session")
             waiter.set_result(True)
 
@@ -247,8 +247,9 @@ def get_first_response(waiter):
         thread.start()
 
         try:
-            waiter.result(timeout=DEFAULT_ATTACH_FIRST_RESP_TIMEOUT)
+            waiter.result(timeout=first_resp_timeout)
         except Exception as e:
+            self._state.reset()
             status_stream.cancel()
             raise e
 
@@ -258,7 +259,7 @@ def get_first_response(waiter):
         threading.Thread(
             target=self._check_session_status_loop,
             args=(status_stream,),
-            name="check session status thread",
+            name="attach stream thread",
             daemon=True,
         ).start()