diff --git a/api/src/api/v1/websocket.py b/api/src/api/v1/websocket.py index 402e6d8..7a0ab90 100644 --- a/api/src/api/v1/websocket.py +++ b/api/src/api/v1/websocket.py @@ -73,6 +73,7 @@ class ConnectionManager: ConnectionKeyEnum.FORM: {}, ConnectionKeyEnum.PROJECT: {}, } + self.cell_locks: dict[Any, int] = {} async def connect( @@ -139,6 +140,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, ) @@ -152,8 +155,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: @@ -182,6 +186,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): @@ -490,6 +603,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, @@ -500,6 +626,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, @@ -517,10 +649,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( @@ -529,7 +746,14 @@ async def process_websocket( **kwargs_process ) 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", diff --git a/api/src/core/exception_handlers.py b/api/src/core/exception_handlers.py index a524dbf..9047683 100644 --- a/api/src/core/exception_handlers.py +++ b/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(