optimization of database transactions/fixing middlewares/ cache of images

2025-10-14 01:05:52 +03:00
parent a1347dff04
commit 5a769f8972
5 changed files with 284 additions and 118 deletions
@@ -1,6 +1,6 @@
 from datetime import datetime

-from sqlalchemy import delete, exists, or_, select, update
+from sqlalchemy import delete, exists, or_, select, update, func
 from sqlalchemy.dialects.postgresql import insert
 from sqlalchemy.exc import SQLAlchemyError
 from sqlalchemy.ext.asyncio import AsyncSession
@@ -16,6 +16,7 @@ from database.models import (
    Referral,
    TemporaryData,
    User,
+    Key,
 )
 from logger import logger

@@ -44,7 +45,6 @@ async def add_user(
            )
            .on_conflict_do_nothing(index_elements=[User.tg_id])
        )
-
        await session.execute(stmt)
        await session.commit()
        logger.info(f"[DB] Новый пользователь добавлен: {tg_id} (source: {source_code})")
@@ -56,12 +56,19 @@ async def add_user(

 async def update_balance(session: AsyncSession, tg_id: int, amount: float) -> None:
    try:
-        result = await session.execute(select(User.balance).where(User.tg_id == tg_id))
-        current = result.scalar_one_or_none() or 0
-        new_balance = current + amount
-        await session.execute(update(User).where(User.tg_id == tg_id).values(balance=new_balance))
+        res = await session.execute(
+            update(User)
+            .where(User.tg_id == tg_id)
+            .values(balance=func.coalesce(User.balance, 0) + amount)
+            .returning(User.balance)
+        )
+        new_balance = res.scalar_one_or_none()
        await session.commit()
-        logger.info(f"[DB] Баланс пользователя {tg_id} обновлён: {current} → {new_balance}")
+        if new_balance is not None:
+            old_balance = new_balance - amount
+            logger.info(f"[DB] Баланс пользователя {tg_id} обновлён: {old_balance} → {new_balance}")
+        else:
+            logger.info(f"[DB] Баланс пользователя {tg_id} не изменён: пользователь не найден")
    except SQLAlchemyError as e:
        logger.error(f"[DB] Ошибка при обновлении баланса пользователя {tg_id}: {e}")
        await session.rollback()
@@ -74,14 +81,17 @@ async def check_user_exists(session: AsyncSession, tg_id: int) -> bool:


 async def get_balance(session: AsyncSession, tg_id: int) -> float:
-    result = await session.execute(select(User.balance).where(User.tg_id == tg_id))
-    balance = result.scalar_one_or_none()
-    return round(balance, 1) if balance is not None else 0.0
+    result = await session.execute(
+        select(func.coalesce(User.balance, 0.0)).where(User.tg_id == tg_id)
+    )
+    return round(float(result.scalar_one()), 1)


 async def set_user_balance(session: AsyncSession, tg_id: int, balance: float) -> None:
    try:
-        await session.execute(update(User).where(User.tg_id == tg_id).values(balance=balance))
+        await session.execute(
+            update(User).where(User.tg_id == tg_id).values(balance=balance)
+        )
        await session.commit()
    except SQLAlchemyError as e:
        logger.error(f"Ошибка при установке баланса для пользователя {tg_id}: {e}")
@@ -90,7 +100,9 @@ async def set_user_balance(session: AsyncSession, tg_id: int, balance: float) ->

 async def update_trial(session: AsyncSession, tg_id: int, status: int):
    try:
-        await session.execute(update(User).where(User.tg_id == tg_id).values(trial=status))
+        await session.execute(
+            update(User).where(User.tg_id == tg_id).values(trial=status)
+        )
        await session.commit()
        logger.info(f"[DB] Триал статус обновлён для пользователя {tg_id}: {status}")
    except SQLAlchemyError as e:
@@ -99,9 +111,10 @@ async def update_trial(session: AsyncSession, tg_id: int, status: int):


 async def get_trial(session: AsyncSession, tg_id: int) -> int:
-    result = await session.execute(select(User.trial).where(User.tg_id == tg_id))
-    trial = result.scalar_one_or_none()
-    return trial or 0
+    result = await session.execute(
+        select(func.coalesce(User.trial, 0)).where(User.tg_id == tg_id)
+    )
+    return int(result.scalar_one())


 async def upsert_user(
@@ -120,7 +133,6 @@ async def upsert_user(
            user = result.scalar_one_or_none()
            if not user:
                return None
-
            await session.execute(
                update(User)
                .where(User.tg_id == tg_id)
@@ -133,8 +145,11 @@ async def upsert_user(
                    updated_at=datetime.utcnow(),
                )
            )
+            await session.commit()
+            result = await session.execute(select(User).where(User.tg_id == tg_id))
+            return dict(result.scalar_one().__dict__)
        else:
-            await session.execute(
+            res = await session.execute(
                insert(User)
                .values(
                    tg_id=tg_id,
@@ -157,10 +172,13 @@ async def upsert_user(
                        "updated_at": datetime.utcnow(),
                    },
                )
+                .returning(User)
            )
-        await session.commit()
-        result = await session.execute(select(User).where(User.tg_id == tg_id))
-        return dict(result.scalar_one().__dict__)
+            obj = res.scalar_one()
+            await session.commit()
+            d = obj.__dict__.copy()
+            d.pop("_sa_instance_state", None)
+            return d
    except SQLAlchemyError as e:
        logger.error(f"[DB] Ошибка при UPSERT пользователя {tg_id}: {e}")
        await session.rollback()
@@ -170,29 +188,28 @@ async def upsert_user(
 async def delete_user_data(session: AsyncSession, tg_id: int):
    try:
        await session.execute(delete(Notification).where(Notification.tg_id == tg_id))
-
-        result = await session.execute(select(Gift.gift_id).where(Gift.sender_tg_id == tg_id))
-        gift_ids = [row[0] for row in result.all()]
-        if gift_ids:
-            await session.execute(delete(GiftUsage).where(GiftUsage.gift_id.in_(gift_ids)))
-
+        await session.execute(
+            delete(GiftUsage).where(
+                GiftUsage.gift_id.in_(select(Gift.gift_id).where(Gift.sender_tg_id == tg_id))
+            )
+        )
        await session.execute(delete(Gift).where(Gift.sender_tg_id == tg_id))
-
-        await session.execute(update(Gift).where(Gift.recipient_tg_id == tg_id).values(recipient_tg_id=None))
-
+        await session.execute(
+            update(Gift).where(Gift.recipient_tg_id == tg_id).values(recipient_tg_id=None)
+        )
        await session.execute(delete(Payment).where(Payment.tg_id == tg_id))
        await session.execute(
-            delete(Referral).where(or_(Referral.referrer_tg_id == tg_id, Referral.referred_tg_id == tg_id))
+            delete(Referral).where(
+                or_(Referral.referrer_tg_id == tg_id, Referral.referred_tg_id == tg_id)
+            )
        )
        await session.execute(delete(CouponUsage).where(CouponUsage.user_id == tg_id))
        await delete_key(session, tg_id)
        await session.execute(delete(TemporaryData).where(TemporaryData.tg_id == tg_id))
        await session.execute(delete(BlockedUser).where(BlockedUser.tg_id == tg_id))
        await session.execute(delete(User).where(User.tg_id == tg_id))
-
        await session.commit()
        logger.info(f"[DB] Данные пользователя {tg_id} полностью удалены")
-
    except SQLAlchemyError as e:
        await session.rollback()
        logger.error(f"[DB] Ошибка при удалении данных пользователя {tg_id}: {e}")
@@ -202,3 +219,33 @@ async def delete_user_data(session: AsyncSession, tg_id: int):
 async def mark_trial_extended(tg_id: int, session: AsyncSession):
    await session.execute(update(User).where(User.tg_id == tg_id).values(trial=-1))
    await session.commit()
+
+
+async def get_user_snapshot(session: AsyncSession, tg_id: int) -> tuple[int, int] | None:
+    res = await session.execute(
+        select(func.coalesce(User.trial, 0), func.count(Key.client_id))
+        .select_from(User)
+        .join(Key, Key.tg_id == User.tg_id, isouter=True)
+        .where(User.tg_id == tg_id)
+        .group_by(User.tg_id, User.trial)
+    )
+    row = res.first()
+    if row is None:
+        return None
+    return int(row[0]), int(row[1])
+
+
+async def upsert_source_if_empty(session: AsyncSession, tg_id: int, source_code: str) -> None:
+    if not source_code:
+        return
+    stmt = (
+        insert(User)
+        .values(tg_id=tg_id, source_code=source_code)
+        .on_conflict_do_update(
+            index_elements=[User.tg_id],
+            set_={"source_code": insert(User).excluded.source_code},
+            where=(User.source_code.is_(None)),
+        )
+    )
+    await session.execute(stmt)
+    await session.commit()