import hashlib from collections import defaultdict from datetime import datetime from sqlalchemy import delete, func, insert, select, update from sqlalchemy.exc import SQLAlchemyError from sqlalchemy.ext.asyncio import AsyncSession from core.cache_config import TARIFF_BY_ID_CACHE_TTL_SEC, TARIFFS_FOR_CLUSTER_CACHE_TTL_SEC from core.redis_cache import cache_delete, cache_delete_pattern, cache_get, cache_key, cache_set from database.models import Server, Tariff from logger import logger def _row_to_cache_dict(row_dict: dict) -> dict: """Делает dict строки БД пригодным для JSON/Redis (datetime → str, убирает _sa_instance_state).""" out = {} for k, v in row_dict.items(): if k.startswith("_"): continue if isinstance(v, datetime): out[k] = v.isoformat() else: out[k] = v return out async def _invalidate_tariff_cache(tariff_id: int | None = None) -> None: """Сброс кэша тарифов при изменении (по id и списков по кластерам).""" if tariff_id is not None: await cache_delete(cache_key("tariff", tariff_id)) await cache_delete_pattern("tariffs_cluster:*") def create_subgroup_hash(subgroup_title: str, group_code: str) -> str: if not subgroup_title: return "" unique_key = f"{subgroup_title}:{group_code}" hash_object = hashlib.md5(unique_key.encode("utf-8")) return hash_object.hexdigest()[:8] async def find_subgroup_by_hash(session: AsyncSession, subgroup_hash: str, group_code: str) -> str | None: result = await session.execute( select(Tariff.subgroup_title) .where(Tariff.group_code == group_code, Tariff.subgroup_title.isnot(None)) .distinct() ) subgroups = [row[0] for row in result.fetchall()] for subgroup_title in subgroups: if create_subgroup_hash(subgroup_title, group_code) == subgroup_hash: return subgroup_title return None async def get_tariffs( session: AsyncSession, tariff_id: int = None, group_code: str = None, with_subgroup_weights: bool = False ): try: if tariff_id: result = await session.execute(select(Tariff).where(Tariff.id == tariff_id)) elif group_code: result = await session.execute( select(Tariff).where(Tariff.group_code == group_code).order_by(Tariff.sort_order, Tariff.id) ) else: result = await session.execute(select(Tariff).order_by(Tariff.sort_order, Tariff.id)) tariffs = [dict(r.__dict__) for r in result.scalars().all()] if with_subgroup_weights and group_code: tariffs_without_order = [t for t in tariffs if t.get("sort_order") is None] if tariffs_without_order: for tariff in tariffs_without_order: tariff["sort_order"] = 1 await session.execute(update(Tariff).where(Tariff.id == tariff["id"]).values(sort_order=1)) await session.commit() grouped = defaultdict(list) for t in tariffs: grouped[t.get("subgroup_title")].append(t) subgroup_weights = {} for subgroup, tariffs_list in grouped.items(): if subgroup: total_weight = sum(t.get("sort_order", 1) for t in tariffs_list) subgroup_weights[subgroup] = total_weight return {"tariffs": tariffs, "subgroup_weights": subgroup_weights} return tariffs except SQLAlchemyError as e: logger.error(f"[TARIFF] Ошибка при получении тарифов: {e}") await session.rollback() return [] async def get_tariff_names_groups_subgroups_durations( session: AsyncSession, tariff_ids: list[int] ) -> tuple[dict[int, str], dict[int, str], dict[int, str | None], dict[int, int]]: """Один запрос: id, name, group_code, subgroup_title, duration_days → четыре словаря.""" if not tariff_ids: return {}, {}, {}, {} result = await session.execute( select( Tariff.id, Tariff.name, Tariff.group_code, Tariff.subgroup_title, Tariff.duration_days, ).where(Tariff.id.in_(tariff_ids)) ) rows = result.all() names = {} groups = {} subgroups = {} durations = {} for tid, name, group_code, subgroup_title, duration_days in rows: names[tid] = name groups[tid] = group_code subgroups[tid] = subgroup_title durations[tid] = duration_days return names, groups, subgroups, durations async def get_tariff_by_id(session: AsyncSession, tariff_id: int): key = cache_key("tariff", tariff_id) cached = await cache_get(key) if isinstance(cached, dict): return cached try: result = await session.execute(select(Tariff).where(Tariff.id == tariff_id)) tariff = result.scalar_one_or_none() if not tariff: return None row = _row_to_cache_dict(dict(tariff.__dict__)) await cache_set(key, row, TARIFF_BY_ID_CACHE_TTL_SEC) return row except SQLAlchemyError as e: logger.error(f"[TARIFF] Ошибка при получении тарифа по ID {tariff_id}: {e}") await session.rollback() return None async def get_tariffs_for_cluster(session: AsyncSession, cluster_name: str): key = cache_key("tariffs_cluster", cluster_name) cached = await cache_get(key) if isinstance(cached, list): return cached try: server_row = await session.execute( select(Server.tariff_group).where(Server.cluster_name == cluster_name).limit(1) ) row = server_row.first() if not row: server_row = await session.execute( select(Server.tariff_group).where(Server.server_name == cluster_name).limit(1) ) row = server_row.first() if not row or not row[0]: return [] group_code = row[0] result = await session.execute( select(Tariff) .where(Tariff.group_code == group_code, Tariff.is_active.is_(True)) .order_by(Tariff.sort_order, Tariff.id) ) rows = [_row_to_cache_dict(dict(r.__dict__)) for r in result.scalars().all()] await cache_set(key, rows, TARIFFS_FOR_CLUSTER_CACHE_TTL_SEC) return rows except SQLAlchemyError as e: logger.error(f"[TARIFF] Ошибка при получении тарифов для кластера {cluster_name}: {e}") return [] async def create_tariff(session: AsyncSession, data: dict): try: data["created_at"] = datetime.utcnow() data["updated_at"] = datetime.utcnow() if "sort_order" not in data: group_code = data.get("group_code") if group_code: result = await session.execute( select(func.max(Tariff.sort_order)).where( Tariff.group_code == group_code, Tariff.sort_order.isnot(None) ) ) max_order = result.scalar() or 0 else: result = await session.execute(select(func.max(Tariff.sort_order)).where(Tariff.sort_order.isnot(None))) max_order = result.scalar() or 0 data["sort_order"] = max_order + 1 stmt = insert(Tariff).values(**data).returning(Tariff) result = await session.execute(stmt) await session.commit() await _invalidate_tariff_cache() return result.scalar_one() except SQLAlchemyError as e: logger.error(f"[TARIFF] Ошибка при создании тарифа: {e}") await session.rollback() return None async def update_tariff(session: AsyncSession, tariff_id: int, updates: dict): if not updates: return False try: updates["updated_at"] = datetime.utcnow() await session.execute(update(Tariff).where(Tariff.id == tariff_id).values(**updates)) await session.commit() await _invalidate_tariff_cache(tariff_id) return True except SQLAlchemyError as e: logger.error(f"[TARIFF] Ошибка при обновлении тарифа ID={tariff_id}: {e}") await session.rollback() return False async def delete_tariff(session: AsyncSession, tariff_id: int): try: await session.execute(delete(Tariff).where(Tariff.id == tariff_id)) await session.commit() await _invalidate_tariff_cache(tariff_id) return True except SQLAlchemyError as e: logger.error(f"[TARIFF] Ошибка при удалении тарифа ID={tariff_id}: {e}") await session.rollback() return False async def check_tariff_exists(session: AsyncSession, tariff_id: int): try: result = await session.execute(select(Tariff).where(Tariff.id == tariff_id, Tariff.is_active.is_(True))) tariff = result.scalar_one_or_none() if tariff: return True logger.warning(f"[TARIFF] Тариф {tariff_id} не найден в БД") return False except SQLAlchemyError as e: logger.error(f"[TARIFF] Ошибка при проверке тарифа {tariff_id}: {e}") await session.rollback() return False async def get_tariff_sort_order(session: AsyncSession, tariff_id: int) -> int: try: result = await session.execute(select(Tariff.sort_order).where(Tariff.id == tariff_id)) sort_order = result.scalar_one_or_none() if sort_order is None: await session.execute(update(Tariff).where(Tariff.id == tariff_id).values(sort_order=1)) await session.commit() await _invalidate_tariff_cache(tariff_id) return 1 return sort_order except SQLAlchemyError as e: logger.error(f"[TARIFF] Ошибка при получении sort_order для тарифа {tariff_id}: {e}") await session.rollback() return None async def move_tariff_up(session: AsyncSession, tariff_id: int) -> bool: try: current_order = await get_tariff_sort_order(session, tariff_id) new_order = max(1, current_order - 1) await session.execute(update(Tariff).where(Tariff.id == tariff_id).values(sort_order=new_order)) await session.commit() await _invalidate_tariff_cache(tariff_id) return True except SQLAlchemyError as e: logger.error(f"[TARIFF] Ошибка при перемещении тарифа {tariff_id} вверх: {e}") await session.rollback() return False async def move_tariff_down(session: AsyncSession, tariff_id: int) -> bool: try: current_order = await get_tariff_sort_order(session, tariff_id) new_order = current_order + 1 await session.execute(update(Tariff).where(Tariff.id == tariff_id).values(sort_order=new_order)) await session.commit() await _invalidate_tariff_cache(tariff_id) return True except SQLAlchemyError as e: logger.error(f"[TARIFF] Ошибка при перемещении тарифа {tariff_id} вниз: {e}") await session.rollback() return False async def initialize_tariff_sort_orders(session: AsyncSession, group_code: str) -> bool: try: result = await session.execute(select(Tariff).where(Tariff.group_code == group_code).order_by(Tariff.id)) tariffs = result.scalars().all() if not tariffs: return True for i, tariff in enumerate(tariffs): new_sort_order = 1 + i await session.execute(update(Tariff).where(Tariff.id == tariff.id).values(sort_order=new_sort_order)) await session.commit() await _invalidate_tariff_cache() return True except SQLAlchemyError as e: logger.error(f"[TARIFF] Ошибка при инициализации sort_order для группы {group_code}: {e}") await session.rollback() return False async def initialize_all_tariff_weights(session: AsyncSession) -> bool: try: result = await session.execute(select(Tariff).where(Tariff.sort_order.is_(None))) tariffs_without_weight = result.scalars().all() if not tariffs_without_weight: return True for tariff in tariffs_without_weight: await session.execute(update(Tariff).where(Tariff.id == tariff.id).values(sort_order=1)) await session.commit() await _invalidate_tariff_cache() return True except SQLAlchemyError as e: logger.error(f"[TARIFF] Ошибка при инициализации весов тарифов: {e}") await session.rollback() return False