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

This commit is contained in:
tsygankoviva 2026-05-29 13:52:20 +03:00
parent 3f6d59ba06
commit 38b0033c4a
2 changed files with 231 additions and 6 deletions

View File

@ -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",

View File

@ -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(