Add public wrappers for describe lock result

Author: slampy97

tests/coordination/test_coordination_client.py

@@ -10,6 +10,8 @@
     ConsistencyMode,
     RateLimiterCountersMode,
     CoordinationClient,
+    CreateSemaphoreResult,
+    DescribeLockResult,
 )
 
 
@@ -119,33 +121,29 @@ async def test_coordination_lock_full_lifecycle(self, aio_connection):
 
         lock = client.lock("test_lock", node_path)
 
-        create_resp = await lock.create(init_limit=1, init_data=b"init-data")
+        create_resp: CreateSemaphoreResult = await lock.create(init_limit=1, init_data=b"init-data")
         assert create_resp.status == StatusCode.SUCCESS
 
-        describe_resp = await lock.describe()
+        describe_resp: DescribeLockResult = await lock.describe()
         assert describe_resp.status == StatusCode.SUCCESS
-
-        sem = describe_resp.semaphore_description
-        assert sem.name == "test_lock"
-        assert sem.data == b"init-data"
-        assert sem.count == 0
-        assert sem.ephemeral is False
-        assert list(sem.owners) == []
-        assert list(sem.waiters) == []
+        assert describe_resp.name == "test_lock"
+        assert describe_resp.data == b"init-data"
+        assert describe_resp.count == 0
+        assert describe_resp.ephemeral is False
+        assert list(describe_resp.owners) == []
+        assert list(describe_resp.waiters) == []
 
         update_resp = await lock.update(new_data=b"updated-data")
         assert update_resp.status == StatusCode.SUCCESS
 
-        describe_resp2 = await lock.describe()
+        describe_resp2: DescribeLockResult = await lock.describe()
         assert describe_resp2.status == StatusCode.SUCCESS
-
-        sem2 = describe_resp2.semaphore_description
-        assert sem2.name == "test_lock"
-        assert sem2.data == b"updated-data"
-        assert sem2.count == 0
-        assert sem2.ephemeral is False
-        assert list(sem2.owners) == []
-        assert list(sem2.waiters) == []
+        assert describe_resp2.name == "test_lock"
+        assert describe_resp2.data == b"updated-data"
+        assert describe_resp2.count == 0
+        assert describe_resp2.ephemeral is False
+        assert list(describe_resp2.owners) == []
+        assert list(describe_resp2.waiters) == []
 
         lock2_started = asyncio.Event()
         lock2_acquired = asyncio.Event()
@@ -160,16 +158,14 @@ async def second_lock_task():
             assert lock1._stream is not None
             assert lock1._stream.session_id is not None
 
-            resp = await lock1.describe()
+            resp: DescribeLockResult = await lock1.describe()
             assert resp.status == StatusCode.SUCCESS
-
-            sem_under_lock = resp.semaphore_description
-            assert sem_under_lock.name == "test_lock"
-            assert sem_under_lock.data == b"updated-data"
-            assert sem_under_lock.count == 1
-            assert sem_under_lock.ephemeral is False
-            assert len(list(sem_under_lock.owners)) == 1
-            assert list(sem_under_lock.waiters) == []
+            assert resp.name == "test_lock"
+            assert resp.data == b"updated-data"
+            assert resp.count == 1
+            assert resp.ephemeral is False
+            assert len(list(resp.owners)) == 1
+            assert list(resp.waiters) == []
 
             t2 = asyncio.create_task(second_lock_task())
             await lock2_started.wait()
@@ -185,13 +181,12 @@ async def second_lock_task():
             assert lock3._stream is not None
             assert lock3._stream.session_id is not None
 
-            resp3 = await lock3.describe()
+            resp3: DescribeLockResult = await lock3.describe()
             assert resp3.status == StatusCode.SUCCESS
-            sem3 = resp3.semaphore_description
-            assert sem3.count == 1
+            assert resp3.count == 1
 
         delete_resp = await lock.delete()
         assert delete_resp.status == StatusCode.SUCCESS
 
-        describe_after_delete = await lock.describe()
+        describe_after_delete: DescribeLockResult = await lock.describe()
         assert describe_after_delete.status == StatusCode.NOT_FOUND

ydb/_grpc/grpcwrapper/ydb_coordination.py

@@ -1,7 +1,6 @@
 import typing
 from dataclasses import dataclass
 
-from .ydb_coordination_public_types import NodeConfig
 
 if typing.TYPE_CHECKING:
     from ..v4.protos import ydb_coordination_pb2
@@ -14,7 +13,7 @@
 @dataclass
 class CreateNodeRequest(IToProto):
     path: str
-    config: typing.Optional[NodeConfig]
+    config: typing.Any
 
     def to_proto(self) -> ydb_coordination_pb2.CreateNodeRequest:
         cfg_proto = self.config.to_proto() if self.config else None
@@ -27,7 +26,7 @@ def to_proto(self) -> ydb_coordination_pb2.CreateNodeRequest:
 @dataclass
 class AlterNodeRequest(IToProto):
     path: str
-    config: NodeConfig
+    config: typing.Any
 
     def to_proto(self) -> ydb_coordination_pb2.AlterNodeRequest:
         cfg_proto = self.config.to_proto() if self.config else None
@@ -204,62 +203,35 @@ def from_proto(resp: ydb_coordination_pb2.SessionResponse) -> "FromServer":
     def __getattr__(self, name: str):
         return getattr(self.raw, name)
 
-    @property
-    def status(self) -> typing.Optional[int]:
-        return getattr(self.raw, "status", None)
-
     @property
     def session_started(self) -> typing.Optional[ydb_coordination_pb2.SessionResponse.SessionStarted]:
-        x = getattr(self.raw, "session_started", None)
-        return x if getattr(x, "session_id", 0) else None
-
-    @property
-    def session_stopped(self) -> typing.Optional[ydb_coordination_pb2.SessionResponse.SessionStopped]:
-        return getattr(self.raw, "session_stopped", None) or None
-
-    @property
-    def failure(self) -> typing.Optional[ydb_coordination_pb2.SessionResponse.Failure]:
-        return getattr(self.raw, "failure", None) or None
-
-    @property
-    def pong(self) -> typing.Optional[ydb_coordination_pb2.SessionResponse.PingPong]:
-        return getattr(self.raw, "pong", None) or None
-
-    @property
-    def ping(self) -> typing.Optional[ydb_coordination_pb2.SessionResponse.PingPong]:
-        return getattr(self.raw, "ping", None) or None
+        s = self.raw.session_started
+        return s if s.session_id else None
 
     @property
     def opaque(self) -> typing.Optional[int]:
-        if getattr(self.raw, "ping") is not None:
-            return getattr(getattr(self.raw, "ping"), "opaque", None)
+        if self.raw.HasField("ping"):
+            return self.raw.ping.opaque
+        return None
 
     @property
-    def acquire_semaphore_result(
-        self,
-    ) -> typing.Optional[ydb_coordination_pb2.SessionResponse.AcquireSemaphoreResult]:
-        return getattr(self.raw, "acquire_semaphore_result", None) or None
+    def acquire_semaphore_result(self) -> typing.Optional[ydb_coordination_pb2.SessionResponse.AcquireSemaphoreResult]:
+        return self.raw.acquire_semaphore_result if self.raw.HasField("acquire_semaphore_result") else None
 
     @property
-    def create_semaphore_result(
-        self,
-    ) -> typing.Optional[ydb_coordination_pb2.SessionResponse.CreateSemaphoreResult]:
-        return getattr(self.raw, "create_semaphore_result", None) or None
+    def create_semaphore_result(self) -> typing.Optional[ydb_coordination_pb2.SessionResponse.CreateSemaphoreResult]:
+        return self.raw.create_semaphore_result if self.raw.HasField("create_semaphore_result") else None
 
     @property
-    def delete_semaphore_result(
-        self,
-    ) -> typing.Optional[ydb_coordination_pb2.SessionResponse.DeleteSemaphoreResult]:
-        return getattr(self.raw, "delete_semaphore_result", None) or None
+    def delete_semaphore_result(self) -> typing.Optional[ydb_coordination_pb2.SessionResponse.DeleteSemaphoreResult]:
+        return self.raw.delete_semaphore_result if self.raw.HasField("delete_semaphore_result") else None
 
     @property
-    def update_semaphore_result(
-        self,
-    ) -> typing.Optional[ydb_coordination_pb2.SessionResponse.UpdateSemaphoreResult]:
-        return getattr(self.raw, "update_semaphore_result", None) or None
+    def update_semaphore_result(self) -> typing.Optional[ydb_coordination_pb2.SessionResponse.UpdateSemaphoreResult]:
+        return self.raw.update_semaphore_result if self.raw.HasField("update_semaphore_result") else None
 
     @property
     def describe_semaphore_result(
         self,
     ) -> typing.Optional[ydb_coordination_pb2.SessionResponse.DescribeSemaphoreResult]:
-        return getattr(self.raw, "describe_semaphore_result", None) or None
+        return self.raw.describe_semaphore_result if self.raw.HasField("describe_semaphore_result") else None

ydb/_grpc/grpcwrapper/ydb_coordination_public_types.py

@@ -2,7 +2,6 @@
 from enum import IntEnum
 import typing
 
-
 if typing.TYPE_CHECKING:
     from ..v4.protos import ydb_coordination_pb2
 else:
@@ -83,3 +82,32 @@ def from_proto(msg: ydb_coordination_pb2.SessionResponse.CreateSemaphoreResult)
             req_id=msg.req_id,
             status=msg.status,
         )
+
+
+@dataclass
+class DescribeLockResult:
+    req_id: int
+    status: int
+    watch_added: bool
+    count: int
+    data: bytes
+    ephemeral: bool
+    limit: int
+    name: str
+    owners: list
+    waiters: list
+
+    @staticmethod
+    def from_proto(msg: ydb_coordination_pb2.SessionResponse.DescribeSemaphoreResult) -> "DescribeLockResult":
+        return DescribeLockResult(
+            req_id=msg.req_id,
+            status=msg.status,
+            watch_added=msg.watch_added,
+            count=msg.semaphore_description.count,
+            data=msg.semaphore_description.data,
+            ephemeral=msg.semaphore_description.ephemeral,
+            limit=msg.semaphore_description.limit,
+            name=msg.semaphore_description.name,
+            owners=msg.semaphore_description.owners,
+            waiters=msg.semaphore_description.waiters,
+        )

ydb/aio/coordination/lock.py

@@ -11,6 +11,7 @@
     DeleteSemaphore,
     FromServer,
 )
+from ydb._grpc.grpcwrapper.ydb_coordination_public_types import CreateSemaphoreResult, DescribeLockResult
 from ydb.aio.coordination.stream import CoordinationStream
 from ydb.aio.coordination.reconnector import CoordinationReconnector
 
@@ -158,7 +159,7 @@ async def create(self, init_limit, init_data):
         await self.send(req)
 
         resp = await self._wait_for_response(req_id, kind="create")
-        return resp
+        return CreateSemaphoreResult.from_proto(resp)
 
     async def delete(self):
         await self._ensure_session()
@@ -192,7 +193,7 @@ async def describe(self):
         await self.send(req)
 
         resp = await self._wait_for_response(req_id, kind="describe")
-        return resp
+        return DescribeLockResult.from_proto(resp)
 
     async def update(self, new_data):
         await self._ensure_session()

ydb/coordination/__init__.py

@@ -5,6 +5,16 @@
     ConsistencyMode,
     RateLimiterCountersMode,
     DescribeResult,
+    CreateSemaphoreResult,
+    DescribeLockResult,
 )
 
-__all__ = ["CoordinationClient", "NodeConfig", "ConsistencyMode", "RateLimiterCountersMode", "DescribeResult"]
+__all__ = [
+    "CoordinationClient",
+    "NodeConfig",
+    "ConsistencyMode",
+    "RateLimiterCountersMode",
+    "DescribeResult",
+    "CreateSemaphoreResult",
+    "DescribeLockResult",
+]