Merge pull request 'блокировки ячеек для websockets' (#23) from websocket_locks into test
Reviewed-on: #23 Reviewed-by: Raykov-MS <RaykovMS@avt.rshb.ru>
This commit is contained in:
commit
30e3426408
@ -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(
|
||||
@ -530,6 +747,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",
|
||||
|
||||
@ -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(
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user