Add missing ability to configure QueryClientSettings

Author: vgvoleg

ydb/aio/query/pool.py

@@ -13,6 +13,7 @@
     RetrySettings,
     retry_operation_async,
 )
+from ...query.base import QueryClientSettings
 from ... import convert
 from ..._grpc.grpcwrapper import common_utils
 
@@ -22,10 +23,16 @@
 class QuerySessionPool:
     """QuerySessionPool is an object to simplify operations with sessions of Query Service."""
 
-    def __init__(self, driver: common_utils.SupportedDriverType, size: int = 100):
+    def __init__(
+        self,
+        driver: common_utils.SupportedDriverType,
+        size: int = 100,
+        query_client_settings: Optional[QueryClientSettings] = None,
+    ):
         """
         :param driver: A driver instance
         :param size: Size of session pool
+        :param query_client_settings: ydb.QueryClientSettings object to configure QueryService behavior
         """
 
         self._driver = driver
@@ -35,9 +42,10 @@ def __init__(self, driver: common_utils.SupportedDriverType, size: int = 100):
         self._current_size = 0
         self._waiters = 0
         self._loop = asyncio.get_running_loop()
+        self._query_client_settings = query_client_settings
 
     async def _create_new_session(self):
-        session = QuerySession(self._driver)
+        session = QuerySession(self._driver, settings=self._query_client_settings)
         await session.create()
         logger.debug(f"New session was created for pool. Session id: {session._state.session_id}")
         return session

ydb/driver.py

@@ -89,6 +89,7 @@ class DriverConfig(object):
         "secure_channel",
         "table_client_settings",
         "topic_client_settings",
+        "query_client_settings",
         "endpoints",
         "primary_user_agent",
         "tracer",
@@ -112,6 +113,7 @@ def __init__(
         grpc_keep_alive_timeout=None,
         table_client_settings=None,
         topic_client_settings=None,
+        query_client_settings=None,
         endpoints=None,
         primary_user_agent="python-library",
         tracer=None,
@@ -159,6 +161,7 @@ def __init__(
         self.grpc_keep_alive_timeout = grpc_keep_alive_timeout
         self.table_client_settings = table_client_settings
         self.topic_client_settings = topic_client_settings
+        self.query_client_settings = query_client_settings
         self.primary_user_agent = primary_user_agent
         self.tracer = tracer if tracer is not None else tracing.Tracer(None)
         self.grpc_lb_policy_name = grpc_lb_policy_name

ydb/query/pool.py

@@ -8,6 +8,7 @@
 import threading
 import queue
 
+from .base import QueryClientSettings
 from .session import (
     QuerySession,
 )
@@ -27,10 +28,16 @@
 class QuerySessionPool:
     """QuerySessionPool is an object to simplify operations with sessions of Query Service."""
 
-    def __init__(self, driver: common_utils.SupportedDriverType, size: int = 100):
+    def __init__(
+        self,
+        driver: common_utils.SupportedDriverType,
+        size: int = 100,
+        query_client_settings: Optional[QueryClientSettings] = None,
+    ):
         """
         :param driver: A driver instance.
         :param size: Max size of Session Pool.
+        :param query_client_settings: ydb.QueryClientSettings object to configure QueryService behavior
         """
 
         self._driver = driver
@@ -39,9 +46,10 @@ def __init__(self, driver: common_utils.SupportedDriverType, size: int = 100):
         self._size = size
         self._should_stop = threading.Event()
         self._lock = threading.RLock()
+        self._query_client_settings = query_client_settings
 
     def _create_new_session(self, timeout: Optional[float]):
-        session = QuerySession(self._driver)
+        session = QuerySession(self._driver, settings=self._query_client_settings)
         session.create(settings=BaseRequestSettings().with_timeout(timeout))
         logger.debug(f"New session was created for pool. Session id: {session._state.session_id}")
         return session

ydb/query/session.py

@@ -134,9 +134,20 @@ class BaseQuerySession:
 
     def __init__(self, driver: common_utils.SupportedDriverType, settings: Optional[base.QueryClientSettings] = None):
         self._driver = driver
-        self._settings = settings if settings is not None else base.QueryClientSettings()
+        self._settings = self._get_client_settings(driver, settings)
         self._state = QuerySessionState(settings)
 
+    def _get_client_settings(
+        self,
+        driver: common_utils.SupportedDriverType,
+        settings: Optional[base.QueryClientSettings] = None,
+    ) -> base.QueryClientSettings:
+        if settings is not None:
+            return settings
+        if driver._driver_config.query_client_settings is not None:
+            return driver._driver_config.query_client_settings
+        return base.QueryClientSettings()
+
     def _create_call(self, settings: Optional[BaseRequestSettings] = None) -> "BaseQuerySession":
         return self._driver(
             _apis.ydb_query.CreateSessionRequest(),