#!/usr/bin/env python3 from __future__ import annotations import argparse import asyncio import logging import re 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 from sqlalchemy import text from sqlalchemy.ext.asyncio import AsyncConnection, AsyncEngine, create_async_engine 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 from app.library.log import get_logger if TYPE_CHECKING: from collections.abc import Iterable, Iterator, Sequence LOG = get_logger() 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() async def connect_db(db_file: Path) -> tuple[AsyncConnection, AsyncEngine]: """Create async SQLAlchemy connection.""" engine = create_async_engine(f"sqlite+aiosqlite:///{db_file}") conn = await engine.connect() await conn.execute(text("PRAGMA journal_mode=wal")) await conn.commit() return conn, engine 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": { "source_name": "Tester", "source_id": "0ddc40f3-9227-4ffa-9f3c-c52bcf33ef2f", "source_handler": "YoutubeHandler", "metadata": { "published": formatdate(timestamp.timestamp()), }, "is_video": True, "is_audio": True, }, "cli": "", "auto_start": True, "is_archivable": True, "is_archived": True, "archive_id": archive_id, "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, } 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) async def insert_batches( conn: AsyncConnection, rows: Iterable[tuple[str, str, str, str, str]], batch_size: int, ) -> int: """Insert rows in batches using SQLAlchemy.""" sql = """ INSERT INTO "history" ("id", "type", "url", "data", "created_at") VALUES (:p0, :p1, :p2, :p3, :p4) 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: # Convert tuples to dicts for SQLAlchemy params = [{"p0": r[0], "p1": r[1], "p2": r[2], "p3": r[3], "p4": r[4]} for r in batch] await conn.execute(text(sql), params) await conn.commit() total_written += len(batch) LOG.info("Inserted %d rows...", total_written) batch.clear() if batch: params = [{"p0": r[0], "p1": r[1], "p2": r[2], "p3": r[3], "p4": r[4]} for r in batch] await conn.execute(text(sql), params) await conn.commit() total_written += len(batch) return total_written async 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, engine = await connect_db(args.db_file) try: written = await insert_batches(conn, rows, args.batch) finally: await conn.close() await engine.dispose() LOG.info("Done. Inserted %d rows into '%s'.", written, store.value) if __name__ == "__main__": asyncio.run(main())