diff --git a/database/db.py b/database/db.py index 0ebd7e2c..f4e77d4a 100644 --- a/database/db.py +++ b/database/db.py @@ -4,13 +4,17 @@ from sqlalchemy import text from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine from sqlalchemy.orm import declarative_base -from config import DATABASE_URL, DB_MAX_OVERFLOW, DB_POOL_SIZE +from config import DATABASE_URL, DB_MAX_OVERFLOW, DB_POOL_SIZE, USE_PGBOUNCER from core.cache_config import UPDATE_STALE_AGE_SEC CONCURRENT_UPDATES_LIMIT = DB_POOL_SIZE + DB_MAX_OVERFLOW MAX_UPDATE_AGE_SEC = UPDATE_STALE_AGE_SEC +_connect_args = {} +if USE_PGBOUNCER and "+asyncpg" in DATABASE_URL: + _connect_args["statement_cache_size"] = 0 + engine = create_async_engine( DATABASE_URL, echo=False, @@ -20,6 +24,7 @@ engine = create_async_engine( pool_timeout=60, pool_pre_ping=True, pool_recycle=300, + connect_args=_connect_args, ) async_session_maker = async_sessionmaker( diff --git a/database/notifications.py b/database/notifications.py index 2b0f1fd7..2dc5d7e9 100644 --- a/database/notifications.py +++ b/database/notifications.py @@ -139,6 +139,21 @@ async def check_hot_lead_discount(session: AsyncSession, tg_id: int) -> dict: return {"available": False} +_BULK_NOTIFICATION_BATCH_SIZE = 250 + + +def _batched_pairs(tg_ids: list[int], emails: list[str], batch_size: int): + """Yield (tg_ids_chunk, emails_chunk) of length <= batch_size. Lists must have same length.""" + for i in range(0, len(tg_ids), batch_size): + yield tg_ids[i : i + batch_size], emails[i : i + batch_size] + + +def _batched_list(items: list, batch_size: int): + """Yield chunks of items of length <= batch_size.""" + for i in range(0, len(items), batch_size): + yield items[i : i + batch_size] + + async def check_notifications_bulk( session: AsyncSession, notification_type: str, @@ -159,49 +174,133 @@ async def check_notifications_bulk( .subquery() ) - stmt = ( - select( - User.tg_id, - Key.email, - User.username, - User.first_name, - User.last_name, - subq_last_notification.c.last_notification_time, - ) - .outerjoin(Key, Key.tg_id == User.tg_id) - .outerjoin(subq_last_notification, subq_last_notification.c.tg_id == User.tg_id) - ) - - if notification_type == "inactive_trial": - stmt = stmt.where( - and_( - User.trial.in_([0, -1]), - ~User.tg_id.in_(select(BlockedUser.tg_id)), - ~User.tg_id.in_(select(Key.tg_id.distinct())), + def make_stmt(tg_ids_batch: list[int] | None, emails_batch: list[str] | None): + stmt = ( + select( + User.tg_id, + Key.email, + User.username, + User.first_name, + User.last_name, + subq_last_notification.c.last_notification_time, ) + .select_from(User) + .outerjoin(Key, Key.tg_id == User.tg_id) + .outerjoin(subq_last_notification, subq_last_notification.c.tg_id == User.tg_id) ) + if notification_type == "inactive_trial": + stmt = stmt.where( + and_( + User.trial.in_([0, -1]), + ~User.tg_id.in_(select(BlockedUser.tg_id)), + ~User.tg_id.in_(select(Key.tg_id.distinct())), + ) + ) + if tg_ids_batch: + stmt = stmt.where(User.tg_id.in_(tg_ids_batch)) + if emails_batch: + stmt = stmt.where(Key.email.in_(emails_batch)) + return stmt - if tg_ids: - stmt = stmt.where(User.tg_id.in_(tg_ids)) - if emails: - stmt = stmt.where(Key.email.in_(emails)) + users: list[dict] = [] + seen: set[tuple[int, str | None]] = set() - result = await session.execute(stmt) - users = [] - - for row in result: - last_time = row.last_notification_time - can_notify = not last_time or (now - last_time > timedelta(hours=hours)) - - if can_notify: - users.append({ - "tg_id": row.tg_id, - "email": row.email, - "username": row.username, - "first_name": row.first_name, - "last_name": row.last_name, - "last_notification_time": int(last_time.timestamp() * 1000) if last_time else None, - }) + if tg_ids and emails and len(tg_ids) == len(emails): + for tg_ids_chunk, emails_chunk in _batched_pairs(tg_ids, emails, _BULK_NOTIFICATION_BATCH_SIZE): + stmt = make_stmt(tg_ids_chunk, emails_chunk) + result = await session.execute(stmt) + for row in result: + key = (row.tg_id, row.email) + if key in seen: + continue + seen.add(key) + last_time = row.last_notification_time + can_notify = not last_time or (now - last_time > timedelta(hours=hours)) + if can_notify: + users.append({ + "tg_id": row.tg_id, + "email": row.email, + "username": row.username, + "first_name": row.first_name, + "last_name": row.last_name, + "last_notification_time": int(last_time.timestamp() * 1000) if last_time else None, + }) + elif tg_ids and emails: + for tg_ids_chunk in _batched_list(tg_ids, _BULK_NOTIFICATION_BATCH_SIZE): + for emails_chunk in _batched_list(emails, _BULK_NOTIFICATION_BATCH_SIZE): + stmt = make_stmt(tg_ids_chunk, emails_chunk) + result = await session.execute(stmt) + for row in result: + key = (row.tg_id, row.email) + if key in seen: + continue + seen.add(key) + last_time = row.last_notification_time + can_notify = not last_time or (now - last_time > timedelta(hours=hours)) + if can_notify: + users.append({ + "tg_id": row.tg_id, + "email": row.email, + "username": row.username, + "first_name": row.first_name, + "last_name": row.last_name, + "last_notification_time": int(last_time.timestamp() * 1000) if last_time else None, + }) + elif tg_ids: + for tg_ids_chunk in _batched_list(tg_ids, _BULK_NOTIFICATION_BATCH_SIZE): + stmt = make_stmt(tg_ids_chunk, None) + result = await session.execute(stmt) + for row in result: + key = (row.tg_id, row.email) + if key in seen: + continue + seen.add(key) + last_time = row.last_notification_time + can_notify = not last_time or (now - last_time > timedelta(hours=hours)) + if can_notify: + users.append({ + "tg_id": row.tg_id, + "email": row.email, + "username": row.username, + "first_name": row.first_name, + "last_name": row.last_name, + "last_notification_time": int(last_time.timestamp() * 1000) if last_time else None, + }) + elif emails: + for emails_chunk in _batched_list(emails, _BULK_NOTIFICATION_BATCH_SIZE): + stmt = make_stmt(None, emails_chunk) + result = await session.execute(stmt) + for row in result: + key = (row.tg_id, row.email) + if key in seen: + continue + seen.add(key) + last_time = row.last_notification_time + can_notify = not last_time or (now - last_time > timedelta(hours=hours)) + if can_notify: + users.append({ + "tg_id": row.tg_id, + "email": row.email, + "username": row.username, + "first_name": row.first_name, + "last_name": row.last_name, + "last_notification_time": int(last_time.timestamp() * 1000) if last_time else None, + }) + else: + stmt = make_stmt(None, None) + result = await session.execute(stmt) + for row in result: + last_time = row.last_notification_time + can_notify = not last_time or (now - last_time > timedelta(hours=hours)) + if can_notify: + users.append({ + "tg_id": row.tg_id, + "email": row.email, + "username": row.username, + "first_name": row.first_name, + "last_name": row.last_name, + "last_notification_time": int(last_time.timestamp() * 1000) if last_time else None, + }) logger.info(f"Найдено {len(users)} пользователей, готовых к уведомлению типа {notification_type}") return users diff --git a/handlers/admin/ads/ads_handler.py b/handlers/admin/ads/ads_handler.py index e2d81294..7cb2c16a 100644 --- a/handlers/admin/ads/ads_handler.py +++ b/handlers/admin/ads/ads_handler.py @@ -12,6 +12,8 @@ from sqlalchemy import delete, select, update from sqlalchemy.ext.asyncio import AsyncSession from config import USERNAME_BOT +from core.cache_config import START_UTM_EXISTS_TTL_SEC +from core.redis_cache import cache_key, cache_delete, cache_set from database import create_tracking_source, get_tracking_source_stats from database.models import TrackingSource, User from filters.admin import IsAdminFilter @@ -83,6 +85,7 @@ async def handle_ads_code_input(message: Message, state: FSMContext, session: As created_by=message.from_user.id, session=session, ) + await cache_set(cache_key("utm_exists", code_with_prefix), True, START_UTM_EXISTS_TTL_SEC) stats = await get_tracking_source_stats(session, code_with_prefix) if not stats: await message.answer("❌ Источник не найден или не содержит данных.") @@ -159,6 +162,7 @@ async def handle_ads_delete( await session.execute(update(User).where(User.source_code == code).values(source_code=None)) await session.execute(delete(TrackingSource).where(TrackingSource.code == code)) await session.commit() + await cache_delete(cache_key("utm_exists", code)) await callback_query.message.edit_text( f"🗑️ Ссылка {code} удалена.", reply_markup=build_ads_kb(), diff --git a/handlers/start.py b/handlers/start.py index 7cbc8afc..ee253d3e 100644 --- a/handlers/start.py +++ b/handlers/start.py @@ -263,11 +263,13 @@ async def prompt_subscription(callback: CallbackQuery): async def handle_utm_link(utm_code: str, message: Message, state: FSMContext, session: AsyncSession, user_data: dict): - is_known = await cache_get(cache_key("utm_exists", utm_code)) + key = cache_key("utm_exists", utm_code) + is_known = await cache_get(key) if is_known is None: - res = await session.execute(select(TrackingSource).where(TrackingSource.code == utm_code)) + stmt = select(1).select_from(TrackingSource).where(TrackingSource.code == utm_code).limit(1) + res = await session.execute(stmt) is_known = res.scalar_one_or_none() is not None - await cache_set(cache_key("utm_exists", utm_code), bool(is_known), START_UTM_EXISTS_TTL_SEC) + await cache_set(key, bool(is_known), START_UTM_EXISTS_TTL_SEC) if not is_known: await message.answer("❌ UTM ссылка не найдена.")