isekai-toolkit/api/model/toolkit_ui/page_title.py

from __future__ import annotations
import time

from typing import Optional
import sqlalchemy
from sqlalchemy import select, update
from sqlalchemy.orm import mapped_column, Mapped

from api.model.base import BaseHelper, BaseModel
from service.database import DatabaseService


class PageTitleModel(BaseModel):
    __tablename__ = "toolkit_ui_page_title"

    id: Mapped[int] = mapped_column(
        sqlalchemy.Integer, primary_key=True, autoincrement=True)
    page_id: Mapped[int] = mapped_column(sqlalchemy.Integer, index=True)
    title: Mapped[str] = mapped_column(sqlalchemy.String(255), nullable=True)
    updated_at: Mapped[int] = mapped_column(sqlalchemy.BigInteger, index=True)


class PageTitleHelper(BaseHelper):
    async def find_by_page_id(self, page_id: int):
        stmt = select(PageTitleModel).where(PageTitleModel.page_id == page_id)
        return await self.session.scalar(stmt)

    async def find_by_title(self, title: str):
        stmt = select(PageTitleModel).where(PageTitleModel.title == title)
        return await self.session.scalar(stmt)

    async def get_page_id_by_title(self, title: str):
        obj = await self.find_by_title(title)
        if obj is None:
            return None
        return obj.page_id

    async def should_update(self, title: str):
        title_info = await self.find_by_title(title)
        if title_info is None:
            return True
        if time.time() - title_info.updated_at > 60:
            return True

    async def add(self, page_id: int, title: Optional[str] = None):
        obj = PageTitleModel(page_id=page_id, title=title, updated_at=int(time.time()))

        self.session.add(obj)
        await self.session.commit()
        await self.session.refresh(obj)
        return obj

    async def set_title(self, page_id: int, title: Optional[str] = None):
        stmt = update(PageTitleModel).where(
            PageTitleModel.page_id == page_id).values(title=title, updated_at=int(time.time()))
        await self.session.execute(stmt)
        await self.session.commit()

    async def update(self, obj: PageTitleModel, ignore_updated_at: bool = False):
        await self.session.merge(obj)
        if not ignore_updated_at:
            obj.updated_at = int(time.time())
        await self.session.commit()
        await self.session.refresh(obj)
        return obj