修复大量导入时的数据库冲突

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

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:

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)

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)