Изменил вебсокет

app/exceptions.py

@@ -59,7 +59,7 @@ class MessageNotFoundException(BlackPhoenixException):
     detail = "Сообщение не найдено"
 
 
-class PasswordsМismatchException(BlackPhoenixException):
+class PasswordsMismatchException(BlackPhoenixException):
     status_code = status.HTTP_409_CONFLICT
     detail = "Пароли не совпадают"
 
@@ -97,3 +97,11 @@ class UserMustConfirmEmailException(BlackPhoenixException):
 class SomethingWentWrongException(BlackPhoenixException):
     status_code = status.HTTP_500_INTERNAL_SERVER_ERROR
     detail = "Что-то пошло не так"
+
+
+class IncorrectDataException(BlackPhoenixException):
+    status_code = status.HTTP_409_CONFLICT
+    detail = "Ты передал какую-то хуйню"
+
+
+

app/users/chat/dao.py

@@ -1,7 +1,7 @@
 from sqlalchemy import insert, select, update, and_, delete
 
 from app.dao.base import BaseDAO
-from app.database import async_session_maker, engine
+from app.database import async_session_maker, engine  # noqa
 from app.exceptions import UserAlreadyInChatException, UserAlreadyPinnedChatException
 from app.users.models import Users
 from app.users.chat.models import Chats, Messages, UsersXChats, PinnedChats, PinnedMessages, Answers
@@ -116,6 +116,14 @@ class ChatDAO(BaseDAO):
             result = [dict(res) for res in result]
             return result
 
+    @staticmethod
+    async def edit_message(message_id: int, new_message: str, new_image_url: str) -> bool:
+        query = update(Messages).where(Messages.id == message_id).values(message=new_message, image_url=new_image_url)
+        async with async_session_maker() as session:
+            await session.execute(query)
+            await session.commit()
+            return True
+
     @staticmethod
     async def delete_chat(chat_id: int) -> bool:
         query = update(Chats).where(Chats.id == chat_id).values(visibility=False)

app/users/chat/router.py

@@ -45,6 +45,7 @@ async def add_message_to_chat(
     )
     return send_message_to_chat
 
+
 @router.delete("/delete_message", status_code=status.HTTP_204_NO_CONTENT)
 async def delete_message_from_chat(
         message_id: int,

app/users/chat/websocket.py

@@ -1,5 +1,6 @@
 from fastapi import WebSocket, WebSocketDisconnect
 
+from app.exceptions import IncorrectDataException, UserDontHavePermissionException
 from app.users.chat.dao import ChatDAO
 from app.users.auth import validate_user_access_to_chat, check_verificated_user_with_exc
 from app.users.chat.router import router
@@ -19,29 +20,46 @@ class ConnectionManager(WebSocket):
         self.active_connections[chat_id].remove(websocket)
 
     async def broadcast(self, user_id: int, chat_id: int, message: dict):
-        new_message = await self.add_message_to_database(
-            user_id=user_id, chat_id=chat_id, message=message['message'], image_url=message['image_url']
-        )
-        new_message = dict(new_message)
-        new_message['created_at'] = new_message['created_at'].isoformat()
+        if "flag" not in message:
+            raise IncorrectDataException
+
+        if message['flag'] == 'send':
+            new_message = await self.add_message_to_database(
+                user_id=user_id, chat_id=chat_id, message=message['message'], image_url=message['image_url']
+            )
+            new_message = dict(new_message)
+            new_message['created_at'] = new_message['created_at'].isoformat()
+            new_message['flag'] = 'send'
+        elif message['flag'] == 'delete':
+            if message["user_id"] != user_id:
+                raise UserDontHavePermissionException
+            deleted_message = await self.delete_message(message["id"])
+            new_message = {'deleted_message': deleted_message, 'id': message['id'], "flag": "delete"}
+        elif message['flag'] == 'edit':
+            if message["user_id"] != user_id:
+                raise UserDontHavePermissionException
+            edited_message = await self.edit_message(message['id'], message['new_message'], message['new_image_url'])
+            new_message = {'edited_message': edited_message, "flag": "edit", "message_id": message["id"],
+                           "new_message": message["new_message"], "new_image_url": message["new_image_url"]}
+        else:
+            raise IncorrectDataException
+
         for websocket in self.active_connections[chat_id]:
             await websocket.send_json(new_message)
 
     @staticmethod
-    async def add_message_to_database(user_id: int, chat_id: int, message: str, image_url: str):
+    async def add_message_to_database(user_id: int, chat_id: int, message: str, image_url: str) -> dict:
         new_message = await ChatDAO.send_message(user_id=user_id, chat_id=chat_id, message=message, image_url=image_url)
         return new_message[0]
 
     @staticmethod
-    async def delete_message(user_id: int, chat_id: int, message: str, image_url: str):
-        result = await ChatDAO.send_message(user_id=user_id, chat_id=chat_id, message=message, image_url=image_url)
-        new_message = await ChatDAO.get_message_by_id(message_id=result)
+    async def delete_message(message_id: int) -> bool:
+        new_message = await ChatDAO.delete_message(message_id)
         return new_message
 
     @staticmethod
-    async def change_message(user_id: int, chat_id: int, message: str, image_url: str):
-        result = await ChatDAO.send_message(user_id=user_id, chat_id=chat_id, message=message, image_url=image_url)
-        new_message = await ChatDAO.get_message_by_id(message_id=result)
+    async def edit_message(message_id: int, new_message: str, image_url: str) -> bool:
+        new_message = await ChatDAO.edit_message(message_id=message_id, new_message=new_message, new_image_url=image_url)
         return new_message
 
 
@@ -61,4 +79,3 @@ async def websocket_endpoint(chat_id: int, user_id: int, websocket: WebSocket):
     except WebSocketDisconnect:
         manager.disconnect(chat_id, websocket)
 
-

app/users/router.py

@@ -3,7 +3,7 @@ from fastapi.responses import RedirectResponse
 
 from app.config import settings
 from app.exceptions import UsernameAlreadyInUseException, \
-    IncorrectPasswordException, PasswordsМismatchException, WrongCodeException, UserNotFoundException, \
+    IncorrectPasswordException, PasswordsMismatchException, WrongCodeException, UserNotFoundException, \
     SomethingWentWrongException
 from app.users.auth import get_password_hash, authenticate_user_by_email, \
     create_access_token, verify_password, VERIFICATED_USER, authenticate_user, \
@@ -118,7 +118,7 @@ async def change_password(new_password: SUserPassword, current_user: Users = Dep
     if not verify_password(new_password.current_password, existing_user.hashed_password):
         raise IncorrectPasswordException
     if new_password.new_password != new_password.new_password2:
-        raise PasswordsМismatchException
+        raise PasswordsMismatchException
     hashed_password = get_password_hash(new_password.new_password)
     await UserDAO.change_data(current_user.id, hashed_password=hashed_password)
     send_password_change_email.delay(current_user.username, current_user.email, MODE=settings.MODE)
@@ -152,7 +152,7 @@ async def confirm_password_recovery(user_code: SUserCode):
 @router.post("/password_recovery", status_code=status.HTTP_200_OK, response_model=SPasswordRecovered)
 async def password_recovery(passwords: SUserPasswordChange):
     if passwords.password1 != passwords.password2:
-        raise PasswordsМismatchException
+        raise PasswordsMismatchException
     hashed_password = get_password_hash(passwords.password1)
     username = await UserDAO.change_data(passwords.user_id, hashed_password=hashed_password)
     user = await UserDAO.find_one_or_none(username=username, id=passwords.user_id)