From e4ffa5d77affa489725516a82ecf5f451a702149 Mon Sep 17 00:00:00 2001 From: arabcoders Date: Mon, 15 Dec 2025 17:22:41 +0300 Subject: [PATCH] Fix: Fix performance regression while loading items. --- app/library/DataStore.py | 8 +- app/library/DownloadQueue.py | 37 +++-- app/library/Utils.py | 6 +- app/routes/socket/connection.py | 2 +- app/scripts/seed_db.py | 240 ++++++++++++++++++++++++++++++++ 5 files changed, 273 insertions(+), 20 deletions(-) create mode 100644 app/scripts/seed_db.py diff --git a/app/library/DataStore.py b/app/library/DataStore.py index 9fcad9d3..7d668189 100644 --- a/app/library/DataStore.py +++ b/app/library/DataStore.py @@ -2,6 +2,7 @@ import copy import json import logging from collections import OrderedDict +from dataclasses import fields from datetime import UTC, datetime from email.utils import formatdate from enum import Enum @@ -14,6 +15,9 @@ from .Utils import init_class LOG: logging.Logger = logging.getLogger("datastore") +ITEM_DTO_FIELDS = {f.name for f in fields(ItemDTO)} +"""Caching ItemDTO fields for performance.""" + class StoreType(str, Enum): HISTORY: str = "done" @@ -149,7 +153,7 @@ class DataStore: rowDate: datetime = 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: ItemDTO = init_class(ItemDTO, data, ITEM_DTO_FIELDS) item._id = row["id"] item.datetime = formatdate(rowDate.replace(tzinfo=UTC).timestamp()) items.append((row["id"], item)) @@ -294,7 +298,7 @@ class DataStore: rowDate: datetime = 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: ItemDTO = init_class(ItemDTO, data, ITEM_DTO_FIELDS) item._id = row["id"] item.datetime = formatdate(rowDate.replace(tzinfo=UTC).timestamp()) items.append((row["id"], item)) diff --git a/app/library/DownloadQueue.py b/app/library/DownloadQueue.py index 245c6d58..544415f6 100644 --- a/app/library/DownloadQueue.py +++ b/app/library/DownloadQueue.py @@ -1003,33 +1003,40 @@ class DownloadQueue(metaclass=Singleton): return status - def get(self) -> dict[str, list[dict[str, ItemDTO]]]: + def get(self, mode: str = "all") -> dict[str, list[dict[str, ItemDTO]]]: """ Get the download queue and the download history. + Args: + mode (str): The mode to get the items. Supported modes are 'all', 'queue', 'done'. + Returns: dict: The download queue and the download history. """ items = {"queue": {}, "done": {}} - for k, v in self.queue.saved_items(): - items["queue"][k] = self._active[k].info if k in self._active else v - - for k, v in self.done.saved_items(): - v.get_file_sidecar() - items["done"][k] = v - - for k, v in self.queue.items(): - if k not in items["queue"]: + if mode in ("all", "queue"): + for k, v in self.queue.saved_items(): items["queue"][k] = self._active[k].info if k in self._active else v - for k, v in self.done.items(): - if k in items["done"]: - continue + if mode in ("all", "done"): + for k, v in self.done.saved_items(): + v.get_file_sidecar() + items["done"][k] = v - v.info.get_file_sidecar() - items["done"][k] = v.info + if mode in ("all", "queue"): + for k, v in self.queue.items(): + if k not in items["queue"]: + items["queue"][k] = self._active[k].info if k in self._active else v + + if mode in ("all", "done"): + for k, v in self.done.items(): + if k in items["done"]: + continue + + v.info.get_file_sidecar() + items["done"][k] = v.info return items diff --git a/app/library/Utils.py b/app/library/Utils.py index 876c6223..4ba05617 100644 --- a/app/library/Utils.py +++ b/app/library/Utils.py @@ -1543,13 +1543,14 @@ def delete_dir(dir: Path) -> bool: return False -def init_class(cls: type[T], data: dict) -> T: +def init_class(cls: type[T], data: dict, _fields=None) -> T: """ Initialize a class instance with data from a dictionary, filtering out keys not present in the class fields. Args: cls (type): The class to initialize. data (dict): The data to use for initialization. + _fields (set, optional): A set of field names to consider. If None, all fields from the class will be used. Returns: T: An instance of the class initialized with the provided data. @@ -1557,7 +1558,8 @@ def init_class(cls: type[T], data: dict) -> T: """ from dataclasses import fields - return cls(**{k: v for k, v in data.items() if k in {f.name for f in fields(cls)}}) + _classMembers = _fields if _fields is not None else {f.name for f in fields(cls)} + return cls(**{k: v for k, v in data.items() if k in _classMembers}) def load_modules(root_path: Path, directory: Path): diff --git a/app/routes/socket/connection.py b/app/routes/socket/connection.py index d16a5e95..64586289 100644 --- a/app/routes/socket/connection.py +++ b/app/routes/socket/connection.py @@ -45,7 +45,7 @@ async def connect(config: Config, queue: DownloadQueue, notify: EventBus, sid: s depth_limit=config.download_path_depth - 1, ), "history_count": queue.done.get_total_count(), - "queue": queue.get()["queue"], + "queue": queue.get("queue")["queue"], }, title="Sending initial download data", message=f"Sending initial download data to client '{sid}'.", diff --git a/app/scripts/seed_db.py b/app/scripts/seed_db.py new file mode 100644 index 00000000..6871118c --- /dev/null +++ b/app/scripts/seed_db.py @@ -0,0 +1,240 @@ +#!/usr/bin/env python3 +from __future__ import annotations + +import argparse +import logging +import re +import sqlite3 +import sys +import time +import uuid +from datetime import UTC, datetime +from email.utils import formatdate +from itertools import cycle, islice +from pathlib import Path +from typing import TYPE_CHECKING + +APP_ROOT = str((Path(__file__).parent / ".." / "..").resolve()) +if APP_ROOT not in sys.path: + sys.path.insert(0, APP_ROOT) + +from app.library.DataStore import StoreType +from app.library.encoder import Encoder + +if TYPE_CHECKING: + from collections.abc import Iterable, Iterator, Sequence + +LOG = logging.getLogger("seed_db") + +USED_IDS: set[str] = set() + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="Seed history/queue with synthetic records for performance testing.") + parser.add_argument( + "--root", + type=Path, + default=Path("/storage/media/youtube"), + help="Root folder to scan for media files", + ) + parser.add_argument("--count", type=int, default=10_000, help="Total records to insert (default: 10000).") + parser.add_argument("--batch", type=int, default=500, help="Insert batch size (default: 500).") + parser.add_argument( + "--store", + choices=[StoreType.HISTORY.value, StoreType.QUEUE.value], + default=StoreType.HISTORY.value, + help="Target store type (done|queue). Default: done.", + ) + parser.add_argument( + "--extension", + default=".mkv", + help="File extension to look for (default: .mkv).", + ) + parser.add_argument( + "--status", + default="finished", + help="Status value to set on seeded items (default: finished).", + ) + parser.add_argument("--preset", default="default", help="Preset to stamp on items (default: default).") + parser.add_argument( + "--dry-run", + action="store_true", + help="Scan and plan only; do not write to the database.", + ) + parser.add_argument("db_file", type=Path, help="Path to the database file.") + return parser.parse_args() + + +def connect_db(db_file: str) -> sqlite3.Connection: + conn = sqlite3.connect(database=db_file, isolation_level=None) + conn.row_factory = sqlite3.Row + conn.execute("PRAGMA journal_mode=wal") + return conn + + +def find_files(root: Path, extension: str) -> list[Path]: + if not root.exists(): + msg = f"Root path does not exist: {root}" + raise FileNotFoundError(msg) + + files = [p for p in root.rglob(f"*.{extension}") if p.is_file()] + if not files: + msg = f"No .{extension} files found under {root!s}" + raise FileNotFoundError(msg) + + return files + + +def _cycle(files: Sequence[Path]) -> Iterator[Path]: + return cycle(files) + + +def _relative_folder(path: Path, root: Path) -> tuple[str, str]: + try: + rel = path.relative_to(root) + folder = "" if rel.parent == Path(".") else rel.parent.as_posix() + filename = rel.name + except ValueError: + folder = "" + filename = path.name + + return folder, filename + + +def _build_row( + path: Path, + root: Path, + idx: int, + status: str, + preset: str, + store: StoreType, +) -> tuple[str, str, str, str, str]: + folder, filename = _relative_folder(path, root) + stat = path.stat() + + video_id = re.search(r"\[youtube-(\S+)\]", path.stem) + video_id = f"seeded_video_{idx:06d}" if not video_id else video_id.group(1) + if video_id in USED_IDS: + video_id = f"{video_id}_{uuid.uuid4().hex[:6]}" + USED_IDS.add(video_id) + + archive_id = f"youtube {video_id}" + record_id = f"seed-{idx:06d}-{uuid.uuid4().hex[:12]}" + timestamp = datetime.now(UTC) + + data = { + "_id": record_id, + "id": video_id, + "title": path.stem, + "description": "", + "url": f"https://www.youtube.com/watch?v={video_id}", + "preset": preset, + "folder": folder, + "download_dir": str(root), + "temp_dir": None, + "status": status, + "cookies": None, + "template": None, + "template_chapter": None, + "timestamp": time.time_ns(), + "is_live": False, + "datetime": formatdate(timestamp.timestamp()), + "live_in": None, + "file_size": stat.st_size, + "options": {}, + "extras": {}, + "cli": "", + "auto_start": True, + "is_archivable": True, + "is_archived": True, + "archive_id": archive_id, + "sidecar": {}, + "tmpfilename": None, + "filename": filename, + "total_bytes": stat.st_size, + "total_bytes_estimate": stat.st_size, + "downloaded_bytes": stat.st_size, + "msg": None, + "percent": 100, + "speed": None, + "eta": None, + "_recomputed": True, + "_archive_file": None, + } + + created_at = timestamp.strftime("%Y-%m-%d %H:%M:%S") + encoded = Encoder(sort_keys=True, indent=4).encode(data) + + return (record_id, str(store), data["url"], encoded, created_at) + + +def generate_rows( + files: Sequence[Path], + root: Path, + total: int, + status: str, + preset: str, + store: StoreType, +) -> Iterable[tuple[str, str, str, str, str]]: + _iter = _cycle(files) + for idx, path in enumerate(islice(_iter, total)): + yield _build_row(path, root, idx, status, preset, store) + + +def insert_batches( + conn: sqlite3.Connection, + rows: Iterable[tuple[str, str, str, str, str]], + batch_size: int, +) -> int: + sql = """ + INSERT INTO "history" ("id", "type", "url", "data", "created_at") + VALUES (?, ?, ?, ?, ?) + ON CONFLICT("id") DO UPDATE SET + "type" = excluded.type, + "url" = excluded.url, + "data" = excluded.data, + "created_at" = excluded.created_at + """ + total_written = 0 + batch: list[tuple[str, str, str, str, str]] = [] + + for row in rows: + batch.append(row) + if len(batch) >= batch_size: + conn.executemany(sql, batch) + total_written += len(batch) + LOG.info("Inserted %d rows...", total_written) + batch.clear() + + if batch: + conn.executemany(sql, batch) + total_written += len(batch) + + return total_written + + +def main() -> None: + args = parse_args() + logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s") + + extension = args.extension.lstrip(".") + store = StoreType.from_value(args.store) + + files = find_files(args.root, extension) + LOG.info("Found %d %s files under %s", len(files), extension, args.root) + + if args.dry_run: + LOG.info("Dry run only. Would insert %d rows into '%s'.", args.count, store.value) + return + + rows = generate_rows(files, args.root, args.count, args.status, args.preset, store) + conn = connect_db(args.db_file) + try: + written = insert_batches(conn, rows, args.batch) + finally: + conn.close() + + LOG.info("Done. Inserted %d rows into '%s'.", written, store.value) + + +if __name__ == "__main__": + main()