add alter node

Author: slampy97

tests/coordination/test_coordination_alter_node.py

@@ -0,0 +1,34 @@
+import ydb
+from ydb import _apis
+
+
+def test_coordination_alter_node(driver_sync: ydb.Driver):
+    client = driver_sync.coordination_client
+    node_path = "/local/test_alter_node"
+
+    try:
+        client.delete_node(node_path)
+    except ydb.SchemeError:
+        pass
+
+    client.create_node(node_path)
+
+    new_config = _apis.ydb_coordination.Config(
+        session_grace_period_millis=12345,
+        attach_consistency_mode=_apis.ydb_coordination.ConsistencyMode.CONSISTENCY_MODE_STRICT,
+        read_consistency_mode=_apis.ydb_coordination.ConsistencyMode.CONSISTENCY_MODE_RELAXED,
+    )
+
+    alter_res = client.alter_node(node_path, new_config)
+
+    assert alter_res.status == ydb.StatusCode.SUCCESS, f"Alter operation failed: {alter_res.status}"
+
+
+    node = client.describe_node(node_path)
+    assert node.config.session_grace_period_millis == 12345, "Session grace period not updated"
+    assert node.config.attach_consistency_mode == _apis.ydb_coordination.ConsistencyMode.CONSISTENCY_MODE_STRICT, \
+        "Attach consistency mode not updated"
+    assert node.config.read_consistency_mode == _apis.ydb_coordination.ConsistencyMode.CONSISTENCY_MODE_RELAXED, \
+        "Read consistency mode not updated"
+
+    client.delete_node(node_path)

tests/coordination/coordination_client.py renamed to tests/coordination/test_coordination_describe_node.py

@@ -28,6 +28,21 @@ def to_proto(self) -> ydb_coordination_pb2.CreateNodeRequest:
             operation_params=self.operation_params,
         )
 
+@dataclass
+class AlterNodeRequest(IToProto):
+    path: str
+    config: typing.Optional[public_types.NodeConfig] = None
+    operation_params: typing.Any = 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,
+            operation_params=self.operation_params,
+        )
+
+
 
 @dataclass
 class DescribeNodeRequest(IToProto):
@@ -53,41 +68,3 @@ def to_proto(self) -> ydb_coordination_pb2.DropNodeRequest:
         )
 
 
-@dataclass
-class CreateNodeResponse(IFromProto):
-    operation : ydb.Operation
-    OPERATION_FIELD_NUMBER : int
-
-    @staticmethod
-    def from_proto(msg: ydb_coordination_pb2.CreateNodeResponse) -> "CreateNodeResponse":
-        return CreateNodeResponse(
-            operation=msg.operation,
-            OPERATION_FIELD_NUMBER=msg.OPERATION_FIELD_NUMBER
-        )
-
-
-@dataclass
-class DescribeNodeResponse(IFromProto):
-    operation : ydb.Operation
-    OPERATION_FIELD_NUMBER : int
-
-    @staticmethod
-    def from_proto(msg: "ydb_coordination_pb2.DescribeNodeResponse") -> "DescribeNodeResponse":
-        return DescribeNodeResponse(
-            operation=msg.operation,
-            OPERATION_FIELD_NUMBER=msg.OPERATION_FIELD_NUMBER
-        )
-
-
-@dataclass
-class DropNodeResponse(IFromProto):
-    operation : ydb.Operation
-    OPERATION_FIELD_NUMBER : int
-
-    @staticmethod
-    def from_proto(msg: ydb_coordination_pb2.DropNodeResponse) -> "DropNodeResponse":
-        return DropNodeResponse(
-            operation=msg.operation,
-            OPERATION_FIELD_NUMBER=msg.OPERATION_FIELD_NUMBER
-        )
-

ydb/_grpc/grpcwrapper/ydb_coordination.py

@@ -0,0 +1,5 @@
+from .coordination_client import CoordinationClient
+
+__all__ = [
+    "CoordinationClient",
+]

ydb/coordination/__init__.py

@@ -4,8 +4,7 @@
 from ydb import _apis, issues
 
 
-from .operations import DescribeNodeOperation, CreateNodeOperation, DropNodeOperation
-
+from .operations import DescribeNodeOperation, CreateNodeOperation, DropNodeOperation, AlterNodeOperation
 
 if typing.TYPE_CHECKING:
     import ydb
@@ -25,6 +24,10 @@ def wrapper_delete_node(rpc_state, response_pb, path, *_args, **_kwargs):
     issues._process_response(response_pb.operation)
     return DropNodeOperation(rpc_state, response_pb, path)
 
+def wrapper_alter_node(rpc_state, response_pb, path, *_args, **_kwargs):
+    issues._process_response(response_pb.operation)
+    return AlterNodeOperation(rpc_state, response_pb, path)
+
 
 class CoordinationClient:
     def __init__(self, driver: "ydb.Driver"):
@@ -103,3 +106,30 @@ def delete_node(
             settings=settings,
         )
 
+    def alter_node(
+            self,
+            path: str,
+            new_config: typing.Optional[typing.Any] = None,
+            operation_params: typing.Optional[typing.Any] = None,
+            settings: Optional["ydb.BaseRequestSettings"] = None,
+    ):
+        """
+        Alter node configuration.
+        """
+        request = _apis.ydb_coordination.AlterNodeRequest(
+            path=path,
+            config=new_config,
+            operation_params=operation_params,
+        )
+
+        return self._call_node(
+            request,
+            _apis.CoordinationService.AlterNode,
+            wrapper_alter_node,
+            wrap_args=(path,),
+            settings=settings,
+        )
+
+    def close(self):
+        pass
+

ydb/coordination/coordination_client.py

@@ -48,3 +48,16 @@ def __init__(self, rpc_state, response, path, driver=None):
 
     def __repr__(self):
         return f"DropNodeOperation<id={self.id}, path={self.path}, status={self.status}>"
+
+class AlterNodeOperation(ydb_op.Operation):
+    def __init__(self, rpc_state, response, path, driver=None):
+        super().__init__(rpc_state, response, driver)
+        self.path = path
+        self.status = response.operation.status
+
+    def __repr__(self):
+        return f"AlterNodeOperation<id={self.id}, path={self.path}, status={self.status}>"
+
+    __str__ = __repr__
+
+

ydb/coordination/operations.py

@@ -8,7 +8,7 @@
 from . import tracing
 from . import iam
 from . import _utilities
-from .coordination.coordination_client import CoordinationClient
+from . import coordination
 
 
 
@@ -289,16 +289,12 @@ def __init__(
         self._credentials = driver_config.credentials
 
         self.scheme_client = scheme.SchemeClient(self)
-        self.coordination_client = CoordinationClient(self)
+        self.coordination_client = coordination.CoordinationClient(self)
         self.table_client = table.TableClient(self, driver_config.table_client_settings)
         self.topic_client = topic.TopicClient(self, driver_config.topic_client_settings)
 
     def stop(self, timeout=10):
         self.table_client._stop_pool_if_needed(timeout=timeout)
         self.topic_client.close()
-        if hasattr(self, "coordination_client"):
-            try:
-                self.coordination_client.close()
-            except Exception as e:
-                logger.warning(f"Failed to close coordination client: {e}")
+        self.coordination_client.close()
         super().stop(timeout=timeout)