add AssetsResolver support [in progress]

Author: bigcat88

alembic_db/versions/0001_assets.py

@@ -1,3 +1,4 @@
+# File: /alembic_db/versions/0001_assets.py
 """initial assets schema + per-asset state cache
 
 Revision ID: 0001_assets
@@ -22,15 +23,12 @@ def upgrade() -> None:
         sa.Column("size_bytes", sa.BigInteger(), nullable=False, server_default="0"),
         sa.Column("mime_type", sa.String(length=255), nullable=True),
         sa.Column("refcount", sa.BigInteger(), nullable=False, server_default="0"),
-        sa.Column("storage_backend", sa.String(length=32), nullable=False, server_default="fs"),
-        sa.Column("storage_locator", sa.Text(), nullable=False),
         sa.Column("created_at", sa.DateTime(timezone=False), nullable=False),
         sa.Column("updated_at", sa.DateTime(timezone=False), nullable=False),
         sa.CheckConstraint("size_bytes >= 0", name="ck_assets_size_nonneg"),
         sa.CheckConstraint("refcount >= 0", name="ck_assets_refcount_nonneg"),
     )
     op.create_index("ix_assets_mime_type", "assets", ["mime_type"])
-    op.create_index("ix_assets_backend_locator", "assets", ["storage_backend", "storage_locator"])
 
     # ASSETS_INFO: user-visible references (mutable metadata)
     op.create_table(
@@ -52,11 +50,12 @@ def upgrade() -> None:
     op.create_index("ix_assets_info_name", "assets_info", ["name"])
     op.create_index("ix_assets_info_created_at", "assets_info", ["created_at"])
     op.create_index("ix_assets_info_last_access_time", "assets_info", ["last_access_time"])
+    op.create_index("ix_assets_info_owner_name", "assets_info", ["owner_id", "name"])
 
     # TAGS: normalized tag vocabulary
     op.create_table(
         "tags",
-        sa.Column("name", sa.String(length=128), primary_key=True),
+        sa.Column("name", sa.String(length=512), primary_key=True),
         sa.Column("tag_type", sa.String(length=32), nullable=False, server_default="user"),
         sa.CheckConstraint("name = lower(name)", name="ck_tags_lowercase"),
     )
@@ -65,8 +64,8 @@ def upgrade() -> None:
     # ASSET_INFO_TAGS: many-to-many for tags on AssetInfo
     op.create_table(
         "asset_info_tags",
-        sa.Column("asset_info_id", sa.BigInteger(), sa.ForeignKey("assets_info.id", ondelete="CASCADE"), nullable=False),
-        sa.Column("tag_name", sa.String(length=128), sa.ForeignKey("tags.name", ondelete="RESTRICT"), nullable=False),
+        sa.Column("asset_info_id", sa.Integer(), sa.ForeignKey("assets_info.id", ondelete="CASCADE"), nullable=False),
+        sa.Column("tag_name", sa.String(length=512), sa.ForeignKey("tags.name", ondelete="RESTRICT"), nullable=False),
         sa.Column("origin", sa.String(length=32), nullable=False, server_default="manual"),
         sa.Column("added_by", sa.String(length=128), nullable=True),
         sa.Column("added_at", sa.DateTime(timezone=False), nullable=False),
@@ -75,15 +74,15 @@ def upgrade() -> None:
     op.create_index("ix_asset_info_tags_tag_name", "asset_info_tags", ["tag_name"])
     op.create_index("ix_asset_info_tags_asset_info_id", "asset_info_tags", ["asset_info_id"])
 
-    # ASSET_LOCATOR_STATE: 1:1 filesystem metadata(for fast integrity checking) for an Asset records
+    # ASSET_CACHE_STATE: 1:1 local cache metadata for an Asset
     op.create_table(
-        "asset_locator_state",
+        "asset_cache_state",
         sa.Column("asset_hash", sa.String(length=256), sa.ForeignKey("assets.hash", ondelete="CASCADE"), primary_key=True),
+        sa.Column("file_path", sa.Text(), nullable=False),  # absolute local path to cached file
         sa.Column("mtime_ns", sa.BigInteger(), nullable=True),
-        sa.Column("etag", sa.String(length=256), nullable=True),
-        sa.Column("last_modified", sa.String(length=128), nullable=True),
-        sa.CheckConstraint("(mtime_ns IS NULL) OR (mtime_ns >= 0)", name="ck_als_mtime_nonneg"),
+        sa.CheckConstraint("(mtime_ns IS NULL) OR (mtime_ns >= 0)", name="ck_acs_mtime_nonneg"),
     )
+    op.create_index("ix_asset_cache_state_file_path", "asset_cache_state", ["file_path"])
 
     # ASSET_INFO_META: typed KV projection of user_metadata for filtering/sorting
     op.create_table(
@@ -102,6 +101,21 @@ def upgrade() -> None:
     op.create_index("ix_asset_info_meta_key_val_num", "asset_info_meta", ["key", "val_num"])
     op.create_index("ix_asset_info_meta_key_val_bool", "asset_info_meta", ["key", "val_bool"])
 
+    # ASSET_LOCATIONS: remote locations per asset
+    op.create_table(
+        "asset_locations",
+        sa.Column("id", sa.Integer(), primary_key=True, autoincrement=True),
+        sa.Column("asset_hash", sa.String(length=256), sa.ForeignKey("assets.hash", ondelete="CASCADE"), nullable=False),
+        sa.Column("provider", sa.String(length=32), nullable=False),  # e.g., "gcs"
+        sa.Column("locator", sa.Text(), nullable=False),             # e.g., "gs://bucket/path/to/blob"
+        sa.Column("expected_size_bytes", sa.BigInteger(), nullable=True),
+        sa.Column("etag", sa.String(length=256), nullable=True),
+        sa.Column("last_modified", sa.String(length=128), nullable=True),
+        sa.UniqueConstraint("asset_hash", "provider", "locator", name="uq_asset_locations_triplet"),
+    )
+    op.create_index("ix_asset_locations_hash", "asset_locations", ["asset_hash"])
+    op.create_index("ix_asset_locations_provider", "asset_locations", ["provider"])
+
     # Tags vocabulary for models
     tags_table = sa.table(
         "tags",
@@ -143,13 +157,18 @@ def upgrade() -> None:
 
 
 def downgrade() -> None:
+    op.drop_index("ix_asset_locations_provider", table_name="asset_locations")
+    op.drop_index("ix_asset_locations_hash", table_name="asset_locations")
+    op.drop_table("asset_locations")
+
     op.drop_index("ix_asset_info_meta_key_val_bool", table_name="asset_info_meta")
     op.drop_index("ix_asset_info_meta_key_val_num", table_name="asset_info_meta")
     op.drop_index("ix_asset_info_meta_key_val_str", table_name="asset_info_meta")
     op.drop_index("ix_asset_info_meta_key", table_name="asset_info_meta")
     op.drop_table("asset_info_meta")
 
-    op.drop_table("asset_locator_state")
+    op.drop_index("ix_asset_cache_state_file_path", table_name="asset_cache_state")
+    op.drop_table("asset_cache_state")
 
     op.drop_index("ix_asset_info_tags_asset_info_id", table_name="asset_info_tags")
     op.drop_index("ix_asset_info_tags_tag_name", table_name="asset_info_tags")
@@ -159,13 +178,13 @@ def downgrade() -> None:
     op.drop_table("tags")
 
     op.drop_constraint("uq_assets_info_hash_owner_name", table_name="assets_info")
+    op.drop_index("ix_assets_info_owner_name", table_name="assets_info")
     op.drop_index("ix_assets_info_last_access_time", table_name="assets_info")
     op.drop_index("ix_assets_info_created_at", table_name="assets_info")
     op.drop_index("ix_assets_info_name", table_name="assets_info")
     op.drop_index("ix_assets_info_asset_hash", table_name="assets_info")
     op.drop_index("ix_assets_info_owner_id", table_name="assets_info")
     op.drop_table("assets_info")
 
-    op.drop_index("ix_assets_backend_locator", table_name="assets")
     op.drop_index("ix_assets_mime_type", table_name="assets")
     op.drop_table("assets")

app/_assets_helpers.py

@@ -105,7 +105,10 @@ def resolve_destination_from_tags(tags: list[str]) -> tuple[str, list[str]]:
     if root == "models":
         if len(tags) < 2:
             raise ValueError("at least two tags required for model asset")
-        bases = folder_paths.folder_names_and_paths[tags[1]][0]
+        try:
+            bases = folder_paths.folder_names_and_paths[tags[1]][0]
+        except KeyError:
+            raise ValueError(f"unknown model category '{tags[1]}'")
         if not bases:
             raise ValueError(f"no base path configured for category '{tags[1]}'")
         base_dir = os.path.abspath(bases[0])

app/api/assets_routes.py

@@ -9,7 +9,7 @@
 import folder_paths
 
 from .. import assets_manager, assets_scanner
-from . import schemas_in
+from . import schemas_in, schemas_out
 
 
 ROUTES = web.RouteTableDef()
@@ -69,7 +69,7 @@ async def download_asset_content(request: web.Request) -> web.Response:
     except FileNotFoundError:
         return _error_response(404, "FILE_NOT_FOUND", "Underlying file not found on disk.")
 
-    quoted = filename.replace('"', "'")
+    quoted = (filename or "").replace("\r", "").replace("\n", "").replace('"', "'")
     cd = f'{disposition}; filename="{quoted}"; filename*=UTF-8\'\'{urllib.parse.quote(filename)}'
 
     resp = web.FileResponse(abs_path)
@@ -182,6 +182,8 @@ async def upload_asset(request: web.Request) -> web.Response:
             client_filename=file_client_name,
         )
         return web.json_response(created.model_dump(mode="json"), status=201)
+    except ValueError:
+        return _error_response(400, "BAD_REQUEST", "Invalid inputs.")
     except Exception:
         try:
             if os.path.exists(tmp_path):
@@ -341,6 +343,7 @@ async def get_asset_scan_status(request: web.Request) -> web.Response:
     states = assets_scanner.current_statuses()
     if root in {"models", "input", "output"}:
         states = [s for s in states.scans if s.root == root]  # type: ignore
+        states = schemas_out.AssetScanStatusResponse(scans=states)
     return web.json_response(states.model_dump(mode="json"), status=200)
 
 

app/assets_fetcher.py

@@ -0,0 +1,130 @@
+from __future__ import annotations
+import asyncio
+import os
+import tempfile
+from typing import Optional
+import aiohttp
+
+from .storage.hashing import blake3_hash_sync
+from .database.db import create_session
+from .database.services import ingest_fs_asset, get_cache_state_by_asset_hash
+from .resolvers import resolve_asset
+from ._assets_helpers import resolve_destination_from_tags, ensure_within_base
+
+_FETCH_LOCKS: dict[str, asyncio.Lock] = {}
+
+
+def _sanitize_filename(name: str) -> str:
+    return os.path.basename((name or "").strip()) or "file"
+
+
+async def ensure_asset_cached(
+    asset_hash: str,
+    *,
+    preferred_name: Optional[str] = None,
+    tags_hint: Optional[list[str]] = None,
+) -> str:
+    """
+    Ensure there is a verified local file for `asset_hash` in the correct Comfy folder.
+    Policy:
+      - Resolver must provide valid tags (root and, for models, category).
+      - If target path already exists:
+          * if hash matches -> reuse & ingest
+          * else -> remove and overwrite with the correct content
+    """
+    lock = _FETCH_LOCKS.setdefault(asset_hash, asyncio.Lock())
+    async with lock:
+        # 1) If we already have a state -> trust the path
+        async with await create_session() as sess:
+            state = await get_cache_state_by_asset_hash(sess, asset_hash=asset_hash)
+            if state and os.path.isfile(state.file_path):
+                return state.file_path
+
+        # 2) Resolve remote location + placement hints (must include valid tags)
+        res = await resolve_asset(asset_hash)
+        if not res:
+            raise FileNotFoundError(f"No resolver/locations for {asset_hash}")
+
+        placement_tags = tags_hint or res.tags
+        if not placement_tags:
+            raise ValueError(f"Resolver did not provide placement tags for {asset_hash}")
+
+        name_hint = res.filename or preferred_name or asset_hash.replace(":", "_")
+        safe_name = _sanitize_filename(name_hint)
+
+        # 3) Map tags -> destination (strict: raises if invalid root or models category)
+        base_dir, subdirs = resolve_destination_from_tags(placement_tags)  # may raise
+        dest_dir = os.path.join(base_dir, *subdirs) if subdirs else base_dir
+        os.makedirs(dest_dir, exist_ok=True)
+
+        final_path = os.path.abspath(os.path.join(dest_dir, safe_name))
+        ensure_within_base(final_path, base_dir)
+
+        # 4) If target path exists, try to reuse; else delete invalid cache
+        if os.path.exists(final_path) and os.path.isfile(final_path):
+            existing_digest = blake3_hash_sync(final_path)
+            if f"blake3:{existing_digest}" == asset_hash:
+                size_bytes = os.path.getsize(final_path)
+                mtime_ns = getattr(os.stat(final_path), "st_mtime_ns", int(os.path.getmtime(final_path) * 1_000_000_000))
+                async with await create_session() as sess:
+                    await ingest_fs_asset(
+                        sess,
+                        asset_hash=asset_hash,
+                        abs_path=final_path,
+                        size_bytes=size_bytes,
+                        mtime_ns=mtime_ns,
+                        mime_type=None,
+                        info_name=None,
+                        tags=(),
+                    )
+                    await sess.commit()
+                return final_path
+            else:
+                # Invalid cache: remove before re-downloading
+                os.remove(final_path)
+
+        # 5) Download to temp next to destination
+        timeout = aiohttp.ClientTimeout(total=60 * 30)
+        async with aiohttp.ClientSession(timeout=timeout) as session:
+            async with session.get(res.download_url, headers=dict(res.headers)) as resp:
+                resp.raise_for_status()
+                cl = resp.headers.get("Content-Length")
+                if res.expected_size and cl and int(cl) != int(res.expected_size):
+                    raise ValueError("server Content-Length does not match expected size")
+                with tempfile.NamedTemporaryFile("wb", delete=False, dir=dest_dir) as tmp:
+                    tmp_path = tmp.name
+                    async for chunk in resp.content.iter_chunked(8 * 1024 * 1024):
+                        if chunk:  # TODO: check do we need "await" here
+                            tmp.write(chunk)
+
+        # 6) Verify content hash
+        digest = blake3_hash_sync(tmp_path)
+        canonical = f"blake3:{digest}"
+        if canonical != asset_hash:
+            try:
+                os.remove(tmp_path)
+            finally:
+                raise ValueError(f"Hash mismatch: expected {asset_hash}, got {canonical}")
+
+        # 7) Atomically move into place (we already removed an invalid file if it existed)
+        if os.path.exists(final_path):
+            os.remove(final_path)
+        os.replace(tmp_path, final_path)
+
+        # 8) Record identity + cache state
+        size_bytes = os.path.getsize(final_path)
+        mtime_ns = getattr(os.stat(final_path), "st_mtime_ns", int(os.path.getmtime(final_path) * 1_000_000_000))
+        async with await create_session() as sess:
+            await ingest_fs_asset(
+                sess,
+                asset_hash=asset_hash,
+                abs_path=final_path,
+                size_bytes=size_bytes,
+                mtime_ns=mtime_ns,
+                mime_type=None,
+                info_name=None,
+                tags=(),
+            )
+            await sess.commit()
+
+        return final_path