Add missing ability to configure QueryClientSettings

Author: vgvoleg

examples/topic_transactions/example.py

@@ -0,0 +1,95 @@
+import ydb
+
+
+def writer_example(driver: ydb.Driver, topic: str):
+    session_pool = ydb.QuerySessionPool(driver)
+
+    def callee(session: ydb.QuerySession):
+        tx = session.transaction()
+        tx.begin()
+
+        tx_writer: ydb.TopicTxWriter = driver.topic_client.tx_writer(topic, tx=tx)
+        # tx_writer: ydb.TopicWriter = driver.topic_client.writer(topic)
+
+        def prepare_messages(result_sets):
+            # some messages based on query result
+            return ["1"]
+
+        with tx.execute(query="select 1") as result_sets:
+            messages = prepare_messages(result_sets)
+
+        tx_writer.write(messages) # wait only on commit
+        # tx_writer.flush() -> maybe
+
+        tx.commit() # <----------------------
+
+    session_pool.retry_operation_sync(callee)
+
+
+# =============================================================================
+
+
+def reader_example(driver: ydb.Driver, reader: ydb.TopicReader):
+    session_pool = ydb.QuerySessionPool(driver)
+
+    def callee(session: ydb.QuerySession):
+        tx = session.transaction().begin()
+
+        batch = reader.receive_batch_with_tx(tx=tx)
+        # this batch will be commited on tx commit
+        # important to commit in correct order
+
+        def prepare_query(batch):
+            # some query based on batch
+            return "select 1"
+
+        with tx.execute(
+            query=prepare_query(batch),
+            commit_tx=True,
+        ) as result_sets:
+            # handle results
+            pass
+
+    session_pool.retry_operation_sync(callee)
+
+# =============================================================================
+
+def writer_example_update(driver: ydb.Driver, topic: str):
+    session_pool = ydb.QuerySessionPool(driver)
+
+    def callee(tx: ydb.QueryTxContext):
+        tx_writer: ydb.TopicWriter = driver.topic_client.tx_writer(topic, tx=tx)
+
+        def prepare_messages(result_sets):
+            # some messages based on query result
+            return ["1"]
+
+        with tx.execute(query="select 1") as result_sets:
+            messages = prepare_messages(result_sets)
+
+        tx_writer.write(messages)
+
+    session_pool.retry_operation_tx(callee)
+
+
+# =============================================================================
+
+
+def reader_example_update(driver: ydb.Driver, reader: ydb.TopicReader):
+    session_pool = ydb.QuerySessionPool(driver)
+
+    def callee(tx: ydb.QueryTxContext):
+        batch = reader.receive_batch_with_tx(tx=tx)
+        # this batch will be commited on tx commit
+
+        def prepare_query(batch):
+            # some query based on patch
+            return "select 1"
+
+        with tx.execute(
+            query=prepare_query(batch)
+        ) as result_sets:
+            # handle results
+            pass
+
+    session_pool.retry_operation_tx(callee)

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(),