Capybara-z
432 lines
15 KiB
Python
432 lines
15 KiB
Python
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
|