блокировки ячеек для websockets

api/src/api/v1/websocket.py

@@ -72,6 +72,7 @@ class ConnectionManager:
             ConnectionKeyEnum.FORM: {},
             ConnectionKeyEnum.PROJECT: {},
         }
+        self.cell_locks: dict[Any, int] = {}
         
 
     async def connect(
@@ -138,6 +139,8 @@ class ConnectionManager:
                 self.connections[con_key][key].remove(con_info)
         if not len(self.connections[con_key][key]):
             del self.connections[con_key][key]
+        
+        self.release_locks_for_user(con_key=con_key, user_id=con_info.user_id, form_key=kwargs)
         await websocket.close(
             code=code, reason=reason,
         )
@@ -151,8 +154,9 @@ class ConnectionManager:
     ):
         key = frozenset(kwargs.items())
 
-        if key not in self.connections:
+        if key not in self.connections[con_key]:
             return
+        
         for con_info in self.connections[con_key][key]:
             if con_info.user_id != user_id:
                 try:
@@ -181,6 +185,115 @@ class ConnectionManager:
     async def get_data(self, websocket: WebSocket) -> dict:
         return loads(await websocket.receive_text())
 
+    def acquire_cell_lock(
+            self, 
+            con_key: ConnectionKeyEnum, 
+            user_id: int,
+            cell_key: dict,
+            form_key: dict,
+    ) -> bool:
+        if not cell_key:
+            return False
+        cell_key = frozenset(cell_key.items())
+        
+        lock_key = (con_key, frozenset(form_key.items()), cell_key)
+        lock_owner_id = self.cell_locks.get(lock_key)
+        if lock_owner_id is not None and lock_owner_id != user_id:
+            return False
+        self.cell_locks[lock_key] = user_id
+        return True
+
+    def release_cell_lock(
+            self, 
+            con_key: ConnectionKeyEnum, 
+            user_id: int,
+            cell_key: dict,
+            form_key: dict,
+    ):
+        if not cell_key:
+            return
+        cell_key = frozenset(cell_key.items())
+
+        lock_key = (con_key, frozenset(form_key.items()), cell_key)
+        lock_owner_id = self.cell_locks.get(lock_key)
+        if lock_owner_id == user_id:
+            self.cell_locks.pop(lock_key, None)
+
+    def release_locks_for_user(
+            self, 
+            con_key: ConnectionKeyEnum, 
+            user_id: int,
+            form_key: dict
+        ):
+        filter_key = (con_key, frozenset(form_key.items()))
+        lock_keys = [
+            lock_key
+            for lock_key, lock_owner_id in self.cell_locks.items()
+            if lock_key[:-1] == filter_key and lock_owner_id == user_id
+        ]
+        for lock_key in lock_keys:
+            self.cell_locks.pop(lock_key, None)
+
+    def is_locked_by_other(
+            self, 
+            con_key: ConnectionKeyEnum, 
+            user_id: int,
+            cell_key: dict,
+            form_key,
+    ) -> bool:
+        if not cell_key:
+            return False
+        cell_key = frozenset(cell_key.items())
+        lock_owner_id = self.cell_locks.get((con_key, frozenset(form_key.items()), cell_key))
+        return lock_owner_id is not None and lock_owner_id != user_id
+
+    def is_row_locked_by_other(
+            self, 
+            con_key: ConnectionKeyEnum, 
+            user_id: int,
+            line_id: int,
+    ) -> bool:
+        if line_id is None:
+            return False
+
+        for (dict_con_key, form_key, cell_key), lock_owner_id in self.cell_locks.items():
+            if dict_con_key != con_key:
+                continue
+            
+            if form_key != form_key or lock_owner_id == user_id:
+                continue
+
+            for k, v in cell_key:
+                if k == "line_id" and v == line_id:
+                    return True
+        return False
+
+    def release_locks_for_row(
+            self, 
+            con_key: ConnectionKeyEnum, 
+            user_id: int,
+            line_id: int,
+    ) -> bool:
+        
+        if line_id is None:
+            return False
+        
+        lock_keys = []
+        for lock_key, lock_owner_id in self.cell_locks.items():
+            dict_con_key, form_key, cell_key = lock_key
+            if dict_con_key != con_key:
+                continue
+            
+            if form_key != form_key or lock_owner_id == user_id:
+                continue
+
+            for k, v in cell_key.items():
+                if k == "line_id" and v == line_id:
+                    lock_keys.append(lock_key)
+        
+        for lock_key in lock_keys:
+            self.cell_locks.pop(lock_key, None)
+
 
 class FormEventProcess:
     def __init__(self, db: AsyncSession):
@@ -477,6 +590,19 @@ async def login(websocket: WebSocket, **kwargs) -> int:
         return user.id
 
 
+
+
+def resolve_cell_key(event_data: dict) -> dict | None:
+    data = event_data.get("data")
+
+    if not data:
+        return None
+    return {
+        "line_id": data["line_id"],
+        "column": data["column"],
+    }
+
+
 async def process_websocket(
     websocket: WebSocket,
     processor_cls,
@@ -487,6 +613,12 @@ async def process_websocket(
         **kwargs,
     )
 
+    
+    kwargs_process = kwargs.copy()
+    if "con_key" in kwargs_process:
+        kwargs_process.pop("con_key")
+    con_key = kwargs.get("con_key")
+
     try:
         user_id = await login(
             websocket=websocket,
@@ -504,10 +636,95 @@ async def process_websocket(
         while True:
             data = loads(await websocket.receive_text())
 
-            kwargs_process = kwargs.copy()
-            if "con_key" in kwargs_process:
-                kwargs_process.pop("con_key")
             try:
+                event_name = data.get("event")
+                cell_key = resolve_cell_key(data)
+
+                match event_name:
+                    case "cell_edit_start":
+                        if not manager.acquire_cell_lock(
+                            con_key=con_key,
+                            user_id=user_id,
+                            cell_key=cell_key,
+                            form_key=kwargs_process,
+                        ):
+                            data["error"] = "Ячейка уже редактируется другим пользователем"
+                            await manager.send_back(
+                                dumps(data, default=_convert_error, ensure_ascii=False), 
+                                websocket,
+                            )
+                            continue
+                        await manager.send_back(
+                            dumps(data, default=_convert_error, ensure_ascii=False), 
+                            websocket,
+                        )
+                        await manager.broadcast_to_other(
+                            message=dumps(data, default=_convert_error, ensure_ascii=False),
+                            user_id=user_id,
+                            con_key=con_key,
+                            **kwargs_process,
+                        )
+                        continue
+                    case "cell_edit_end":
+                        manager.release_cell_lock(
+                            con_key=con_key,
+                            user_id=user_id,
+                            cell_key=cell_key,
+                            form_key=kwargs_process,
+                        )
+                        await manager.send_back(
+                            dumps(data, default=_convert_error, ensure_ascii=False), 
+                            websocket,
+                        )
+                        await manager.broadcast_to_other(
+                            message=dumps(data, default=_convert_error, ensure_ascii=False),
+                            user_id=user_id,
+                            con_key=con_key,
+                            **kwargs_process,
+                        )
+                        continue
+                    case "cell_updated":
+                        if manager.is_locked_by_other(
+                                con_key=con_key,
+                                user_id=user_id,
+                                cell_key=cell_key,
+                                form_key=kwargs_process,
+                        ):
+                            data["error"] = "Ячейка уже редактируется другим пользователем"
+                            await manager.send_back(
+                                dumps(data, default=_convert_error, ensure_ascii=False), 
+                                websocket,
+                            )
+                            continue
+                        if cell_key and not manager.acquire_cell_lock(
+                            con_key=con_key,
+                            user_id=user_id,
+                            cell_key=cell_key,
+                            form_key=kwargs_process,
+                        ):
+                            data["error"] = "Ячейка уже редактируется другим пользователем"
+                            await manager.send_back(
+                                dumps(data, default=_convert_error, ensure_ascii=False), 
+                                websocket,
+                            )
+                            continue
+                    case "row_deleted":
+                        if manager.is_row_locked_by_other(
+                            con_key=con_key,
+                            user_id=user_id,
+                            line_id=data["data"]["line_id"]
+                        ):
+                            data["error"] = (
+                                "Нельзя удалить строку: в этой строке есть редактируемые ячейки"
+                            )
+                            await manager.send_back(
+                                dumps(data, default=_convert_error, ensure_ascii=False), 
+                                websocket,
+                            )
+                            continue
+                    case _:
+                        pass
+
                 async with get_db_session() as db:
                     processor = processor_cls(db)
                     data["result"] = await processor.process(
@@ -517,6 +734,13 @@ async def process_websocket(
                     )
                     await db.commit()
                 
+                if data.get("event") == "row_deleted":
+                    manager.release_locks_for_row(
+                        con_key=con_key,
+                        user_id=user_id,
+                        line_id=data["data"]["line_id"]
+                    )
+                
                 if data.get("event") in [
                     "cell_updated",
                     "row_added",

api/src/core/exception_handlers.py

@@ -1,6 +1,6 @@
 from fastapi import FastAPI, HTTPException, status
 from fastapi.responses import JSONResponse
-from sqlalchemy.exc import IntegrityError, OperationalError, SQLAlchemyError
+from sqlalchemy.exc import IntegrityError, OperationalError, SQLAlchemyError, DBAPIError
 
 from src.core.errors import (
     AccessDeniedException,
@@ -136,6 +136,7 @@ def register_exception_handlers(app: FastAPI) -> None:
         )
 
     @app.exception_handler(SQLAlchemyError)
+    @app.exception_handler(DBAPIError)
     async def sqlalchemy_exception_handler(request, exc: SQLAlchemyError):
         status_code, code, message, field = map_sqlalchemy_error(exc)
         return JSONResponse(