wangbo/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2025-12-18 18:43:05 +08:00
Code Issues Actions 10 Packages Projects Releases Wiki Activity

refactor(4): use one query to init DB with all tags for assets

Browse Source
This commit is contained in:
bigcat88 2025-09-16 21:18:18 +03:00
parent 677a0e2508
commit d0aa64d57b
No known key found for this signature in database
GPG Key ID: 1F0BF0EC3CF22721
2 changed files with 50 additions and 34 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

39
app/assets_scanner.py
View File

@ -23,6 +23,7 @@ from .api import schemas_in, schemas_out
from .database.db import create_session from .database.db import create_session
from .database.helpers import ( from .database.helpers import (
add_missing_tag_for_asset_id, add_missing_tag_for_asset_id,
ensure_tags_exist,
escape_like_prefix, escape_like_prefix,
fast_asset_file_check, fast_asset_file_check,
remove_missing_tag_for_asset_id, remove_missing_tag_for_asset_id,
@ -118,38 +119,52 @@ async def sync_seed_assets(roots: list[schemas_in.RootType]) -> None:
if "output" in roots: if "output" in roots:
paths.extend(list_tree(folder_paths.get_output_directory())) paths.extend(list_tree(folder_paths.get_output_directory()))
async with await create_session() as sess: new_specs: list[tuple[str, int, int, str, list[str]]] = []
tag_pool: set[str] = set()
for p in paths: for p in paths:
try: ap = os.path.abspath(p)
if os.path.abspath(p) in existing_paths: if ap in existing_paths:
skipped_existing += 1 skipped_existing += 1
continue continue
try:
st = os.stat(p, follow_symlinks=True) st = os.stat(p, follow_symlinks=True)
except OSError:
continue
if not int(st.st_size or 0): if not int(st.st_size or 0):
continue continue
size_bytes = int(st.st_size) name, tags = get_name_and_tags_from_asset_path(ap)
mtime_ns = getattr(st, "st_mtime_ns", int(st.st_mtime * 1_000_000_000)) new_specs.append((
name, tags = get_name_and_tags_from_asset_path(p) ap,
int(st.st_size),
getattr(st, "st_mtime_ns", int(st.st_mtime * 1_000_000_000)),
name,
tags,
))
for t in tags:
tag_pool.add(t)
async with await create_session() as sess:
if tag_pool:
await ensure_tags_exist(sess, tag_pool, tag_type="user")
for ap, sz, mt, name, tags in new_specs:
await ensure_seed_for_path( await ensure_seed_for_path(
sess, sess,
abs_path=p, abs_path=ap,
size_bytes=size_bytes, size_bytes=sz,
mtime_ns=mtime_ns, mtime_ns=mt,
info_name=name, info_name=name,
tags=tags, tags=tags,
owner_id="", owner_id="",
skip_tag_ensure=True,
) )
created += 1 created += 1
if created % 500 == 0: if created % 500 == 0:
await sess.commit() await sess.commit()
except OSError:
continue
await sess.commit() await sess.commit()
finally: finally:
LOGGER.info( LOGGER.info(
"Assets scan(roots=%s) completed in %.3f s (created=%d, skipped_existing=%d, total_seen=%d)", "Assets scan(roots=%s) completed in %.3fs (created=%d, skipped_existing=%d, total_seen=%d)",
roots, roots,
time.perf_counter() - t_total, time.perf_counter() - t_total,
created, created,

13
app/database/services/content.py
View File

@ -65,6 +65,7 @@ async def ensure_seed_for_path(
info_name: str, info_name: str,
tags: Sequence[str], tags: Sequence[str],
owner_id: str = "", owner_id: str = "",
skip_tag_ensure: bool = False,
) -> str: ) -> str:
"""Ensure: Asset(hash=NULL), AssetCacheState(file_path), and AssetInfo exist for the path. Returns asset_id.""" """Ensure: Asset(hash=NULL), AssetCacheState(file_path), and AssetInfo exist for the path. Returns asset_id."""
locator = os.path.abspath(abs_path) locator = os.path.abspath(abs_path)
@ -81,20 +82,20 @@ async def ensure_seed_for_path(
if state: if state:
state_row: AssetCacheState = state[0] state_row: AssetCacheState = state[0]
asset_row: Asset = state[1] asset_row: Asset = state[1]
changed = state_row.mtime_ns is None or int(state_row.mtime_ns) != int(mtime_ns) changed = state_row.mtime_ns is None or int(state_row.mtime_ns) != mtime_ns
if changed: if changed:
state_row.mtime_ns = int(mtime_ns) state_row.mtime_ns = mtime_ns
state_row.needs_verify = True state_row.needs_verify = True
if asset_row.size_bytes == 0 and size_bytes > 0: if asset_row.size_bytes == 0 and size_bytes > 0:
asset_row.size_bytes = int(size_bytes) asset_row.size_bytes = size_bytes
await session.flush() await session.flush()
return asset_row.id return asset_row.id
asset = Asset(hash=None, size_bytes=int(size_bytes), mime_type=None, created_at=now) asset = Asset(hash=None, size_bytes=size_bytes, mime_type=None, created_at=now)
session.add(asset) session.add(asset)
await session.flush() # to get id await session.flush() # to get id
cs = AssetCacheState(asset_id=asset.id, file_path=locator, mtime_ns=int(mtime_ns), needs_verify=False) cs = AssetCacheState(asset_id=asset.id, file_path=locator, mtime_ns=mtime_ns, needs_verify=False)
session.add(cs) session.add(cs)
info = AssetInfo( info = AssetInfo(
@ -120,12 +121,12 @@ async def ensure_seed_for_path(
want = normalize_tags(tags) want = normalize_tags(tags)
if want: if want:
if not skip_tag_ensure:
await ensure_tags_exist(session, want, tag_type="user") await ensure_tags_exist(session, want, tag_type="user")
session.add_all([ session.add_all([
AssetInfoTag(asset_info_id=info.id, tag_name=t, origin="automatic", added_at=now) AssetInfoTag(asset_info_id=info.id, tag_name=t, origin="automatic", added_at=now)
for t in want for t in want
]) ])
await session.flush() await session.flush()
return asset.id return asset.id