100 lines
3.4 KiB
Python
100 lines
3.4 KiB
Python
from fastapi import WebSocket, WebSocketDisconnect
|
|
|
|
from app.exceptions import IncorrectDataException, UserDontHavePermissionException
|
|
from app.services.message_service import MessageService
|
|
from app.users.auth import AuthService
|
|
from app.chat.router import router
|
|
from app.chat.shemas import SSendMessage, SMessage, SDeleteMessage, SEditMessage
|
|
|
|
|
|
class ConnectionManager(WebSocket):
|
|
def __init__(self):
|
|
self.active_connections: dict[int, list[WebSocket]] = {}
|
|
|
|
async def connect(self, chat_id: int, websocket: WebSocket):
|
|
await websocket.accept()
|
|
if chat_id not in self.active_connections:
|
|
self.active_connections[chat_id] = []
|
|
self.active_connections[chat_id].append(websocket)
|
|
|
|
def disconnect(self, chat_id: int, websocket: WebSocket):
|
|
self.active_connections[chat_id].remove(websocket)
|
|
|
|
async def broadcast(self, user_id: int, chat_id: int, message: dict):
|
|
if "flag" not in message:
|
|
raise IncorrectDataException
|
|
|
|
if message["flag"] == "send":
|
|
message = SSendMessage.model_validate(message)
|
|
|
|
new_message = await self.add_message_to_database(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"
|
|
|
|
elif message["flag"] == "delete":
|
|
message = SDeleteMessage.model_validate(message)
|
|
|
|
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":
|
|
message = SEditMessage.model_validate(message)
|
|
|
|
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",
|
|
"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: SSendMessage) -> SMessage:
|
|
new_message = await MessageService.send_message(
|
|
user_id=user_id, chat_id=chat_id, message=message.message, image_url=message.image_url
|
|
)
|
|
if message.answer:
|
|
new_message = await MessageService.add_answer(self_id=new_message.id, answer_id=message.answer)
|
|
return new_message
|
|
|
|
@staticmethod
|
|
async def delete_message(message_id: int) -> bool:
|
|
new_message = await MessageService.delete_message(message_id)
|
|
return new_message
|
|
|
|
@staticmethod
|
|
async def edit_message(message_id: int, new_message: str, image_url: str) -> bool:
|
|
new_message = await MessageService.edit_message(
|
|
message_id=message_id, new_message=new_message, new_image_url=image_url
|
|
)
|
|
return new_message
|
|
|
|
|
|
manager = ConnectionManager()
|
|
|
|
|
|
@router.websocket("/ws/{chat_id}")
|
|
async def websocket_endpoint(chat_id: int, user_id: int, websocket: WebSocket):
|
|
await AuthService.check_verificated_user_with_exc(user_id=user_id)
|
|
await AuthService.validate_user_access_to_chat(user_id=user_id, chat_id=chat_id)
|
|
await manager.connect(chat_id, websocket)
|
|
try:
|
|
while True:
|
|
data = await websocket.receive_json()
|
|
|
|
await manager.broadcast(user_id=user_id, chat_id=chat_id, message=data)
|
|
except WebSocketDisconnect:
|
|
manager.disconnect(chat_id, websocket)
|