support session balancer feature

Author: gridnevvvit

CHANGELOG.md

@@ -2,6 +2,7 @@
 
 * allow to refer endpoints by node id
 * support null type in queries
+* support session balancer feature
 
 ## 2.1.0 ##
 

ydb/connection.py

@@ -146,6 +146,7 @@ def _construct_metadata(driver_config, settings):
             metadata.append((YDB_TRACE_ID_HEADER, settings.trace_id))
         if settings.request_type is not None:
             metadata.append((YDB_REQUEST_TYPE_HEADER, settings.request_type))
+        metadata.extend(getattr(settings, "headers", []))
 
     metadata.append(_utilities.x_ydb_sdk_build_info_header())
     return metadata

ydb/settings.py

@@ -10,6 +10,7 @@ class BaseRequestSettings(object):
         "operation_timeout",
         "tracer",
         "compression",
+        "headers",
     )
 
     def __init__(self):
@@ -22,11 +23,38 @@ def __init__(self):
         self.cancel_after = None
         self.operation_timeout = None
         self.compression = None
+        self.headers = []
+
+    def make_copy(self):
+        return (
+            BaseRequestSettings()
+            .with_trace_id(self.trace_id)
+            .with_request_type(self.request_type)
+            .with_timeout(self.timeout)
+            .with_cancel_after(self.cancel_after)
+            .with_operation_timeout(self.operation_timeout)
+            .with_compression(self.compression)
+        )
 
     def with_compression(self, compression):
+        """
+        Enables compression for the specific RPC
+        :param compression: An RPCCompression enum value.
+        :return The self instance.
+        """
         self.compression = compression
         return self
 
+    def with_header(self, key, value):
+        """
+        Adds a key-value pair to the request headers.
+        :param key: A string with a header key.
+        :param value: A string with a header value.
+        :return The self instance.
+        """
+        self.headers.append((key, value))
+        return self
+
     def with_trace_id(self, trace_id):
         """
         Includes trace id for RPC headers

ydb/table.py

@@ -1602,12 +1602,18 @@ def keep_alive(self, settings=None):
     def create(self, settings=None):
         if self._state.session_id is not None:
             return self
+        create_settings = settings_impl.BaseRequestSettings()
+        if settings is not None:
+            create_settings = settings.make_copy()
+        create_settings = create_settings.with_header(
+            "x-ydb-client-capabilities", "session-balancer"
+        )
         return self._driver(
             _apis.ydb_table.CreateSessionRequest(),
             _apis.TableService.Stub,
             _apis.TableService.CreateSession,
             _session_impl.initialize_session,
-            settings,
+            create_settings,
             (self._state, self),
             self._state.endpoint,
         )
@@ -1860,12 +1866,18 @@ def async_keep_alive(self, settings=None):
     def async_create(self, settings=None):
         if self._state.session_id is not None:
             return _utilities.wrap_result_in_future(self)
+        create_settings = settings_impl.BaseRequestSettings()
+        if settings is not None:
+            create_settings = settings.make_copy()
+        create_settings = create_settings.with_header(
+            "x-ydb-client-capabilities", "session-balancer"
+        )
         return self._driver.future(
             _apis.ydb_table.CreateSessionRequest(),
             _apis.TableService.Stub,
             _apis.TableService.CreateSession,
             _session_impl.initialize_session,
-            settings,
+            create_settings,
             (self._state, self),
             self._state.endpoint,
         )