Files
remnawave-bedolaga-telegram…/app/database/crud/user_promo_group.py
T
Fringg 3ec9e71de7 fix: propagate exceptions from get_primary_user_promo_group
Consistent with has_user_promo_group and get_user_promo_groups which
now propagate exceptions instead of masking them with default returns.
2026-03-20 05:30:42 +03:00

287 lines
9.5 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
"""CRUD операции для связи пользователей с промогруппами (Many-to-Many)."""
from datetime import UTC, datetime
import structlog
from sqlalchemy import and_, desc, func, select
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.orm import selectinload
from app.database.models import PromoGroup, User, UserPromoGroup
logger = structlog.get_logger(__name__)
async def _sync_user_primary_promo_group(
db: AsyncSession,
user_id: int,
) -> None:
"""Синхронизирует колонку users.promo_group_id с primary промогруппой."""
try:
result = await db.execute(
select(UserPromoGroup.promo_group_id)
.join(PromoGroup, UserPromoGroup.promo_group_id == PromoGroup.id)
.where(UserPromoGroup.user_id == user_id)
.order_by(desc(PromoGroup.priority), desc(PromoGroup.id))
)
first = result.first()
new_primary_id = first[0] if first else None
user = await db.get(User, user_id)
if not user:
return
if user.promo_group_id != new_primary_id:
user.promo_group_id = new_primary_id
user.updated_at = datetime.now(UTC)
except Exception as error:
logger.error('Ошибка синхронизации primary промогруппы пользователя', user_id=user_id, error=error)
raise
async def sync_user_primary_promo_group(
db: AsyncSession,
user_id: int,
) -> None:
"""Публичная обертка для синхронизации primary промогруппы пользователя."""
await _sync_user_primary_promo_group(db, user_id)
async def add_user_to_promo_group(
db: AsyncSession,
user_id: int,
promo_group_id: int,
assigned_by: str = 'admin',
*,
commit: bool = True,
) -> UserPromoGroup | None:
"""
Добавляет пользователю промогруппу.
Args:
db: Сессия БД
user_id: ID пользователя
promo_group_id: ID промогруппы
assigned_by: Кто назначил ('admin', 'system', 'auto', 'promocode')
commit: Коммитить транзакцию (False для батчевых операций)
Returns:
UserPromoGroup или None если уже существует
"""
try:
# Проверяем существование связи
existing = await has_user_promo_group(db, user_id, promo_group_id)
if existing:
logger.info('Пользователь уже имеет промогруппу', user_id=user_id, promo_group_id=promo_group_id)
return None
# Создаем новую связь
user_promo_group = UserPromoGroup(
user_id=user_id,
promo_group_id=promo_group_id,
assigned_by=assigned_by,
)
db.add(user_promo_group)
await db.flush()
await _sync_user_primary_promo_group(db, user_id)
if commit:
await db.commit()
await db.refresh(user_promo_group)
logger.info(
'Пользователю добавлена промогруппа',
user_id=user_id,
promo_group_id=promo_group_id,
assigned_by=assigned_by,
)
return user_promo_group
except Exception as error:
logger.error('Ошибка добавления промогруппы пользователю', error=error)
if commit:
await db.rollback()
return None
raise
async def remove_user_from_promo_group(
db: AsyncSession,
user_id: int,
promo_group_id: int,
*,
commit: bool = True,
) -> bool:
"""
Удаляет промогруппу у пользователя.
Args:
db: Сессия БД
user_id: ID пользователя
promo_group_id: ID промогруппы
commit: Коммитить транзакцию (False для батчевых операций)
Returns:
True если удалено, False если связи не было
"""
try:
result = await db.execute(
select(UserPromoGroup).where(
and_(UserPromoGroup.user_id == user_id, UserPromoGroup.promo_group_id == promo_group_id)
)
)
user_promo_group = result.scalar_one_or_none()
if not user_promo_group:
logger.warning(
'Связь пользователя с промогруппой не найдена', user_id=user_id, promo_group_id=promo_group_id
)
return False
await db.delete(user_promo_group)
await db.flush()
await _sync_user_primary_promo_group(db, user_id)
if commit:
await db.commit()
logger.info('У пользователя удалена промогруппа', user_id=user_id, promo_group_id=promo_group_id)
return True
except Exception as error:
logger.error('Ошибка удаления промогруппы у пользователя', error=error)
if commit:
await db.rollback()
return False
raise
async def get_user_promo_groups(db: AsyncSession, user_id: int) -> list[UserPromoGroup]:
"""
Получает все промогруппы пользователя, отсортированные по приоритету.
Args:
db: Сессия БД
user_id: ID пользователя
Returns:
Список UserPromoGroup с загруженными PromoGroup, отсортированный по приоритету DESC
"""
result = await db.execute(
select(UserPromoGroup)
.options(selectinload(UserPromoGroup.promo_group))
.where(UserPromoGroup.user_id == user_id)
.join(PromoGroup, UserPromoGroup.promo_group_id == PromoGroup.id)
.order_by(desc(PromoGroup.priority), desc(PromoGroup.id))
)
return list(result.scalars().all())
async def get_primary_user_promo_group(db: AsyncSession, user_id: int) -> PromoGroup | None:
"""
Получает промогруппу пользователя с максимальным приоритетом.
Args:
db: Сессия БД
user_id: ID пользователя
Returns:
PromoGroup с максимальным приоритетом или None
"""
user_promo_groups = await get_user_promo_groups(db, user_id)
if not user_promo_groups:
return None
# Первая в списке имеет максимальный приоритет (список уже отсортирован)
return user_promo_groups[0].promo_group or None
async def has_user_promo_group(db: AsyncSession, user_id: int, promo_group_id: int) -> bool:
"""
Проверяет наличие промогруппы у пользователя.
Args:
db: Сессия БД
user_id: ID пользователя
promo_group_id: ID промогруппы
Returns:
True если пользователь уже имеет эту промогруппу
"""
result = await db.execute(
select(UserPromoGroup).where(
and_(UserPromoGroup.user_id == user_id, UserPromoGroup.promo_group_id == promo_group_id)
)
)
return result.scalar_one_or_none() is not None
async def count_user_promo_groups(db: AsyncSession, user_id: int) -> int:
"""
Подсчитывает количество промогрупп у пользователя.
Args:
db: Сессия БД
user_id: ID пользователя
Returns:
Количество промогрупп
"""
try:
result = await db.execute(
select(func.count()).select_from(UserPromoGroup).where(UserPromoGroup.user_id == user_id)
)
return result.scalar_one()
except Exception as error:
logger.error('Ошибка подсчета промогрупп пользователя', error=error)
return 0
async def replace_user_promo_groups(
db: AsyncSession, user_id: int, promo_group_ids: list[int], assigned_by: str = 'admin'
) -> bool:
"""
Заменяет все промогруппы пользователя на новый список.
Args:
db: Сессия БД
user_id: ID пользователя
promo_group_ids: Список ID промогрупп
assigned_by: Кто назначил
Returns:
True если успешно
"""
try:
# Удаляем все текущие промогруппы
result = await db.execute(select(UserPromoGroup).where(UserPromoGroup.user_id == user_id))
for upg in result.scalars().all():
await db.delete(upg)
await db.flush()
# Добавляем новые
for promo_group_id in promo_group_ids:
user_promo_group = UserPromoGroup(user_id=user_id, promo_group_id=promo_group_id, assigned_by=assigned_by)
db.add(user_promo_group)
await db.flush()
await _sync_user_primary_promo_group(db, user_id)
await db.commit()
logger.info('Промогруппы пользователя заменены на', user_id=user_id, promo_group_ids=promo_group_ids)
return True
except Exception as error:
logger.error('Ошибка замены промогрупп пользователя', error=error)
await db.rollback()
return False