diff --git a/backend/community_manager/handlers/bot.py b/backend/community_manager/handlers/bot.py index 9d2613a5..c447e256 100644 --- a/backend/community_manager/handlers/bot.py +++ b/backend/community_manager/handlers/bot.py @@ -1,5 +1,8 @@ from telethon import events, types +from core.services.db import DBService +from core.services.user import UserService + START_MESSAGE = """**Access bot** What I can do: @@ -14,6 +17,11 @@ async def handle_start_message(event: events.NewMessage()) -> None: + # Mark user as writable since they interacted with the bot + with DBService().db_session() as db_session: + user_service = UserService(db_session) + user_service.mark_as_writable(event.sender_id) + await event.respond( START_MESSAGE, buttons=types.ReplyInlineMarkup( diff --git a/backend/core/src/core/services/user.py b/backend/core/src/core/services/user.py index fc44ad3a..7c9b9b57 100644 --- a/backend/core/src/core/services/user.py +++ b/backend/core/src/core/services/user.py @@ -25,22 +25,30 @@ def get_all( return query.all() - def get_all_telegram_ids(self) -> list[int]: + def get_all_writable_telegram_ids(self) -> list[int]: return [ u.telegram_id for u in self.db_session.query(User.telegram_id) .filter(User.telegram_id.is_not(None)) - .filter(User.is_blocked.is_(False)) + .filter(User.allows_write_to_pm.is_(True)) .all() ] - def mark_as_blocked(self, telegram_ids: list[int]) -> None: + def mark_as_not_writable(self, telegram_ids: list[int]) -> None: if not telegram_ids: return - logger.info(f"Marking {len(telegram_ids)} users as blocked.") + logger.info( + f"Marking {len(telegram_ids)} users as not writable (allows_write_to_pm=False)." + ) self.db_session.query(User).filter(User.telegram_id.in_(telegram_ids)).update( - {"is_blocked": True}, synchronize_session=False + {"allows_write_to_pm": False}, synchronize_session=False + ) + self.db_session.flush() + + def mark_as_writable(self, telegram_id: int) -> None: + self.db_session.query(User).filter(User.telegram_id == telegram_id).update( + {"allows_write_to_pm": True} ) self.db_session.flush() @@ -88,7 +96,6 @@ def update(self, user: User, telegram_user: TelegramUserDTO) -> User: user.last_name = telegram_user.last_name user.username = telegram_user.username user.is_premium = bool(telegram_user.is_premium) - user.allows_write_to_pm = telegram_user.allow_write_to_pm # TODO add photo_url self.db_session.add(user) self.db_session.flush() diff --git a/backend/indexer_gifts/actions/user.py b/backend/indexer_gifts/actions/user.py index 43110f8d..9a8f7a18 100644 --- a/backend/indexer_gifts/actions/user.py +++ b/backend/indexer_gifts/actions/user.py @@ -27,7 +27,8 @@ async def sync_single_user_gifts(self, telegram_user_id: int) -> None: await self._sync_user_gifts(telegram_user_id, indexer=indexer) except Exception as e: if "user not found" in str(e).lower(): - self.user_service.mark_as_blocked([telegram_user_id]) + self.user_service.mark_as_not_writable([telegram_user_id]) + return raise async def sync_users_batch( @@ -36,7 +37,7 @@ async def sync_users_batch( existing_collection_ids: set[int] | None = None, ) -> None: """Public method for batch sync, manages a single indexer across multiple users.""" - blocked_user_ids = [] + not_writable_user_ids = [] async with BotApiGiftIndexer() as indexer: for user_id in telegram_user_ids: try: @@ -48,10 +49,10 @@ async def sync_users_batch( except Exception as e: logger.error(f"Error refreshing gifts for user {user_id}: {e}") if "user not found" in str(e).lower(): - blocked_user_ids.append(user_id) + not_writable_user_ids.append(user_id) - if blocked_user_ids: - self.user_service.mark_as_blocked(blocked_user_ids) + if not_writable_user_ids: + self.user_service.mark_as_not_writable(not_writable_user_ids) async def _sync_user_gifts( self, @@ -95,20 +96,12 @@ async def _sync_user_gifts( ) updated_count += 1 else: - # Ensure collection exists using in-memory cache collection_id = int(ug.gift_id) if collection_id not in existing_collection_ids: - logger.info( - f"Collection {collection_id} missing, creating stub collection." + logger.debug( + f"Skipping gift {ug.name}: collection {collection_id} not yet synced." ) - self.gift_collection_service.create( - id=collection_id, - title=ug.base_name or f"Gift {collection_id}", - preview_url=None, - supply=0, - upgraded_count=0, - ) - existing_collection_ids.add(collection_id) + continue # Extract string identifiers for visual attributes if available model_name = ( diff --git a/backend/indexer_gifts/indexers/bot_api.py b/backend/indexer_gifts/indexers/bot_api.py index 3732e0ea..f97ccd75 100644 --- a/backend/indexer_gifts/indexers/bot_api.py +++ b/backend/indexer_gifts/indexers/bot_api.py @@ -1,15 +1,18 @@ -import logging import asyncio -from typing import AsyncGenerator +import logging +from typing import Any, AsyncGenerator from aiogram import Bot from aiogram.client.default import DefaultBotProperties +from aiogram.exceptions import TelegramRetryAfter from aiogram.types import OwnedGifts, OwnedGiftUnique from core.settings import core_settings logger = logging.getLogger(__name__) +MAX_FLOOD_RETRIES = 3 + class BotApiGiftIndexer: def __init__(self): @@ -19,48 +22,63 @@ def __init__(self): default=DefaultBotProperties(parse_mode="HTML"), ) + async def _safe_request(self, func, *args, **kwargs) -> Any: + """ + Wraps a request with retry logic for Telegram flood control (429). + """ + attempt = 1 + while True: + try: + return await func(*args, **kwargs) + except TelegramRetryAfter as e: + if attempt >= MAX_FLOOD_RETRIES: + logger.error( + f"Flood control: exceeded {MAX_FLOOD_RETRIES} retries, giving up." + ) + raise + logger.warning( + f"Flood control: sleeping {e.retry_after}s " + f"(attempt {attempt}/{MAX_FLOOD_RETRIES})" + ) + await asyncio.sleep(e.retry_after) + attempt += 1 + async def iter_user_gifts( self, telegram_user_id: int ) -> AsyncGenerator[OwnedGiftUnique, None]: """ Yields OwnedGiftUnique objects owned by the user via pagination. - Handles rate limits or pagination as needed. + Handles flood control by sleeping and retrying up to MAX_FLOOD_RETRIES times. """ offset = "" page_num = 1 logger.info(f"Starting to fetch gifts for user {telegram_user_id}...") while True: - try: - logger.debug( - f"Fetching page {page_num} for user {telegram_user_id} (offset: {offset!r})" - ) - # get_user_gifts is available in aiogram 3.24.0 - user_gifts: OwnedGifts = await self.bot.get_user_gifts( - user_id=telegram_user_id, offset=offset, limit=100 - ) + user_gifts: OwnedGifts = await self._safe_request( + self.bot.get_user_gifts, + user_id=telegram_user_id, + offset=offset, + limit=100, + ) - logger.info( - f"Fetched {len(user_gifts.gifts)} gifts for user {telegram_user_id} on page {page_num}." - ) + logger.info( + f"Fetched {len(user_gifts.gifts)} gifts for user " + f"{telegram_user_id} on page {page_num}." + ) - # Yield only unique gifts from the current page - for gift in user_gifts.gifts: - if isinstance(gift, OwnedGiftUnique): - yield gift + for gift in user_gifts.gifts: + if isinstance(gift, OwnedGiftUnique): + yield gift - # Check if there are more pages - if not user_gifts.next_offset: - logger.info( - f"Finished fetching all gifts for user {telegram_user_id}. Total count: {user_gifts.total_count}" - ) - break - - offset = user_gifts.next_offset - page_num += 1 + if not user_gifts.next_offset: + logger.info( + f"Finished fetching all gifts for user {telegram_user_id}. " + f"Total count: {user_gifts.total_count}" + ) + break - except Exception as e: - logger.error(f"Error fetching gifts for user {telegram_user_id}: {e}") - raise + offset = user_gifts.next_offset + page_num += 1 async def __aenter__(self): return self diff --git a/backend/indexer_gifts/tasks.py b/backend/indexer_gifts/tasks.py index 5c4cd6a0..0e738818 100644 --- a/backend/indexer_gifts/tasks.py +++ b/backend/indexer_gifts/tasks.py @@ -78,7 +78,7 @@ def refresh_all_user_gifts() -> None: """Master dispatcher: chunks all registered users and enqueues batch tasks.""" with DBService().db_session() as db_session: user_service = UserService(db_session) - all_telegram_ids = user_service.get_all_telegram_ids() + all_telegram_ids = user_service.get_all_writable_telegram_ids() if not all_telegram_ids: logger.info("No registered Telegram users found for gift refresh.") diff --git a/backend/scheduler/celery_app.py b/backend/scheduler/celery_app.py index 99c4c601..88aa3870 100644 --- a/backend/scheduler/celery_app.py +++ b/backend/scheduler/celery_app.py @@ -49,19 +49,15 @@ def create_app() -> Celery: "schedule": crontab(minute="*/10"), # Every 10 minutes "options": {"queue": CELERY_STICKER_FETCH_QUEUE_NAME}, }, - # "fetch-gift-ownerships": { - # "task": "fetch-gift-ownership-details", - # "schedule": crontab(hour="*/1", minute="0"), # Every hour - # "options": {"queue": CELERY_GIFT_COLLECTIONS_SYNC_QUEUE}, - # }, "sync-gift-collections-from-api": { "task": "sync-gift-collections-from-api", - "schedule": crontab(hour="1", minute="0"), # Daily at 1:00 AM + # "schedule": crontab(hour="1", minute="0"), # Daily at 01:00 + "schedule": crontab(minute="*/10"), # Every 15 minutes "options": {"queue": CELERY_GIFT_COLLECTIONS_SYNC_QUEUE}, }, "refresh-all-user-gifts": { "task": "refresh-all-user-gifts", - "schedule": crontab(hour="*/1", minute="15"), # Hourly at :15 + "schedule": crontab(hour="*/4", minute="2"), # Every 4 hours at :02 "options": {"queue": CELERY_GIFT_USER_BATCH_QUEUE}, }, "refresh-prices": { diff --git a/backend/tests/unit/core/services/test_user.py b/backend/tests/unit/core/services/test_user.py index 04856a90..511d96fe 100644 --- a/backend/tests/unit/core/services/test_user.py +++ b/backend/tests/unit/core/services/test_user.py @@ -2,19 +2,52 @@ from tests.factories.user import UserFactory -def test_get_all_telegram_ids(db_session): +def test_get_all_writable_telegram_ids(db_session): user_service = UserService(db_session) - # Create some users with and without telegram_ids - UserFactory.with_session(db_session).create(telegram_id=10001) - UserFactory.with_session(db_session).create(telegram_id=10002) - UserFactory.with_session(db_session).create(telegram_id=None) + # Create some users with and without telegram_ids and allows_write_to_pm + UserFactory.with_session(db_session).create( + telegram_id=10001, allows_write_to_pm=True + ) + UserFactory.with_session(db_session).create( + telegram_id=10002, allows_write_to_pm=True + ) + UserFactory.with_session(db_session).create( + telegram_id=10003, allows_write_to_pm=False + ) + UserFactory.with_session(db_session).create( + telegram_id=None, allows_write_to_pm=True + ) db_session.flush() - telegram_ids = user_service.get_all_telegram_ids() + telegram_ids = user_service.get_all_writable_telegram_ids() assert 10001 in telegram_ids assert 10002 in telegram_ids + assert 10003 not in telegram_ids assert None not in telegram_ids - assert set([10001, 10002]).issubset(set(telegram_ids)) + +def test_mark_as_not_writable_and_writable(db_session): + user_service = UserService(db_session) + + user = UserFactory.with_session(db_session).create( + telegram_id=20001, allows_write_to_pm=True + ) + db_session.flush() + + assert user.allows_write_to_pm is True + + # Mark as not writable + user_service.mark_as_not_writable([20001]) + db_session.refresh(user) + assert user.allows_write_to_pm is False + + # Mark back as writable + user_service.mark_as_writable(20001) + db_session.refresh(user) + assert user.allows_write_to_pm is True + + # Calling mark_as_writable on a non-existent user should be a safe no-op + user_service.mark_as_writable(999999999) + user_service.mark_as_not_writable([999999999])