Skip to content

Commit

Permalink
add base draft dynconfig sdk impl
Browse files Browse the repository at this point in the history
  • Loading branch information
Enjection committed Jan 25, 2024
1 parent f2f869e commit 5ba14ea
Show file tree
Hide file tree
Showing 4 changed files with 212 additions and 0 deletions.
1 change: 1 addition & 0 deletions ydb/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@
from .import_client import * # noqa
from .tracing import * # noqa
from .topic import * # noqa
import ydb.draft # noqa

try:
import ydb.aio as aio # noqa
Expand Down
26 changes: 26 additions & 0 deletions ydb/_apis.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,10 @@
ydb_topic_v1_pb2_grpc,
)

from ._grpc.v4.draft import (
ydb_dynamic_config_v1_pb2_grpc,
)

from ._grpc.v4.protos import (
ydb_status_codes_pb2,
ydb_discovery_pb2,
Expand All @@ -21,6 +25,10 @@
ydb_operation_pb2,
ydb_common_pb2,
)

from ._grpc.v4.draft.protos import (
ydb_dynamic_config_pb2,
)
else:
from ._grpc.common import (
ydb_cms_v1_pb2_grpc,
Expand All @@ -31,6 +39,10 @@
ydb_topic_v1_pb2_grpc,
)

from ._grpc.common.draft import (
ydb_dynamic_config_v1_pb2_grpc,
)

from ._grpc.common.protos import (
ydb_status_codes_pb2,
ydb_discovery_pb2,
Expand All @@ -41,6 +53,10 @@
ydb_common_pb2,
)

from ._grpc.common.draft.protos import (
ydb_dynamic_config_pb2,
)


StatusIds = ydb_status_codes_pb2.StatusIds
FeatureFlag = ydb_common_pb2.FeatureFlag
Expand All @@ -50,6 +66,7 @@
ydb_table = ydb_table_pb2
ydb_discovery = ydb_discovery_pb2
ydb_operation = ydb_operation_pb2
ydb_dynamic_config = ydb_dynamic_config_pb2


class CmsService(object):
Expand Down Expand Up @@ -109,3 +126,12 @@ class TopicService(object):
DropTopic = "DropTopic"
StreamRead = "StreamRead"
StreamWrite = "StreamWrite"


class DynamicConfigService(object):
Stub = ydb_dynamic_config_v1_pb2_grpc.DynamicConfigServiceStub

ReplaceConfig = "ReplaceConfig"
SetConfig = "SetConfig"
GetConfig = "GetConfig"
GetNodeLabels = "GetNodeLabels"
1 change: 1 addition & 0 deletions ydb/draft/__init__.py
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
from .dynamic_config import * # noqa
184 changes: 184 additions & 0 deletions ydb/draft/dynamic_config.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,184 @@
import abc
import enum
from abc import abstractmethod
from .. import issues, operation, settings as settings_impl, _apis


class IDynamicConfigClient(abc.ABC):
@abstractmethod
def __init__(self, driver):
pass

@abstractmethod
def replace_config(self, config, dry_run, allow_unknown_fields, settings):
pass

@abstractmethod
def set_config(self, config, dry_run, allow_unknown_fields, settings):
pass

@abstractmethod
def get_config(self, settings):
pass

@abstractmethod
def get_node_labels(self, node_id, settings):
pass


class DynamicConfig(object):
__slots__ = ("version", "cluster", "config")

def __init__(self, version, cluster, config, *args, **kwargs):
self.version = version
self.cluster = cluster
self.config = config


class NodeLabels(object):
__slots__ = ("labels")

def __init__(self, labels, *args, **kwargs):
self.labels = labels


def _replace_config_request_factory(config, dry_run, allow_unknown_fields):
request = _apis.ydb_dynamic_config.ReplaceConfigRequest()
request.config = config
request.dry_run = dry_run
request.allow_unknown_fields = allow_unknown_fields
return request


def _set_config_request_factory(config, dry_run, allow_unknown_fields):
request = _apis.ydb_dynamic_config.SetConfigRequest()
request.config = config
request.dry_run = dry_run
request.allow_unknown_fields = allow_unknown_fields
return request


def _get_config_request_factory():
request = _apis.ydb_dynamic_config.GetConfigRequest()
return request


def _get_node_labels_request_factory(node_id):
request = _apis.ydb_dynamic_config.GetNodeLabelsRequest()
request.node_id = node_id
return request


def _wrap_dynamic_config(config_pb, dynamic_config_cls=None, *args, **kwargs):
dynamic_config_cls = DynamicConfig if dynamic_config_cls is None else dynamic_config_cls
return dynamic_config_cls(
config_pb.identity.version,
config_pb.identity.cluster,
config_pb.config,
*args,
**kwargs
)


def _wrap_get_config_response(rpc_state, response):
issues._process_response(response.operation)
message = _apis.ydb_dynamic_config.GetConfigResult()
response.operation.result.Unpack(message)
return _wrap_dynamic_config(message)


def _wrap_node_labels(labels_pb, node_labels_cls=None, *args, **kwargs):
node_labels_cls = NodeLabels if node_labels_cls is None else node_labels_cls
return node_labels_cls(
dict([(entry.label, entry.value) for entry in labels_pb.labels]),
*args,
**kwargs
)


def _wrap_get_node_labels_response(rpc_state, response):
issues._process_response(response.operation)
message = _apis.ydb_dynamic_config.GetNodeLabelsResult()
response.operation.result.Unpack(message)
return _wrap_node_labels(message)


class BaseDynamicConfigClient(IDynamicConfigClient):
__slots__ = ("_driver",)

def __init__(self, driver):
self._driver = driver

def replace_config(self, config, dry_run, allow_unknown_fields, settings=None):
return self._driver(
_replace_config_request_factory(config, dry_run, allow_unknown_fields),
_apis.DynamicConfigService.Stub,
_apis.DynamicConfigService.ReplaceConfig,
operation.Operation,
settings,
)

def set_config(self, config, dry_run, allow_unknown_fields, settings=None):
return self._driver(
_set_config_request_factory(config, dry_run, allow_unknown_fields),
_apis.DynamicConfigService.Stub,
_apis.DynamicConfigService.SetConfig,
operation.Operation,
settings,
)

def get_config(self, settings=None):
return self._driver(
_get_config_request_factory(),
_apis.DynamicConfigService.Stub,
_apis.DynamicConfigService.GetConfig,
_wrap_get_config_response,
settings,
)

def get_node_labels(self, node_id, settings=None):
return self._driver(
_get_node_labels_request_factory(node_id),
_apis.DynamicConfigService.Stub,
_apis.DynamicConfigService.GetNodeLabels,
_wrap_get_node_labels_response,
settings,
)


class DynamicConfigClient(BaseDynamicConfigClient):
def async_replace_config(self, config, dry_run, allow_unknown_fields, settings=None):
return self._driver.future(
_replace_config_request_factory(config, dry_run, allow_unknown_fields),
_apis.DynamicConfigService.Stub,
_apis.DynamicConfigService.ReplaceConfig,
operation.Operation,
settings,
)

def async_set_config(self, config, dry_run, allow_unknown_fields, settings=None):
return self._driver.future(
_set_config_request_factory(config, dry_run, allow_unknown_fields),
_apis.DynamicConfigService.Stub,
_apis.DynamicConfigService.SetConfig,
operation.Operation,
settings,
)

def async_get_config(self, settings=None):
return self._driver.future(
_get_config_request_factory(),
_apis.DynamicConfigService.Stub,
_apis.DynamicConfigService.GetConfig,
_wrap_get_config_response,
settings,
)

def async_get_node_labels(self, node_id, settings=None):
return self._driver.future(
_get_node_labels_request_factory(node_id),
_apis.DynamicConfigService.Stub,
_apis.DynamicConfigService.GetNodeLabels,
_wrap_get_node_labels_response,
settings,
)

0 comments on commit 5ba14ea

Please sign in to comment.