PVideoDl/app/core/storage.py

"""Абстракция хранилища загрузок.

Интерфейс Storage + реализация на aiosqlite. Сегодня SQLite, завтра Postgres —
сервисы и роуты зовут только методы Storage и не знают о бэкенде.
"""

from __future__ import annotations

import json
from abc import ABC, abstractmethod
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

import aiosqlite

from app.models import Download, DownloadStatus


def _now() -> datetime:
    return datetime.now(timezone.utc)


class Storage(ABC):
    @abstractmethod
    async def init(self) -> None: ...

    @abstractmethod
    async def close(self) -> None: ...

    @abstractmethod
    async def create(self, download: Download) -> None: ...

    @abstractmethod
    async def update(self, download_id: str, **fields: Any) -> Download | None: ...

    @abstractmethod
    async def get(self, download_id: str) -> Download | None: ...

    @abstractmethod
    async def list(self) -> list[Download]: ...

    @abstractmethod
    async def delete(self, download_id: str) -> bool: ...


_COLUMNS = (
    "id",
    "url",
    "filename",
    "status",
    "progress",
    "size_bytes",
    "downloaded_bytes",
    "speed",
    "eta",
    "error",
    "created_at",
    "finished_at",
)


def _row_to_download(row: aiosqlite.Row) -> Download:
    data = dict(row)
    data["status"] = DownloadStatus(data["status"])
    data["created_at"] = datetime.fromisoformat(data["created_at"])
    if data.get("finished_at"):
        data["finished_at"] = datetime.fromisoformat(data["finished_at"])
    return Download.model_validate(data)


def _to_db_value(key: str, value: Any) -> Any:
    if isinstance(value, DownloadStatus):
        return value.value
    if isinstance(value, datetime):
        return value.isoformat()
    return value


class SqliteStorage(Storage):
    def __init__(self, db_path: Path) -> None:
        self._db_path = db_path
        self._db: aiosqlite.Connection | None = None

    @property
    def _conn(self) -> aiosqlite.Connection:
        if self._db is None:
            raise RuntimeError("Storage не инициализирован — вызовите init()")
        return self._db

    async def init(self) -> None:
        self._db_path.parent.mkdir(parents=True, exist_ok=True)
        self._db = await aiosqlite.connect(self._db_path)
        self._db.row_factory = aiosqlite.Row
        await self._db.execute(
            """
            CREATE TABLE IF NOT EXISTS downloads (
                id               TEXT PRIMARY KEY,
                url              TEXT NOT NULL,
                filename         TEXT,
                status           TEXT NOT NULL,
                progress         REAL NOT NULL DEFAULT 0,
                size_bytes       INTEGER,
                downloaded_bytes INTEGER NOT NULL DEFAULT 0,
                speed            REAL,
                eta              REAL,
                error            TEXT,
                created_at       TEXT NOT NULL,
                finished_at      TEXT
            )
            """
        )
        # Миграция старых БД, созданных до появления finished_at.
        await self._migrate_add_column("finished_at", "TEXT")
        # На старте всё, что осталось "качающимся" после прошлого запуска,
        # помечаем упавшим — воркеры этого процесса о них не знают.
        await self._db.execute(
            "UPDATE downloads SET status = ?, error = ?, finished_at = ? "
            "WHERE status = ?",
            (
                DownloadStatus.FAILED.value,
                "Прервано при перезапуске приложения",
                _now().isoformat(),
                DownloadStatus.DOWNLOADING.value,
            ),
        )
        await self._db.commit()

    async def _migrate_add_column(self, name: str, decl: str) -> None:
        """Добавить колонку, если её ещё нет (idempotent-миграция старых БД)."""
        async with self._conn.execute("PRAGMA table_info(downloads)") as cursor:
            cols = {row["name"] for row in await cursor.fetchall()}
        if name not in cols:
            await self._conn.execute(
                f"ALTER TABLE downloads ADD COLUMN {name} {decl}"
            )

    async def close(self) -> None:
        if self._db is not None:
            await self._db.close()
            self._db = None

    async def create(self, download: Download) -> None:
        await self._conn.execute(
            f"INSERT INTO downloads ({', '.join(_COLUMNS)}) "
            f"VALUES ({', '.join('?' for _ in _COLUMNS)})",
            tuple(_to_db_value(c, getattr(download, c)) for c in _COLUMNS),
        )
        await self._conn.commit()

    async def update(self, download_id: str, **fields: Any) -> Download | None:
        if not fields:
            return await self.get(download_id)
        allowed = {k: v for k, v in fields.items() if k in _COLUMNS and k != "id"}
        if not allowed:
            return await self.get(download_id)
        assignments = ", ".join(f"{k} = ?" for k in allowed)
        params = [_to_db_value(k, v) for k, v in allowed.items()]
        params.append(download_id)
        await self._conn.execute(
            f"UPDATE downloads SET {assignments} WHERE id = ?", params
        )
        await self._conn.commit()
        return await self.get(download_id)

    async def get(self, download_id: str) -> Download | None:
        async with self._conn.execute(
            "SELECT * FROM downloads WHERE id = ?", (download_id,)
        ) as cursor:
            row = await cursor.fetchone()
        return _row_to_download(row) if row else None

    async def list(self) -> list[Download]:
        async with self._conn.execute(
            "SELECT * FROM downloads ORDER BY created_at DESC"
        ) as cursor:
            rows = await cursor.fetchall()
        return [_row_to_download(r) for r in rows]

    async def delete(self, download_id: str) -> bool:
        cursor = await self._conn.execute(
            "DELETE FROM downloads WHERE id = ?", (download_id,)
        )
        await self._conn.commit()
        return cursor.rowcount > 0