rumine-api-wrapper/db/database.py

from sqlalchemy import select, delete
from sqlalchemy.ext.asyncio import create_async_engine
from sqlalchemy.orm import sessionmaker
from db import LuckpermsPlayer, LitebansBan, LuckpermsUserPermission, TabUser, Subscribe, Whitelist
from db.schemas import MySQLConfig, UserSchema, RequestSchema, TabSchema
from sqlalchemy.ext.asyncio import AsyncSession
from aiocache import cached, SimpleMemoryCache
from aiocache.serializers import PickleSerializer
import datetime
import httpx
from config import configs


class Database:
    def __init__(self, uri: MySQLConfig):
        self.engine = create_async_engine(str(uri))
        self.AsyncSessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=self.engine, class_=AsyncSession)
        self.uri = uri

    @cached(ttl=3600, cache=SimpleMemoryCache, serializer=PickleSerializer())
    async def get_uuid_by_username(self, request: RequestSchema):
        async with self.AsyncSessionLocal() as session:
            async with session.begin():
                player = await session.execute(select(LuckpermsPlayer).filter(LuckpermsPlayer.username == request.username))
                player = player.scalar_one_or_none()
                if player:
                    return player.uuid
                else:
                    return None

    async def unban_by_uuid(self, user: UserSchema):
        async with self.AsyncSessionLocal() as session:
            bans = await session.execute(
                select(LitebansBan).filter(
                    LitebansBan.uuid == user.uuid,
                    LitebansBan.active == True
                )
            )
            bans = bans.scalars().all()
            for ban in bans:
                ban.active = False
                ban.removed_by_uuid = "fa87b50a-791b-4c61-a56f-d6cc03df1582"
                ban.removed_by_name = "The_MrKroll"
                ban.removed_by_reason = "Куплен разбан на сайте"
            await session.commit()

    async def grant_permissions_by_uuid(self, user: UserSchema):
        async with self.AsyncSessionLocal() as session:
            permission = LuckpermsUserPermission(
                uuid=str(user.uuid),
                permission="group.subscribe",
                value="1",
                server="global",
                world="global",
                expiry=str(user.expiry),
                contexts="{}"
            )
            session.add(permission)
            await session.commit()

    async def grant_tab_by_username(self, user: UserSchema):
        async with self.AsyncSessionLocal() as session:
            prefix = TabUser(
                user=user.username,
                property=user.tab.property,
                value=user.tab.value,
                expiry=user.expiry
            )
            session.add(prefix)
            await session.commit()

    async def check_subscription_by_uuid_old(self, user: UserSchema) -> bool:
        async with self.AsyncSessionLocal() as session:
            subscription = await session.execute(
                select(LuckpermsUserPermission).filter(
                    LuckpermsUserPermission.uuid == user.uuid,
                    LuckpermsUserPermission.permission == "group.subscribe"
                )
            )
            if subscription.scalar():
                return True
            else:
                return False
    
    async def check_subscription_by_nickname(self, user: UserSchema) -> bool:
        async with self.AsyncSessionLocal() as session:
            subscription = await session.execute(
                select(Subscribe).filter(
                    Subscribe.nickname == user.username,
                    Subscribe.expiry > datetime.datetime.now()
                )
            )
            if subscription.scalar():
                return True
            else:
                return False
            
    @cached(ttl=2, cache=SimpleMemoryCache, serializer=PickleSerializer())
    async def check_ban_status_by_uuid(self, user: UserSchema) -> bool:
        async with self.AsyncSessionLocal() as session:
            ban = await session.execute(
                select(LitebansBan).filter(
                    LitebansBan.uuid == user.uuid,
                    LitebansBan.active == True
                )
            )
            if not ban.scalar():
                return True
            else:
                return False
            
    async def delete_expired_tab_users(self):
        async with self.AsyncSessionLocal() as session:
            async with session.begin():
                current_timestamp = int((datetime.datetime.utcnow()).timestamp())
                await session.execute(
                    delete(TabUser)
                    .where(TabUser.expiry < current_timestamp)
                    .where(TabUser.expiry != 0) 
                )
                await session.commit()

    async def get_discord(self, user: UserSchema):
        discord = None
        async with self.AsyncSessionLocal() as session:
            try:
                result = await session.execute(select(Whitelist).where(Whitelist.player == user.username))
                record = result.fetchone()
                if record:
                    discord = record[0].discord
                else:
                    discord = None
            except Exception as e:
                print("Ошибка при выполнении запроса:", e)

        if isinstance(discord, str) and len(discord) > 8:
            return discord
        else:
            return None
        
    async def update_subscription(self, user: UserSchema):
        async with self.AsyncSessionLocal() as session:
            try:
                await self.delete_expired_subscriptions()
                await self.delete_expired_tab_users()
                expiry_datetime = datetime.datetime.fromtimestamp(user.expiry, datetime.timezone.utc)
                if user.has_sub:
                    subscription = await session.execute(select(Subscribe).where(Subscribe.discord == user.discord_id and Subscribe.nickname == user.username))
                    subscription = subscription.fetchone()[0]
                    if subscription:
                        updated_expiry = subscription.expiry + (expiry_datetime - datetime.datetime.now(datetime.timezone.utc))
                        subscription.expiry = updated_expiry
                    else:
                        new_subscription = Subscribe(discord=user.discord_id, nickname=user.username, expiry=expiry_datetime)
                        session.add(new_subscription)
                else:
                    new_subscription = Subscribe(discord=user.discord_id, nickname=user.username, expiry=expiry_datetime)
                    session.add(new_subscription)
                await session.commit()
                await self.grant_tab_by_username(user=user)
                await self.grant_permissions_by_uuid(user=user)
                await discord_role(str(user.discord_id))
            except Exception as e:
                print("Ошибка при выполнении запроса:", e)
                await session.rollback()

    async def delete_expired_subscriptions(self):
        async with self.AsyncSessionLocal() as session:
            try:
                await session.execute(delete(Subscribe).where(Subscribe.expiry < datetime.datetime.now(datetime.timezone.utc)))
                await session.commit()
            except Exception as e:
                print("Ошибка при удалении устаревших подписок:", e)
                await session.rollback()
    
async def discord_role(discord_id: int):
    params = {"user_id": discord_id, "token": configs.TOKEN_DS}
    
    async with httpx.AsyncClient() as client:
        try:
            response = await client.get(configs.DS_URL, params=params)
            data = response.json()
            return data
        except httpx.HTTPError as e:
            print(f"Ошибка при отправке запроса: {e}")
            return {"error": "Произошла ошибка при выполнении запроса"}