mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2026-02-06 19:42:34 +08:00
4e02245012
- Fix unused imports in routes.py, asset.py, manager.py, asset_management.py, ingest.py - Fix whitespace issues in upload.py, asset_info.py, ingest.py - Fix typo in manager.py (stray character after result["asset"]) - Fix broken import in test_metadata.py (project_kv moved to asset_info.py) - Add fixture override in queries/conftest.py for unit test isolation Add 48 new tests covering all previously untested query functions: - asset.py: upsert_asset, bulk_insert_assets - cache_state.py: upsert_cache_state, delete_cache_states_outside_prefixes, get_orphaned_seed_asset_ids, delete_assets_by_ids, get_cache_states_for_prefixes, bulk_set_needs_verify, delete_cache_states_by_ids, delete_orphaned_seed_asset, bulk_insert_cache_states_ignore_conflicts, get_cache_states_by_paths_and_asset_ids - asset_info.py: insert_asset_info, get_or_create_asset_info, update_asset_info_timestamps, replace_asset_info_metadata_projection, bulk_insert_asset_infos_ignore_conflicts, get_asset_info_ids_by_ids - tags.py: bulk_insert_tags_and_meta Total: 119 tests pass (up from 71) Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
95 lines
2.4 KiB
Python
95 lines
2.4 KiB
Python
|
|
import sqlalchemy as sa
|
|
from sqlalchemy import select
|
|
from sqlalchemy.dialects import sqlite
|
|
from sqlalchemy.orm import Session
|
|
|
|
from app.assets.database.models import Asset
|
|
|
|
MAX_BIND_PARAMS = 800
|
|
|
|
|
|
def _rows_per_stmt(cols: int) -> int:
|
|
return max(1, MAX_BIND_PARAMS // max(1, cols))
|
|
|
|
|
|
def _iter_chunks(seq, n: int):
|
|
for i in range(0, len(seq), n):
|
|
yield seq[i : i + n]
|
|
|
|
|
|
def asset_exists_by_hash(
|
|
session: Session,
|
|
asset_hash: str,
|
|
) -> bool:
|
|
"""
|
|
Check if an asset with a given hash exists in database.
|
|
"""
|
|
row = (
|
|
session.execute(
|
|
select(sa.literal(True)).select_from(Asset).where(Asset.hash == asset_hash).limit(1)
|
|
)
|
|
).first()
|
|
return row is not None
|
|
|
|
|
|
def get_asset_by_hash(
|
|
session: Session,
|
|
asset_hash: str,
|
|
) -> Asset | None:
|
|
return (
|
|
session.execute(select(Asset).where(Asset.hash == asset_hash).limit(1))
|
|
).scalars().first()
|
|
|
|
|
|
def upsert_asset(
|
|
session: Session,
|
|
asset_hash: str,
|
|
size_bytes: int,
|
|
mime_type: str | None = None,
|
|
) -> tuple[Asset, bool, bool]:
|
|
"""Upsert an Asset by hash. Returns (asset, created, updated)."""
|
|
vals = {"hash": asset_hash, "size_bytes": int(size_bytes)}
|
|
if mime_type:
|
|
vals["mime_type"] = mime_type
|
|
|
|
ins = (
|
|
sqlite.insert(Asset)
|
|
.values(**vals)
|
|
.on_conflict_do_nothing(index_elements=[Asset.hash])
|
|
)
|
|
res = session.execute(ins)
|
|
created = int(res.rowcount or 0) > 0
|
|
|
|
asset = session.execute(
|
|
select(Asset).where(Asset.hash == asset_hash).limit(1)
|
|
).scalars().first()
|
|
if not asset:
|
|
raise RuntimeError("Asset row not found after upsert.")
|
|
|
|
updated = False
|
|
if not created:
|
|
changed = False
|
|
if asset.size_bytes != int(size_bytes) and int(size_bytes) > 0:
|
|
asset.size_bytes = int(size_bytes)
|
|
changed = True
|
|
if mime_type and asset.mime_type != mime_type:
|
|
asset.mime_type = mime_type
|
|
changed = True
|
|
if changed:
|
|
updated = True
|
|
|
|
return asset, created, updated
|
|
|
|
|
|
def bulk_insert_assets(
|
|
session: Session,
|
|
rows: list[dict],
|
|
) -> None:
|
|
"""Bulk insert Asset rows. Each dict should have: id, hash, size_bytes, mime_type, created_at."""
|
|
if not rows:
|
|
return
|
|
ins = sqlite.insert(Asset)
|
|
for chunk in _iter_chunks(rows, _rows_per_stmt(5)):
|
|
session.execute(ins, chunk)
|