anicli/db/engine.py - master - dig/anicli_ru

anicli/db/engine.py
92 строки3 KB
An0nX
refactor: modernize type hints, improve aniskip merging and history fallback
09 янв 2026, 00:43
09 янв 2026, 00:43925b7f8
from pathlib import Path
from typing import Any

from sqlalchemy import text
from sqlalchemy.ext.asyncio import (
    AsyncEngine,
    AsyncSession,
    async_sessionmaker,
    create_async_engine,
)

from anicli.db.models import Base

DB_PATH = Path("anicli.db")
DATABASE_URL = f"sqlite+aiosqlite:///{DB_PATH}"


class DBManager:
    """Manages the async database connection and migrations."""

    def __init__(self) -> None:
        """Initializes the DBManager."""
        self.engine: AsyncEngine = create_async_engine(DATABASE_URL, echo=False)
        self.session_factory = async_sessionmaker(
            self.engine, expire_on_commit=False, class_=AsyncSession
        )

    async def init_db(self) -> None:
        """Initializes database tables and performs simple migrations."""
        async with self.engine.begin() as conn:
            await conn.run_sync(Base.metadata.create_all)

            # --- Simple Migration Logic for anime_progress table ---
            progress_columns: list[tuple[str, str]] = [
                ("shikimori_id", "INTEGER"),
                ("shikimori_rate_id", "INTEGER"),
                ("shikimori_status", "VARCHAR"),
                ("score", "INTEGER DEFAULT 0"),
                ("total_episodes", "INTEGER DEFAULT 0"),
                ("rewatches", "INTEGER DEFAULT 0"),
                ("needs_correction", "BOOLEAN DEFAULT 0"),
                ("shikimori_title", "VARCHAR"),
                ("bound_title", "VARCHAR"),
                ("bound_similarity", "FLOAT"),
            ]
            await self._add_columns(conn, "anime_progress", progress_columns)

            # --- Simple Migration Logic for predicted_skip_time table ---
            skip_time_columns: list[tuple[str, str]] = [
                ("episode_length", "FLOAT DEFAULT 0.0"),
            ]
            await self._add_columns(conn, "predicted_skip_time", skip_time_columns)

            # --- Drop old search_cache table if it exists ---
            try:
                await conn.execute(text("DROP TABLE IF EXISTS search_cache"))
            except Exception:
                # Table likely doesn't exist, which is fine.
                pass

    @staticmethod
    async def _add_columns(
        conn: Any, table: str, columns: list[tuple[str, str]]
    ) -> None:
        """
        Safely adds multiple columns to a table if they don't exist.

        Args:
            conn: SQLAlchemy async connection object.
            table: The name of the table to alter.
            columns: A list of (column_name, column_type) tuples.
        """
        for col_name, col_type in columns:
            try:
                await conn.execute(
                    text(f"ALTER TABLE {table} ADD COLUMN {col_name} {col_type}")
                )
            except Exception:
                # Column likely exists or another error, ignore safely.
                pass

    def get_session(self) -> AsyncSession:
        """
        Returns a new async session.

        Returns:
            AsyncSession: A new database session.
        """
        return self.session_factory()


db_manager = DBManager()