working coord client plus not working lock/ conenction problemsg

Author: slampy97

ydb/coordination/client.py

@@ -0,0 +1,56 @@
+import ydb
+from ydb import operation
+
+from ydb._grpc.v5.ydb_coordination_v1_pb2_grpc import CoordinationServiceStub
+from ydb._grpc.v5.protos import ydb_coordination_pb2
+from ydb.coordination.coordination_lock import CoordinationLock
+from ydb.coordination.сoordination_session import CoordinationSession
+
+
+class CoordinationClient:
+    def __init__(self, driver: "ydb.Driver"):
+        self._driver = driver
+
+    def session(self):
+        return CoordinationSession(self._driver)
+
+
+    def create_node(self, path: str, config=None, operation_params=None):
+        request = ydb_coordination_pb2.CreateNodeRequest(
+            path=path,
+            config=config,
+            operation_params=operation_params,
+        )
+        return self._driver(
+            request,
+            CoordinationServiceStub,
+            "CreateNode",
+            operation.Operation,
+        )
+
+    def describe_node(self, path: str, operation_params=None):
+        request = ydb_coordination_pb2.DescribeNodeRequest(
+            path=path,
+            operation_params=operation_params,
+        )
+        return self._driver(
+            request,
+            CoordinationServiceStub,
+            "DescribeNode",
+            operation.Operation,
+        )
+
+    def delete_node(self, path: str, operation_params=None):
+        request = ydb_coordination_pb2.DropNodeRequest(
+            path=path,
+            operation_params=operation_params,
+        )
+        return self._driver(
+            request,
+            CoordinationServiceStub,
+            "DropNode",
+            operation.Operation,
+        )
+
+    def lock(self, path: str, timeout: int = 5000, count: int = 1):
+        return CoordinationLock(self.session(), path, timeout, count)

ydb/coordination/coordination_lock.py

@@ -0,0 +1,22 @@
+from ydb.coordination.сoordination_session import CoordinationSession
+
+
+class CoordinationLock:
+    def __init__(self, session: CoordinationSession, path: str, timeout: int = 5000, count: int = 1):
+        self._session = session
+        self._path = path
+        self._timeout = timeout
+        self._count = count
+
+    def acquire(self):
+        self._session.acquire_semaphore(self._path, self._count, self._timeout)
+
+    def release(self):
+        self._session.release_semaphore(self._path)
+
+    def __enter__(self):
+        self.acquire()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.release()

ydb/coordination/exceptions.py

@@ -0,0 +1,19 @@
+
+class CoordinationError(Exception):
+    """Базовое исключение для всех ошибок координации."""
+
+
+class NodeAlreadyExists(CoordinationError):
+    """Узел координации уже существует."""
+
+
+class NodeNotFound(CoordinationError):
+    """Узел координации не найден."""
+
+
+class NodeLocked(CoordinationError):
+    """Узел координации уже захвачен."""
+
+
+class NodeTimeout(CoordinationError):
+    """Истекло время ожидания при захвате узла."""

ydb/coordination/tests/test_coordination_minimal.py

@@ -0,0 +1,93 @@
+import pytest
+import ydb
+from ydb.coordination.client import CoordinationClient
+import time
+
+
+
+@pytest.mark.integration
+def test_coordination_client_local():
+    driver_config = ydb.DriverConfig(
+        endpoint="grpc://localhost:2136",
+        database="/local",
+    )
+
+    with ydb.Driver(driver_config) as driver:
+        for _ in range(10):
+            try:
+                driver.wait(timeout=5)
+                break
+            except Exception:
+                time.sleep(1)
+
+
+        scheme = ydb.SchemeClient(driver)
+        base_path = "/local/coordination"
+
+        try:
+            scheme.describe_path(base_path)
+        except ydb.issues.SchemeError:
+            scheme.make_directory(base_path)
+            desc = scheme.describe_path(base_path)
+            assert desc is not None, f"Directory {base_path} was not created"
+
+
+        node_path = f"{base_path}/test_node"
+
+        client = CoordinationClient(driver)
+
+
+        create_future = client.create_node(path=node_path)
+        assert create_future is not None
+
+
+        res_desc = client.describe_node(path=node_path)
+        assert res_desc is not None
+
+
+        res_delete = client.delete_node(path=node_path)
+        assert res_delete is not None
+
+@pytest.mark.integration
+def test_coordination_lock_lifecycle():
+    driver_config = ydb.DriverConfig(
+        endpoint="grpc://localhost:2136",
+        database="/local",
+    )
+
+    with ydb.Driver(driver_config) as driver:
+        for _ in range(10):
+            try:
+                driver.wait(timeout=5)
+                break
+            except Exception:
+                time.sleep(1)
+
+        scheme = driver.scheme_client
+        base_path = "/local/coordination"
+        try:
+            scheme.describe_path(base_path)
+        except ydb.SchemeError:
+            scheme.make_directory(base_path)
+            desc = scheme.describe_path(base_path)
+            assert desc is not None, f"Directory {base_path} was not created"
+
+        # создаём client
+        client = CoordinationClient(driver)
+
+        lock_path = f"{base_path}/test_lock"
+
+
+        with client.lock(lock_path, timeout=2000, count=1) as lock:
+            assert lock._session_id is not None, "Lock должен иметь session_id после acquire"
+
+
+            sem_state = client.describe_node(lock_path)
+            assert sem_state is not None, "Семафор должен существовать после acquire"
+
+
+        assert lock._session_id is None, "Lock должен быть освобождён после выхода из with"
+
+
+        sem_state_after = client.describe_node(lock_path)
+        assert sem_state_after is not None, "Семафор всё ещё существует"

ydb/coordination/сoordination_session.py

@@ -1,41 +1,59 @@
 import time
 
+import ydb
 from ydb._grpc.v5.protos.ydb_coordination_pb2 import SessionRequest
 from ydb._grpc.v5.ydb_coordination_v1_pb2_grpc import CoordinationServiceStub
+from ydb._utilities import SyncResponseIterator
 
 
 class CoordinationSession:
     def __init__(self, driver: "ydb.Driver"):
         self._driver = driver
-
-    def acquire_semaphore(self, name: str, count: int = 1, timeout: int = 5000):
-        req = SessionRequest(
+        self._session_id = None
+
+    def _ensure_session(self):
+        if self._session_id is None:
+            req = SessionRequest(session_start=SessionRequest.SessionStart())
+            stream_it = self._driver(
+                req,
+                CoordinationServiceStub,
+                "Session",
+            )
+            iterator = SyncResponseIterator(stream_it, lambda resp: resp)
+            first_resp = next(iterator)
+            self._session_id = first_resp.session_started.session_id
+        return self._session_id
+
+    def acquire_semaphore(self, path: str, count: int = 1, timeout_millis: int = 5000):
+        session_id = self._ensure_session()
+        acquire_req = SessionRequest(
             acquire_semaphore=SessionRequest.AcquireSemaphore(
-                name=name,
+                name=path,
                 count=count,
-                timeout_millis=timeout,
+                timeout_millis=timeout_millis,
                 req_id=int(time.time() * 1000),
                 data=b"",
                 ephemeral=True,
             ),
-            session_start=SessionRequest.SessionStart()
+            session_start=SessionRequest.SessionStart(session_id=session_id)
         )
-
-        res_iter = self._driver(req, CoordinationServiceStub, "Session")
-        res = next(res_iter)
-        acquire_result = getattr(res, "acquire_semaphore_result", None)
-
-        if not acquire_result or not acquire_result.acquired:
-            raise RuntimeError(f"Failed to acquire semaphore {name}")
-
-        return res.session_started.session_id
-
-    def release_semaphore(self, name: str, session_id: int):
-        req = SessionRequest(
+        stream_it = self._driver(acquire_req, CoordinationServiceStub, "Session")
+        iterator = SyncResponseIterator(stream_it, lambda resp: resp)
+        resp = next(iterator)
+        result = getattr(resp, "acquire_semaphore_result", None)
+        if not result or not result.acquired:
+            raise RuntimeError(f"Failed to acquire semaphore {path}")
+
+    def release_semaphore(self, path: str):
+        if self._session_id is None:
+            return
+        release_req = SessionRequest(
             release_semaphore=SessionRequest.ReleaseSemaphore(
-                name=name,
+                name=path,
                 req_id=int(time.time() * 1000),
             ),
             session_stop=SessionRequest.SessionStop()
         )
-        self._driver(req, CoordinationServiceStub, "Session")
+        stream_it = self._driver(release_req, CoordinationServiceStub, "Session")
+        SyncResponseIterator(stream_it, lambda resp: resp)
+        self._session_id = None