feat: re-schedule welcome message deletion on bot restart

On startup, load persisted welcome messages from the database and
re-schedule their deletion jobs so messages are cleaned up even
after a bot restart.

Author: 0xMett

src/python_italy_bot/db/base.py

@@ -1,6 +1,7 @@
 """Abstract repository interface for persistence."""
 
 from abc import ABC, abstractmethod
+from datetime import datetime
 
 from .models import Chat, KnownUser
 
@@ -421,6 +422,13 @@ async def get_welcome_message_user(
         """Get the user_id associated with a welcome message, or None."""
         ...
 
+    @abstractmethod
+    async def get_all_welcome_messages(
+        self,
+    ) -> list[tuple[int, int, int, datetime]]:
+        """Return all tracked welcome messages as (chat_id, message_id, user_id, created_at)."""
+        ...
+
     @abstractmethod
     async def delete_welcome_message(self, chat_id: int, message_id: int) -> None:
         """Remove a welcome message mapping."""

src/python_italy_bot/db/in_memory.py

@@ -239,5 +239,14 @@ async def get_welcome_message_user(
     ) -> int | None:
         return self._welcome_message_map.get((chat_id, message_id))
 
+    async def get_all_welcome_messages(
+        self,
+    ) -> list[tuple[int, int, int, datetime]]:
+        now = datetime.now(timezone.utc)
+        return [
+            (chat_id, message_id, user_id, now)
+            for (chat_id, message_id), user_id in self._welcome_message_map.items()
+        ]
+
     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

@@ -1,5 +1,7 @@
 """PostgreSQL implementation of the repository using psycopg async pool."""
 
+from datetime import datetime
+
 from psycopg_pool import AsyncConnectionPool
 
 from .base import AsyncRepository
@@ -427,6 +429,18 @@ async def get_welcome_message_user(
                 row = await cur.fetchone()
                 return row[0] if row else None
 
+    async def get_all_welcome_messages(
+        self,
+    ) -> list[tuple[int, int, int, datetime]]:
+        async with self._pool.connection() as conn:
+            async with conn.cursor() as cur:
+                await cur.execute(
+                    "SELECT chat_id, message_id, user_id, created_at "
+                    "FROM welcome_messages ORDER BY created_at"
+                )
+                rows = await cur.fetchall()
+                return [(r[0], r[1], r[2], r[3]) for r in rows]
+
     async def delete_welcome_message(self, chat_id: int, message_id: int) -> None:
         async with self._pool.connection() as conn:
             await conn.execute(

src/python_italy_bot/main.py

@@ -1,6 +1,7 @@
 """Entry point for the Python Italy Telegram Bot."""
 
 import logging
+from datetime import datetime, timezone
 
 from telegram.ext import ApplicationBuilder
 
@@ -11,9 +12,14 @@
 from .handlers.moderation import create_moderation_handlers
 from .handlers.ping import create_ping_handlers
 from .handlers.settings import create_settings_handlers
+
 # from .handlers.spam import create_spam_handler
 from .handlers.utils import create_user_tracking_handler
-from .handlers.welcome import create_welcome_handlers
+from .handlers.welcome import (
+    DEFAULT_WELCOME_DELAY_MINUTES,
+    _delete_welcome_message,
+    create_welcome_handlers,
+)
 from .services.captcha import CaptchaService
 from .services.moderation import ModerationService
 # from .services.spam_detector import SpamDetector
@@ -61,6 +67,51 @@ async def _post_init(application) -> None:
         application.add_handler(handler)
     # application.add_handler(create_spam_handler(spam_detector))
 
+    # Re-schedule deletion for welcome messages persisted before a restart
+    await _reschedule_welcome_deletions(application, captcha_service)
+
+
+async def _reschedule_welcome_deletions(
+    application, captcha_service: CaptchaService
+) -> None:
+    """Re-schedule deletion jobs for welcome messages that survived a bot restart."""
+    job_queue = application.job_queue
+    if job_queue is None:
+        return
+
+    messages = await captcha_service.get_all_welcome_messages()
+    if not messages:
+        return
+
+    now = datetime.now(timezone.utc)
+    welcome_map = application.bot_data.setdefault("welcome_message_map", {})
+    scheduled = 0
+
+    for chat_id, message_id, user_id, created_at in messages:
+        # Restore in-memory mapping for ban-by-reply
+        welcome_map[(chat_id, message_id)] = user_id
+
+        delay_minutes = await captcha_service.get_welcome_delay(chat_id)
+        if delay_minutes is None:
+            delay_minutes = DEFAULT_WELCOME_DELAY_MINUTES
+        if delay_minutes <= 0:
+            continue
+
+        # Calculate remaining time; delete immediately if overdue
+        elapsed = (now - created_at).total_seconds()
+        remaining = max(delay_minutes * 60 - elapsed, 0)
+
+        job_queue.run_once(
+            _delete_welcome_message,
+            when=remaining,
+            data=(chat_id, message_id),
+            name=f"del_welcome_{chat_id}_{message_id}",
+        )
+        scheduled += 1
+
+    if scheduled:
+        logger.info("Re-scheduled deletion for %d welcome message(s)", scheduled)
+
 
 async def _post_shutdown(application) -> None:
     """Clean up resources on shutdown."""

src/python_italy_bot/services/captcha.py

@@ -1,6 +1,7 @@
 """Captcha verification logic (file + secret command flow)."""
 
 import re
+from datetime import datetime
 from pathlib import Path
 
 from telegram import (
@@ -204,3 +205,9 @@ async def get_welcome_message_user(
     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)
+
+    async def get_all_welcome_messages(
+        self,
+    ) -> list[tuple[int, int, int, datetime]]:
+        """Return all tracked welcome messages as (chat_id, message_id, user_id, created_at)."""
+        return await self._repo.get_all_welcome_messages()