diff --git a/CHANGELOG.md b/CHANGELOG.md index 16975ee4..e39d0ae2 100644 --- a/CHANGELOG.md +++ b/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 ## diff --git a/ydb/connection.py b/ydb/connection.py index a5173672..ed86de44 100644 --- a/ydb/connection.py +++ b/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 diff --git a/ydb/settings.py b/ydb/settings.py index e1e1f0f2..f55b1abb 100644 --- a/ydb/settings.py +++ b/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 diff --git a/ydb/table.py b/ydb/table.py index c49c4660..97b1fba8 100644 --- a/ydb/table.py +++ b/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, )