Skip to content

Commit 41323e6

Browse files
slampy97slampy97
committed
fix review remarks
1 parent c3324cf commit 41323e6

File tree

2 files changed

+87
-104
lines changed

2 files changed

+87
-104
lines changed

‎ydb/_grpc/grpcwrapper/ydb_coordination_public_types.py‎

Lines changed: 20 additions & 104 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
from dataclasses import dataclass
2-
from enum import Enum
2+
from enum import IntEnum
33
import typing
44
import ydb
55

@@ -9,50 +9,16 @@
99
from ..common.protos import ydb_coordination_pb2
1010

1111

12-
class ConsistencyMode(Enum):
13-
UNSET = 0
14-
STRICT = 1
15-
RELAXED = 2
12+
class ConsistencyMode(IntEnum):
13+
UNSET = ydb_coordination_pb2.ConsistencyMode.CONSISTENCY_MODE_UNSET
14+
STRICT = ydb_coordination_pb2.ConsistencyMode.CONSISTENCY_MODE_STRICT
15+
RELAXED = ydb_coordination_pb2.ConsistencyMode.CONSISTENCY_MODE_RELAXED
1616

17-
@classmethod
18-
def from_proto(cls, proto_val: int) -> "ConsistencyMode":
19-
mapping = {
20-
ydb_coordination_pb2.ConsistencyMode.CONSISTENCY_MODE_UNSET: cls.UNSET,
21-
ydb_coordination_pb2.ConsistencyMode.CONSISTENCY_MODE_STRICT: cls.STRICT,
22-
ydb_coordination_pb2.ConsistencyMode.CONSISTENCY_MODE_RELAXED: cls.RELAXED,
23-
}
24-
return mapping[proto_val]
2517

26-
def to_proto(self) -> int:
27-
mapping = {
28-
self.UNSET: ydb_coordination_pb2.ConsistencyMode.CONSISTENCY_MODE_UNSET,
29-
self.STRICT: ydb_coordination_pb2.ConsistencyMode.CONSISTENCY_MODE_STRICT,
30-
self.RELAXED: ydb_coordination_pb2.ConsistencyMode.CONSISTENCY_MODE_RELAXED,
31-
}
32-
return mapping[self]
33-
34-
35-
class RateLimiterCountersMode(Enum):
36-
UNSET = 0
37-
AGGREGATED = 1
38-
DETAILED = 2
39-
40-
@classmethod
41-
def from_proto(cls, proto_val: int) -> "RateLimiterCountersMode":
42-
mapping = {
43-
ydb_coordination_pb2.RateLimiterCountersMode.RATE_LIMITER_COUNTERS_MODE_UNSET: cls.UNSET,
44-
ydb_coordination_pb2.RateLimiterCountersMode.RATE_LIMITER_COUNTERS_MODE_AGGREGATED: cls.AGGREGATED,
45-
ydb_coordination_pb2.RateLimiterCountersMode.RATE_LIMITER_COUNTERS_MODE_DETAILED: cls.DETAILED,
46-
}
47-
return mapping[proto_val]
48-
49-
def to_proto(self) -> int:
50-
mapping = {
51-
self.UNSET: ydb_coordination_pb2.RateLimiterCountersMode.RATE_LIMITER_COUNTERS_MODE_UNSET,
52-
self.AGGREGATED: ydb_coordination_pb2.RateLimiterCountersMode.RATE_LIMITER_COUNTERS_MODE_AGGREGATED,
53-
self.DETAILED: ydb_coordination_pb2.RateLimiterCountersMode.RATE_LIMITER_COUNTERS_MODE_DETAILED,
54-
}
55-
return mapping[self]
18+
class RateLimiterCountersMode(IntEnum):
19+
UNSET = ydb_coordination_pb2.RateLimiterCountersMode.RATE_LIMITER_COUNTERS_MODE_UNSET
20+
AGGREGATED = ydb_coordination_pb2.RateLimiterCountersMode.RATE_LIMITER_COUNTERS_MODE_AGGREGATED
21+
DETAILED = ydb_coordination_pb2.RateLimiterCountersMode.RATE_LIMITER_COUNTERS_MODE_DETAILED
5622

5723

5824
@dataclass
@@ -66,18 +32,18 @@ class NodeConfig:
6632
@staticmethod
6733
def from_proto(msg: ydb_coordination_pb2.Config) -> "NodeConfig":
6834
return NodeConfig(
69-
attach_consistency_mode=ConsistencyMode.from_proto(msg.attach_consistency_mode),
70-
rate_limiter_counters_mode=RateLimiterCountersMode.from_proto(msg.rate_limiter_counters_mode),
71-
read_consistency_mode=ConsistencyMode.from_proto(msg.read_consistency_mode),
35+
attach_consistency_mode=msg.attach_consistency_mode,
36+
rate_limiter_counters_mode=msg.rate_limiter_counters_mode,
37+
read_consistency_mode=msg.read_consistency_mode,
7238
self_check_period_millis=msg.self_check_period_millis,
7339
session_grace_period_millis=msg.session_grace_period_millis,
7440
)
7541

7642
def to_proto(self) -> ydb_coordination_pb2.Config:
7743
return ydb_coordination_pb2.Config(
78-
attach_consistency_mode=self.attach_consistency_mode.to_proto(),
79-
rate_limiter_counters_mode=self.rate_limiter_counters_mode.to_proto(),
80-
read_consistency_mode=self.read_consistency_mode.to_proto(),
44+
attach_consistency_mode=self.attach_consistency_mode,
45+
rate_limiter_counters_mode=self.rate_limiter_counters_mode,
46+
read_consistency_mode=self.read_consistency_mode,
8147
self_check_period_millis=self.self_check_period_millis,
8248
session_grace_period_millis=self.session_grace_period_millis,
8349
)
@@ -88,58 +54,8 @@ class NodeDescription:
8854
path: str
8955
config: NodeConfig
9056

91-
92-
class CoordinationClientSettings:
93-
def __init__(self):
94-
self._trace_id = None
95-
self._request_type = None
96-
self._timeout = None
97-
self._cancel_after = None
98-
self._operation_timeout = None
99-
self._compression = None
100-
self._need_rpc_auth = True
101-
self._headers = []
102-
103-
def with_trace_id(self, trace_id: str) -> "CoordinationClientSettings":
104-
self._trace_id = trace_id
105-
return self
106-
107-
def with_request_type(self, request_type: str) -> "CoordinationClientSettings":
108-
self._request_type = request_type
109-
return self
110-
111-
def with_timeout(self, timeout: float) -> "CoordinationClientSettings":
112-
self._timeout = timeout
113-
return self
114-
115-
def with_cancel_after(self, timeout: float) -> "CoordinationClientSettings":
116-
self._cancel_after = timeout
117-
return self
118-
119-
def with_operation_timeout(self, timeout: float) -> "CoordinationClientSettings":
120-
self._operation_timeout = timeout
121-
return self
122-
123-
def with_compression(self, compression) -> "CoordinationClientSettings":
124-
self._compression = compression
125-
return self
126-
127-
def with_need_rpc_auth(self, need_rpc_auth: bool) -> "CoordinationClientSettings":
128-
self._need_rpc_auth = need_rpc_auth
129-
return self
130-
131-
def with_header(self, key: str, value: str) -> "CoordinationClientSettings":
132-
self._headers.append((key, value))
133-
return self
134-
135-
def to_base_request_settings(self) -> "ydb.BaseRequestSettings":
136-
brs = ydb.BaseRequestSettings()
137-
brs.trace_id = self._trace_id
138-
brs.request_type = self._request_type
139-
brs.timeout = self._timeout
140-
brs.cancel_after = self._cancel_after
141-
brs.operation_timeout = self._operation_timeout
142-
brs.compression = self._compression
143-
brs.need_rpc_auth = self._need_rpc_auth
144-
brs.headers.extend(self._headers)
145-
return brs
57+
@staticmethod
58+
def from_proto(path: str, response_pb: ydb_coordination_pb2.DescribeNodeResponse) -> "NodeDescription":
59+
result = ydb_coordination_pb2.DescribeNodeResult()
60+
response_pb.operation.result.Unpack(result)
61+
return NodeDescription(path=path, config=NodeConfig.from_proto(result.config))

‎ydb/coordination/base_coordination_client.py‎

Lines changed: 67 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,67 @@
1+
import typing
2+
from typing import Optional
3+
4+
from ydb import _apis, issues
5+
from ydb._grpc.grpcwrapper.ydb_coordination_public_types import NodeConfig, NodeDescription
6+
7+
8+
def wrapper_create_node(rpc_state, response_pb, client: "BaseCoordinationClient"):
9+
issues._process_response(response_pb.operation)
10+
11+
12+
def wrapper_describe_node(rpc_state, response_pb, client: "BaseCoordinationClient", path: str) -> NodeDescription:
13+
issues._process_response(response_pb.operation)
14+
return NodeDescription.from_proto(path, response_pb)
15+
16+
17+
def wrapper_delete_node(rpc_state, response_pb, client: "BaseCoordinationClient"):
18+
issues._process_response(response_pb.operation)
19+
20+
21+
def wrapper_alter_node(rpc_state, response_pb, client: "BaseCoordinationClient"):
22+
issues._process_response(response_pb.operation)
23+
24+
25+
class BaseCoordinationClient:
26+
def __init__(self, driver):
27+
self._driver = driver
28+
29+
def _call_create(self, request, settings=None, wrap_args=None):
30+
return self._driver(
31+
request,
32+
_apis.CoordinationService.Stub,
33+
_apis.CoordinationService.CreateNode,
34+
wrap_result=wrapper_create_node,
35+
wrap_args=wrap_args,
36+
settings=settings,
37+
)
38+
39+
def _call_describe(self, request, settings=None, wrap_args=None):
40+
return self._driver(
41+
request,
42+
_apis.CoordinationService.Stub,
43+
_apis.CoordinationService.DescribeNode,
44+
wrap_result=wrapper_describe_node,
45+
wrap_args=wrap_args,
46+
settings=settings,
47+
)
48+
49+
def _call_alter(self, request, settings=None, wrap_args=None):
50+
return self._driver(
51+
request,
52+
_apis.CoordinationService.Stub,
53+
_apis.CoordinationService.AlterNode,
54+
wrap_result=wrapper_alter_node,
55+
wrap_args=wrap_args,
56+
settings=settings,
57+
)
58+
59+
def _call_delete(self, request, settings=None, wrap_args=None):
60+
return self._driver(
61+
request,
62+
_apis.CoordinationService.Stub,
63+
_apis.CoordinationService.DropNode,
64+
wrap_result=wrapper_delete_node,
65+
wrap_args=wrap_args,
66+
settings=settings,
67+
)

0 commit comments

Comments
 (0)