query session pool timeout refactor

Author: vgvoleg

ydb/aio/query/session.py

@@ -8,6 +8,7 @@
 from .transaction import QueryTxContextAsync
 from .. import _utilities
 from ... import issues
+from ...settings import BaseRequestSettings
 from ..._grpc.grpcwrapper import common_utils
 from ..._grpc.grpcwrapper import ydb_query_public_types as _ydb_query_public
 
@@ -62,7 +63,7 @@ async def _check_session_status_loop(self) -> None:
                 self._state.reset()
                 self._state._change_state(QuerySessionStateEnum.CLOSED)
 
-    async def delete(self) -> None:
+    async def delete(self, settings: Optional[BaseRequestSettings] = None) -> None:
         """WARNING: This API is experimental and could be changed.
 
         Deletes a Session of Query Service on server side and releases resources.
@@ -73,10 +74,10 @@ async def delete(self) -> None:
             return
 
         self._state._check_invalid_transition(QuerySessionStateEnum.CLOSED)
-        await self._delete_call()
+        await self._delete_call(settings=settings)
         self._stream.cancel()
 
-    async def create(self) -> "QuerySessionAsync":
+    async def create(self, settings: Optional[BaseRequestSettings] = None) -> "QuerySessionAsync":
         """WARNING: This API is experimental and could be changed.
 
         Creates a Session of Query Service on server side and attaches it.
@@ -87,7 +88,7 @@ async def create(self) -> "QuerySessionAsync":
             return
 
         self._state._check_invalid_transition(QuerySessionStateEnum.CREATED)
-        await self._create_call()
+        await self._create_call(settings=settings)
         await self._attach()
 
         return self
@@ -110,6 +111,7 @@ async def execute(
         syntax: base.QuerySyntax = None,
         exec_mode: base.QueryExecMode = None,
         concurrent_result_sets: bool = False,
+        settings: Optional[BaseRequestSettings] = None,
     ) -> AsyncResponseContextIterator:
         """WARNING: This API is experimental and could be changed.
 
@@ -132,6 +134,7 @@ async def execute(
             exec_mode=exec_mode,
             parameters=parameters,
             concurrent_result_sets=concurrent_result_sets,
+            settings=settings,
         )
 
         return AsyncResponseContextIterator(

ydb/query/pool.py

@@ -17,6 +17,7 @@
 )
 from .. import issues
 from .. import convert
+from ..settings import BaseRequestSettings
 from .._grpc.grpcwrapper import common_utils
 
 
@@ -39,9 +40,9 @@ def __init__(self, driver: common_utils.SupportedDriverType, size: int = 100):
         self._should_stop = threading.Event()
         self._lock = threading.RLock()
 
-    def _create_new_session(self):
+    def _create_new_session(self, timeout: float):
         session = QuerySessionSync(self._driver)
-        session.create()
+        session.create(settings=BaseRequestSettings().with_timeout(timeout))
         logger.debug(f"New session was created for pool. Session id: {session._state.session_id}")
         return session
 
@@ -73,12 +74,9 @@ def acquire(self, timeout: float) -> QuerySessionSync:
                     logger.debug(f"Acquired dead session from queue: {session._state.session_id}")
 
             logger.debug(f"Session pool is not large enough: {self._current_size} < {self._size}, will create new one.")
-            session = self._create_new_session()
-
             finish = time.monotonic()
-            if finish - start > timeout:
-                session.delete()
-                raise issues.SessionPoolEmpty("Timeout on acquire session")
+            time_left = timeout - (finish - start)
+            session = self._create_new_session(time_left)
 
             self._current_size += 1
             return session

ydb/query/session.py

@@ -1,6 +1,7 @@
 import abc
 import enum
 import logging
+import time
 import threading
 from typing import (
     Iterable,
@@ -10,6 +11,7 @@
 from . import base
 
 from .. import _apis, issues, _utilities
+from ..settings import BaseRequestSettings
 from ..connection import _RpcState as RpcState
 from .._grpc.grpcwrapper import common_utils
 from .._grpc.grpcwrapper import ydb_query as _ydb_query
@@ -136,29 +138,32 @@ def __init__(self, driver: common_utils.SupportedDriverType, settings: Optional[
         self._settings = settings if settings is not None else base.QueryClientSettings()
         self._state = QuerySessionState(settings)
 
-    def _create_call(self) -> "BaseQuerySession":
+    def _create_call(self, settings: Optional[BaseRequestSettings] = None) -> "BaseQuerySession":
         return self._driver(
             _apis.ydb_query.CreateSessionRequest(),
             _apis.QueryService.Stub,
             _apis.QueryService.CreateSession,
             wrap_result=wrapper_create_session,
             wrap_args=(self._state, self),
+            settings=settings,
         )
 
-    def _delete_call(self) -> "BaseQuerySession":
+    def _delete_call(self, settings: Optional[BaseRequestSettings] = None) -> "BaseQuerySession":
         return self._driver(
             _apis.ydb_query.DeleteSessionRequest(session_id=self._state.session_id),
             _apis.QueryService.Stub,
             _apis.QueryService.DeleteSession,
             wrap_result=wrapper_delete_session,
             wrap_args=(self._state, self),
+            settings=settings,
         )
 
-    def _attach_call(self) -> Iterable[_apis.ydb_query.SessionState]:
+    def _attach_call(self, settings: Optional[BaseRequestSettings] = None) -> Iterable[_apis.ydb_query.SessionState]:
         return self._driver(
             _apis.ydb_query.AttachSessionRequest(session_id=self._state.session_id),
             _apis.QueryService.Stub,
             _apis.QueryService.AttachSession,
+            settings=settings,
         )
 
     def _execute_call(
@@ -169,6 +174,7 @@ def _execute_call(
         exec_mode: base.QueryExecMode = None,
         parameters: dict = None,
         concurrent_result_sets: bool = False,
+        settings: Optional[BaseRequestSettings] = None,
     ) -> Iterable[_apis.ydb_query.ExecuteQueryResponsePart]:
         request = base.create_execute_query_request(
             query=query,
@@ -186,6 +192,7 @@ def _execute_call(
             request.to_proto(),
             _apis.QueryService.Stub,
             _apis.QueryService.ExecuteQuery,
+            settings=settings,
         )
 
 
@@ -196,8 +203,8 @@ class QuerySessionSync(BaseQuerySession):
 
     _stream = None
 
-    def _attach(self) -> None:
-        self._stream = self._attach_call()
+    def _attach(self, settings: Optional[BaseRequestSettings] = None) -> None:
+        self._stream = self._attach_call(settings=settings)
         status_stream = _utilities.SyncResponseIterator(
             self._stream,
             lambda response: common_utils.ServerStatus.from_proto(response),
@@ -228,7 +235,7 @@ def _check_session_status_loop(self, status_stream: _utilities.SyncResponseItera
                 self._state.reset()
                 self._state._change_state(QuerySessionStateEnum.CLOSED)
 
-    def delete(self) -> None:
+    def delete(self, settings: Optional[BaseRequestSettings] = None) -> None:
         """WARNING: This API is experimental and could be changed.
 
         Deletes a Session of Query Service on server side and releases resources.
@@ -239,10 +246,10 @@ def delete(self) -> None:
             return
 
         self._state._check_invalid_transition(QuerySessionStateEnum.CLOSED)
-        self._delete_call()
+        self._delete_call(settings=settings)
         self._stream.cancel()
 
-    def create(self) -> "QuerySessionSync":
+    def create(self, settings: Optional[BaseRequestSettings] = None) -> "QuerySessionSync":
         """WARNING: This API is experimental and could be changed.
 
         Creates a Session of Query Service on server side and attaches it.
@@ -253,7 +260,8 @@ def create(self) -> "QuerySessionSync":
             return
 
         self._state._check_invalid_transition(QuerySessionStateEnum.CREATED)
-        self._create_call()
+
+        self._create_call(settings=settings)
         self._attach()
 
         return self
@@ -289,6 +297,7 @@ def execute(
         syntax: base.QuerySyntax = None,
         exec_mode: base.QueryExecMode = None,
         concurrent_result_sets: bool = False,
+        settings: Optional[BaseRequestSettings] = None,
     ) -> base.SyncResponseContextIterator:
         """WARNING: This API is experimental and could be changed.
 
@@ -311,6 +320,7 @@ def execute(
             exec_mode=exec_mode,
             parameters=parameters,
             concurrent_result_sets=concurrent_result_sets,
+            settings=settings,
         )
 
         return base.SyncResponseContextIterator(