Merge pull request #386 from Enjection/feature/KIKIMR-20854/dynconfig-sdk Add base draft dynconfig sdk impl

Author: rekby

ydb/__init__.py

@@ -18,6 +18,7 @@
 from .import_client import *  # noqa
 from .tracing import *  # noqa
 from .topic import *  # noqa
+from .draft import *  # noqa
 
 try:
     import ydb.aio as aio  # noqa

ydb/_apis.py

@@ -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,
@@ -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,
@@ -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,
@@ -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
@@ -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):
@@ -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"

ydb/draft/__init__.py

@@ -0,0 +1 @@
+from .dynamic_config import *  # noqa

ydb/draft/dynamic_config.py

@@ -0,0 +1,173 @@
+import abc
+from abc import abstractmethod
+from .. import issues, operation, _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,
+        )