PGboucer/UTM cache

2026-02-21 18:32:52 +03:00
parent aed63de100
commit 4b1a9a2567
4 changed files with 153 additions and 43 deletions
@@ -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(
@@ -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
@@ -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>{code}</code> удалена.",
            reply_markup=build_ads_kb(),
@@ -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 ссылка не найдена.")