diff --git a/py/services/download_queue_service.py b/py/services/download_queue_service.py index 4c87699d..7eae08e5 100644 --- a/py/services/download_queue_service.py +++ b/py/services/download_queue_service.py @@ -12,6 +12,15 @@ from ..utils.cache_paths import get_cache_base_dir logger = logging.getLogger(__name__) +# SQL fragment extracting the CivitAI file id from the JSON ``file_params`` +# column (#1058). ``json_valid`` guards against NULL and legacy/unparseable +# values, yielding NULL for rows without a file identity; NULL keys group +# together so such rows keep the old version-level dedup behavior. +_FILE_ID_SQL = ( + "CASE WHEN json_valid(file_params) " + "THEN json_extract(file_params, '$.id') END" +) + def _resolve_database_path() -> str: base_dir = get_cache_base_dir(create=True) @@ -866,33 +875,44 @@ class DownloadQueueService: async with self._lock: conn = self._get_conn() - # 1. History: for each (model_id, model_version_id, status) triplet - # keep only the row with the highest id (most recently inserted). - conn.execute(""" + # 1. History: for each (model_id, model_version_id, file_id, + # status) group keep only the row with the highest id (most + # recently inserted). file_id comes from file_params (#1058) + # so distinct files of the same version never collapse. + conn.execute(f""" DELETE FROM download_history WHERE id NOT IN ( SELECT MAX(id) FROM download_history - GROUP BY model_id, model_version_id, status + GROUP BY model_id, model_version_id, status, + {_FILE_ID_SQL} ) """) result["removed_history"] = conn.execute( "SELECT changes()" ).fetchone()[0] - # 2. Cross-status dedup: for each (model_id, model_version_id), - # keep only the entry with the highest-priority terminal status. + # 2. Cross-status dedup: for each (model_id, model_version_id, + # file_id), keep only the entry with the highest-priority + # terminal status. # Priority: completed (3) > failed (2) > canceled (1). - # This prevents the same model version from having both a - # 'failed' and a 'canceled' entry (or a 'completed' alongside - # either) after the bug-created duplicates are removed. - conn.execute(""" + # This prevents the same file of a model version from having + # both a 'failed' and a 'canceled' entry (or a 'completed' + # alongside either) after the bug-created duplicates are + # removed. ``IS`` matches NULL file ids against each other so + # rows without file identity keep the old behavior. + conn.execute(f""" DELETE FROM download_history WHERE id NOT IN ( SELECT dh.id - FROM download_history dh + FROM ( + SELECT id, model_id, model_version_id, status, + {_FILE_ID_SQL} AS file_id + FROM download_history + ) dh INNER JOIN ( SELECT model_id, model_version_id, + {_FILE_ID_SQL} AS file_id, MAX(CASE status WHEN 'completed' THEN 3 WHEN 'failed' THEN 2 @@ -900,17 +920,18 @@ class DownloadQueueService: ELSE 0 END) AS best_prio FROM download_history - GROUP BY model_id, model_version_id + GROUP BY model_id, model_version_id, {_FILE_ID_SQL} ) best ON dh.model_id = best.model_id AND dh.model_version_id = best.model_version_id + AND dh.file_id IS best.file_id AND CASE dh.status WHEN 'completed' THEN 3 WHEN 'failed' THEN 2 WHEN 'canceled' THEN 1 ELSE 0 END = best.best_prio - GROUP BY dh.model_id, dh.model_version_id + GROUP BY dh.model_id, dh.model_version_id, dh.file_id HAVING dh.id = MAX(dh.id) ) """) @@ -918,15 +939,17 @@ class DownloadQueueService: "SELECT changes()" ).fetchone()[0] - # 3. Queue: for each (model_id, model_version_id) keep only the - # row with the latest added_at (most recently enqueued). - conn.execute(""" + # 3. Queue: for each (model_id, model_version_id, file_id) keep + # only the row with the latest added_at (most recently + # enqueued). file_id comes from file_params (#1058) so + # distinct files of the same version never collapse. + conn.execute(f""" DELETE FROM download_queue WHERE rowid NOT IN ( SELECT MAX(rowid) FROM download_queue WHERE status IN ('queued', 'downloading', 'paused', 'waiting') - GROUP BY model_id, model_version_id + GROUP BY model_id, model_version_id, {_FILE_ID_SQL} ) AND status IN ('queued', 'downloading', 'paused', 'waiting') """) diff --git a/tests/services/test_download_queue_service.py b/tests/services/test_download_queue_service.py index 7b8c50b6..48b2f711 100644 --- a/tests/services/test_download_queue_service.py +++ b/tests/services/test_download_queue_service.py @@ -7,6 +7,7 @@ compatibility with ``id``. import json import sqlite3 +import time from pathlib import Path import pytest @@ -321,3 +322,176 @@ async def test_legacy_history_db_gains_file_params_column(tmp_path: Path) -> Non item = await svc.retry_from_history(download_id="dl-legacy-fp") assert item is not None assert json.loads(item["file_params"]) == {"id": 5} + + +# --------------------------------------------------------------------------- +# deduplicate() — file identity in the dedup key (#1058) +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_dedup_queue_keeps_distinct_files_of_same_version(tmp_path: Path) -> None: + """Two queued downloads of different files of one version both survive.""" + svc = _make_service(tmp_path) + await svc.add_to_queue( + download_id="dl-a", + model_id=1, + model_version_id=100, + file_params={"id": 1001, "type": "Model"}, + ) + await svc.add_to_queue( + download_id="dl-b", + model_id=1, + model_version_id=100, + file_params={"id": 1002, "type": "Model"}, + ) + + result = await svc.deduplicate() + + assert result["removed_queue"] == 0 + queue = await svc.get_queue() + assert sorted(item["download_id"] for item in queue) == ["dl-a", "dl-b"] + + +@pytest.mark.asyncio +async def test_dedup_queue_collapses_same_file_and_legacy_rows(tmp_path: Path) -> None: + """Same version + same/absent file id still collapses to the latest row.""" + svc = _make_service(tmp_path) + # Same file id -> only the most recently enqueued row survives + await svc.add_to_queue( + download_id="dl-old", + model_id=1, + model_version_id=100, + file_params={"id": 1001}, + ) + await svc.add_to_queue( + download_id="dl-new", + model_id=1, + model_version_id=100, + file_params={"id": 1001}, + ) + # Rows without file identity keep the old per-version behavior + await svc.add_to_queue( + download_id="dl-legacy-old", model_id=2, model_version_id=200 + ) + await svc.add_to_queue( + download_id="dl-legacy-new", model_id=2, model_version_id=200 + ) + + result = await svc.deduplicate() + + assert result["removed_queue"] == 2 + remaining = {item["download_id"] for item in await svc.get_queue()} + assert remaining == {"dl-new", "dl-legacy-new"} + + +@pytest.mark.asyncio +async def test_dedup_queue_does_not_mix_null_and_file_id(tmp_path: Path) -> None: + """A row with a file id never dedups against a row without one.""" + svc = _make_service(tmp_path) + await svc.add_to_queue( + download_id="dl-file", + model_id=1, + model_version_id=100, + file_params={"id": 1001}, + ) + await svc.add_to_queue( + download_id="dl-nofile", model_id=1, model_version_id=100 + ) + + result = await svc.deduplicate() + + assert result["removed_queue"] == 0 + assert len(await svc.get_queue()) == 2 + + +@pytest.mark.asyncio +async def test_dedup_queue_unparseable_file_params_treated_as_none(tmp_path: Path) -> None: + """Corrupt file_params JSON falls back to the NULL file identity.""" + svc = _make_service(tmp_path) + await svc.add_to_queue( + download_id="dl-plain", model_id=1, model_version_id=100 + ) + conn = sqlite3.connect(str(tmp_path / "queue.sqlite")) + conn.execute( + "INSERT INTO download_queue (download_id, model_id, model_version_id, " + "file_params, status, added_at) VALUES (?, ?, ?, ?, 'queued', ?)", + ("dl-corrupt", 1, 100, "{not json", time.time()), + ) + conn.commit() + conn.close() + + result = await svc.deduplicate() + + assert result["removed_queue"] == 1 + assert len(await svc.get_queue()) == 1 + + +@pytest.mark.asyncio +async def test_dedup_history_keeps_distinct_files_of_same_version(tmp_path: Path) -> None: + """History rows of different files never collapse, even across statuses.""" + svc = _make_service(tmp_path) + await svc.add_to_history( + download_id="dl-h1", + model_id=1, + model_version_id=100, + status="completed", + file_params={"id": 1001}, + ) + await svc.add_to_history( + download_id="dl-h2", + model_id=1, + model_version_id=100, + status="completed", + file_params={"id": 1002}, + ) + # A failed entry for file 1002 must not remove file 1001's completed row + await svc.add_to_history( + download_id="dl-h3", + model_id=1, + model_version_id=100, + status="failed", + file_params={"id": 1002}, + ) + + result = await svc.deduplicate() + + assert result["removed_history"] == 1 # only dl-h3 collapses (same file) + history = await svc.get_history() + remaining = {item["download_id"] for item in history["items"]} + assert remaining == {"dl-h1", "dl-h2"} + + +@pytest.mark.asyncio +async def test_dedup_history_collapses_same_file_and_legacy_rows(tmp_path: Path) -> None: + """Same file id / absent file id history rows still dedup as before.""" + svc = _make_service(tmp_path) + # Same file id + same status -> keep the most recent row + await svc.add_to_history( + download_id="dl-x1", + model_id=1, + model_version_id=100, + status="completed", + file_params={"id": 1001}, + ) + await svc.add_to_history( + download_id="dl-x2", + model_id=1, + model_version_id=100, + status="completed", + file_params={"id": 1001}, + ) + # Legacy rows without file identity: cross-status dedup still applies + await svc.add_to_history( + download_id="dl-y1", model_id=2, model_version_id=200, status="failed" + ) + await svc.add_to_history( + download_id="dl-y2", model_id=2, model_version_id=200, status="completed" + ) + + result = await svc.deduplicate() + + assert result["removed_history"] == 2 + history = await svc.get_history() + remaining = {item["download_id"] for item in history["items"]} + assert remaining == {"dl-x2", "dl-y2"}