Рефакторинг

app/services/message_service.py

@@ -0,0 +1,25 @@
+from app.users.chat.dao import ChatDAO
+from app.users.chat.shemas import SMessage
+
+
+class MessageService:
+	@staticmethod
+	async def send_message(user_id: int, chat_id: int, message: str, image_url: str | None = None) -> SMessage:
+		new_message = await ChatDAO.send_message(user_id=user_id, chat_id=chat_id, message=message, image_url=image_url)
+		return new_message
+
+	@staticmethod
+	async def add_answer(self_id: int, answer_id: int) -> SMessage:
+		new_message = await ChatDAO.add_answer(self_id=self_id, answer_id=answer_id)
+		return new_message
+
+	@staticmethod
+	async def delete_message(message_id: int) -> bool:
+		new_message = await ChatDAO.delete_message(message_id=message_id)
+		return new_message
+
+	@staticmethod
+	async def edit_message(message_id: int, new_message: str, new_image_url: str) -> bool:
+		new_message = await ChatDAO.edit_message(message_id=message_id, new_message=new_message,
+												 new_image_url=new_image_url)
+		return new_message

app/users/chat/dao.py

@@ -3,6 +3,7 @@ from sqlalchemy import insert, select, update, and_, delete
 from app.dao.base import BaseDAO
 from app.database import async_session_maker, engine  # noqa
 from app.exceptions import UserAlreadyInChatException, UserAlreadyPinnedChatException
+from app.users.chat.shemas import SMessage
 from app.users.models import Users
 from app.users.chat.models import Chats, Messages, UsersXChats, PinnedChats, PinnedMessages, Answers
 
@@ -33,7 +34,7 @@ class ChatDAO(BaseDAO):
 		return True
 
 	@staticmethod
-	async def send_message(user_id: int, chat_id: int, message: str, image_url: str | None = None) -> list[dict]:
+	async def send_message(user_id: int, chat_id: int, message: str, image_url: str | None = None) -> SMessage:
 		inserted_image = (
 			insert(Messages).values(chat_id=chat_id, user_id=user_id, message=message, image_url=image_url)
 			.returning(Messages.id, Messages.message, Messages.image_url, Messages.chat_id,
@@ -49,7 +50,8 @@ class ChatDAO(BaseDAO):
 		async with async_session_maker() as session:
 			result = await session.execute(query)
 			await session.commit()
-		return result.mappings().all()
+			result = result.mappings().one()
+		return SMessage.model_validate(result, from_attributes=True)
 
 	@staticmethod
 	async def get_message_by_id(message_id: int):
@@ -127,9 +129,10 @@ class ChatDAO(BaseDAO):
 			return True
 
 	@staticmethod
-	async def add_answer(self_id: int, answer_id: int) -> list[dict]:
+	async def add_answer(self_id: int, answer_id: int) -> SMessage:
 		answer = (insert(Answers).values(self_id=self_id, answer_id=answer_id)
 				  .returning(Answers.self_id, Answers.answer_id).cte("answer"))
+
 		query = (select(Messages.id, Messages.message, Messages.image_url,
 						Messages.chat_id, Messages.user_id, Messages.created_at,
 						Users.avatar_image, Users.username, Users.avatar_hex, answer.c.self_id, answer.c.answer_id)
@@ -137,10 +140,12 @@ class ChatDAO(BaseDAO):
 				 .join(Users, Users.id == Messages.user_id)
 				 .join(answer, answer.c.answer_id == Messages.id, isouter=True)
 				 .where(Messages.id == answer_id))
+
 		async with async_session_maker() as session:
 			result = await session.execute(query)
 			await session.commit()
-		return result.mappings().all()
+		result = result.scalar_one()
+		return SMessage.model_validate(result)
 
 	@staticmethod
 	async def delete_chat(chat_id: int) -> bool:
@@ -256,7 +261,8 @@ class ChatDAO(BaseDAO):
 		query = (select(Messages.id, Messages.message, Messages.image_url,
 						Messages.chat_id, Messages.user_id,
 						Messages.created_at, Users.avatar_image,
-						Users.username, Users.avatar_hex, Answers.self_id, Answers.answer_id).select_from(PinnedMessages)
+						Users.username, Users.avatar_hex, Answers.self_id, Answers.answer_id).select_from(
+			PinnedMessages)
 				 .join(Messages, PinnedMessages.message_id == Messages.id, isouter=True)
 				 .join(Users, PinnedMessages.user_id == Users.id, isouter=True)
 				 .join(Answers, Answers.self_id == Messages.id, isouter=True)

app/users/chat/shemas.py

@@ -1,3 +1,5 @@
+from datetime import datetime
+
 from pydantic import BaseModel, ConfigDict
 
 
@@ -10,7 +12,7 @@ class SMessage(BaseModel):
 	chat_id: int
 	user_id: int
 	username: str
-	created_at: str
+	created_at: datetime
 	avatar_image: str
 	avatar_hex: str
 	self_id: int | None
@@ -57,5 +59,19 @@ class SChat(BaseModel):
 class SSendMessage(BaseModel):
 	flag: str
 	message: str
-	image_url: str
+	image_url: str | None
 	answer: int | None
+
+
+class SDeleteMessage(BaseModel):
+	flag: str
+	user_id: int
+	id: int
+
+
+class SEditMessage(BaseModel):
+	flag: str
+	user_id: int
+	id: int
+	new_message: str
+	new_image_url: str

app/users/chat/websocket.py

@@ -1,10 +1,10 @@
 from fastapi import WebSocket, WebSocketDisconnect
 
 from app.exceptions import IncorrectDataException, UserDontHavePermissionException
-from app.users.chat.dao import ChatDAO
+from app.services.message_service import MessageService
 from app.users.auth import validate_user_access_to_chat, check_verificated_user_with_exc
 from app.users.chat.router import router
-from app.users.chat.shemas import SSendMessage
+from app.users.chat.shemas import SSendMessage, SMessage, SDeleteMessage, SEditMessage
 
 
 class ConnectionManager(WebSocket):
@@ -28,20 +28,28 @@ class ConnectionManager(WebSocket):
 			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 = dict(new_message)
+			new_message = new_message.model_dump()
 			new_message["created_at"] = new_message["created_at"].isoformat()
 			new_message["flag"] = "send"
+
 		elif message["flag"] == "delete":
-			if message["user_id"] != user_id:
+			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"}
+
+			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:
+			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"]}
+
+			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
 
@@ -49,21 +57,23 @@ class ConnectionManager(WebSocket):
 			await websocket.send_json(new_message)
 
 	@staticmethod
-	async def add_message_to_database(user_id: int, chat_id: int, message: SSendMessage) -> dict:
-		new_message = await ChatDAO.send_message(
-			user_id=user_id, chat_id=chat_id, message=message.message, image_url=message.image_url)
+	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 ChatDAO.add_answer(self_id=new_message[0]["id"], answer_id=message.answer)
-		return new_message[0]
+			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 ChatDAO.delete_message(message_id)
+		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 ChatDAO.edit_message(message_id=message_id, new_message=new_message, new_image_url=image_url)
+		new_message = await MessageService.edit_message(message_id=message_id, new_message=new_message,
+														new_image_url=image_url)
 		return new_message
 
 
@@ -82,4 +92,3 @@ async def websocket_endpoint(chat_id: int, user_id: int, websocket: WebSocket):
 			await manager.broadcast(user_id=user_id, chat_id=chat_id, message=data)
 	except WebSocketDisconnect:
 		manager.disconnect(chat_id, websocket)
-