mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2026-02-19 06:30:07 +00:00
8ff4d38ad1
This change solves the basename collision bug by using UNIQUE(file_path) on the unified asset_references table. Key changes: Database: - Migration 0005 merges asset_cache_states and asset_infos into asset_references - AssetReference now contains: cache state fields (file_path, mtime_ns, needs_verify, is_missing, enrichment_level) plus info fields (name, owner_id, preview_id, etc.) - AssetReferenceMeta replaces AssetInfoMeta - AssetReferenceTag replaces AssetInfoTag - UNIQUE constraint on file_path prevents duplicate entries for same file Code: - New unified query module: asset_reference.py (replaces asset_info.py, cache_state.py) - Updated scanner, seeder, and services to use AssetReference - Updated API routes to use reference_id instead of asset_info_id Tests: - All 175 unit tests updated and passing - Integration tests require server environment (not run here) Amp-Thread-ID: https://ampcode.com/threads/T-019c4fe8-9dcb-75ce-bea8-ea786343a581 Co-authored-by: Amp <amp@ampcode.com>
140 lines
3.6 KiB
Python
140 lines
3.6 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
|
|
from app.assets.database.queries.common import calculate_rows_per_statement, iter_chunks
|
|
|
|
|
|
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 with ON CONFLICT DO NOTHING on hash."""
|
|
if not rows:
|
|
return
|
|
ins = sqlite.insert(Asset).on_conflict_do_nothing(index_elements=[Asset.hash])
|
|
for chunk in iter_chunks(rows, calculate_rows_per_statement(5)):
|
|
session.execute(ins, chunk)
|
|
|
|
|
|
def get_existing_asset_ids(
|
|
session: Session,
|
|
asset_ids: list[str],
|
|
) -> set[str]:
|
|
"""Return the subset of asset_ids that exist in the database."""
|
|
if not asset_ids:
|
|
return set()
|
|
rows = session.execute(
|
|
select(Asset.id).where(Asset.id.in_(asset_ids))
|
|
).fetchall()
|
|
return {row[0] for row in rows}
|
|
|
|
|
|
def update_asset_hash_and_mime(
|
|
session: Session,
|
|
asset_id: str,
|
|
asset_hash: str | None = None,
|
|
mime_type: str | None = None,
|
|
) -> bool:
|
|
"""Update asset hash and/or mime_type. Returns True if asset was found."""
|
|
asset = session.get(Asset, asset_id)
|
|
if not asset:
|
|
return False
|
|
if asset_hash is not None:
|
|
asset.hash = asset_hash
|
|
if mime_type is not None:
|
|
asset.mime_type = mime_type
|
|
return True
|
|
|
|
|
|
def reassign_asset_references(
|
|
session: Session,
|
|
from_asset_id: str,
|
|
to_asset_id: str,
|
|
reference_id: str,
|
|
) -> None:
|
|
"""Reassign a reference from one asset to another.
|
|
|
|
Used when merging a stub asset into an existing asset with the same hash.
|
|
"""
|
|
from app.assets.database.models import AssetReference
|
|
|
|
ref = session.get(AssetReference, reference_id)
|
|
if ref:
|
|
ref.asset_id = to_asset_id
|
|
|
|
session.flush()
|