From 1d22cfd3dbfe2a80dc32c806017bad608a777d33 Mon Sep 17 00:00:00 2001
From: h <h@kotikot.com>
Date: Mon, 25 Aug 2025 00:49:06 +0300
Subject: [PATCH] feat(bot): add basic in-db message storage (no attachments
 yet)

---
 src/bot/handlers/__init__.py       |  5 ++++
 src/bot/middlewares/__init__.py    |  1 +
 src/bot/middlewares/saver.py       | 46 ++++++++++++++++++++++++++++++
 src/utils/db/__init__.py           |  2 ++
 src/utils/db/models/__init__.py    |  1 +
 src/utils/db/models/message_log.py | 44 ++++++++++++++++++++++++++++
 6 files changed, 99 insertions(+)
 create mode 100644 src/bot/middlewares/__init__.py
 create mode 100644 src/bot/middlewares/saver.py
 create mode 100644 src/utils/db/models/message_log.py

diff --git a/src/bot/handlers/__init__.py b/src/bot/handlers/__init__.py
index 93e9d54..cb5d1b6 100644
--- a/src/bot/handlers/__init__.py
+++ b/src/bot/handlers/__init__.py
@@ -1,5 +1,7 @@
 from aiogram import Router
 
+from bot.middlewares import MessageSaverMiddleware
+
 from . import (
     initialize,
     message,
@@ -9,6 +11,9 @@ from . import (
 router = Router()
 
 
+router.message.outer_middleware(MessageSaverMiddleware())
+
+
 router.include_routers(
     start.router,
     initialize.router,
diff --git a/src/bot/middlewares/__init__.py b/src/bot/middlewares/__init__.py
new file mode 100644
index 0000000..adad009
--- /dev/null
+++ b/src/bot/middlewares/__init__.py
@@ -0,0 +1 @@
+from .saver import MessageSaverMiddleware
diff --git a/src/bot/middlewares/saver.py b/src/bot/middlewares/saver.py
new file mode 100644
index 0000000..57c13ca
--- /dev/null
+++ b/src/bot/middlewares/saver.py
@@ -0,0 +1,46 @@
+from typing import Any, Awaitable, Callable, Dict, TypeVar
+
+from aiogram import BaseMiddleware, types
+
+from utils.db.models.message_log import MessageAttachment, MessageLog
+
+
+class MessageSaverMiddleware(BaseMiddleware):
+    async def __call__(
+        self,
+        handler: Callable[[types.TelegramObject, Dict[str, Any]], Awaitable[Any]],
+        event: types.Message,
+        data: Dict[str, Any],
+    ):
+        await MessageLog(
+            chat_id=event.chat.id,
+            message_id=event.message_id,
+            user_id=event.from_user.id,
+            username=event.from_user.username,
+            first_name=event.from_user.first_name,
+            last_name=event.from_user.last_name,
+            text=event.text or event.caption,
+            timestamp=event.date,
+            reply_to_message_id=(
+                event.reply_to_message.message_id if event.reply_to_message else None
+            ),
+            forward_from_type=(
+                event.forward_origin.type if event.forward_origin else None
+            ),
+            original_message_date=(
+                event.forward_origin.date if event.forward_origin else None
+            ),
+            original_message_author=(
+                event.forward_origin.sender_user.full_name
+                if (event.forward_origin and event.forward_origin.sender_user)
+                else None
+            ),
+            original_message_author_id=(
+                event.forward_origin.sender_user.id
+                if (event.forward_origin and event.forward_origin.sender_user)
+                else None
+            ),
+            # attachment_file_ids=sum([] for attachment in [event.photo])
+        ).insert()
+
+        return await handler(event, data)
diff --git a/src/utils/db/__init__.py b/src/utils/db/__init__.py
index 9d529e9..00e8e22 100644
--- a/src/utils/db/__init__.py
+++ b/src/utils/db/__init__.py
@@ -9,6 +9,7 @@ client = AsyncMongoClient(env.db.connection_url)
 async def init_db():
     from .models import (
         DynamicConfig,
+        MessageLog,
         RespondSession,
         ReviewSession,
     )
@@ -19,5 +20,6 @@ async def init_db():
             DynamicConfig,
             RespondSession,
             ReviewSession,
+            MessageLog,
         ],
     )
diff --git a/src/utils/db/models/__init__.py b/src/utils/db/models/__init__.py
index 2255ab5..89761e4 100644
--- a/src/utils/db/models/__init__.py
+++ b/src/utils/db/models/__init__.py
@@ -1,2 +1,3 @@
 from .config import DynamicConfig
+from .message_log import MessageLog
 from .session import RespondSession, ReviewSession
diff --git a/src/utils/db/models/message_log.py b/src/utils/db/models/message_log.py
new file mode 100644
index 0000000..e592201
--- /dev/null
+++ b/src/utils/db/models/message_log.py
@@ -0,0 +1,44 @@
+import datetime
+from typing import Annotated, Optional
+
+from aiogram.enums.message_origin_type import MessageOriginType
+from beanie import Document, Indexed
+from pydantic import BaseModel, Field
+
+
+class MessageAttachment(BaseModel):
+    file_id: str
+    file_size: int
+    file_name: str
+    mime_type: str
+
+
+class MessageLogBase(BaseModel):
+    chat_id: Annotated[int, Indexed()]
+    message_id: Annotated[int, Indexed()]
+
+    user_id: Annotated[int, Indexed()]
+    username: Optional[str] = None
+    first_name: str
+    last_name: Optional[str] = None
+
+    text: Optional[str]
+    timestamp: datetime.datetime
+    reply_to_message_id: Optional[int] = None
+
+    forward_from_type: Optional[MessageOriginType] = None
+    original_message_date: Optional[datetime.datetime] = None
+    original_message_author: Optional[str] = None
+    original_message_author_id: Optional[int] = None
+
+    attachment_file_ids: list[str] = Field(default_factory=list)
+
+
+class MessageLog(MessageLogBase, Document):
+    class Settings:
+        name = "message_logs"
+        indexes = [[("chat_id", 1), ("timestamp", -1)]]
+
+    @classmethod
+    def get_n_messages(cls, chat_id: int, n: int):
+        return cls.find(cls.chat_id == chat_id).sort("timestamp", "-").limit(n)