ytptube/app/library/DataStore.py

import asyncio
import copy
import json
import logging
from collections import OrderedDict
from datetime import UTC, datetime
from email.utils import formatdate
from enum import Enum
from sqlite3 import Connection

from .Download import Download
from .ItemDTO import ItemDTO
from .Utils import init_class

LOG = logging.getLogger("datastore")


class StoreType(str, Enum):
    HISTORY = "done"
    QUEUE = "queue"

    @classmethod
    def all(cls) -> list[str]:
        return [member.value for member in cls]

    @classmethod
    def from_value(cls, value: str) -> "StoreType":
        """
        Returns the StoreType enum member corresponding to the given value.

        Args:
            value (str): The value to match against the enum members.

        Returns:
            StoreType: The enum member that matches the value.

        Raises:
            ValueError: If the value does not match any member.

        """
        for member in cls:
            if member.value == value:
                return member

        msg = f"Invalid StoreType value: {value}"
        raise ValueError(msg)

    def __str__(self) -> str:
        return self.value


class DataStore:
    """
    Persistent queue.
    """

    _type: StoreType = None
    """Type of the store, e.g., DONE, QUEUE, PENDING."""

    _dict: OrderedDict[str, Download] = None
    """Ordered dictionary to store Download objects."""

    _connection: Connection
    """SQLite connection to the database."""

    def __init__(self, type: StoreType, connection: Connection):
        self._dict = OrderedDict()
        self._type = type
        self._connection = connection

    def load(self) -> None:
        for id, item in self.saved_items():
            self._dict.update({id: Download(info=item)})

    def exists(self, key: str | None = None, url: str | None = None) -> bool:
        if not key and not url:
            msg = "key or url must be provided."
            raise KeyError(msg)

        if key and key in self._dict:
            return True

        return any(
            (key and self._dict[i].info._id == key) or (url and self._dict[i].info.url == url) for i in self._dict
        )

    def get(self, key: str | None = None, url: str | None = None) -> Download:
        if not key and not url:
            msg = "key or url must be provided."
            raise KeyError(msg)

        for i in self._dict:
            if (key and self._dict[i].info._id == key) or (url and self._dict[i].info.url == url):
                return self._dict[i]

        msg: str = f"{key=} or {url=} not found."
        raise KeyError(msg)

    def get_by_id(self, id: str) -> Download | None:
        return self._dict.get(id, None)

    def items(self) -> list[tuple[str, Download]]:
        return self._dict.items()

    def saved_items(self) -> list[tuple[str, ItemDTO]]:
        items: list[tuple[str, ItemDTO]] = []

        cursor = self._connection.execute(
            'SELECT "id", "data", "created_at" FROM "history" WHERE "type" = ? ORDER BY "created_at" ASC',
            (str(self._type),),
        )

        for row in cursor:
            rowDate = datetime.strptime(row["created_at"], "%Y-%m-%d %H:%M:%S")  # noqa: DTZ007
            data: dict = json.loads(row["data"])
            data.pop("_id", None)
            item: ItemDTO = init_class(ItemDTO, data)
            item._id = row["id"]
            item.datetime = formatdate(rowDate.replace(tzinfo=UTC).timestamp())
            items.append((row["id"], item))

        return items

    def put(self, value: Download) -> Download:
        if "error" == value.info.status:
            from app.library.Events import EventBus, Events

            asyncio.create_task(EventBus.get_instance().emit(Events.ITEM_ERROR, value.info), name="emit_item_error")

        self._dict.update({value.info._id: value})
        self._update_store_item(self._type, value.info)

        return self._dict[value.info._id]

    def delete(self, key: str) -> None:
        self._dict.pop(key, None)
        self._delete_store_item(key)

    def next(self) -> tuple[str, Download]:
        return next(iter(self._dict.items()))

    def empty(self):
        return not bool(self._dict)

    def has_downloads(self):
        if 0 == len(self._dict):
            return False

        return any(self._dict[key].info.auto_start and self._dict[key].started() is False for key in self._dict)

    def get_next_download(self) -> Download:
        for key in self._dict:
            ref = self._dict[key]
            if ref.info.auto_start and ref.started() is False and ref.is_cancelled() is False:
                return ref

        return None

    async def test(self) -> bool:
        self._connection.execute('SELECT "id" FROM "history" LIMIT 1').fetchone()
        return True

    def _update_store_item(self, type: StoreType, item: ItemDTO) -> None:
        sqlStatement = """
        INSERT INTO "history" ("id", "type", "url", "data")
        VALUES (?, ?, ?, ?)
        ON CONFLICT DO UPDATE SET "type" = ?, "url" = ?, "data" = ?, created_at = ?
        """

        stored: ItemDTO = copy.deepcopy(item)

        if hasattr(stored, "datetime"):
            try:
                delattr(stored, "datetime")
            except AttributeError:
                pass

        if hasattr(stored, "live_in") and stored.status == "finished":
            try:
                delattr(stored, "live_in")
            except AttributeError:
                pass

        self._connection.execute(
            sqlStatement.strip(),
            (
                stored._id,
                str(type),
                stored.url,
                stored.json(),
                str(type),
                stored.url,
                stored.json(),
                datetime.now(UTC).strftime("%Y-%m-%d %H:%M:%S"),
            ),
        )

    def _delete_store_item(self, key: str) -> None:
        self._connection.execute('DELETE FROM "history" WHERE "type" = ? AND "id" = ?', (str(self._type), key))