from datetime import datetime from sqlalchemy import delete, func, or_, select, update from sqlalchemy.dialects.postgresql import insert from sqlalchemy.ext.asyncio import AsyncSession from core.cache_config import ( BALANCE_CACHE_TTL_SEC, USER_EXISTS_CACHE_TTL_SEC, USER_SNAPSHOT_CACHE_TTL_SEC, ) from core.redis_cache import cache_delete, cache_get, cache_key, cache_set from database.access.resolution import resolve_user_optional from database.models import ( BlockedUser, CouponUsage, Gift, GiftUsage, Identity, Key, Notification, Payment, Referral, ScheduledBroadcast, TemporaryData, User, WebNotification, WebPushSubscription, ) from logger import logger def invalidate_user_snapshot(tg_id: int) -> None: import asyncio try: loop = asyncio.get_running_loop() loop.create_task(cache_delete(cache_key("user_snapshot", tg_id))) except RuntimeError: return async def add_user( session: AsyncSession, tg_id: int, username: str = None, first_name: str = None, last_name: str = None, language_code: str = None, is_bot: bool = False, source_code: str = None, ) -> int | None: stmt = ( insert(User) .values( tg_id=tg_id, username=username, first_name=first_name, last_name=last_name, language_code=language_code, is_bot=is_bot, source_code=source_code, ) .on_conflict_do_nothing(index_elements=["tg_id"]) .returning(User.id) ) res = await session.execute(stmt) inserted_id = res.scalar_one_or_none() if inserted_id is None: return None await cache_set(cache_key("user_exists", tg_id), True, USER_EXISTS_CACHE_TTL_SEC) logger.info(f"[DB] Новый пользователь добавлен: tg_id={tg_id} id={inserted_id} (source: {source_code})") return int(inserted_id) async def invalidate_balance_cache(tg_id: int) -> None: await cache_delete(cache_key("balance", tg_id)) async def invalidate_profile_cache(tg_id: int) -> None: await cache_delete(cache_key("profile_data", tg_id)) async def update_balance( session: AsyncSession, legacy_user_ref: int, amount: float, ) -> None: u = await resolve_user_optional(session, legacy_user_ref) if u is None: logger.info(f"[DB] Баланс не изменён: пользователь {legacy_user_ref} не найден") return uid = u.id amount = float(amount) res = await session.execute( update(User) .where(User.id == uid) .values(balance=func.coalesce(User.balance, 0) + amount) .returning(User.balance) ) new_balance = res.scalar_one_or_none() if new_balance is not None: old_balance = new_balance - amount logger.info(f"[DB] Баланс пользователя id={uid} обновлён: {old_balance} → {new_balance}") else: logger.info(f"[DB] Баланс пользователя id={uid} не изменён: пользователь не найден") await invalidate_balance_cache(uid) await invalidate_profile_cache(uid) if u.tg_id is not None: await invalidate_balance_cache(u.tg_id) await invalidate_profile_cache(u.tg_id) async def check_user_exists(session: AsyncSession, legacy_user_ref: int) -> bool: cached = await cache_get(cache_key("user_exists", legacy_user_ref)) if isinstance(cached, bool): return cached u = await resolve_user_optional(session, legacy_user_ref) value = u is not None await cache_set(cache_key("user_exists", legacy_user_ref), value, USER_EXISTS_CACHE_TTL_SEC) return value async def get_balance(session: AsyncSession, legacy_user_ref: int) -> float: u = await resolve_user_optional(session, legacy_user_ref) if u is None: return 0.0 uid = u.id cached = await cache_get(cache_key("balance", uid)) if cached is not None: try: return round(float(cached), 1) except (TypeError, ValueError): pass result = await session.execute(select(func.coalesce(User.balance, 0.0)).where(User.id == uid)) balance = result.scalar_one_or_none() value = round(float(balance or 0.0), 1) await cache_set(cache_key("balance", uid), value, BALANCE_CACHE_TTL_SEC) return value async def set_user_balance( session: AsyncSession, legacy_user_ref: int, balance: float, ) -> None: u = await resolve_user_optional(session, legacy_user_ref) if u is None: return uid = u.id balance = float(balance) await session.execute(update(User).where(User.id == uid).values(balance=balance)) await invalidate_balance_cache(uid) await invalidate_profile_cache(uid) async def get_user_preferred_currency(session: AsyncSession, tg_id: int) -> str | None: """Предпочитаемая валюта пользователя по ``tg_id``, если установлена.""" result = await session.execute(select(User.preferred_currency).where(User.tg_id == int(tg_id))) return result.scalar() async def mark_trial_started_if_eligible(session: AsyncSession, tg_id: int) -> None: """Переводит `trial` в 1, только если текущее значение in [0, -1] (пользователь ещё не использовал триал). Условный update без пред-чтения — атомарно на уровне БД. Используется в `services.operations.creation.create_key_on_cluster` после успешного создания ключа. """ await session.execute(update(User).where(User.tg_id == tg_id, User.trial.in_([0, -1])).values(trial=1)) async def update_trial(session: AsyncSession, legacy_user_ref: int, status: int): u = await resolve_user_optional(session, legacy_user_ref) if u is None: return uid = u.id await session.execute(update(User).where(User.id == uid).values(trial=status)) await invalidate_profile_cache(uid) invalidate_user_snapshot(uid) if u.tg_id is not None: await invalidate_profile_cache(u.tg_id) invalidate_user_snapshot(u.tg_id) logger.info(f"[DB] Триал статус обновлён для пользователя id={uid}: {status}") async def get_trial(session: AsyncSession, legacy_user_ref: int) -> int: u = await resolve_user_optional(session, legacy_user_ref) if u is None: return 0 result = await session.execute(select(func.coalesce(User.trial, 0)).where(User.id == u.id)) trial = result.scalar_one_or_none() return int(trial or 0) async def get_balance_and_trial(session: AsyncSession, legacy_user_ref: int) -> tuple[float, int]: """Один запрос к БД для баланса и триала (профиль при промахе кэша).""" u = await resolve_user_optional(session, legacy_user_ref) if u is None: return 0.0, 0 result = await session.execute( select( func.coalesce(User.balance, 0.0), func.coalesce(User.trial, 0), ).where(User.id == u.id) ) row = result.one_or_none() if row is None: return 0.0, 0 balance, trial = row return round(float(balance or 0.0), 1), int(trial or 0) async def get_balance_trial_key_count(session: AsyncSession, legacy_user_ref: int) -> tuple[float, int, int]: """ Один запрос: баланс, триал и число ключей пользователя (для профиля при промахе кэша). Возвращает (balance_rub, trial_status, key_count). """ u = await resolve_user_optional(session, legacy_user_ref) if u is None: return 0.0, 0, 0 key_count_subq = select(func.count()).select_from(Key).where(Key.user_id == User.id).scalar_subquery() result = await session.execute( select( func.coalesce(User.balance, 0.0), func.coalesce(User.trial, 0), key_count_subq, ).where(User.id == u.id) ) row = result.one_or_none() if row is None: return 0.0, 0, 0 balance, trial, key_count = row return ( round(float(balance or 0.0), 1), int(trial or 0), int(key_count or 0), ) async def upsert_user( session: AsyncSession, tg_id: int, username: str = None, first_name: str = None, last_name: str = None, language_code: str = None, is_bot: bool = False, only_if_exists: bool = False, ) -> dict | None: """Создаёт пользователя или обновляет поля профиля.""" now = datetime.utcnow() returning_cols = list(User.__table__.c) if only_if_exists: username_value = username if username else User.username first_name_value = first_name if first_name else User.first_name last_name_value = last_name if last_name else User.last_name language_code_value = language_code if language_code else User.language_code res = await session.execute( update(User) .where(User.tg_id == tg_id) .values( username=username_value, first_name=first_name_value, last_name=last_name_value, language_code=language_code_value, is_bot=is_bot, updated_at=now, ) .returning(*returning_cols) ) row = res.mappings().one_or_none() if row is None: return None await cache_set(cache_key("user_exists", tg_id), True, USER_EXISTS_CACHE_TTL_SEC) return dict(row) res = await session.execute( insert(User) .values( tg_id=tg_id, username=username, first_name=first_name, last_name=last_name, language_code=language_code, is_bot=is_bot, created_at=now, updated_at=now, ) .on_conflict_do_update( index_elements=[User.tg_id], set_={ "username": username, "first_name": first_name, "last_name": last_name, "language_code": language_code, "is_bot": is_bot, "updated_at": now, }, ) .returning(*returning_cols) ) row = res.mappings().one() await cache_set(cache_key("user_exists", tg_id), True, USER_EXISTS_CACHE_TTL_SEC) return dict(row) async def delete_user_data(session: AsyncSession, legacy_user_ref: int): from database.keys import delete_key u = await resolve_user_optional(session, legacy_user_ref) if u is None: return uid = u.id identity_id = u.identity_id await session.execute(delete(Notification).where(Notification.user_id == uid)) await session.execute( delete(GiftUsage).where(GiftUsage.gift_id.in_(select(Gift.gift_id).where(Gift.sender_user_id == uid))) ) await session.execute(delete(GiftUsage).where(GiftUsage.user_id == uid)) await session.execute(delete(Gift).where(Gift.sender_user_id == uid)) await session.execute( update(Gift).where(Gift.recipient_user_id == uid).values(recipient_user_id=None, recipient_tg_id=None) ) if u.tg_id is not None: await session.execute( update(Gift).where(Gift.recipient_tg_id == u.tg_id).values(recipient_tg_id=None) ) await session.execute( update(Gift).where(Gift.sender_tg_id == u.tg_id).values(sender_tg_id=None) ) await session.execute( update(Payment).where(Payment.tg_id == u.tg_id).values(tg_id=None) ) await session.execute(delete(Payment).where(Payment.user_id == uid)) await session.execute( delete(Referral).where(or_(Referral.referrer_user_id == uid, Referral.referred_user_id == uid)) ) await session.execute(delete(CouponUsage).where(CouponUsage.user_id == uid)) if u.tg_id is not None: await session.execute(update(Key).where(Key.tg_id == u.tg_id).values(tg_id=None)) await delete_key(session, uid) await session.execute( delete(TemporaryData).where( or_( TemporaryData.user_id == uid, TemporaryData.tg_id == u.tg_id, ) ) ) await session.execute( delete(BlockedUser).where( or_( BlockedUser.user_id == uid, BlockedUser.tg_id == u.tg_id, ) ) ) await session.execute(delete(WebPushSubscription).where(WebPushSubscription.user_id == uid)) await session.execute(delete(WebNotification).where(WebNotification.user_id == uid)) await session.execute( update(ScheduledBroadcast).where(ScheduledBroadcast.created_by_user_id == uid).values(created_by_user_id=None) ) await session.execute(delete(User).where(User.id == uid)) if identity_id: still_linked = await session.scalar( select(func.count()).select_from(User).where(User.identity_id == identity_id) ) if not still_linked: await session.execute(delete(Identity).where(Identity.id == identity_id)) logger.info(f"[DB] Данные пользователя id={uid} полностью удалены") async def mark_trial_extended(legacy_user_ref: int, session: AsyncSession): u = await resolve_user_optional(session, legacy_user_ref) if u is None: return await session.execute(update(User).where(User.id == u.id).values(trial=-1)) invalidate_user_snapshot(u.id) if u.tg_id is not None: invalidate_user_snapshot(u.tg_id) async def get_user_snapshot(session: AsyncSession, legacy_user_ref: int) -> tuple[int, int] | None: u = await resolve_user_optional(session, legacy_user_ref) if u is None: return None uid = u.id cached = await cache_get(cache_key("user_snapshot", uid)) if isinstance(cached, list) and len(cached) == 2: return (int(cached[0]), int(cached[1])) if isinstance(cached, tuple) and len(cached) == 2: return (int(cached[0]), int(cached[1])) keys_count_sq = select(func.count(Key.client_id)).where(Key.user_id == uid).scalar_subquery() res = await session.execute(select(func.coalesce(User.trial, 0), keys_count_sq).where(User.id == uid)) row = res.first() if row is None: return None value = (int(row[0]), int(row[1])) await cache_set(cache_key("user_snapshot", uid), [value[0], value[1]], USER_SNAPSHOT_CACHE_TTL_SEC) return value async def upsert_source_if_empty( session: AsyncSession, tg_id: int, source_code: str, ) -> bool: if not source_code: return False stmt = ( insert(User) .values(tg_id=tg_id, source_code=source_code) .on_conflict_do_update( index_elements=["tg_id"], set_={"source_code": insert(User).excluded.source_code}, where=(User.source_code.is_(None)), ) .returning(User.tg_id) ) res = await session.execute(stmt) changed_tg_id = res.scalar_one_or_none() return changed_tg_id is not None