79 lines
2.9 KiB
Python
79 lines
2.9 KiB
Python
from __future__ import annotations
|
|
|
|
from typing import Any
|
|
|
|
from sqlalchemy import text
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from database import bulk_add_notifications, bulk_delete_notifications
|
|
from database.models import Key
|
|
from logger import logger
|
|
|
|
|
|
def create_bulk_updates() -> dict[str, Any]:
|
|
return {
|
|
"balance_changes": {},
|
|
"key_expiry_updates": [],
|
|
"key_tariff_updates": [],
|
|
"notifications_to_add": [],
|
|
"notifications_to_delete": [],
|
|
}
|
|
|
|
|
|
async def execute_bulk_updates(session: AsyncSession, bulk_updates: dict[str, Any]) -> None:
|
|
try:
|
|
balance_changes = bulk_updates.get("balance_changes") or {}
|
|
if balance_changes:
|
|
tg_ids = list(balance_changes.keys())
|
|
changes = [balance_changes[tg_id] for tg_id in tg_ids]
|
|
await session.execute(
|
|
text(
|
|
"UPDATE users SET balance = balance + v.change FROM "
|
|
"(SELECT unnest(CAST(:tg_ids AS bigint[])) AS tg_id, unnest(CAST(:changes AS double precision[])) AS change) AS v "
|
|
"WHERE users.tg_id = v.tg_id"
|
|
),
|
|
{"tg_ids": tg_ids, "changes": changes},
|
|
)
|
|
logger.info(f"Bulk: обновлено {len(balance_changes)} балансов")
|
|
|
|
key_expiry = bulk_updates.get("key_expiry_updates") or []
|
|
if key_expiry:
|
|
await session.run_sync(
|
|
lambda sync_sess: sync_sess.bulk_update_mappings(
|
|
Key,
|
|
[{"client_id": cid, "expiry_time": exp} for cid, exp in key_expiry],
|
|
)
|
|
)
|
|
logger.info(f"Bulk: обновлено {len(key_expiry)} сроков действия ключей")
|
|
|
|
key_tariff = bulk_updates.get("key_tariff_updates") or []
|
|
if key_tariff:
|
|
await session.run_sync(
|
|
lambda sync_sess: sync_sess.bulk_update_mappings(
|
|
Key,
|
|
[{"client_id": cid, "tariff_id": tid} for cid, tid in key_tariff],
|
|
)
|
|
)
|
|
logger.info(f"Bulk: обновлено {len(key_tariff)} тарифов ключей")
|
|
|
|
to_add = bulk_updates.get("notifications_to_add") or []
|
|
if to_add:
|
|
await bulk_add_notifications(session, to_add, commit=False)
|
|
|
|
to_delete = bulk_updates.get("notifications_to_delete") or []
|
|
if to_delete:
|
|
await bulk_delete_notifications(session, to_delete, commit=False)
|
|
|
|
if to_add or to_delete:
|
|
logger.info(f"Bulk: {len(to_add)} добавлений, {len(to_delete)} удалений уведомлений")
|
|
|
|
await session.commit()
|
|
|
|
except Exception as error:
|
|
logger.error(f"Ошибка в bulk-обновлениях: {error}")
|
|
try:
|
|
await session.rollback()
|
|
except Exception:
|
|
pass
|
|
raise
|