feat(db): persist welcome message map for ban-by-reply across restarts

Store welcome-message-to-user mappings in a new welcome_messages table
so that /ban reply works even after a bot restart. The in-memory cache
is still used as a fast path; the database is a fallback. Cleanup
happens both in memory and in the DB when welcome messages are
auto-deleted.

Author: 0xMett

schema.sql

@@ -92,3 +92,11 @@ ALTER TABLE group_settings
 -- Add title column to bot_chats for group name lookup (idempotent).
 ALTER TABLE bot_chats
     ADD COLUMN IF NOT EXISTS title TEXT;
+
+CREATE TABLE IF NOT EXISTS welcome_messages (
+    chat_id BIGINT NOT NULL,
+    message_id BIGINT NOT NULL,
+    user_id BIGINT NOT NULL,
+    created_at TIMESTAMPTZ DEFAULT NOW(),
+    PRIMARY KEY (chat_id, message_id)
+);

src/python_italy_bot/db/base.py

@@ -195,6 +195,25 @@ def set_welcome_delay(self, chat_id: int, minutes: int | None) -> None:
         """Set welcome message auto-delete delay. None to reset to default."""
         ...
 
+    # -- Welcome message tracking (ban-by-reply) --
+
+    @abstractmethod
+    def store_welcome_message(
+        self, chat_id: int, message_id: int, user_id: int
+    ) -> None:
+        """Store mapping from a welcome message to the user who triggered it."""
+        ...
+
+    @abstractmethod
+    def get_welcome_message_user(self, chat_id: int, message_id: int) -> int | None:
+        """Get the user_id associated with a welcome message, or None."""
+        ...
+
+    @abstractmethod
+    def delete_welcome_message(self, chat_id: int, message_id: int) -> None:
+        """Remove a welcome message mapping."""
+        ...
+
 
 class AsyncRepository(ABC):
     """Abstract interface for data persistence (async)."""
@@ -386,6 +405,27 @@ async def set_welcome_delay(self, chat_id: int, minutes: int | None) -> None:
         """Set welcome message auto-delete delay. None to reset to default."""
         ...
 
+    # -- Welcome message tracking (ban-by-reply) --
+
+    @abstractmethod
+    async def store_welcome_message(
+        self, chat_id: int, message_id: int, user_id: int
+    ) -> None:
+        """Store mapping from a welcome message to the user who triggered it."""
+        ...
+
+    @abstractmethod
+    async def get_welcome_message_user(
+        self, chat_id: int, message_id: int
+    ) -> int | None:
+        """Get the user_id associated with a welcome message, or None."""
+        ...
+
+    @abstractmethod
+    async def delete_welcome_message(self, chat_id: int, message_id: int) -> None:
+        """Remove a welcome message mapping."""
+        ...
+
     async def close(self) -> None:
         """Close any resources (override if needed)."""
         pass

src/python_italy_bot/db/in_memory.py

@@ -23,6 +23,7 @@ def __init__(self) -> None:
         self._username_to_user_id: dict[str, int] = {}
         self._welcomed: set[tuple[int, int]] = set()
         self._welcome_delays: dict[int, int] = {}
+        self._welcome_message_map: dict[tuple[int, int], int] = {}
 
     async def add_pending_verification(self, user_id: int, chat_id: int) -> None:
         self._pending.add((user_id, chat_id))
@@ -225,3 +226,18 @@ async def set_welcome_delay(self, chat_id: int, minutes: int | None) -> None:
             self._welcome_delays.pop(chat_id, None)
         else:
             self._welcome_delays[chat_id] = minutes
+
+    # -- Welcome message tracking --
+
+    async def store_welcome_message(
+        self, chat_id: int, message_id: int, user_id: int
+    ) -> None:
+        self._welcome_message_map[(chat_id, message_id)] = user_id
+
+    async def get_welcome_message_user(
+        self, chat_id: int, message_id: int
+    ) -> int | None:
+        return self._welcome_message_map.get((chat_id, message_id))
+
+    async def delete_welcome_message(self, chat_id: int, message_id: int) -> None:
+        self._welcome_message_map.pop((chat_id, message_id), None)

src/python_italy_bot/db/postgres.py

@@ -398,5 +398,41 @@ async def set_welcome_delay(self, chat_id: int, minutes: int | None) -> None:
                     (chat_id, minutes),
                 )
 
+    # -- Welcome message tracking (ban-by-reply) --
+
+    async def store_welcome_message(
+        self, chat_id: int, message_id: int, user_id: int
+    ) -> None:
+        async with self._pool.connection() as conn:
+            await conn.execute(
+                """
+                INSERT INTO welcome_messages (chat_id, message_id, user_id)
+                VALUES (%s, %s, %s)
+                ON CONFLICT (chat_id, message_id)
+                DO UPDATE SET user_id = EXCLUDED.user_id
+                """,
+                (chat_id, message_id, user_id),
+            )
+
+    async def get_welcome_message_user(
+        self, chat_id: int, message_id: int
+    ) -> int | None:
+        async with self._pool.connection() as conn:
+            async with conn.cursor() as cur:
+                await cur.execute(
+                    "SELECT user_id FROM welcome_messages "
+                    "WHERE chat_id = %s AND message_id = %s",
+                    (chat_id, message_id),
+                )
+                row = await cur.fetchone()
+                return row[0] if row else None
+
+    async def delete_welcome_message(self, chat_id: int, message_id: int) -> None:
+        async with self._pool.connection() as conn:
+            await conn.execute(
+                "DELETE FROM welcome_messages WHERE chat_id = %s AND message_id = %s",
+                (chat_id, message_id),
+            )
+
     async def close(self) -> None:
         await self._pool.close()

src/python_italy_bot/handlers/moderation.py

@@ -9,6 +9,7 @@
 
 from .. import strings
 from ..db.base import AsyncRepository
+from ..services.captcha import CaptchaService
 from ..services.moderation import ModerationService
 
 logger = logging.getLogger(__name__)
@@ -98,6 +99,12 @@ async def _handle_ban(update: Update, context: ContextTypes.DEFAULT_TYPE) -> Non
                 "welcome_message_map", {}
             )
             user_id = welcome_map.get((chat.id, reply.message_id))
+            # Fall back to database if not in memory (e.g. after bot restart)
+            if user_id is None:
+                captcha_service: CaptchaService = context.bot_data["captcha_service"]
+                user_id = await captcha_service.get_welcome_message_user(
+                    chat.id, reply.message_id
+                )
         reason = " ".join(args) if args else None
     elif args:
         target = args[0]

src/python_italy_bot/handlers/welcome.py

@@ -52,6 +52,20 @@ async def _delete_welcome_message(context: ContextTypes.DEFAULT_TYPE) -> None:
             "Could not delete welcome message %s in chat %s: %s", message_id, chat_id, e
         )
 
+    # Clean up in-memory mapping
+    welcome_map: dict[tuple[int, int], int] = context.bot_data.get(
+        "welcome_message_map", {}
+    )
+    welcome_map.pop((chat_id, message_id), None)
+
+    # Clean up database mapping
+    captcha_service: CaptchaService | None = context.bot_data.get("captcha_service")
+    if captcha_service is not None:
+        try:
+            await captcha_service.delete_welcome_message(chat_id, message_id)
+        except Exception as e:
+            logger.debug("Could not delete welcome message mapping from DB: %s", e)
+
 
 async def _handle_new_member(
     update: Update,
@@ -149,6 +163,12 @@ async def _handle_new_member(
     welcome_map = context.bot_data.setdefault("welcome_message_map", {})
     welcome_map[(chat.id, sent.message_id)] = user.id
 
+    # Persist to database so the mapping survives bot restarts
+    try:
+        await captcha_service.store_welcome_message(chat.id, sent.message_id, user.id)
+    except Exception as e:
+        logger.warning("Could not persist welcome message mapping: %s", e)
+
     # Schedule auto-deletion of the welcome message
     delay_minutes = await captcha_service.get_welcome_delay(chat.id)
     if delay_minutes is None:

src/python_italy_bot/services/captcha.py

@@ -186,3 +186,21 @@ async def get_welcome_delay(self, chat_id: int) -> int | None:
     async def set_welcome_delay(self, chat_id: int, minutes: int | None) -> None:
         """Set welcome message auto-delete delay. None to reset to default."""
         await self._repo.set_welcome_delay(chat_id, minutes)
+
+    # -- Welcome message tracking (ban-by-reply) --
+
+    async def store_welcome_message(
+        self, chat_id: int, message_id: int, user_id: int
+    ) -> None:
+        """Persist mapping from a welcome message to the user who triggered it."""
+        await self._repo.store_welcome_message(chat_id, message_id, user_id)
+
+    async def get_welcome_message_user(
+        self, chat_id: int, message_id: int
+    ) -> int | None:
+        """Get the user_id associated with a welcome message, or None."""
+        return await self._repo.get_welcome_message_user(chat_id, message_id)
+
+    async def delete_welcome_message(self, chat_id: int, message_id: int) -> None:
+        """Remove a welcome message mapping."""
+        await self._repo.delete_welcome_message(chat_id, message_id)