diff --git a/app/chat/router.py b/app/chat/router.py index 6372bb3..0244868 100644 --- a/app/chat/router.py +++ b/app/chat/router.py @@ -211,7 +211,7 @@ async def unpinn_chat(chat_id: int, user: SUser = Depends(check_verificated_user ) async def get_pinned_chats(user: SUser = Depends(check_verificated_user_with_exc), uow=Depends(UnitOfWork)): async with uow: - return await uow.chat.get_pinned_chats(user_id=user.id) + return {"pinned_chats": await uow.chat.get_pinned_chats(user_id=user.id)} @router.get( diff --git a/app/chat/shemas.py b/app/chat/shemas.py index 8880596..fe34ec6 100644 --- a/app/chat/shemas.py +++ b/app/chat/shemas.py @@ -52,7 +52,7 @@ class SAllowedChats(BaseModel): class SPinnedChats(BaseModel): - pinned_chat: list[SChat] + pinned_chats: list[SChat] class SPinnedMessages(BaseModel): diff --git a/app/chat/websocket.py b/app/chat/websocket.py index 8b6d73e..2e06bcd 100644 --- a/app/chat/websocket.py +++ b/app/chat/websocket.py @@ -13,6 +13,13 @@ from app.users.schemas import SUser class ConnectionManager: def __init__(self): self.active_connections: dict[int, list[WebSocket]] = {} + self.message_methods = { + "send": self.send, + "delete": self.delete, + "edit": self.edit, + "pin": self.pin, + "unpin": self.unpin, + } async def connect(self, chat_id: int, websocket: WebSocket): await websocket.accept() @@ -24,60 +31,65 @@ class ConnectionManager: self.active_connections[chat_id].remove(websocket) async def broadcast(self, uow: UnitOfWork, user_id: int, chat_id: int, message: dict): - if "flag" not in message: + try: + new_message = await self.message_methods[message["flag"]](uow, user_id, chat_id, message) + for websocket in self.active_connections[chat_id]: + await websocket.send_json(new_message) + except KeyError: raise IncorrectDataException - if message["flag"] == "send": - message = SSendMessage.model_validate(message) + async def send(self, uow: UnitOfWork, user_id: int, chat_id: int, message: dict) -> dict: + message = SSendMessage.model_validate(message) - new_message = await self.add_message_to_database(uow=uow, user_id=user_id, chat_id=chat_id, message=message) - new_message = new_message.model_dump() - new_message["created_at"] = new_message["created_at"].isoformat() - new_message["flag"] = "send" + new_message = await self.add_message_to_database(uow=uow, user_id=user_id, chat_id=chat_id, message=message) + new_message = new_message.model_dump() + new_message["created_at"] = new_message["created_at"].isoformat() + new_message["flag"] = "send" + return new_message - elif message["flag"] == "delete": - message = SDeleteMessage.model_validate(message) + async def delete(self, uow: UnitOfWork, user_id: int, _: int, message: dict) -> dict: + message = SDeleteMessage.model_validate(message) - if message.user_id != user_id: - raise UserDontHavePermissionException + if message.user_id != user_id: + raise UserDontHavePermissionException - deleted_message = await self.delete_message(uow=uow, message_id=message.id) - new_message = {"deleted_message": deleted_message, "id": message.id, "flag": "delete"} + deleted_message = await self.delete_message(uow=uow, message_id=message.id) + new_message = {"deleted_message": deleted_message, "id": message.id, "flag": "delete"} + return new_message - elif message["flag"] == "edit": - message = SEditMessage.model_validate(message) + async def edit(self, uow: UnitOfWork, user_id: int, _: int, message: dict) -> dict: + message = SEditMessage.model_validate(message) - if message.user_id != user_id: - raise UserDontHavePermissionException + if message.user_id != user_id: + raise UserDontHavePermissionException - edited_message = await self.edit_message( - uow=uow, message_id=message.id, new_message=message.new_message, image_url=message.new_image_url - ) - new_message = { - "flag": "edit", - "id": message.id, - "edited_message": edited_message, - "new_message": message.new_message, - "new_image_url": message.new_image_url, - } + edited_message = await self.edit_message( + uow=uow, message_id=message.id, new_message=message.new_message, image_url=message.new_image_url + ) + new_message = { + "flag": "edit", + "id": message.id, + "edited_message": edited_message, + "new_message": message.new_message, + "new_image_url": message.new_image_url, + } + return new_message - elif message["flag"] == "pin": - message = SPinMessage.model_validate(message) - pinned_message = await self.pin_message(uow=uow, chat_id=chat_id, user_id=message.user_id, message_id=message.id) - new_message = pinned_message.model_dump() - new_message["created_at"] = new_message["created_at"].isoformat() - new_message["flag"] = "pin" + async def pin(self, uow: UnitOfWork, _: int, chat_id: int, message: dict) -> dict: + message = SPinMessage.model_validate(message) + pinned_message = await self.pin_message( + uow=uow, chat_id=chat_id, user_id=message.user_id, message_id=message.id + ) + new_message = pinned_message.model_dump() + new_message["created_at"] = new_message["created_at"].isoformat() + new_message["flag"] = "pin" + return new_message - elif message["flag"] == "unpin": - message = SUnpinMessage.model_validate(message) - unpinned_message = await self.unpin_message(uow=uow, chat_id=chat_id, message_id=message.id) - new_message = {"flag": "pin", "id": unpinned_message} - - else: - raise IncorrectDataException - - for websocket in self.active_connections[chat_id]: - await websocket.send_json(new_message) + async def unpin(self, uow: UnitOfWork, _: int, chat_id: int, message: dict) -> dict: + message = SUnpinMessage.model_validate(message) + unpinned_message = await self.unpin_message(uow=uow, chat_id=chat_id, message_id=message.id) + new_message = {"flag": "unpin", "id": unpinned_message} + return new_message @staticmethod async def add_message_to_database(uow: UnitOfWork, user_id: int, chat_id: int, message: SSendMessage) -> SMessage: