diff --git a/backend/app/db/engine.py b/backend/app/db/engine.py index efb1d46..1e933e8 100644 --- a/backend/app/db/engine.py +++ b/backend/app/db/engine.py @@ -7,7 +7,11 @@ from sqlalchemy import event from sqlalchemy.engine.url import make_url from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine -SQLITE_BUSY_TIMEOUT_MS = 5_000 +# Under concurrent writers (API requests + worker jobs), large imports/backfills can +# legitimately hold the SQLite writer lock for several seconds. A higher default +# busy timeout makes the system much more resilient under load (still overrideable +# via SQLITE_BUSY_TIMEOUT_MS env). +SQLITE_BUSY_TIMEOUT_MS = 30_000 SQLITE_POOL_SIZE = 10 SQLITE_MAX_OVERFLOW = 10 SQLITE_POOL_TIMEOUT_S = 5 diff --git a/backend/app/jobs/handlers/import_images.py b/backend/app/jobs/handlers/import_images.py index 2dbdf67..66322d3 100644 --- a/backend/app/jobs/handlers/import_images.py +++ b/backend/app/jobs/handlers/import_images.py @@ -8,12 +8,12 @@ from pathlib import Path from typing import Any import sqlalchemy as sa +from sqlalchemy.dialects.sqlite import insert as sqlite_insert from sqlalchemy.ext.asyncio import AsyncEngine from app.core.config import load_settings from app.core.data_files import get_sqlite_db_dir, resolve_file_ref from app.core.pixiv_urls import parse_pixiv_original_url -from app.db.images_upsert import upsert_image_by_illust_page from app.db.models.images import Image from app.db.models.imports import Import from app.db.models.jobs import JobRow @@ -118,154 +118,218 @@ def build_import_images_handler(engine: AsyncEngine): hydrate_on_import = _as_bool(payload.get("hydrate_on_import"), default=False) file_path = _resolve_payload_file(payload) - async def _op() -> None: + async def _ensure_import_exists() -> None: async with Session() as session: imp = await session.get(Import, import_id) if imp is None: raise JobPermanentError("Import not found") - total = 0 - accepted = 0 - success = 0 - deduped = 0 - error_total = 0 - errors: list[ImportLineError] = [] - seen: set[tuple[int, int]] = set() - illust_ids: set[int] = set() + await with_sqlite_busy_retry(_ensure_import_exists) - chunk: list[tuple[str, Any]] = [] + total = 0 + accepted = 0 + success = 0 + deduped = 0 + error_total = 0 + errors: list[ImportLineError] = [] + seen: set[tuple[int, int]] = set() + illust_ids: set[int] = set() - for line_no, raw in _iter_lines(payload, file_path=file_path): - url = raw.strip() - if not url: - continue - total += 1 - try: - parsed = parse_pixiv_original_url(url) - except Exception as exc: - error_total += 1 - if len(errors) < _MAX_ERRORS: - errors.append( - ImportLineError( - line=int(line_no), - url=url, - code="unsupported_url", - message=str(exc) or "unsupported_url", - ) - ) - continue + chunk_rows: list[dict[str, Any]] = [] + chunk_keys: list[tuple[int, int]] = [] - key = (parsed.illust_id, parsed.page_index) - if key in seen: - deduped += 1 - continue - seen.add(key) + async def _persist_chunk( + rows: list[dict[str, Any]], + keys: list[tuple[int, int]], + *, + total_v: int, + accepted_v: int, + success_v: int, + failed_v: int, + ) -> None: + if not rows: + return - accepted += 1 - if hydrate_on_import: - illust_ids.add(int(parsed.illust_id)) + async def _op() -> None: + async with Session() as session: + now_expr = sa.text("(strftime('%Y-%m-%dT%H:%M:%fZ','now'))") - chunk.append((url, parsed)) + stmt = sqlite_insert(Image).values(rows) + stmt = stmt.on_conflict_do_update( + index_elements=["illust_id", "page_index"], + set_={ + "ext": stmt.excluded.ext, + "original_url": stmt.excluded.original_url, + "proxy_path": sa.case( + (sa.func.length(stmt.excluded.proxy_path) > 0, stmt.excluded.proxy_path), + else_=Image.proxy_path, + ), + "created_import_id": stmt.excluded.created_import_id, + "updated_at": now_expr, + }, + ) + await session.execute(stmt) - if len(chunk) < _CHUNK_SIZE: - continue - - for url_v, parsed_v in chunk: - image_id = await upsert_image_by_illust_page( - session, - illust_id=int(parsed_v.illust_id), - page_index=int(parsed_v.page_index), - ext=str(parsed_v.ext), - original_url=url_v, - proxy_path="", - random_key=random.random(), - created_import_id=int(import_id), - ) - proxy_path = f"/i/{image_id}.{parsed_v.ext}" + # Fill proxy_path for (newly inserted) rows that are still empty. + if keys: await session.execute( - sa.update(Image).where(Image.id == image_id).values(proxy_path=proxy_path) + sa.update(Image) + .where(Image.created_import_id == int(import_id)) + .where(Image.proxy_path == "") + .where(sa.tuple_(Image.illust_id, Image.page_index).in_(keys)) + .values(proxy_path=sa.text("'/i/' || id || '.' || ext")) ) - success += len(chunk) - imp.total = max(int(imp.total or 0), int(total)) - imp.accepted = max(int(imp.accepted or 0), int(accepted)) - imp.success = max(int(imp.success or 0), int(success)) - imp.failed = max(int(imp.failed or 0), int(error_total)) - await session.commit() - chunk.clear() - if chunk: - for url_v, parsed_v in chunk: - image_id = await upsert_image_by_illust_page( - session, - illust_id=int(parsed_v.illust_id), - page_index=int(parsed_v.page_index), - ext=str(parsed_v.ext), - original_url=url_v, - proxy_path="", - random_key=random.random(), - created_import_id=int(import_id), + await session.execute( + sa.update(Import) + .where(Import.id == int(import_id)) + .values( + total=sa.func.max(Import.total, int(total_v)), + accepted=sa.func.max(Import.accepted, int(accepted_v)), + success=sa.func.max(Import.success, int(success_v)), + failed=sa.func.max(Import.failed, int(failed_v)), ) - proxy_path = f"/i/{image_id}.{parsed_v.ext}" - await session.execute(sa.update(Image).where(Image.id == image_id).values(proxy_path=proxy_path)) - success += len(chunk) - imp.total = max(int(imp.total or 0), int(total)) - imp.accepted = max(int(imp.accepted or 0), int(accepted)) - imp.success = max(int(imp.success or 0), int(success)) - imp.failed = max(int(imp.failed or 0), int(error_total)) + ) await session.commit() - imp.total = max(int(imp.total or 0), int(total)) - imp.accepted = max(int(imp.accepted or 0), int(accepted)) - imp.success = max(int(imp.success or 0), int(success)) - imp.failed = max(int(imp.failed or 0), int(error_total)) - imp.detail_json = json.dumps( + await with_sqlite_busy_retry(_op) + + for line_no, raw in _iter_lines(payload, file_path=file_path): + url = raw.strip() + if not url: + continue + total += 1 + try: + parsed = parse_pixiv_original_url(url) + except Exception as exc: + error_total += 1 + if len(errors) < _MAX_ERRORS: + errors.append( + ImportLineError( + line=int(line_no), + url=url, + code="unsupported_url", + message=str(exc) or "unsupported_url", + ) + ) + continue + + key = (int(parsed.illust_id), int(parsed.page_index)) + if key in seen: + deduped += 1 + continue + seen.add(key) + + accepted += 1 + if hydrate_on_import: + illust_ids.add(int(parsed.illust_id)) + + chunk_keys.append(key) + chunk_rows.append( + { + "illust_id": int(parsed.illust_id), + "page_index": int(parsed.page_index), + "ext": str(parsed.ext), + "original_url": url, + "proxy_path": "", + "random_key": float(random.random()), + "created_import_id": int(import_id), + } + ) + + if len(chunk_rows) < _CHUNK_SIZE: + continue + + success_after = int(success + len(chunk_rows)) + await _persist_chunk( + chunk_rows, + chunk_keys, + total_v=int(total), + accepted_v=int(accepted), + success_v=int(success_after), + failed_v=int(error_total), + ) + success = success_after + chunk_rows.clear() + chunk_keys.clear() + + if chunk_rows: + success_after = int(success + len(chunk_rows)) + await _persist_chunk( + chunk_rows, + chunk_keys, + total_v=int(total), + accepted_v=int(accepted), + success_v=int(success_after), + failed_v=int(error_total), + ) + success = success_after + chunk_rows.clear() + chunk_keys.clear() + + async def _persist_detail() -> None: + async with Session() as session: + detail_json = json.dumps( {"deduped": int(deduped), "errors": [asdict(e) for e in errors[:_MAX_ERRORS]]}, ensure_ascii=False, separators=(",", ":"), ) + await session.execute( + sa.update(Import) + .where(Import.id == int(import_id)) + .values( + total=sa.func.max(Import.total, int(total)), + accepted=sa.func.max(Import.accepted, int(accepted)), + success=sa.func.max(Import.success, int(success)), + failed=sa.func.max(Import.failed, int(error_total)), + detail_json=detail_json, + ) + ) await session.commit() - if not hydrate_on_import or not illust_ids: - return + await with_sqlite_busy_retry(_persist_detail) - existing = set( - ( - await session.execute( - sa.select(JobRow.ref_id) - .where(JobRow.type == "hydrate_metadata") - .where(JobRow.ref_type == "import") - .where(JobRow.ref_id.like(f"{import_id}:%")) + if hydrate_on_import and illust_ids: + async def _enqueue_hydrate_jobs() -> None: + async with Session() as session: + existing = set( + ( + await session.execute( + sa.select(JobRow.ref_id) + .where(JobRow.type == "hydrate_metadata") + .where(JobRow.ref_type == "import") + .where(JobRow.ref_id.like(f"{import_id}:%")) + ) ) + .scalars() + .all() ) - .scalars() - .all() - ) - added = 0 - for illust_id in sorted(illust_ids): - ref_id = f"{import_id}:{int(illust_id)}" - if ref_id in existing: - continue - session.add( - JobRow( - type="hydrate_metadata", - status="pending", - payload_json=json.dumps( - {"illust_id": int(illust_id), "reason": "import"}, - ensure_ascii=False, - separators=(",", ":"), - ), - ref_type="import", - ref_id=ref_id, + added = 0 + for illust_id in sorted(illust_ids): + ref_id = f"{import_id}:{int(illust_id)}" + if ref_id in existing: + continue + session.add( + JobRow( + type="hydrate_metadata", + status="pending", + payload_json=json.dumps( + {"illust_id": int(illust_id), "reason": "import"}, + ensure_ascii=False, + separators=(",", ":"), + ), + ref_type="import", + ref_id=ref_id, + ) ) - ) - added += 1 - if added % 500 == 0: + added += 1 + if added % 500 == 0: + await session.commit() + if added: await session.commit() - if added: - await session.commit() - await with_sqlite_busy_retry(_op) + await with_sqlite_busy_retry(_enqueue_hydrate_jobs) if file_path is not None: try: diff --git a/backend/tests/test_job_handler_import_images.py b/backend/tests/test_job_handler_import_images.py index a06bfc8..09ba342 100644 --- a/backend/tests/test_job_handler_import_images.py +++ b/backend/tests/test_job_handler_import_images.py @@ -121,6 +121,81 @@ def test_job_handler_import_images_happy_path_and_enqueues_hydrate(tmp_path: Pat asyncio.run(_run()) +def test_job_handler_import_images_retries_on_sqlite_busy(tmp_path: Path, monkeypatch) -> None: + monkeypatch.setenv("SQLITE_BUSY_TIMEOUT_MS", "100") + monkeypatch.setenv("SQLITE_BUSY_RETRIES", "20") + + db_path = tmp_path / "handler_import_images_busy.db" + engine = create_engine(_sqlite_url(db_path)) + + async def _run() -> None: + async with engine.begin() as conn: + await conn.run_sync(Base.metadata.create_all) + + Session = create_sessionmaker(engine) + async with Session() as session: + imp = Import(created_by="admin", source="manual") + session.add(imp) + await session.commit() + await session.refresh(imp) + + payload = { + "import_id": int(imp.id), + "hydrate_on_import": False, + "text_lines": [ + "https://i.pximg.net/img-original/img/2020/01/01/00/00/00/111_p0.jpg", + "https://i.pximg.net/img-original/img/2020/01/01/00/00/00/222_p0.png", + ], + } + job = JobRow( + type="import_images", + status="pending", + payload_json=json.dumps(payload, ensure_ascii=False, separators=(",", ":")), + ref_type="import", + ref_id=str(int(imp.id)), + ) + session.add(job) + await session.commit() + + dispatcher = JobDispatcher() + dispatcher.register("import_images", build_import_images_handler(engine)) + + claimed = await claim_next_job(engine, worker_id="w1") + assert claimed is not None + assert claimed["type"] == "import_images" + + lock_acquired = asyncio.Event() + + async def _hold_write_lock() -> None: + async with engine.connect() as conn: + await conn.exec_driver_sql("BEGIN IMMEDIATE") + await conn.exec_driver_sql( + "UPDATE imports SET source=source WHERE id=:id", + {"id": int(payload["import_id"])}, + ) + lock_acquired.set() + await asyncio.sleep(0.6) + await conn.exec_driver_sql("COMMIT") + + async def _run_job() -> None: + await lock_acquired.wait() + transition = await execute_claimed_job(engine, dispatcher, job_row=claimed, worker_id="w1") + assert transition is not None + assert transition.status.value == "completed" + + await asyncio.gather(_hold_write_lock(), _run_job()) + + async with Session() as session: + imgs = ((await session.execute(sa.select(Image))).scalars().all()) + assert len(imgs) == 2 + for img in imgs: + assert img.proxy_path.startswith("/i/") + + await engine.dispose() + + asyncio.run(_run()) + + def test_job_handler_import_images_supports_file_ref(tmp_path: Path, monkeypatch) -> None: db_path = tmp_path / "handler_import_images_file_ref.db" db_url = _sqlite_url(db_path) diff --git a/backend/tests/test_pixiv_urls.py b/backend/tests/test_pixiv_urls.py index dcc1739..c3ccaed 100644 --- a/backend/tests/test_pixiv_urls.py +++ b/backend/tests/test_pixiv_urls.py @@ -21,6 +21,22 @@ def test_parse_pixiv_original_url_ok_master1200_suffix() -> None: assert parsed.ext == "png" +def test_parse_pixiv_original_url_ok_square1200_suffix() -> None: + u = "https://i.pximg.net/img-master/img/2023/01/01/00/00/00/12345678_p2_square1200.jpg" + parsed = parse_pixiv_original_url(u) + assert parsed.illust_id == 12345678 + assert parsed.page_index == 2 + assert parsed.ext == "jpg" + + +def test_parse_pixiv_original_url_ok_mirror_host_pixiv_re() -> None: + u = "https://i.pixiv.re/img-original/img/2023/01/01/00/00/00/12345678_p0.jpg" + parsed = parse_pixiv_original_url(u) + assert parsed.illust_id == 12345678 + assert parsed.page_index == 0 + assert parsed.ext == "jpg" + + def test_parse_pixiv_original_url_ok_ugoira0() -> None: u = "https://i.pximg.net/img-original/img/2014/07/01/21/17/59/44439242_ugoira0.jpg" parsed = parse_pixiv_original_url(u) @@ -29,6 +45,14 @@ def test_parse_pixiv_original_url_ok_ugoira0() -> None: assert parsed.ext == "jpg" +def test_parse_pixiv_original_url_ok_ugoira_zip() -> None: + u = "https://i.pximg.net/img-zip-ugoira/img/2023/01/01/00/00/00/12345678_ugoira600x600.zip" + parsed = parse_pixiv_original_url(u) + assert parsed.illust_id == 12345678 + assert parsed.page_index == 0 + assert parsed.ext == "zip" + + def test_parse_pixiv_original_url_ignores_query_and_whitespace() -> None: u = " https://i.pximg.net/img-original/img/2023/01/01/00/00/00/12345678_p12.PNG?foo=bar \n" parsed = parse_pixiv_original_url(u)