Coordination service impementation - 1 step

tests/coordination/test_coordination_client.py

@@ -0,0 +1,56 @@
+import ydb
+
+from ydb.coordination import NodeConfig, ConsistencyMode, RateLimiterCountersMode, CoordinationClient
+
+
+class TestCoordination:
+    def test_coordination_alter_node(self, driver_sync: ydb.Driver):
+        client = CoordinationClient(driver_sync)
+        node_path = "/local/test_alter_node"
+
+        try:
+            client.delete_node(node_path)
+        except ydb.SchemeError:
+            pass
+
+        client.create_node(node_path)
+
+        new_config = NodeConfig(
+            session_grace_period_millis=12345,
+            attach_consistency_mode=ConsistencyMode.STRICT,
+            read_consistency_mode=ConsistencyMode.RELAXED,
+            rate_limiter_counters_mode=RateLimiterCountersMode.UNSET,
+            self_check_period_millis=0,
+        )
+
+        client.alter_node(node_path, new_config)
+
+        node_desc = client.describe_node(node_path)
+        node_config = node_desc.config
+        path = node_desc.path
+
+        assert node_path == path
+        assert node_config.session_grace_period_millis == 12345
+        assert node_config.attach_consistency_mode == ConsistencyMode.STRICT
+        assert node_config.read_consistency_mode == ConsistencyMode.RELAXED
+
+        client.delete_node(node_path)
+
+    def test_coordination_nodes(self, driver_sync: ydb.Driver):
+        client = CoordinationClient(driver_sync)
+        node_path = "/local/test_node"
+
+        try:
+            client.delete_node(node_path)
+        except ydb.SchemeError:
+            pass
+
+        client.create_node(node_path)
+
+        node_descr = client.describe_node(node_path)
+
+        node_descr_path = node_descr.path
+
+        assert node_descr_path == node_path
+
+        client.delete_node(node_path)

ydb/_apis.py

@@ -11,6 +11,7 @@
         ydb_operation_v1_pb2_grpc,
         ydb_topic_v1_pb2_grpc,
         ydb_query_v1_pb2_grpc,
+        ydb_coordination_v1_pb2_grpc,
     )
 
     from ._grpc.v4.protos import (
@@ -22,6 +23,7 @@
         ydb_operation_pb2,
         ydb_common_pb2,
         ydb_query_pb2,
+        ydb_coordination_pb2,
     )
 
 else:
@@ -33,6 +35,7 @@
         ydb_operation_v1_pb2_grpc,
         ydb_topic_v1_pb2_grpc,
         ydb_query_v1_pb2_grpc,
+        ydb_coordination_v1_pb2_grpc,
     )
 
     from ._grpc.common.protos import (
@@ -44,6 +47,7 @@
         ydb_operation_pb2,
         ydb_common_pb2,
         ydb_query_pb2,
+        ydb_coordination_pb2,
     )
 
 
@@ -56,6 +60,7 @@
 ydb_discovery = ydb_discovery_pb2
 ydb_operation = ydb_operation_pb2
 ydb_query = ydb_query_pb2
+ydb_coordination = ydb_coordination_pb2
 
 
 class CmsService(object):
@@ -134,3 +139,13 @@ class QueryService(object):
     ExecuteQuery = "ExecuteQuery"
     ExecuteScript = "ExecuteScript"
     FetchScriptResults = "FetchScriptResults"
+
+
+class CoordinationService(object):
+    Stub = ydb_coordination_v1_pb2_grpc.CoordinationServiceStub
+
+    Session = "Session"
+    CreateNode = "CreateNode"
+    AlterNode = "AlterNode"
+    DropNode = "DropNode"
+    DescribeNode = "DescribeNode"

ydb/_grpc/grpcwrapper/ydb_coordination.py

@@ -0,0 +1,57 @@
+import typing
+from dataclasses import dataclass
+
+
+if typing.TYPE_CHECKING:
+    from ..v4.protos import ydb_coordination_pb2
+else:
+    from ..common.protos import ydb_coordination_pb2
+
+from .common_utils import IToProto
+from ydb.coordination import NodeConfig
+
+
+@dataclass
+class CreateNodeRequest(IToProto):
+    path: str
+    config: typing.Optional[NodeConfig] = None
+
+    def to_proto(self) -> ydb_coordination_pb2.CreateNodeRequest:
+        cfg_proto = self.config.to_proto() if self.config else None
+        return ydb_coordination_pb2.CreateNodeRequest(
+            path=self.path,
+            config=cfg_proto,
+        )
+
+
+@dataclass
+class AlterNodeRequest(IToProto):
+    path: str
+    config: typing.Optional[NodeConfig] = None
+
+    def to_proto(self) -> ydb_coordination_pb2.AlterNodeRequest:
+        cfg_proto = self.config.to_proto() if self.config else None
+        return ydb_coordination_pb2.AlterNodeRequest(
+            path=self.path,
+            config=cfg_proto,
+        )
+
+
+@dataclass
+class DescribeNodeRequest(IToProto):
+    path: str
+
+    def to_proto(self) -> ydb_coordination_pb2.DescribeNodeRequest:
+        return ydb_coordination_pb2.DescribeNodeRequest(
+            path=self.path,
+        )
+
+
+@dataclass
+class DropNodeRequest(IToProto):
+    path: str
+
+    def to_proto(self) -> ydb_coordination_pb2.DropNodeRequest:
+        return ydb_coordination_pb2.DropNodeRequest(
+            path=self.path,
+        )

ydb/_grpc/grpcwrapper/ydb_coordination_public_types.py

@@ -0,0 +1,145 @@
+from dataclasses import dataclass
+from enum import Enum
+import typing
+import ydb
+
+if typing.TYPE_CHECKING:
+    from ..v4.protos import ydb_coordination_pb2
+else:
+    from ..common.protos import ydb_coordination_pb2
+
+
+class ConsistencyMode(Enum):
+    UNSET = 0
+    STRICT = 1
+    RELAXED = 2
+
+    @classmethod
+    def from_proto(cls, proto_val: int) -> "ConsistencyMode":
+        mapping = {
+            ydb_coordination_pb2.ConsistencyMode.CONSISTENCY_MODE_UNSET: cls.UNSET,
+            ydb_coordination_pb2.ConsistencyMode.CONSISTENCY_MODE_STRICT: cls.STRICT,
+            ydb_coordination_pb2.ConsistencyMode.CONSISTENCY_MODE_RELAXED: cls.RELAXED,
+        }
+        return mapping[proto_val]
+
+    def to_proto(self) -> int:
+        mapping = {
+            self.UNSET: ydb_coordination_pb2.ConsistencyMode.CONSISTENCY_MODE_UNSET,
+            self.STRICT: ydb_coordination_pb2.ConsistencyMode.CONSISTENCY_MODE_STRICT,
+            self.RELAXED: ydb_coordination_pb2.ConsistencyMode.CONSISTENCY_MODE_RELAXED,
+        }
+        return mapping[self]
+
+
+class RateLimiterCountersMode(Enum):
+    UNSET = 0
+    AGGREGATED = 1
+    DETAILED = 2
+
+    @classmethod
+    def from_proto(cls, proto_val: int) -> "RateLimiterCountersMode":
+        mapping = {
+            ydb_coordination_pb2.RateLimiterCountersMode.RATE_LIMITER_COUNTERS_MODE_UNSET: cls.UNSET,
+            ydb_coordination_pb2.RateLimiterCountersMode.RATE_LIMITER_COUNTERS_MODE_AGGREGATED: cls.AGGREGATED,
+            ydb_coordination_pb2.RateLimiterCountersMode.RATE_LIMITER_COUNTERS_MODE_DETAILED: cls.DETAILED,
+        }
+        return mapping[proto_val]
+
+    def to_proto(self) -> int:
+        mapping = {
+            self.UNSET: ydb_coordination_pb2.RateLimiterCountersMode.RATE_LIMITER_COUNTERS_MODE_UNSET,
+            self.AGGREGATED: ydb_coordination_pb2.RateLimiterCountersMode.RATE_LIMITER_COUNTERS_MODE_AGGREGATED,
+            self.DETAILED: ydb_coordination_pb2.RateLimiterCountersMode.RATE_LIMITER_COUNTERS_MODE_DETAILED,
+        }
+        return mapping[self]
+
+
+@dataclass
+class NodeConfig:
+    attach_consistency_mode: ConsistencyMode
+    rate_limiter_counters_mode: RateLimiterCountersMode
+    read_consistency_mode: ConsistencyMode
+    self_check_period_millis: int
+    session_grace_period_millis: int
+
+    @staticmethod
+    def from_proto(msg: ydb_coordination_pb2.Config) -> "NodeConfig":
+        return NodeConfig(
+            attach_consistency_mode=ConsistencyMode.from_proto(msg.attach_consistency_mode),
+            rate_limiter_counters_mode=RateLimiterCountersMode.from_proto(msg.rate_limiter_counters_mode),
+            read_consistency_mode=ConsistencyMode.from_proto(msg.read_consistency_mode),
+            self_check_period_millis=msg.self_check_period_millis,
+            session_grace_period_millis=msg.session_grace_period_millis,
+        )
+
+    def to_proto(self) -> ydb_coordination_pb2.Config:
+        return ydb_coordination_pb2.Config(
+            attach_consistency_mode=self.attach_consistency_mode.to_proto(),
+            rate_limiter_counters_mode=self.rate_limiter_counters_mode.to_proto(),
+            read_consistency_mode=self.read_consistency_mode.to_proto(),
+            self_check_period_millis=self.self_check_period_millis,
+            session_grace_period_millis=self.session_grace_period_millis,
+        )
+
+
+@dataclass
+class NodeDescription:
+    path: str
+    config: NodeConfig
+
+
+class CoordinationClientSettings:
+    def __init__(self):
+        self._trace_id = None
+        self._request_type = None
+        self._timeout = None
+        self._cancel_after = None
+        self._operation_timeout = None
+        self._compression = None
+        self._need_rpc_auth = True
+        self._headers = []
+
+    def with_trace_id(self, trace_id: str) -> "CoordinationClientSettings":
+        self._trace_id = trace_id
+        return self
+
+    def with_request_type(self, request_type: str) -> "CoordinationClientSettings":
+        self._request_type = request_type
+        return self
+
+    def with_timeout(self, timeout: float) -> "CoordinationClientSettings":
+        self._timeout = timeout
+        return self
+
+    def with_cancel_after(self, timeout: float) -> "CoordinationClientSettings":
+        self._cancel_after = timeout
+        return self
+
+    def with_operation_timeout(self, timeout: float) -> "CoordinationClientSettings":
+        self._operation_timeout = timeout
+        return self
+
+    def with_compression(self, compression) -> "CoordinationClientSettings":
+        self._compression = compression
+        return self
+
+    def with_need_rpc_auth(self, need_rpc_auth: bool) -> "CoordinationClientSettings":
+        self._need_rpc_auth = need_rpc_auth
+        return self
+
+    def with_header(self, key: str, value: str) -> "CoordinationClientSettings":
+        self._headers.append((key, value))
+        return self
+
+    def to_base_request_settings(self) -> "ydb.BaseRequestSettings":
+        brs = ydb.BaseRequestSettings()
+        brs.trace_id = self._trace_id
+        brs.request_type = self._request_type
+        brs.timeout = self._timeout
+        brs.cancel_after = self._cancel_after
+        brs.operation_timeout = self._operation_timeout
+        brs.compression = self._compression
+        brs.need_rpc_auth = self._need_rpc_auth
+        brs.headers.extend(self._headers)
+        return brs

ydb/coordination/__init__.py

@@ -0,0 +1,18 @@
+from .coordination_client import CoordinationClient
+
+from ydb._grpc.grpcwrapper.ydb_coordination_public_types import (
+    NodeConfig,
+    NodeDescription,
+    ConsistencyMode,
+    RateLimiterCountersMode,
+    CoordinationClientSettings,
+)
+
+__all__ = [
+    "CoordinationClient",
+    "NodeConfig",
+    "NodeDescription",
+    "ConsistencyMode",
+    "RateLimiterCountersMode",
+    "CoordinationClientSettings",
+]