Add small improvements and query optimization.
Some checks failed
Build package / Build Test (3.13) (push) Has been cancelled
Python Linting / Run Ruff (push) Has been cancelled
Python Linting / Run Pylint (push) Has been cancelled
Build package / Build Test (3.10) (push) Has been cancelled
Build package / Build Test (3.11) (push) Has been cancelled
Build package / Build Test (3.12) (push) Has been cancelled
Build package / Build Test (3.14) (push) Has been cancelled

This commit is contained in:
Talmaj Marinc 2026-07-04 21:43:22 +02:00
parent c32fb2a077
commit 8fc848b763
3 changed files with 16 additions and 12 deletions

View File

@ -4,10 +4,7 @@ import shutil
from app.logger import log_startup_warning from app.logger import log_startup_warning
from utils.install_util import get_missing_requirements_message from utils.install_util import get_missing_requirements_message
from filelock import FileLock, Timeout from filelock import FileLock, Timeout
# NOTE: import the module (not `from ... import args`) so we always read the # Import the module so tests that reload comfy.cli_args see the live object.
# live `args` object. Tests reload `comfy.cli_args`, which replaces the module
# global; a bound `args` reference would go stale and point at the default
# database URL instead of the one configured for the test.
import comfy.cli_args import comfy.cli_args
_DB_AVAILABLE = False _DB_AVAILABLE = False

View File

@ -50,6 +50,19 @@ def list_downloads() -> list[Download]:
return rows return rows
def has_live_download_for_model(
model_id: str, live_statuses: tuple[str, ...], exclude_id: Optional[str] = None
) -> bool:
with create_session() as session:
stmt = select(Download.id).where(
Download.model_id == model_id,
Download.status.in_(live_statuses),
).limit(1)
if exclude_id is not None:
stmt = stmt.where(Download.id != exclude_id)
return session.execute(stmt).first() is not None
def list_segments(download_id: str) -> list[DownloadSegment]: def list_segments(download_id: str) -> list[DownloadSegment]:
with create_session() as session: with create_session() as session:
rows = list( rows = list(

View File

@ -191,12 +191,8 @@ class DownloadManager:
async def _has_live_download( async def _has_live_download(
self, model_id: str, *, exclude_id: Optional[str] = None self, model_id: str, *, exclude_id: Optional[str] = None
) -> bool: ) -> bool:
rows = await asyncio.to_thread(queries.list_downloads) return await asyncio.to_thread(
return any( queries.has_live_download_for_model, model_id, _LIVE_STATUSES, exclude_id
r.model_id == model_id
and r.id != exclude_id
and r.status in _LIVE_STATUSES
for r in rows
) )
# ----- control ----- # ----- control -----
@ -249,8 +245,6 @@ class DownloadManager:
if row is None: if row is None:
raise DownloadError("NOT_FOUND", "No such download.", status=404) raise DownloadError("NOT_FOUND", "No such download.", status=404)
if row.status in _LIVE_STATUSES: if row.status in _LIVE_STATUSES:
import os
try: try:
os.remove(row.temp_path) os.remove(row.temp_path)
except OSError: except OSError: