move wait logic to separate function

Author: vgvoleg

ydb/_utilities.py

@@ -182,3 +182,21 @@ def inc_and_get(self) -> int:
         with self._lock:
             self._value += 1
             return self._value
+
+
+def get_first_message_with_timeout(status_stream: SyncResponseIterator, timeout: int):
+    waiter = future()
+
+    def get_first_response(waiter):
+        first_response = next(status_stream)
+        waiter.set_result(first_response)
+
+    thread = threading.Thread(
+        target=get_first_response,
+        args=(waiter,),
+        name="first response attach stream thread",
+        daemon=True,
+    )
+    thread.start()
+
+    return waiter.result(timeout=timeout)

ydb/aio/_utilities.py

@@ -1,3 +1,6 @@
+import asyncio
+
+
 class AsyncResponseIterator(object):
     def __init__(self, it, wrapper):
         self.it = it.__aiter__()
@@ -21,3 +24,10 @@ async def next(self):
 
     async def __anext__(self):
         return await self._next()
+
+
+async def get_first_message_with_timeout(stream: AsyncResponseIterator, timeout: int):
+    async def get_first_response():
+        return await stream.next()
+
+    return await asyncio.wait_for(get_first_response(), timeout)

ydb/aio/query/session.py

@@ -44,13 +44,13 @@ async def _attach(self) -> None:
             lambda response: common_utils.ServerStatus.from_proto(response),
         )
 
-        async def get_first_response():
-            first_response = await self._status_stream.next()
+        try:
+            first_response = await _utilities.get_first_message_with_timeout(
+                self._status_stream,
+                DEFAULT_ATTACH_FIRST_RESP_TIMEOUT,
+            )
             if first_response.status != issues.StatusCode.SUCCESS:
                 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()

ydb/query/session.py

@@ -230,24 +230,13 @@ def _attach(self, first_resp_timeout: int = DEFAULT_ATTACH_FIRST_RESP_TIMEOUT) -
             lambda response: common_utils.ServerStatus.from_proto(response),
         )
 
-        waiter = _utilities.future()
-
-        def get_first_response(waiter):
-            first_response = next(status_stream)
+        try:
+            first_response = _utilities.get_first_message_with_timeout(
+                status_stream,
+                first_resp_timeout,
+            )
             if first_response.status != issues.StatusCode.SUCCESS:
                 raise RuntimeError("Failed to attach session")
-            waiter.set_result(True)
-
-        thread = threading.Thread(
-            target=get_first_response,
-            args=(waiter,),
-            name="first response attach stream thread",
-            daemon=True,
-        )
-        thread.start()
-
-        try:
-            waiter.result(timeout=first_resp_timeout)
         except Exception as e:
             self._state.reset()
             status_stream.cancel()