add describe -> start making crud lock object

Author: slampy97

tests/coordination/test_coordination_client.py

@@ -3,7 +3,7 @@
 import pytest
 
 import ydb
-from ydb import aio
+from ydb import aio, StatusCode
 
 from ydb.coordination import (
     NodeConfig,
@@ -96,10 +96,10 @@ async def test_coordination_node_lifecycle_async(self, aio_connection):
         with pytest.raises(ydb.SchemeError):
             await client.describe_node(node_path)
 
-    async def test_coordination_lock_context_exclusive(self, aio_connection):
+    async def test_coordination_lock_full_lifecycle(self, aio_connection):
         client = aio.CoordinationClient(aio_connection)
 
-        node_path = "/local/test_lock_context_exclusive"
+        node_path = "/local/test_lock_full_lifecycle"
 
         try:
             await client.delete_node(node_path)
@@ -117,8 +117,34 @@ async def test_coordination_lock_context_exclusive(self, aio_connection):
             ),
         )
 
+        lock2_started = asyncio.Event()
+        lock2_acquired = asyncio.Event()
+
+        async def second_lock_task():
+            lock2_started.set()
+            async with client.lock("test_lock", node_path):
+                lock2_acquired.set()
+                await asyncio.sleep(0.5)
+
         async with client.lock("test_lock", node_path) as lock1:
-            print("Lock1 acquired")
-            await asyncio.sleep(10)
-            print("Lock1 still holding after 10 seconds")
+
+            assert lock1._stream is not None
+            assert lock1._stream.session_id is not None
+            resp = await lock1.describe()
+            assert resp.status == StatusCode.SUCCESS
+
+            t2 = asyncio.create_task(second_lock_task())
+            await lock2_started.wait()
+
+
+            await asyncio.sleep(0.5)
+
+            assert lock1._stream is not None
+
+        await asyncio.wait_for(lock2_acquired.wait(), timeout=5)
+        await asyncio.wait_for(t2, timeout=5)
+
+        async with client.lock("test_lock", node_path) as lock3:
+            assert lock3._stream is not None
+            assert lock3._stream.session_id is not None
 

ydb/_grpc/grpcwrapper/ydb_coordination.py

@@ -161,6 +161,27 @@ def to_proto(self) -> ydb_coordination_pb2.SessionRequest:
                 req_id=self.req_id, name=self.name
             ))
 
+@dataclass
+class DescribeSemaphore(IToProto):
+    include_owners: bool
+    include_waiters: bool
+    name: str
+    req_id: int
+    watch_data: bool
+    watch_owners: bool
+
+    def to_proto(self) -> ydb_coordination_pb2.SessionRequest:
+        return ydb_coordination_pb2.SessionRequest(
+            describe_semaphore=ydb_coordination_pb2.SessionRequest.DescribeSemaphore(
+                include_owners=self.include_owners,
+                include_waiters=self.include_waiters,
+                name=self.name,
+                req_id=self.req_id,
+                watch_data=self.watch_data,
+                watch_owners=self.watch_owners
+            )
+        )
+
 
 @dataclass
 class FromServer:

ydb/aio/coordination/lock.py

@@ -5,6 +5,8 @@
 from ydb._grpc.grpcwrapper.ydb_coordination import (
     AcquireSemaphore,
     ReleaseSemaphore,
+    UpdateSemaphore,
+    DescribeSemaphore,
     FromServer,
 )
 from ydb.aio.coordination.stream import CoordinationStream
@@ -79,6 +81,21 @@ async def _wait_for_acquire_response(self):
                 f"Timeout waiting for lock {self._name} acquisition"
             )
 
+    async def _wait_for_describe_response(self, req_id: int):
+        try:
+            while True:
+                resp = await asyncio.wait_for(
+                    self._stream._incoming_queue.get(),
+                    timeout=self._wait_timeout,
+                )
+                describe_resp = FromServer.from_proto(resp).describe_semaphore_result
+                if describe_resp and describe_resp.req_id == req_id:
+                    return describe_resp
+        except asyncio.TimeoutError:
+            raise issues.Error(
+                f"Timeout waiting for lock {self._name} describe"
+            )
+
     async def __aenter__(self):
         await self._ensure_session()
 
@@ -117,3 +134,26 @@ async def acquire(self):
 
     async def release(self):
         await self.__aexit__(None, None, None)
+
+
+    async def describe(self):
+        await self._ensure_session()
+
+        req_id = self.next_req_id()
+
+        req = DescribeSemaphore(
+            req_id=req_id,
+            name=self._name,
+            include_owners=True,
+            include_waiters=True,
+            watch_data=False,
+            watch_owners=False,
+        ).to_proto()
+
+        await self.send(req)
+
+        resp = await self._wait_for_describe_response(req_id)
+        return resp
+
+    async def update(self, *, limit: Optional[int] = None, settings=None):
+        pass