misskey2telegram/models/services/user.py

from typing import cast, Optional

from sqlalchemy import select
from sqlmodel.ext.asyncio.session import AsyncSession

from init import sqlite
from models.models.user import User, TokenStatusEnum


class UserAction:
    @staticmethod
    async def add_user(user: User):
        async with sqlite.session() as session:
            session = cast(AsyncSession, session)
            session.add(user)
            await session.commit()

    @staticmethod
    async def get_user_by_id(user_id: int) -> Optional[User]:
        async with sqlite.session() as session:
            session = cast(AsyncSession, session)
            statement = select(User).where(User.user_id == user_id)
            results = await session.exec(statement)
            return user[0] if (user := results.first()) else None

    @staticmethod
    async def get_user_if_ok(user_id: int) -> Optional[User]:
        async with sqlite.session() as session:
            session = cast(AsyncSession, session)
            statement = select(User).where(
                User.user_id == user_id
            ).where(
                User.status == TokenStatusEnum.STATUS_SUCCESS
            ).where(
                User.chat_id != 0
            ).where(
                User.timeline_topic != 0
            ).where(
                User.notice_topic != 0
            ).where(
                User.token != ""
            )
            results = await session.exec(statement)
            return user[0] if (user := results.first()) else None

    @staticmethod
    async def get_all_token_ok_users() -> list[User]:
        async with sqlite.session() as session:
            session = cast(AsyncSession, session)
            statement = select(User).where(
                User.status == TokenStatusEnum.STATUS_SUCCESS
            ).where(
                User.chat_id != 0
            ).where(
                User.timeline_topic != 0
            ).where(
                User.notice_topic != 0
            ).where(
                User.token != ""
            )
            results = await session.exec(statement)
            users = results.all()
            return [user[0] for user in users]

    @staticmethod
    async def update_user(user: User):
        async with sqlite.session() as session:
            session = cast(AsyncSession, session)
            session.add(user)
            await session.commit()
            await session.refresh(user)

    @staticmethod
    async def add_or_update_user(user: User):
        if await UserAction.get_user_by_id(user.user_id):
            await UserAction.update_user(user)
        else:
            await UserAction.add_user(user)

    @staticmethod
    async def set_user_status(user_id: int, status: TokenStatusEnum) -> bool:
        user = await UserAction.get_user_by_id(user_id)
        if not user:
            return False
        user.status = status
        await UserAction.update_user(user)
        return True

    @staticmethod
    async def change_user_token(user_id: int, token: str) -> bool:
        user = await UserAction.get_user_by_id(user_id)
        if not user:
            user = User(user_id=user_id, token=token, status=TokenStatusEnum.STATUS_SUCCESS)
        user.token = token
        user.status = TokenStatusEnum.STATUS_SUCCESS
        await UserAction.update_user(user)
        return True

    @staticmethod
    async def change_user_group_id(user_id: int, chat_id: int) -> bool:
        user = await UserAction.get_user_by_id(user_id)
        if not user:
            return False
        user.chat_id = chat_id
        await UserAction.update_user(user)
        return True

    @staticmethod
    async def change_user_timeline(user_id: int, timeline: int) -> bool:
        user = await UserAction.get_user_by_id(user_id)
        if not user:
            return False
        if user.notice_topic == timeline:
            return False
        user.timeline_topic = timeline
        await UserAction.update_user(user)
        return True

    @staticmethod
    async def change_user_notice(user_id: int, notice: int) -> bool:
        user = await UserAction.get_user_by_id(user_id)
        if not user:
            return False
        if user.timeline_topic == notice:
            return False
        user.notice_topic = notice
        await UserAction.update_user(user)
        return True